diff --git a/.github/workflows/_build-wheel.yaml b/.github/workflows/_build-wheel.yaml new file mode 100644 index 0000000000..4cf5da693a --- /dev/null +++ b/.github/workflows/_build-wheel.yaml @@ -0,0 +1,184 @@ +name: _build-wheel + +# Shared build for one wheel variant (matrixed over Python version), used by the +# Release and Pre-Release workflows so both build identically. + +on: + workflow_call: + inputs: + runner: + type: string + required: true + container: + # x86 builds run in manylinux2_28 so the wheel's libstdc++/glibc floor + # stays low enough to import on RHEL/Rocky/Alma 8+. '' = bare runner. + type: string + default: '' + python-versions: + type: string + default: '["3.10", "3.11", "3.12", "3.13"]' + cuda: + # none | container | sbsa-12.8 | sbsa-13.0 + type: string + default: none + cmake-args: + # Space-separated -D flags. No semicolons; use ep-torch-versions for those. + type: string + required: true + cmake-generator: + type: string + default: '' + ep-torch-versions: + type: string + default: '' + build-with-ep: + type: string + default: '0' + variant-flag: + # build_wheel.sh variant set to 1, e.g. CU13_BUILD, NON_CUDA_BUILD + type: string + default: '' + torch-cuda-arch-list: + type: string + default: '' + build-nvlink-allocator: + type: boolean + default: false + artifact-prefix: + type: string + required: true + +env: + SCCACHE_GHA_ENABLED: "true" + +jobs: + build: + runs-on: ${{ inputs.runner }} + container: ${{ inputs.container }} + permissions: + contents: read + strategy: + matrix: + python-version: ${{ fromJSON(inputs.python-versions) }} + env: + BUILD_WITH_EP: ${{ inputs.build-with-ep }} + TORCH_CUDA_ARCH_LIST: ${{ inputs.torch-cuda-arch-list }} + CMAKE_ARGS: ${{ inputs.cmake-args }} + CMAKE_GEN: ${{ inputs.cmake-generator }} + EP_TORCH_VERSIONS_INPUT: ${{ inputs.ep-torch-versions }} + VARIANT_FLAG: ${{ inputs.variant-flag }} + steps: + - name: Checkout source + uses: actions/checkout@v4 + + - name: Mark workspace safe for git (container runs as root) + if: ${{ inputs.container != '' }} + run: git config --global --add safe.directory '*' + + - name: Set version from tag + run: echo "VERSION=${GITHUB_REF_NAME#v}" >> "$GITHUB_ENV" + + - name: Select Python ${{ matrix.python-version }} from manylinux image + if: ${{ inputs.container != '' }} + 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: Set up Python ${{ matrix.python-version }} + if: ${{ inputs.container == '' }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Free up disk space + if: ${{ inputs.container == '' }} + run: | + sudo rm -rf /usr/share/dotnet /opt/ghc /opt/hostedtoolcache/CodeQL /usr/local/lib/android + df -h + + - name: Install CUDA Toolkit (arm64 SBSA) + if: ${{ startsWith(inputs.cuda, 'sbsa-') }} + run: | + ver="${{ inputs.cuda }}"; ver="${ver#sbsa-}" + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb + sudo dpkg -i cuda-keyring_1.1-1_all.deb + sudo apt-get update + sudo apt-get install -y "cuda-toolkit-${ver/./-}" + echo "/usr/local/cuda/bin" >> "$GITHUB_PATH" + /usr/local/cuda/bin/nvcc --version + + - 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 + 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: ${{ inputs.build-nvlink-allocator }} + 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: 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/ci.yml b/.github/workflows/ci.yml index 431dbd5895..345f31638f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -182,6 +182,8 @@ jobs: 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 -lyaml-cpp" # 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 + # KV events publisher (optional; linked when libzmq is installed) + if ldconfig -p 2>/dev/null | grep -q libzmq; then export CGO_LDFLAGS="$CGO_LDFLAGS -lzmq"; fi ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/... kill $MASTER_PID 2>/dev/null || true shell: bash @@ -228,7 +230,7 @@ jobs: } echo "=== Processing coverage data ===" - lcov --remove coverage.info '/usr/*' '*/test/*' '*/third_party/*' --output-file coverage.filtered.info 2>&1 || true + 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..." @@ -729,26 +731,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 @@ -925,12 +907,15 @@ jobs: 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: @@ -948,6 +933,12 @@ jobs: - 'dependencies.sh' - 'scripts/**' - '.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] @@ -985,6 +976,107 @@ jobs: 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 + - name: cuda-off + cmake_flags: '-DUSE_CUDA=OFF' + need_cuda: false + 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 + 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 + 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 }} + 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() @@ -1000,6 +1092,7 @@ jobs: - test-wheel-ubuntu - build-wheel-cu13 - build-wheel-efa + - tent-ci - ascend-test - integration-test runs-on: ubuntu-latest diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml index 0d976f5624..4988dc3179 100644 --- a/.github/workflows/integration-test.yml +++ b/.github/workflows/integration-test.yml @@ -37,7 +37,7 @@ jobs: 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) + artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | select(.name | contains(\"mooncake\") ) | select(.name | contains(\"cu130\") ) | select(.workflow_run.head_sha == \"$SHA\" ) | .id" artifact.json | head -n 1) else echo "Failed to download artifact list. Retrying..." fi diff --git a/.github/workflows/pre-release.yaml b/.github/workflows/pre-release.yaml index 7561c21b97..279d48876a 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,484 +14,75 @@ 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.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 + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '1' + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: 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 + artifact-prefix: pre-release-cuda 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 -E 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 + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '0' + variant-flag: NON_CUDA_BUILD + cmake-args: >- + -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF + -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-non-cuda build-cuda13: - name: Build (CUDA 13) - 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" - 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.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 + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda13.0 + cuda: container + build-with-ep: '1' + variant-flag: CU13_BUILD + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: 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 + artifact-prefix: pre-release-cuda13 build-cuda-arm64: - name: Build (CUDA 12, arm64) - runs-on: ubuntu-22.04-arm - permissions: - contents: read - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - 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: Install CUDA Toolkit 12.8 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-12-8 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --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_arm64 - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export CUDA_HOME=/usr/local/cuda - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: pre-release-cuda-arm64-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-12.8 + cmake-generator: Ninja + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-cuda-arm64 build-cuda13-arm64: - name: Build (CUDA 13, arm64) - runs-on: ubuntu-22.04-arm - permissions: - contents: read - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - 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: Install CUDA Toolkit 13.0 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-13-0 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --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_arm64 - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - shell: bash - - - name: Build Python wheel - run: | - export CUDA_HOME=/usr/local/cuda - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} ./scripts/build_wheel.sh - env: - VERSION: ${{ env.VERSION }} - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: pre-release-cuda13-arm64-py${{ steps.generate_tag_arm64.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_arm64.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-13.0 + cmake-generator: Ninja + variant-flag: CU13_BUILD + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: pre-release-cuda13-arm64 validate-release: name: Validate release artifacts diff --git a/.github/workflows/release-cuda13.yaml b/.github/workflows/release-cuda13.yaml index 80d2f61192..929d9f9fa6 100644 --- a/.github/workflows/release-cuda13.yaml +++ b/.github/workflows/release-cuda13.yaml @@ -5,225 +5,45 @@ on: tags: - 'v*' -env: - SCCACHE_GHA_ENABLED: "true" jobs: build: - runs-on: ubuntu-22.04 - container: pytorch/manylinux2_28-builder:cuda13.0 - permissions: - contents: write - strategy: - 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: Mark workspace safe for git (container runs as root) - run: git config --global --add safe.directory '*' - - - 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: | - bash -x dependencies.sh -y - echo "/usr/local/go/bin" >> "$GITHUB_PATH" - 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.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release -DPython3_EXECUTABLE="$(which python3)" - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - cd build - make -j - 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 + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda13.0 + cuda: container + build-with-ep: '1' + variant-flag: CU13_BUILD + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: 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 + artifact-prefix: mooncake-wheel-cuda13 build-arm64: if: ${{ !contains(github.ref_name, '-') }} - runs-on: ubuntu-22.04-arm - permissions: - contents: write - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - CU13_BUILD: "1" - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - 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: Install CUDA Toolkit 13.0 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-13-0 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --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 CUDA_HOME=/usr/local/cuda - 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 - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-cuda13-arm64-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-13.0 + cmake-generator: Ninja + variant-flag: CU13_BUILD + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel-cuda13-arm64 publish-release: if: ${{ !contains(github.ref_name, '-') }} needs: [build, build-arm64] - 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* - - - 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-non-cuda.yaml b/.github/workflows/release-non-cuda.yaml index 1852717366..a3838b7937 100644 --- a/.github/workflows/release-non-cuda.yaml +++ b/.github/workflows/release-non-cuda.yaml @@ -5,120 +5,30 @@ 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 - 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 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 + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '0' + variant-flag: NON_CUDA_BUILD + cmake-args: >- + -DUSE_HTTP=ON -DUSE_ETCD=ON -DUSE_CUDA=OFF -DWITH_EP=OFF + -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel-non-cuda 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.yaml b/.github/workflows/release.yaml index d0da9e5552..007e4a5c72 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -5,223 +5,44 @@ 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 - container: pytorch/manylinux2_28-builder:cuda12.8 - permissions: - contents: write - 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: Mark workspace safe for git (container runs as root) - run: git config --global --add safe.directory '*' - - - 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: | - bash -x dependencies.sh -y - echo "/usr/local/go/bin" >> "$GITHUB_PATH" - 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.11.0;2.12.0;2.12.1" -DSTORE_USE_ETCD=ON -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release -DPython3_EXECUTABLE="$(which python3)" - shell: bash - - - name: Build project - run: | - export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH - cd build - make -j - 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 + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04 + container: pytorch/manylinux2_28-builder:cuda12.8 + cuda: container + build-with-ep: '1' + torch-cuda-arch-list: '8.0;9.0' + ep-torch-versions: '2.11.0;2.12.0;2.12.1' + build-nvlink-allocator: 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 + artifact-prefix: mooncake-wheel build-arm64: if: ${{ !contains(github.ref_name, '-') }} - runs-on: ubuntu-22.04-arm - permissions: - contents: write - strategy: - matrix: - python-version: ['3.10', '3.11', '3.12', '3.13'] - env: - TORCH_CUDA_ARCH_LIST: "9.0" - CUDA_HOME: "/usr/local/cuda" - 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: Install CUDA Toolkit 12.8 (arm64 SBSA) - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y cuda-toolkit-12-8 - echo "/usr/local/cuda/bin" >> $GITHUB_PATH - /usr/local/cuda/bin/nvcc --version - shell: bash - - - name: Run sccache-cache - uses: mozilla-actions/sccache-action@v0.0.9 - - - name: Configure sccache - uses: actions/github-script@v7 - with: - script: | - core.exportVariable('ACTIONS_RESULTS_URL', process.env.ACTIONS_RESULTS_URL || ''); - core.exportVariable('ACTIONS_RUNTIME_TOKEN', process.env.ACTIONS_RUNTIME_TOKEN || ''); - - - name: Configure project - run: | - sudo apt update -y - sudo bash -x dependencies.sh -y - mkdir build - cd build - cmake -G Ninja .. -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release - shell: bash - - - name: Build project - run: | - export CUDA_HOME=/usr/local/cuda - for dir in /usr/local/cuda/lib64/stubs /usr/local/cuda/targets/*/lib/stubs; do - if [ -d "$dir" ]; then - export LIBRARY_PATH="$dir:${LIBRARY_PATH:-}" - fi - done - cd build - cmake --build . - sudo cmake --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 CUDA_HOME=/usr/local/cuda - 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 - shell: bash - - - name: Upload Python wheel artifact - uses: actions/upload-artifact@v4 - with: - name: mooncake-wheel-arm64-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl + uses: ./.github/workflows/_build-wheel.yaml + with: + runner: ubuntu-22.04-arm + cuda: sbsa-12.8 + cmake-generator: Ninja + torch-cuda-arch-list: '9.0' + cmake-args: >- + -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF + -DWITH_STORE_RUST=OFF -DENABLE_SCCACHE=ON -DCMAKE_BUILD_TYPE=Release + artifact-prefix: mooncake-wheel-arm64 publish-release: needs: [build, build-arm64] - 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 4a3ec467a5..35a253a732 100644 --- a/.gitignore +++ b/.gitignore @@ -198,9 +198,6 @@ 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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3fefbf6af2..d6df05b60d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -47,7 +47,7 @@ repos: hooks: - id: codespell exclude: '^(extern/|FAST25-release/)' - args: ['--ignore-words-list=te,mooncake,KVCache,cann'] + args: ['--ignore-words-list=te,mooncake,KVCache,cann,hsa'] - repo: https://github.com/pre-commit/mirrors-clang-format rev: v20.1.8 diff --git a/.typos.toml b/.typos.toml index 320972c8fa..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,9 +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", + # 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/CONTRIBUTING.md b/CONTRIBUTING.md index 5ce88ac0f3..4116d4bf02 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. diff --git a/README.md b/README.md index b08dbfee34..d52af0e4fa 100644 --- a/README.md +++ b/README.md @@ -32,10 +32,11 @@ Under real workloads, Mooncake’s innovative architecture enables Kimi to handl

🔄 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.
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..68557d0aab 100644 --- a/benchmarks/storage_benchmark_v1/benchmark.py +++ b/benchmarks/storage_benchmark_v1/benchmark.py @@ -11,6 +11,8 @@ 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 @@ -103,7 +105,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 +121,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 +137,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 +166,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 +194,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 +427,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 +439,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 +478,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 +495,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 +570,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 +579,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 +676,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 +695,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 +711,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 fa490aadff..8b53be660a 100644 --- a/benchmarks/storage_benchmark_v1/doc/README.md +++ b/benchmarks/storage_benchmark_v1/doc/README.md @@ -28,12 +28,47 @@ python benchmark.py --scenario conversation \ | `--max-pages` | `2000` | Maximum number of pages (creates modulo mapping if trace is larger) | | `--fsync-mode` | `none` | When to fsync: `none`, `batch`, `always`, or `end` | | `--fsync-batch-size` | `100` | Number of writes between fsync in batch mode | +| `--threads` | `1` | Number of benchmark client worker threads | +| `--replay-scales` | `0` | Comma-separated trace fast-forward speeds; `0` means unpaced | +| `--progress-interval` | `100` | Print progress every N requests; `0` disables per-request progress | + +### Replay Scale + +Use `--replay-scales` to run the same trace at different fast-forward speeds: + +```bash +python benchmark.py --scenario toolagent \ + --trace-dir /path/to/Mooncake/FAST25-release/traces \ + --storage-dir /path/to/test/drive \ + --replay-scales 1,2,4,8 +``` + +For example, `2` means 2x fast-forward and `8` means 8x fast-forward. `0` +preserves the old unpaced behavior. + +### Client Threads + +Use `--threads` to add benchmark client worker threads: + +```bash +python benchmark.py --scenario toolagent \ + --trace-dir /path/to/Mooncake/FAST25-release/traces \ + --storage-dir /path/to/test/drive \ + --threads 4 +``` + +With `--threads > 1`, each benchmark client thread uses an independent storage +file under `thread_N/data.bin`, similar to running multiple clients at the same +time. Final results aggregate the per-thread counters and latency samples. For +strict single-client trace-order read/write and hit-rate accounting, use +`--threads 1`. ## Output Format ### Progress Output -During execution, 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) @@ -56,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 @@ -90,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..64089a6b44 100644 --- a/benchmarks/storage_benchmark_v1/storage/disk.py +++ b/benchmarks/storage_benchmark_v1/storage/disk.py @@ -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 bbbd82aa32..955d0494c0 100755 --- a/dependencies.sh +++ b/dependencies.sh @@ -175,6 +175,7 @@ if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then liburing-dev \ libjemalloc-dev \ libmsgpack-dev \ + libzmq3-dev \ libzstd-dev \ libasio-dev \ libxxhash-dev \ 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/source/deployment/kubernetes-deployment-guide/index.md b/docs/source/deployment/kubernetes-deployment-guide/index.md new file mode 100644 index 0000000000..c87126f436 --- /dev/null +++ b/docs/source/deployment/kubernetes-deployment-guide/index.md @@ -0,0 +1,45 @@ +# Kubernetes Deployment Guide + +Run Mooncake on Kubernetes as a shared **Store** cluster paired 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. + +--- + +## 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. + +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 +::: 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..ecd6fc5590 --- /dev/null +++ b/docs/source/deployment/kubernetes-deployment-guide/rbg-integration.md @@ -0,0 +1,280 @@ +# 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/v1alpha1`): 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](../../getting_started/examples/sglang-integration/hicache-quick-start.md). + +```{note} +The inline manifests below target `workloads.x-k8s.io/v1alpha1`. For the latest +API version `v1alpha2`, please refer to the upstream examples linked under +[RBG example](#rbg-example) above. +``` + +| 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/v1alpha1 +kind: RoleBasedGroup +metadata: + name: qwen3-0 + namespace: default + labels: { app.kubernetes.io/part-of: mooncake } +spec: + roles: + - name: master + replicas: 1 + 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 } + workload: { apiVersion: apps/v1, kind: StatefulSet } + + - name: store-1000gb + replicas: 3 + 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 } } + workload: { apiVersion: apps/v1, kind: StatefulSet } +``` + +### 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/v1alpha1 +kind: RoleBasedGroup +metadata: + name: sglang-workers-0 + namespace: default +spec: + roles: + - name: prefill + replicas: 4 + template: + metadata: + labels: + app: sglang-worker + rolebasedgroup.workloads.x-k8s.io/name: sglang-workers-0 + rolebasedgroup.workloads.x-k8s.io/role: 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 } } + workload: { apiVersion: apps/v1, kind: StatefulSet } + + - 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 + template: { } # fill in a real PodSpec to deploy + workload: { apiVersion: apps/v1, kind: StatefulSet } +``` + +### 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/mooncake-store-deployment-guide.md b/docs/source/deployment/mooncake-store-deployment-guide.md index 5333d7d3b2..653c3b373d 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -199,11 +199,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 @@ -278,7 +279,17 @@ mooncake_master \ --tenant_quota_connector_uri=/etc/mooncake/tenant_quotas.yaml ``` -The v1 connector is a writable YAML file. The file 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: +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 @@ -475,8 +486,8 @@ mooncake_master \ | 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 | -| `--tenant_quota_connector_uri` | empty | Connector URI; for `file`, the writable YAML policy path | +| `--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 @@ -544,6 +555,8 @@ 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 @@ -560,9 +573,11 @@ 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} @@ -637,9 +652,11 @@ Arguments of `MooncakeDistributedStore.setup(...)`: | `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 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. +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) @@ -665,6 +682,9 @@ The store service CLI only accepts `--config`, `-D/--define`, `--port`, and `--m | `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} @@ -695,12 +715,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`) @@ -711,21 +735,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_*`) @@ -794,6 +852,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-offload.md b/docs/source/deployment/ssd-offload.md index 439f5f29b5..edf58f7008 100644 --- a/docs/source/deployment/ssd-offload.md +++ b/docs/source/deployment/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 @@ -92,7 +94,7 @@ 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` | diff --git a/docs/source/design/conductor/indexer-api-design.md b/docs/source/design/conductor/indexer-api-design.md index e79367e6cd..c44af602fc 100644 --- a/docs/source/design/conductor/indexer-api-design.md +++ b/docs/source/design/conductor/indexer-api-design.md @@ -360,3 +360,133 @@ 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_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/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/mooncake-store.md b/docs/source/design/mooncake-store.md index 4ebe9751af..bb34a1ad34 100644 --- a/docs/source/design/mooncake-store.md +++ b/docs/source/design/mooncake-store.md @@ -95,7 +95,7 @@ To reduce cache warm-up time after a master restart, the Master Service supports 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. The v1 connector is a writable YAML file configured by `tenant_quota_connector_type=file` and `tenant_quota_connector_uri=`. 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`. +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`: @@ -494,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. @@ -588,7 +588,7 @@ Valid values are: `random` (default), `free_ratio_first`, `ssd_free_ratio_first` **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/allocation-strategy-benchmark-result.md). +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 @@ -731,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. diff --git a/docs/source/design/ssd-offload.md b/docs/source/design/ssd-offload.md index d9430f5dd1..c55d1cd7e0 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). --- diff --git a/docs/source/design/tent/cpp-api.md b/docs/source/design/tent/cpp-api.md index fc68a22c62..3acead95ea 100644 --- a/docs/source/design/tent/cpp-api.md +++ b/docs/source/design/tent/cpp-api.md @@ -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 | @@ -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/tent/qos.md b/docs/source/design/tent/qos.md index 06ecfc27db..a468b4145a 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: diff --git a/docs/source/design/tent/tebench.md b/docs/source/design/tent/tebench.md index c25745b5bd..f3d061d276 100644 --- a/docs/source/design/tent/tebench.md +++ b/docs/source/design/tent/tebench.md @@ -91,6 +91,68 @@ 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. + ## 5. Runtime Configuration This section summarizes the key runtime options that control workload behavior, @@ -187,8 +249,22 @@ gpu_id + thread_id **Transport (TENT only)** * `--xport_type` : `rdma | shm | mnnvl | gds | iouring` +* `--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 +* `--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. 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/index.md b/docs/source/design/transfer-engine/index.md index f089747be9..d4df0e2106 100644 --- a/docs/source/design/transfer-engine/index.md +++ b/docs/source/design/transfer-engine/index.md @@ -471,9 +471,12 @@ For advanced users, TransferEngine provides the following advanced runtime optio - `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 @@ -482,6 +485,7 @@ 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_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. @@ -495,6 +499,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 796dbaee7a..12b9d22564 100644 --- a/docs/source/getting_started/build.md +++ b/docs/source/getting_started/build.md @@ -121,6 +121,20 @@ 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 passed to `cmake ..`. diff --git a/docs/source/getting_started/examples/sglang-integration-v1.md b/docs/source/getting_started/examples/sglang-integration-v1.md index 4338001dda..a169558bec 100644 --- a/docs/source/getting_started/examples/sglang-integration-v1.md +++ b/docs/source/getting_started/examples/sglang-integration-v1.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/sglang-integration/hicache-integration-v1.md b/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md index 4a6e011cb0..6e9383062b 100644 --- a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md +++ b/docs/source/getting_started/examples/sglang-integration/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):** diff --git a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md b/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md index a338ce2903..a4004f280f 100644 --- a/docs/source/getting_started/examples/sglang-integration/hicache-quick-start.md +++ b/docs/source/getting_started/examples/sglang-integration/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/getting_started/examples/sglang-integration/index.md index 885c444f2d..b927695bc7 100644 --- a/docs/source/getting_started/examples/sglang-integration/index.md +++ b/docs/source/getting_started/examples/sglang-integration/index.md @@ -17,7 +17,7 @@ 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. -**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. --- diff --git a/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md b/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md index c8f0ab246d..49b81849f8 100644 --- a/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md +++ b/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md @@ -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 diff --git a/docs/source/getting_started/examples/vllm-integration/index.md b/docs/source/getting_started/examples/vllm-integration/index.md index d6107e502e..3c6d443c4f 100644 --- a/docs/source/getting_started/examples/vllm-integration/index.md +++ b/docs/source/getting_started/examples/vllm-integration/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/getting_started/examples/vllm-integration/kv-cache-storage.md index c25427adef..c4d9128a6b 100644 --- a/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md +++ b/docs/source/getting_started/examples/vllm-integration/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). --- @@ -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/getting_started/examples/vllm-integration/vllm-integration-v0.2.md index c9421957fe..e4f41ce372 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md +++ b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v0.2.md @@ -10,7 +10,7 @@ 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. diff --git a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md index c1b2c31af7..86a6aedd5f 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md +++ b/docs/source/getting_started/examples/vllm-integration/vllm-integration-v1.0.md @@ -54,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 @@ -127,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/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md index bb0a5f52fc..8517cb5de8 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md +++ b/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md @@ -137,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/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/supported-protocols.md b/docs/source/getting_started/supported-protocols.md index f98edb2fa8..d74ad0d4a3 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) diff --git a/docs/source/http-api-reference/http-service.md b/docs/source/http-api-reference/http-service.md index cbd79e007d..8c8ee90fdd 100644 --- a/docs/source/http-api-reference/http-service.md +++ b/docs/source/http-api-reference/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/index.md b/docs/source/index.md index 288b864f1c..94cc70ad9b 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -99,6 +99,7 @@ getting_started/quick-start :maxdepth: 1 deployment/mooncake-store-deployment-guide +deployment/kubernetes-deployment-guide/index getting_started/examples/sglang-integration/index getting_started/examples/vllm-integration/index Mooncake x LMCache Integration @@ -111,10 +112,9 @@ Mooncake x LMDeploy Integration95% 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/mooncake-common/FindUrma.cmake b/mooncake-common/FindUrma.cmake index 0af8d1a7ca..1a069ddc41 100644 --- a/mooncake-common/FindUrma.cmake +++ b/mooncake-common/FindUrma.cmake @@ -1,20 +1,30 @@ include(FetchContent) -# UMDK 头文件库 -FetchContent_Declare( - urma - GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git - GIT_TAG v25.12.0.B081 -) +# Allow callers to supply headers without downloading UMDK, e.g.: cmake +# -DURMA_INCLUDE_DIR=/usr/include ... cmake +# -DFETCHCONTENT_SOURCE_DIR_URMA=/path/to/umdk ... +if(DEFINED URMA_INCLUDE_DIR AND URMA_INCLUDE_DIR) + set(urma_INCLUDE_DIR ${URMA_INCLUDE_DIR}) + message(STATUS "Using provided URMA_INCLUDE_DIR: ${urma_INCLUDE_DIR}") +elseif(DEFINED FETCHCONTENT_SOURCE_DIR_URMA AND FETCHCONTENT_SOURCE_DIR_URMA) + set(urma_SOURCE_DIR ${FETCHCONTENT_SOURCE_DIR_URMA}) + set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include) + message(STATUS "Using FETCHCONTENT_SOURCE_DIR_URMA: ${urma_SOURCE_DIR}") +else() + FetchContent_Declare( + urma + GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git + GIT_TAG v25.12.0.B081 + GIT_SHALLOW TRUE) -FetchContent_MakeAvailable(urma) + FetchContent_GetProperties(urma) + if(NOT urma_POPULATED) + FetchContent_Populate(urma) + endif() -# 输出实际路径，确认位置 -message(STATUS "URMA source dir: ${urma_SOURCE_DIR}") -message(STATUS "URMA binary dir: ${urma_BINARY_DIR}") + set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include) + message(STATUS "URMA source dir: ${urma_SOURCE_DIR}") + message(STATUS "URMA binary dir: ${urma_BINARY_DIR}") +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 +message(STATUS "urma_INCLUDE_DIR: ${urma_INCLUDE_DIR}") diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index 2426a6ebe1..6af779a249 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -91,6 +91,9 @@ 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) +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) if(USE_UB) add_compile_definitions(USE_UB) @@ -203,6 +206,20 @@ if(USE_CUDA) link_directories(/usr/local/cuda/lib /usr/local/cuda/lib64) 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 diff --git a/mooncake-common/include/environ.h b/mooncake-common/include/environ.h index 78f8f45028..ec860500ae 100644 --- a/mooncake-common/include/environ.h +++ b/mooncake-common/include/environ.h @@ -52,8 +52,30 @@ class Environ { bool GetWithNvidiaPeermem() const { return with_nvidia_peermem_; } int GetEfaCqThreads() const { return efa_cq_threads_; } + // 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_; } + // 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); // 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") @@ -103,6 +125,19 @@ class Environ { bool path_roundrobin_; bool with_nvidia_peermem_; int efa_cq_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/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/src/environ.cpp b/mooncake-common/src/environ.cpp index 46aa1e9a9c..a920414fd9 100644 --- a/mooncake-common/src/environ.cpp +++ b/mooncake-common/src/environ.cpp @@ -31,6 +31,23 @@ int Environ::GetInt(const char* name, int default_value) { return default_value; } +int64_t Environ::GetInt64(const char* name, int64_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) { + std::cerr << "[Mooncake] Warning: invalid value '" << val + << "' for env " << name << ", using default " + << default_value << std::endl; + return default_value; + } + return static_cast(result); + } + return default_value; +} + size_t Environ::GetSizeT(const char* name, size_t default_value) { const char* val = std::getenv(name); if (val) { @@ -108,6 +125,26 @@ Environ::Environ() { path_roundrobin_ = GetBool("MC_PATH_ROUNDROBIN", false); with_nvidia_peermem_ = GetBool("WITH_NVIDIA_PEERMEM", true); efa_cq_threads_ = GetInt("MC_EFA_CQ_THREADS", 1); + + // AWS / S3 client configuration (consumed by s3_helper.cpp) + aws_region_ = GetString("MOONCAKE_AWS_REGION", ""); + aws_s3_endpoint_ = GetString("MOONCAKE_AWS_S3_ENDPOINT", ""); + aws_bucket_name_ = GetString("MOONCAKE_AWS_BUCKET_NAME", ""); + aws_access_key_id_ = GetString("MOONCAKE_AWS_ACCESS_KEY_ID", ""); + aws_secret_access_key_ = GetString("MOONCAKE_AWS_SECRET_ACCESS_KEY", ""); + aws_use_virtual_addressing_ = + GetBool("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING", true); + aws_use_https_ = GetBool("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_ = + GetString("MOONCAKE_AWS_REQUEST_CHECKSUM_CALCULATION", ""); + aws_response_checksum_validation_ = + GetString("MOONCAKE_AWS_RESPONSE_CHECKSUM_VALIDATION", ""); + aws_connect_timeout_ms_ = + GetInt64("MOONCAKE_AWS_CONNECT_TIMEOUT_MS", 10000); + aws_request_timeout_ms_ = + GetInt64("MOONCAKE_AWS_REQUEST_TIMEOUT_MS", 30000); } } // namespace mooncake diff --git a/mooncake-common/tests/environ_test.cpp b/mooncake-common/tests/environ_test.cpp index 9a0b534063..b859d59359 100644 --- a/mooncake-common/tests/environ_test.cpp +++ b/mooncake-common/tests/environ_test.cpp @@ -28,9 +28,22 @@ class EnvironTest : public ::testing::Test { void clearTestEnvVars() { unsetenv("MC_TEST_INT"); + unsetenv("MC_TEST_INT64"); 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"); } }; @@ -85,6 +98,70 @@ TEST_F(EnvironTest, GetIntMinValue) { EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), INT_MIN); } +// --- 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); +} + +// --- 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); + + 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); +} + // --- GetSizeT --- TEST_F(EnvironTest, GetSizeTValidValue) { 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/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_device.h b/mooncake-ep/include/mooncake_ep_device.h index 9322aea49b..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) diff --git a/mooncake-ep/include/mooncake_ep_exception.cuh b/mooncake-ep/include/mooncake_ep_exception.cuh index 097421f3a2..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) \ diff --git a/mooncake-ep/setup.py b/mooncake-ep/setup.py index 955af759bd..79a1625167 100644 --- a/mooncake-ep/setup.py +++ b/mooncake-ep/setup.py @@ -62,7 +62,10 @@ def existing_dirs(*paths): 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 = [ @@ -113,7 +116,9 @@ def existing_dirs(*paths): 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 361df2f96f..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,12 +14,57 @@ 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("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); py::class_(m, "Buffer") @@ -39,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 79aa4fbe7d..b587e29103 100644 --- a/mooncake-ep/src/mooncake_ep_buffer.cpp +++ b/mooncake-ep/src/mooncake_ep_buffer.cpp @@ -1,11 +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, @@ -40,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, @@ -209,6 +236,9 @@ 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_SPLIT_SEND_RECV @@ -247,7 +277,7 @@ 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); @@ -367,6 +397,9 @@ 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_SPLIT_SEND_RECV @@ -405,7 +438,7 @@ 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); 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 0625490b99..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(); @@ -283,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 @@ -352,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); } @@ -465,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 @@ -502,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); @@ -526,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); @@ -610,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 @@ -621,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); } @@ -722,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; @@ -751,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-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 735015a341..c7fe22818b 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -214,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 ec97258201..cf03539601 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -2117,7 +2117,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; @@ -2129,14 +2131,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) { @@ -2167,7 +2172,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, diff --git a/mooncake-p2p-store/src/p2pstore/go.mod b/mooncake-p2p-store/src/p2pstore/go.mod index 5a60aa938a..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.24.0 - -toolchain go1.24.1 +go 1.25.0 require go.etcd.io/etcd/client/v3 v3.5.15 @@ -16,9 +14,9 @@ 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.48.0 // indirect - golang.org/x/sys v0.39.0 // indirect - golang.org/x/text v0.32.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 diff --git a/mooncake-store/CMakeLists.txt b/mooncake-store/CMakeLists.txt index 5cf0f8a4ea..d03d82c67b 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}) diff --git a/mooncake-store/benchmarks/CMakeLists.txt b/mooncake-store/benchmarks/CMakeLists.txt index 76d23cda09..98769bd937 100644 --- a/mooncake-store/benchmarks/CMakeLists.txt +++ b/mooncake-store/benchmarks/CMakeLists.txt @@ -39,3 +39,9 @@ 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) diff --git a/mooncake-store/benchmarks/allocation_strategy_bench.cpp b/mooncake-store/benchmarks/allocation_strategy_bench.cpp index 10e45fe7f7..adf8facb6f 100644 --- a/mooncake-store/benchmarks/allocation_strategy_bench.cpp +++ b/mooncake-store/benchmarks/allocation_strategy_bench.cpp @@ -15,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" 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..463fe959db --- /dev/null +++ b/mooncake-store/benchmarks/batch_evict_bench.cpp @@ -0,0 +1,463 @@ +#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* kTenantId = "default"; + 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 != kTenantId) { + 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, kTenantId, + 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, kTenantId, 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/conf/master.json b/mooncake-store/conf/master.json index a376e72fcb..75e64d8c09 100644 --- a/mooncake-store/conf/master.json +++ b/mooncake-store/conf/master.json @@ -2,12 +2,12 @@ "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, @@ -17,8 +17,16 @@ "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 2b772e3ab8..8f43dace79 100644 --- a/mooncake-store/conf/master.yaml +++ b/mooncake-store/conf/master.yaml @@ -1,16 +1,18 @@ 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 diff --git a/mooncake-store/go/build.sh b/mooncake-store/go/build.sh index ac190e8b1a..7a751c9bc2 100755 --- a/mooncake-store/go/build.sh +++ b/mooncake-store/go/build.sh @@ -50,6 +50,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/allocator.h b/mooncake-store/include/allocator.h index e28b7d5c54..de74e7ffea 100644 --- a/mooncake-store/include/allocator.h +++ b/mooncake-store/include/allocator.h @@ -7,7 +7,7 @@ #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; 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..1ba40614b4 100644 --- a/mooncake-store/include/client_metric.h +++ b/mooncake-store/include/client_metric.h @@ -23,11 +23,12 @@ 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}; + 50000, 100000, 200000, 500000, 1000000, 2000000, 5000000, 10000000, + 20000000}; static inline std::string get_env_or_default( const char* env_var, const std::string& default_val = "") { diff --git a/mooncake-store/include/dummy_client.h b/mooncake-store/include/dummy_client.h index d538708828..bf808abbf1 100644 --- a/mooncake-store/include/dummy_client.h +++ b/mooncake-store/include/dummy_client.h @@ -30,7 +30,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; }; diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index 6aa2ac3d56..4a2f996310 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" 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/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..9f255975d3 --- /dev/null +++ b/mooncake-store/include/kv_event/kv_event_config.h @@ -0,0 +1,35 @@ +#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..39fc4af146 --- /dev/null +++ b/mooncake-store/include/kv_event/kv_event_publisher.h @@ -0,0 +1,124 @@ +#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" + +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. + // tenant_id empty defaults to "default" on the wire. + void PublishStored(const std::string& object_key, const std::string& medium, + const std::string& tenant_id = "", + const std::string& group_id = ""); + void PublishRemoved(const std::string& object_key, + const std::string& medium, + const std::string& tenant_id = "", + 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; + std::string 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 std::string& = "", const std::string& = "") {} + void PublishRemoved(const std::string&, const std::string&, + const std::string& = "", 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 1898cb43ef..5ced911819 100644 --- a/mooncake-store/include/master_admin_service.h +++ b/mooncake-store/include/master_admin_service.h @@ -93,6 +93,8 @@ 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, diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index 0275c57e0f..fd893953f8 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -132,6 +132,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 { @@ -211,6 +225,18 @@ class MasterServiceSupervisorConfig { 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; @@ -253,6 +279,18 @@ class MasterServiceSupervisorConfig { 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); @@ -403,6 +441,18 @@ class WrappedMasterServiceConfig { 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; std::string cluster_id = DEFAULT_CLUSTER_ID; @@ -476,6 +526,18 @@ class WrappedMasterServiceConfig { 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, @@ -573,6 +635,18 @@ class WrappedMasterServiceConfig { 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, @@ -983,6 +1057,18 @@ class MasterServiceConfig { 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; std::string cluster_id = DEFAULT_CLUSTER_ID; @@ -1052,6 +1138,18 @@ class MasterServiceConfig { 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; cluster_id = config.cluster_id; diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index 1479c96fdc..0b2769a194 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -36,11 +36,20 @@ #include "ha/ha_types.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "task_manager.h" +#include "kv_event/kv_event_publisher.h" namespace mooncake { + +// Forward declaration for MasterSnapshotManager +class MasterSnapshotManager; +class MasterSnapshotRepository; + namespace ha { class SnapshotCatalogStore; -} +class MasterSnapshotCodec; +struct MasterSnapshotPayloads; +class MasterSnapshotCodecTest; // test fixture, needs private state access +} // namespace ha class EtcdOpLogStore; @@ -62,6 +71,9 @@ class SnapshotChildProcessTest; class PromotionOnHitTest; class MasterServiceTenantQuotaTest; } // namespace test +namespace benchmarks { +class BatchEvictBench; +} // namespace benchmarks /* * @brief MasterService is the main class for the master server. @@ -82,7 +94,13 @@ class MasterService { friend class test::MasterServiceSnapshotTestBase; friend class test::SnapshotChildProcessTest; friend class test::PromotionOnHitTest; + friend class benchmarks::BatchEvictBench; friend class test::MasterServiceTenantQuotaTest; + 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 public: using NoFProbeFn = @@ -284,6 +302,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. @@ -322,6 +343,15 @@ class MasterService { auto GetReplicaList(const std::string& key, const std::string& 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 std::string& tenant_id) + -> tl::expected; + /** * @brief Get replica lists for a batch of objects. */ @@ -329,6 +359,15 @@ class MasterService { BatchGetReplicaList(const std::vector& keys, const std::string& 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 std::string& tenant_id); + /** * @brief Start a put operation for an object * @param[out] replica_list Vector to store replica information for the @@ -769,44 +808,21 @@ class MasterService { void setHttpMetadataRemoteUrl(const std::string& metadata_connstring); 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); - 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 @@ -1553,10 +1569,9 @@ class MasterService { static constexpr uint64_t kEvictionThreadSleepMs = 10; // 10 ms sleep between eviction checks - std::thread snapshot_thread_; - std::atomic snapshot_running_{false}; - std::mutex snapshot_thread_mutex_; - std::condition_variable snapshot_thread_cv_; + // 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}; @@ -1999,11 +2014,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_; @@ -2108,6 +2121,26 @@ 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 std::string& tenant_id); + void PublishKvRemoved(const std::string& key, + const ObjectMetadata& metadata, + const std::string& tenant_id); + void PublishKvRemoved(const std::string& key, const std::string& medium, + const std::string& 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 std::string& tenant_id); }; } // 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..434a3e6eb6 --- /dev/null +++ b/mooncake-store/include/master_snapshot_manager.h @@ -0,0 +1,133 @@ +#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 + +#ifdef STORE_USE_ETCD +class EtcdOpLogStore; +#endif + +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; + +#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); + + 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_; + +#ifdef STORE_USE_ETCD + mutable std::mutex snapshot_boundary_oplog_store_mutex_; + mutable std::unique_ptr snapshot_boundary_oplog_store_; +#endif + + 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/offset_allocator/offset_allocator.hpp b/mooncake-store/include/offset_allocator/offset_allocator.h similarity index 99% rename from mooncake-store/include/offset_allocator/offset_allocator.hpp rename to mooncake-store/include/offset_allocator/offset_allocator.h index 20052ee1ca..91da68c281 100644 --- a/mooncake-store/include/offset_allocator/offset_allocator.hpp +++ b/mooncake-store/include/offset_allocator/offset_allocator.h @@ -8,7 +8,7 @@ #include #include "mutex.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" namespace mooncake::offset_allocator { typedef unsigned char uint8; diff --git a/mooncake-store/include/pyclient.h b/mooncake-store/include/pyclient.h index c2e2201c5d..b24b5bcebe 100644 --- a/mooncake-store/include/pyclient.h +++ b/mooncake-store/include/pyclient.h @@ -11,7 +11,7 @@ #include #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "mutex.h" #include "utils.h" #include "file_storage.h" @@ -225,7 +225,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, diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index feaa5c72ce..6b3382b3b6 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -14,7 +14,7 @@ #include "pyclient.h" #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "mutex.h" #include "utils.h" #include "rpc_types.h" @@ -89,7 +89,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, @@ -511,7 +513,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); @@ -909,7 +913,7 @@ class RealClient : public PyClient { int stop_ipc_server(); // 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(UdsConnection &connection); 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 bf93f386f4..8e2cfa4895 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -12,6 +12,7 @@ #include "types.h" #include "rpc_types.h" #include "master_config.h" +#include "kv_event/kv_event_publisher.h" #include "segment.h" namespace mooncake { @@ -65,6 +66,15 @@ 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, @@ -290,6 +300,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/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/storage_backend.h b/mooncake-store/include/storage_backend.h index 5fd071b59b..952c751bc8 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -13,7 +13,7 @@ #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 +200,44 @@ 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) +}; + +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(); +}; + struct FileStorageConfig { // type of the storage backend StorageBackendType storage_backend_type = StorageBackendType::kBucket; @@ -881,6 +919,21 @@ class BucketStorageBackend : public StorageBackendInterface { */ 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. @@ -992,7 +1045,8 @@ 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 = {}); /** * @brief Initializes the offset allocator storage backend. @@ -1062,6 +1116,15 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { test_failure_predicate_ = std::move(predicate); } + // 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); + } + private: // On-disk record header: [u32 key_len][u32 value_len] (8 bytes total) struct RecordHeader { @@ -1132,12 +1195,19 @@ 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) {} }; // Returns full path to data file: {storage_path_}/kv_cache.data @@ -1201,6 +1271,39 @@ 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). + void EvictToMakeRoom(int64_t required_bytes, size_t min_victims, + const std::unordered_set& batch_keys, + std::vector& out_evicted); + // Test-only: Predicate to determine which keys should fail in BatchOffload. // Used for deterministic testing of partial success behavior. std::function test_failure_predicate_; @@ -1209,4 +1312,4 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { tl::expected, ErrorCode> CreateStorageBackend(const FileStorageConfig& config); -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/include/tenant_quota_policy_store.h b/mooncake-store/include/tenant_quota_policy_store.h index 30f3d85d41..ad6d2ae9c0 100644 --- a/mooncake-store/include/tenant_quota_policy_store.h +++ b/mooncake-store/include/tenant_quota_policy_store.h @@ -45,7 +45,24 @@ class YamlTenantQuotaPolicyStore final : public TenantQuotaPolicyStore { 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); +CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri, + const std::string& cluster_id); } // namespace mooncake diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index bee73e2b1e..527b0b43b2 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -83,14 +83,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,12 +216,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"; +static constexpr int DEFAULT_CLIENT_HTTP_PORT = 9300; // Original: returns a new string (copies when tenant_id is non-empty). // Kept for backward compatibility with callers that need an owned string. diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 1d1b531d87..ecbdc3c32e 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -377,6 +377,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 +412,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 +447,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}) diff --git a/mooncake-store/rust/build.rs b/mooncake-store/rust/build.rs index abcd95b06c..f8225e6818 100644 --- a/mooncake-store/rust/build.rs +++ b/mooncake-store/rust/build.rs @@ -203,19 +203,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,8 +238,8 @@ 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"); @@ -241,7 +248,8 @@ fn main() { // ----------------------------------------------------------------------- // 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); @@ -355,6 +363,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}"); @@ -365,8 +374,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/src/store.rs b/mooncake-store/rust/src/store.rs index cb59439c7f..fdc0941d7e 100644 --- a/mooncake-store/rust/src/store.rs +++ b/mooncake-store/rust/src/store.rs @@ -57,7 +57,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 @@ -277,8 +281,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 +348,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 +453,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 +507,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 +528,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 +537,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 +661,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()); @@ -706,7 +700,9 @@ mod tests { #[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()); } @@ -718,10 +714,7 @@ 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(_)))); } #[test] @@ -732,9 +725,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 9caf538b7b..1ffff9ac3d 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 @@ -46,6 +48,7 @@ set(MOONCAKE_STORE_SOURCES ha/leadership/master_service_supervisor.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 @@ -114,6 +117,25 @@ 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}") @@ -246,6 +268,9 @@ 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() target_include_directories(mooncake_store PUBLIC ${XXHASH_INCLUDE_DIR}) if(USE_NOF) target_include_directories(mooncake_store PRIVATE ${SPDK_INCLUDE_DIR} diff --git a/mooncake-store/src/aligned_client_buffer.cpp b/mooncake-store/src/aligned_client_buffer.cpp index 55d83e5a1d..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 diff --git a/mooncake-store/src/client_buffer.cpp b/mooncake-store/src/client_buffer.cpp index db102c3b16..6aaea4af10 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 diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 7c078ad7d9..446503c91a 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -35,7 +35,7 @@ #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" diff --git a/mooncake-store/src/device/accelerator_registry.cpp b/mooncake-store/src/device/accelerator_registry.cpp index 1781fcadca..6375f4c2f0 100644 --- a/mooncake-store/src/device/accelerator_registry.cpp +++ b/mooncake-store/src/device/accelerator_registry.cpp @@ -7,6 +7,12 @@ 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); @@ -90,6 +96,10 @@ void RegisterStaticAcceleratorDevice(const AcceleratorDevice& 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(); } diff --git a/mooncake-store/src/device/cuda_like_accelerator_device.cpp b/mooncake-store/src/device/cuda_like_accelerator_device.cpp index 54e7dcf335..385b548542 100644 --- a/mooncake-store/src/device/cuda_like_accelerator_device.cpp +++ b/mooncake-store/src/device/cuda_like_accelerator_device.cpp @@ -8,6 +8,9 @@ namespace mooncake { namespace device { + +void EnsureCudaLikeAcceleratorDeviceLinked() {} + namespace { void FreeCudaLikePinnedHostBuffer(void* addr) { cudaFreeHost(addr); } diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 00e7db4bcd..05d7ef2e02 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -4,7 +4,7 @@ #include #include -#include "aligned_client_buffer.hpp" +#include "aligned_client_buffer.h" #include "storage_backend.h" #include "client_metric.h" #include "utils.h" @@ -446,8 +446,8 @@ tl::expected FileStorage::OffloadObjects( } std::unordered_map> user_batch_object; - auto query_result = BatchQuerySegmentSlices(user_keys, tenant_id, - user_batch_object); + [[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. 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..dc5c8f6fc1 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 { 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/kv_event/kv_event_publisher.cpp b/mooncake-store/src/kv_event/kv_event_publisher.cpp new file mode 100644 index 0000000000..a819a9bf0f --- /dev/null +++ b/mooncake-store/src/kv_event/kv_event_publisher.cpp @@ -0,0 +1,379 @@ +#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 std::string& 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 std::string& 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.empty() ? "default" : item.pending.tenant_id; + + 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 9174bd97ee..66daea1bbf 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -30,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 { @@ -114,7 +114,7 @@ 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"); @@ -225,6 +225,30 @@ 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, + "Deprecated; ignored (event queue is unbounded)"); DEFINE_string(ha_backend_type, "etcd", "HA backend type, e.g. etcd | redis | k8s"); DEFINE_string(ha_backend_connstring, "", @@ -482,6 +506,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", @@ -608,7 +667,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 @@ -624,7 +683,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."; @@ -795,6 +854,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 @@ -1254,6 +1377,9 @@ int main(int argc, char* argv[]) { << master_config.eviction_high_watermark_ratio << ", enable_ha=" << master_config.enable_ha << ", 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 diff --git a/mooncake-store/src/master_admin_service.cpp b/mooncake-store/src/master_admin_service.cpp index 0898dadb8b..798ca55346 100644 --- a/mooncake-store/src/master_admin_service.cpp +++ b/mooncake-store/src/master_admin_service.cpp @@ -562,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) { @@ -920,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}; @@ -935,63 +978,81 @@ 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; - } - item.values->emplace_back( - replica.get_memory_descriptor().buffer_descriptor); + if (results.size() != keys.size()) { + LOG(WARNING) << "BatchGetReplicaListForAdmin size mismatch: keys=" + << keys.size() << " results=" << results.size(); } - payload.data.emplace(keys[i], std::move(item)); - } - - if (results.size() != keys.size()) { - LOG(WARNING) << "BatchGetReplicaList size mismatch: keys=" - << keys.size() << " results=" << results.size(); - } - WriteJsonResponse(resp, coro_http::status_type::ok, payload); + WriteJsonResponse(resp, coro_http::status_type::ok, payload); + }); } void MasterAdminServer::HandleGetTenantQuotas( @@ -1120,6 +1181,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); diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index d9c57cd4e0..ac4a38bfa4 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -1937,7 +1937,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(); diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index c051fe1682..13beddbe37 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -38,31 +38,21 @@ #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 "utils.h" +#include "kv_event/kv_event_config.h" +#include "master_snapshot_manager.h" +#include "master_snapshot_repository.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 @@ -90,12 +80,6 @@ tl::expected ParseSnapshotCatalogKind( std::string(store_type)); } -int64_t CurrentTimeMs() { - return std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()) - .count(); -} - 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); @@ -234,12 +218,12 @@ 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), offloading_queue_limit_(config.offloading_queue_limit), - offload_cap_ratio_(config.offload_cap_ratio) { + 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) { @@ -269,11 +253,18 @@ 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_); + tenant_quota_connector_uri_, + cluster_id_); if (!store) { throw std::invalid_argument(store.error()); } @@ -408,6 +399,9 @@ MasterService::MasterService(const MasterServiceConfig& config) << ")"; } + kv_event_publisher_ = + std::make_unique(BuildKvEventConfig(config)); + eviction_running_ = true; eviction_thread_ = std::thread(&MasterService::EvictionThreadFunc, this); VLOG(1) << "action=start_eviction_thread"; @@ -453,9 +447,30 @@ MasterService::MasterService(const MasterServiceConfig& config) if (enable_snapshot_) { 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(); } } @@ -506,10 +521,12 @@ MasterService::~MasterService() { // Stop and join the threads eviction_running_ = false; client_monitor_running_ = false; - { - std::lock_guard lk(snapshot_thread_mutex_); - 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; @@ -519,7 +536,6 @@ MasterService::~MasterService() { #endif // Wake sleepers so join() doesn't block for long sleep intervals. - snapshot_thread_cv_.notify_all(); task_cleanup_cv_.notify_all(); http_metadata_cleanup_cv_.notify_all(); @@ -534,9 +550,6 @@ MasterService::~MasterService() { nof_heartbeat_thread_.join(); } #endif - if (snapshot_thread_.joinable()) { - snapshot_thread_.join(); - } if (task_cleanup_thread_.joinable()) { task_cleanup_thread_.join(); } @@ -546,6 +559,12 @@ MasterService::~MasterService() { 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(); + } } void MasterService::SetNoFProbeFnForTesting(NoFProbeFn fn) { @@ -2567,6 +2586,35 @@ auto MasterService::GetReplicaList(const std::string& key, return resp; } +auto MasterService::GetReplicaListForAdmin(const std::string& key, + const std::string& tenant_id) + -> tl::expected { + 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_); +} + std::vector> MasterService::BatchGetReplicaList(const std::vector& keys, const std::string& tenant_id) { @@ -2688,6 +2736,93 @@ MasterService::BatchGetReplicaList(const std::vector& keys, return results; } +std::vector> +MasterService::BatchGetReplicaListForAdmin(const std::vector& keys, + const std::string& tenant_id) { + using GetResult = tl::expected; + + std::vector results( + keys.size(), tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND)); + if (keys.empty()) { + return results; + } + + const auto normalized_tenant = NormalizeTenantId(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 = + MakeTenantScopedKey(normalized_tenant, 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_); + } + } + } + + return results; +} + auto MasterService::AllocateAndInsertMetadata( MetadataShardAccessorRW& shard, const UUID& client_id, const std::string& key, uint64_t value_length, @@ -3171,6 +3306,7 @@ 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, tenant_id); return {}; } @@ -3760,6 +3896,7 @@ auto MasterService::EvictDiskReplica(const UUID& client_id, } if (!metadata.IsValid()) { + PublishKvRemoved(key, metadata, tenant_id); accessor.Erase(); } return {}; @@ -4378,7 +4515,8 @@ 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(); + PublishKvRemoved(key, metadata, tenant_id); + auto& tenant_state [[maybe_unused]] = accessor.GetTenantState(); accessor.Erase(); return {}; } @@ -5751,686 +5889,6 @@ uint64_t MasterService::ReleaseExpiredDiscardedReplicas( return released_cnt; } -void MasterService::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(snapshot_interval_seconds_), - [&] { return !snapshot_running_.load(); }); - } - - if (!snapshot_running_) { - break; - } - - 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; - } - - 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; - } - - 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; - } - } - }; - - // 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 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)); - - // 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 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()); - } - - uint64_t sequence_id = 0; - auto err = oplog_store.value()->GetLatestSequenceId(sequence_id); - if (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 snapshot sequence boundary: {}", - toString(err)))); - } - - return static_cast(sequence_id); -#endif -} - -#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)))); - } - - 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)))); - } - - 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")); - } - - 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_); - - 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()); - } - } - - 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); - } - - 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; - - if (old_state_dir == snapshot_id) { - SNAP_LOG_WARN( - "[Snapshot] Skipping deletion of current snapshot " - "directory {}, " - "snapshot_id={}", - old_state_dir, snapshot_id); - continue; - } - - 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); - } - } - } -} - void MasterService::RestoreState() { auto* snapshot_catalog_store = GetSnapshotCatalogStore(); if (!snapshot_catalog_store) { @@ -6442,382 +5900,228 @@ void MasterService::RestoreState() { LOG(INFO) << "[Restore] Backend info: " << snapshot_object_store_->GetConnectionInfo(); - std::vector restore_candidates; - std::unordered_set candidate_ids; - std::optional latest_snapshot_id; - - auto latest_result = snapshot_catalog_store->GetLatest(); + // 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 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; - } - if (!candidate_ids.emplace(snapshot.snapshot_id).second) { - continue; - } - restore_candidates.push_back(snapshot); - } + latest_snapshot = latest_result.value(); } - if (restore_candidates.empty()) { + auto candidates_result = + snapshot_repository_->LoadRestoreCandidates(latest_snapshot); + if (!candidates_result || candidates_result->empty()) { LOG(ERROR) << "[Restore] No previous snapshot found, starting fresh"; return; } + // Phase 2 & 3: Try each candidate const auto now = std::chrono::system_clock::now(); - for (const auto& snapshot : restore_candidates) { + for (const auto& snapshot : candidates_result.value()) { ResetStateAfterFailedRestoreAttempt(); - if (TryRestoreStateFromSnapshot(snapshot, now)) { + + 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; + } + + // 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; + } + + // 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; + } + + 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; } } ResetStateAfterFailedRestoreAttempt(); LOG(ERROR) << "[Restore] Failed to restore from all candidate snapshots " - << "(count=" << restore_candidates.size() << "), starting fresh"; + << "(count=" << candidates_result->size() << "), starting fresh"; } -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 + "/"; - } +void MasterService::ResetStateAfterFailedRestoreAttempt() { + SegmentSerializer segment_serializer(&segment_manager_); + MetadataSerializer metadata_serializer(this); + TaskManagerSerializer task_manager_serializer(&task_manager_); - std::string manifest_path = snapshot.manifest_key; - if (manifest_path.empty()) { - manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; - } + task_manager_serializer.Reset(); + metadata_serializer.Reset(); + segment_serializer.Reset(); - auto fail_restore = [&](const std::string& message) { - LOG(WARNING) << "[Restore] Snapshot candidate " << state_id - << " is unusable: " << message; - ResetStateAfterFailedRestoreAttempt(); - return false; - }; + { + std::unique_lock lock(client_mutex_); + ok_client_.clear(); + } + PodUUID pod_uuid; + while (client_ping_queue_.pop(pod_uuid)) { + } - 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)); - } - LOG(INFO) << "[Restore] Deserialize metadata success"; + MasterMetricManager::instance().reset_allocated_mem_size(); + MasterMetricManager::instance().reset_total_mem_capacity(); + MasterMetricManager::instance().reset_cache_total_nums(); +} - 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)); - } - LOG(INFO) << "[Restore] Deserialize task manager success"; +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. - std::vector segment_names; - { - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - segment_access.GetAllSegmentNames(segment_names); - } + std::vector segment_names; + { + ScopedSegmentAccess segment_access = + segment_manager_.getSegmentAccess(); + segment_access.GetAllSegmentNames(segment_names); + } - { - 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); + // 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 { - ++tenant_it; + ++it; } } + if (tenant_state.Empty()) { + tenant_it = shard.tenants.erase(tenant_it); + } else { + ++tenant_it; + } } } + } - MasterMetricManager::instance().reset_allocated_mem_size(); - RebuildCacheTotalAccounting(); - for (auto& segment_name : segment_names) { - MasterMetricManager::instance() - .reset_segment_allocated_mem_size(segment_name); - } + // 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); + } - 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_)); + 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; } - ++it; + 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; } } - - 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); - } - - 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); - } - } + LOG(INFO) << "[Restore] Total allocated size after restore: " + << MasterMetricManager::instance().get_allocated_mem_size(); + } - std::vector> all_segments; - auto err = segment_access.GetAllSegments(all_segments); + // 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 (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; + 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); } } - LOG(INFO) << "[Restore] Successfully restored state from snapshot: " - << state_id; - 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"); - } -} - -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::vector> all_segments; + auto err = segment_access.GetAllSegments(all_segments); - { - std::unique_lock lock(client_mutex_); - ok_client_.clear(); - } - PodUUID pod_uuid; - while (client_ping_queue_.pop(pod_uuid)) { + 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; + } } - MasterMetricManager::instance().reset_allocated_mem_size(); - MasterMetricManager::instance().reset_total_mem_capacity(); - MasterMetricManager::instance().reset_cache_total_nums(); + return {}; } ha::SnapshotCatalogStore* MasterService::GetSnapshotCatalogStore() { @@ -7227,6 +6531,8 @@ void MasterService::BatchEvict(double evict_ratio_target, result.freed_bytes += freed; if (freed > 0) { result.evicted_objects++; + PublishKvRemovedAfterEvict(member_key, freed, "cpu", + member_metadata, tenant_id); } if (member_key != key && !member_metadata.IsValid()) { EraseMetadata(tenant_state, member_it, tenant_id, @@ -7384,6 +6690,10 @@ void MasterService::BatchEvict(double evict_ratio_target, deferred_replicas, /*allow_soft_pinned=*/false); total_freed_size += evict_result.freed_bytes; + if (!it->second.IsGrouped()) { + PublishKvRemovedAfterEvict(c.key, evict_result.freed_bytes, + "cpu", it->second, c.tenant_id); + } if (!it->second.IsValid()) { EraseMetadata(tenant_state, it, c.tenant_id, QuotaEraseMode::kFull, &shard); @@ -7444,6 +6754,11 @@ 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.IsGrouped()) { + PublishKvRemovedAfterEvict( + it->first, evict_result.freed_bytes, + "cpu", it->second, tenant_it->first); + } if (!it->second.IsValid()) { it = EraseMetadata( tenant_state, it, tenant_it->first, @@ -7504,6 +6819,11 @@ 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.IsGrouped()) { + PublishKvRemovedAfterEvict( + it->first, evict_result.freed_bytes, + "cpu", it->second, tenant_it->first); + } if (!it->second.IsValid()) { it = EraseMetadata( tenant_state, it, tenant_it->first, @@ -7663,6 +6983,8 @@ 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, QuotaEraseMode::kFull, &shard); @@ -8624,23 +7946,6 @@ 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::vector& targets) { @@ -9392,6 +8697,108 @@ 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 std::string& 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 std::string& 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 std::string& 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 std::string& 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) { diff --git a/mooncake-store/src/master_snapshot_manager.cpp b/mooncake-store/src/master_snapshot_manager.cpp new file mode 100644 index 0000000000..2116611c36 --- /dev/null +++ b/mooncake-store/src/master_snapshot_manager.cpp @@ -0,0 +1,680 @@ +#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 "serialize/serializer.h" +#include "segment.h" +#include "task_manager.h" +#include "utils/file_util.h" +#include "utils/zstd_util.h" + +#ifdef STORE_USE_ETCD +#include "ha/oplog/etcd_oplog_store.h" +#include "etcd_helper.h" +#endif + +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 || options_.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()); + } + + uint64_t sequence_id = 0; + auto err = oplog_store.value()->GetLatestSequenceId(sequence_id); + if (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 snapshot sequence boundary: {}", + toString(err)))); + } + + return static_cast(sequence_id); +#endif +} + +#ifdef STORE_USE_ETCD +tl::expected +MasterSnapshotManager::GetSnapshotBoundaryOpLogStore() const { + if (options_.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( + options_.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)))); + } + + auto oplog_store = std::make_unique(options_.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)))); + } + + snapshot_boundary_oplog_store_ = std::move(oplog_store); + return snapshot_boundary_oplog_store_.get(); +} +#endif + +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 {}; +} + +tl::expected +MasterSnapshotManager::UploadSnapshotPayloadFile( + const std::vector& data, const std::string& path, + const std::string& local_filename, const std::string& snapshot_id) { + return repository_->UploadPayloadFile(data, path, local_filename, + snapshot_id); +} + +void MasterSnapshotManager::CleanupOldSnapshot(size_t keep_count, + const std::string& snapshot_id) { + repository_->CleanupOldSnapshots(keep_count, snapshot_id); +} + +} // 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..dd4868514c --- /dev/null +++ b/mooncake-store/src/master_snapshot_repository.cpp @@ -0,0 +1,335 @@ +#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) { + 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) { + 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/offset_allocator.cpp b/mooncake-store/src/offset_allocator.cpp index 68b44025f6..60b2de3d9c 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 @@ -683,4 +683,4 @@ std::ostream& operator<<(std::ostream& os, return os; } -} // namespace mooncake::offset_allocator \ No newline at end of file +} // namespace mooncake::offset_allocator diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index dc544b7872..314d7cf511 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -12,13 +12,15 @@ #include // for atexit #include #include +#include #include #include #include #include #include "real_client.h" -#include "client_buffer.hpp" +#include "client_buffer.h" +#include "replica_selection.h" #include "common.h" #include "config.h" #include "mutex.h" @@ -277,47 +279,22 @@ inline tl::expected scatter_host_to_maybe_device( 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. @@ -626,7 +603,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 = @@ -816,6 +794,9 @@ 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; @@ -841,7 +822,8 @@ 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); } else { ptr = allocate_buffer_allocator_memory(segment_size, this->protocol); @@ -874,6 +856,19 @@ tl::expected RealClient::setup_internal( } else { segment_ptrs_.emplace_back(ptr); } + + // 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 mount_result = client_->MountSegment(ptr, mapped_size, protocol, seg_location); if (!mount_result.has_value()) { @@ -935,10 +930,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"; } } @@ -952,12 +957,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 { @@ -984,6 +990,60 @@ inline std::optional get_config_size(const ConfigDict &config, } return static_cast(parsed_size_opt.value()); } + +inline std::string trim(const std::string &value) { + auto start = value.find_first_not_of(" \t\r\n"); + if (start == std::string::npos) { + return ""; + } + auto end = value.find_last_not_of(" \t\r\n"); + return value.substr(start, end - start + 1); +} + +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; + } + + std::string value = trim(it->second); + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char c) { return std::tolower(c); }); + if (value == "true" || value == "1" || value == "yes" || value == "on" || + value == "enable") { + return true; + } + if (value == "false" || value == "0" || value == "no" || value == "off" || + value == "disable") { + return false; + } + + 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; + } + + std::string value = trim(it->second); + int parsed_value = 0; + const char *begin = value.data(); + const char *end = begin + value.size(); + auto [ptr, ec] = std::from_chars(begin, end, parsed_value); + if (ec != std::errc{} || ptr != end) { + LOG(ERROR) << "Invalid integer value for config key '" << key + << "': " << it->second; + return std::nullopt; + } + return parsed_value; +} } // namespace tl::expected RealClient::setup_internal( @@ -1025,7 +1085,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 @@ -1035,26 +1105,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)) { + 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( @@ -1582,11 +1657,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) { @@ -1652,11 +1731,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; } @@ -1691,7 +1770,7 @@ tl::expected RealClient::put_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - auto scatter_result = scatter_host_to_maybe_device( + 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()); @@ -1773,9 +1852,9 @@ tl::expected RealClient::put_batch_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - auto scatter_result = scatter_host_to_maybe_device( - buffer_handle.ptr(), value.data(), value.size_bytes(), - "put_batch:" + key); + 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()); } @@ -1882,7 +1961,7 @@ tl::expected RealClient::put_parts_internal( // Copy all parts into the contiguous buffer size_t offset = 0; for (const auto &value : values) { - auto scatter_result = scatter_host_to_maybe_device( + 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) { @@ -2613,7 +2692,7 @@ std::shared_ptr RealClient::get_buffer_internal( } // 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(); @@ -3200,9 +3279,42 @@ 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); + 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 get_result = client_->Get(key, filtered_qr, slices); @@ -3286,8 +3398,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) { @@ -3635,17 +3778,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; @@ -3685,15 +3818,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) { @@ -3740,7 +3865,7 @@ tl::expected RealClient::upsert_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - auto scatter_result = scatter_host_to_maybe_device( + 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()); @@ -3800,14 +3925,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) { @@ -3853,15 +3971,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); @@ -3990,7 +4101,7 @@ tl::expected RealClient::upsert_parts_internal( auto &buffer_handle = *alloc_result; size_t offset = 0; for (const auto &value : values) { - auto scatter_result = scatter_host_to_maybe_device( + 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) { @@ -4078,7 +4189,7 @@ tl::expected RealClient::upsert_batch_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - auto scatter_result = scatter_host_to_maybe_device( + auto scatter_result = gather_maybe_device_to_host( buffer_handle.ptr(), value.data(), value.size_bytes(), "upsert_batch:" + key); if (!scatter_result) { @@ -4397,7 +4508,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"; @@ -4666,8 +4777,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( @@ -4684,10 +4796,11 @@ RealClient::batch_get_into_internal(const std::vector &keys, } 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(); @@ -4738,25 +4851,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: " diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index 15c7202f33..8743eaa5e1 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 " @@ -102,6 +105,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 +113,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 +129,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/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index 493a496c23..d5de362f54 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -230,6 +230,21 @@ WrappedMasterService::BatchGetReplicaList(const std::vector& keys, return results; } +std::vector> +WrappedMasterService::BatchGetReplicaListForAdmin( + const std::vector& keys, const std::string& tenant_id) { + return master_service_.BatchGetReplicaListForAdmin(keys, tenant_id); +} + +tl::expected +WrappedMasterService::GetReplicaListForAdmin(const std::string& key, + const std::string& tenant_id) { + return execute_rpc( + "GetReplicaListForAdmin", + [&] { return master_service_.GetReplicaListForAdmin(key, 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, @@ -1243,6 +1258,14 @@ 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 RegisterRpcService( coro_rpc::coro_rpc_server& server, mooncake::WrappedMasterService& wrapped_master_service) { diff --git a/mooncake-store/src/serialize/serializer.cpp b/mooncake-store/src/serialize/serializer.cpp index 00b46d47d5..acfc5feb01 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" diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 3612624a3f..8ba979f918 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -86,6 +87,93 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() { return config; } +bool OffsetAllocatorBackendConfig::Validate() const { + 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) { + std::string s(pol); + if (s == "fifo" || s == "FIFO" || s == "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 = GetEnvOr( + "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 GetEnvOr. + auto max_evict_raw = + GetEnvOr("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; + } + + return cfg; +} + StorageBackendInterface::StorageBackendInterface( const FileStorageConfig& config) : file_storage_config_(config) {} @@ -1330,23 +1418,20 @@ tl::expected BucketStorageBackend::BatchOffload( LOG(ERROR) << "Failed to write bucket with id: " << bucket_id; 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_); // Pre-check for duplicates before modifying any state - for (const auto& key : bucket->keys) { + for (const auto& key : bucket_keys) { if (object_bucket_map_.find(key) != object_bucket_map_.end()) { LOG(WARNING) << "Duplicate key detected in BatchOffload: " << key @@ -1361,18 +1446,40 @@ tl::expected BucketStorageBackend::BatchOffload( // 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])}); + 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]}); if (!inserted) { LOG(ERROR) << "Unexpected duplicate key after pre-check: " - << bucket->keys[i] << ", bucket_id=" << bucket_id; + << bucket_keys[i] << ", bucket_id=" << bucket_id; } } buckets_.emplace(bucket_id, std::move(bucket)); lru_index_.emplace(0LL, bucket_id); } + // 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); + } + } return bucket_id; } @@ -2153,6 +2260,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(); @@ -2177,6 +2292,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). @@ -2716,9 +2922,11 @@ 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; } @@ -2812,16 +3020,128 @@ tl::expected OffsetAllocatorStorageBackend::Init() { fd_guard.release()); } - // Create allocator with base=0, size=capacity - allocator_ = offset_allocator::OffsetAllocator::create(0, capacity_); + // Resolve watermark thresholds from config + 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 (e.g. limit=5, ratio + // 0.95->4, ratio 0.90->4 => low==high==4). Only applies to + // auto-derived values; explicit config values are validated strictly. + 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); + } + + // 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); + } + 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(); @@ -2831,6 +3151,80 @@ tl::expected OffsetAllocatorStorageBackend::Init() { return {}; } +//----------------------------------------------------------------------------- +// EvictToMakeRoom +//----------------------------------------------------------------------------- + +void OffsetAllocatorStorageBackend::EvictToMakeRoom( + int64_t required_bytes, size_t min_victims, + const std::unordered_set& batch_keys, + std::vector& out_evicted) { + 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); + ++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); + 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); + } + fifo_index_.erase(oldest); + out_evicted.push_back(std::move(vkey)); + ++n; + } +} + //----------------------------------------------------------------------------- tl::expected OffsetAllocatorStorageBackend::BatchOffload( @@ -2838,8 +3232,21 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( std::function& keys, std::vector& metadatas)> complete_handler, - std::function& /*evicted_keys*/)> - /*eviction_handler*/) { + std::function& evicted_keys)> + 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 (check + // watermark → evict → update counters) 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."; @@ -2858,33 +3265,47 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT); } + const bool eviction_on = + (cfg_.eviction_policy != OffsetEvictionPolicy::NONE) && + (eviction_handler != nullptr); + + // Warn if eviction policy is set but caller omitted the handler. + // IsEnableOffloading() will still return true in this mode, but + // eviction is effectively disabled (degraded to NONE behavior). + 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 (they aren't committed yet). + std::unordered_set batch_keys; + if (eviction_on) { + for (const auto& [k, _] : batch_object) batch_keys.insert(k); + } + std::vector keys; std::vector metadatas; keys.reserve(batch_object.size()); metadatas.reserve(batch_object.size()); - // Process each object in the batch; continue on individual failures to - // support partial success + // Accumulated evicted keys across the per-key loop. + // Flushed to eviction_handler before each allocate() that may reuse + // freed space, and again after the loop for any leftover victims. + std::vector evicted_keys; + 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; } - // Calculate total value size. RecordHeader stores value_len as a - // uint32_t (RecordHeader::SIZE is 8 bytes), so accumulating directly - // into a uint32_t would silently overflow for an object larger than - // 4 GiB: the record would be under-allocated and then its full slices - // written past the allocation. Sum in 64 bits and reject oversized - // objects instead of corrupting the storage arena. uint64_t total_value_size = 0; for (const auto& slice : slices) { total_value_size += slice.size; @@ -2895,48 +3316,116 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( "in 4 GiB) for key: " << key << ", size: " << total_value_size << " - skipping this key"; - continue; // partial-success model: keep processing other keys + continue; } uint32_t value_size = static_cast(total_value_size); - // Prepare record header RecordHeader header{.key_len = static_cast(key.size()), .value_len = value_size}; - - // 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; - // Step 1: Allocate space (allocator is thread-safe, ensures unique - // offsets) No locks held during allocation + // Guard against record_size exceeding what the on-disk format + // can represent (ObjectEntry::total_size is uint32_t). + if (record_size > UINT32_MAX) { + LOG(ERROR) << "Record too large for key: " << key + << ", record_size=" << record_size; + continue; + } + + // ---- (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) { + // When triggered by key-count overflow, force at least + // fallback_evict_batch victims even if bytes are low. + size_t min_v = over_keys ? cfg_.fallback_evict_batch : 0; + EvictToMakeRoom(static_cast(record_size), min_v, + batch_keys, evicted_keys); + } + } + + // ---- (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. + if (eviction_on && eviction_handler && !evicted_keys.empty()) { + eviction_handler(evicted_keys); + evicted_keys.clear(); + } + + // ---- (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) { + size_t before = evicted_keys.size(); + EvictToMakeRoom(static_cast(record_size), + cfg_.fallback_evict_batch, batch_keys, + evicted_keys); + size_t evicted_this_turn = evicted_keys.size() - before; + fallback_total_evicted += evicted_this_turn; + + // Notify master of fallback victims before retrying. + if (eviction_handler && !evicted_keys.empty()) { + eviction_handler(evicted_keys); + evicted_keys.clear(); + } + + 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); + } + } + + // ---- Handle allocation failure ---- 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; // eviction enabled: try next key + } else { + LOG(ERROR) << "Failed to allocate " << record_size + << " bytes for key: " << key + << " - stopping processing for this batch"; + break; // eviction disabled: preserve old break semantics + } } uint64_t offset = allocation->address(); - // Step 2: Write data to disk (no metadata locks held during I/O) + // ---- (E) Disk write (unchanged from original) ---- std::vector iovs; - iovs.reserve(2 + slices.size()); - - // Header + iovs.reserve(2 + 1 + slices.size()); 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)}); - - // Key iovs.push_back({const_cast(key.data()), static_cast(header.key_len)}); - - // Value slices for (const auto& slice : slices) { iovs.push_back({slice.ptr, slice.size}); } @@ -2945,78 +3434,81 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( 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 + // Lock order: eviction_mutex_ -> shard.mutex. + // Both insert and evict paths obey this order, preventing + // the overwrite-vs-evict race on fifo_index_. + 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()); - - if (!is_new_key) { - // Overwrite: subtract old size + uint64_t seq = 0; + + if (eviction_on) { + seq = insert_seq_.fetch_add(1, std::memory_order_relaxed); + if (!is_new_key) { + // Overwrite: drop old size and remove old FIFO slot. + size_delta -= static_cast(it->second.total_size); + fifo_index_.erase(it->second.fifo_seq); + } + } else if (!is_new_key) { size_delta -= static_cast(it->second.total_size); - // Old AllocationPtr will be dropped, refcount decremented - // Physical extent freed when last reader releases it } - // 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)); - // Update total size atomically (lock-free, separate from map - // updates) - total_size_.fetch_add(size_delta, std::memory_order_relaxed); + if (eviction_on) fifo_index_.emplace(seq, key); - // Update total keys only if inserting a new key + total_size_.fetch_add(size_delta, std::memory_order_relaxed); 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. + if (eviction_on && eviction_handler && !evicted_keys.empty()) { + eviction_handler(evicted_keys); + evicted_keys.clear(); } - // Invoke complete handler only if we have successful keys to report 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); } } @@ -3179,15 +3671,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; } @@ -3301,7 +3795,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/tenant_quota_policy_store.cpp b/mooncake-store/src/tenant_quota_policy_store.cpp index 23e10b7099..8edef856ce 100644 --- a/mooncake-store/src/tenant_quota_policy_store.cpp +++ b/mooncake-store/src/tenant_quota_policy_store.cpp @@ -18,6 +18,9 @@ #include #include +#ifdef STORE_USE_ETCD +#include "etcd_helper.h" +#endif #include "types.h" namespace mooncake { @@ -31,6 +34,34 @@ 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(); @@ -345,17 +376,91 @@ tl::expected YamlTenantQuotaPolicyStore::Save( 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) { - if (type != "file") { - return tl::make_unexpected("unsupported tenant quota connector type '" + - type + "'"); +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 (uri.empty()) { + 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 file connector requires a non-empty uri"); + "tenant quota etcd connector requires STORE_USE_ETCD"); +#endif } - return std::make_unique(uri); + return tl::make_unexpected("unsupported tenant quota connector type '" + + type + "'"); } } // namespace mooncake diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index 7dbaa4b0d2..37f8adba25 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -17,15 +17,18 @@ #include #include -#include -#include #include #include +#include #include -#include +#include +#include #include #include -#include +#include +#include +#include +#include // 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 @@ -237,7 +240,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; @@ -266,7 +395,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 = @@ -290,6 +424,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; @@ -337,12 +476,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; } @@ -363,10 +511,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=[" << 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/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index b661521adf..e7e66c9fbb 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -36,8 +36,31 @@ endfunction() add_store_test(buffer_allocator_test buffer_allocator_test.cpp) add_store_test(runtime_accelerator_test runtime_accelerator_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(master_service_tenant_quota_test master_service_tenant_quota_test.cpp) @@ -72,6 +95,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) @@ -86,6 +110,8 @@ 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) diff --git a/mooncake-store/tests/client_buffer_test.cpp b/mooncake-store/tests/client_buffer_test.cpp index e48069eca3..8a091f9b5a 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 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/ha/snapshot/master_service_ssd_test_for_snapshot.cpp b/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp index 73c6a1734a..d82b62c1ab 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 @@ -311,7 +311,14 @@ TEST_F(MasterServiceSSDSnapshotTest, RemoveKey) { } 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 @@ -365,8 +372,7 @@ TEST_F(MasterServiceSSDSnapshotTest, EvictObject) { } 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(); } 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..6899639175 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) { 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..090e510deb --- /dev/null +++ b/mooncake-store/tests/ha/snapshot/master_snapshot_codec_test.cpp @@ -0,0 +1,200 @@ +#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 "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 std::string kTenant = "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 8902329a39..464db29043 100644 --- a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp +++ b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp @@ -1,4 +1,5 @@ #include "master_service.h" +#include "master_snapshot_manager.h" #include "master_metric_manager.h" #include "ha/snapshot/catalog/snapshot_catalog_store.h" #include "ha/snapshot/object/snapshot_object_store.h" @@ -115,28 +116,57 @@ 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_->CleanupOldSnapshot(keep_count, + snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + temp_manager->CleanupOldSnapshot(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_->UploadSnapshotPayloadFile( + data, path, local_filename, snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->UploadSnapshotPayloadFile( + data, path, local_filename, snapshot_id); + } } SnapshotObjectStore* GetSnapshotObjectStore() { @@ -149,12 +179,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() { @@ -206,6 +246,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_; }; diff --git a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h index a160c1433a..01262f671c 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" 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/kv_event_publisher_test.cpp b/mooncake-store/tests/kv_event_publisher_test.cpp new file mode 100644 index 0000000000..fec8292004 --- /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", "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/master_admin_server_test.cpp b/mooncake-store/tests/master_admin_server_test.cpp index 5ef0cd62a7..20249277cd 100644 --- a/mooncake-store/tests/master_admin_server_test.cpp +++ b/mooncake-store/tests/master_admin_server_test.cpp @@ -1206,6 +1206,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/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/offset_allocator_test.cpp b/mooncake-store/tests/offset_allocator_test.cpp index d0e5519e44..b544798e99 100644 --- a/mooncake-store/tests/offset_allocator_test.cpp +++ b/mooncake-store/tests/offset_allocator_test.cpp @@ -1,4 +1,4 @@ -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "mutex.h" #include "serializer.h" #include "types.h" diff --git a/mooncake-store/tests/promotion_on_hit_test.cpp b/mooncake-store/tests/promotion_on_hit_test.cpp index 2743c4e252..bc2840970c 100644 --- a/mooncake-store/tests/promotion_on_hit_test.cpp +++ b/mooncake-store/tests/promotion_on_hit_test.cpp @@ -283,6 +283,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}, "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}, "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}, + "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) { diff --git a/mooncake-store/tests/pybind_client_test.cpp b/mooncake-store/tests/pybind_client_test.cpp index 27dc31c4a5..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) { 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/serializer_test.cpp b/mooncake-store/tests/serializer_test.cpp index 5d9d6c7318..538d2ee05c 100644 --- a/mooncake-store/tests/serializer_test.cpp +++ b/mooncake-store/tests/serializer_test.cpp @@ -2,7 +2,7 @@ #include #include "serializer.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "segment.h" namespace mooncake::test { diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index 1c66f376e4..bc5e3a3b19 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -507,6 +507,82 @@ TEST_F(StorageBackendTest, OrphanedBucketFileCleanup) { ASSERT_TRUE(is_exist.value()); } +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; @@ -2711,4 +2787,413 @@ TEST_F(StorageBackendTest, AdaptorBatchOffload_EvictionHandlerCalled) { //----------------------------------------------------------------------------- +//----------------------------------------------------------------------------- +// 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 = 4 * 1024; // 4KB — 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); + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // Write A, B, C one by one. The arena is 4 KB; each 1 KB record + // (8 + key + value) ≈ 1040 bytes. 4 records 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 = 4 * 1024; + 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); + }; + 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 = 4 * 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; + } + }; + 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); + }; + + 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_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 = 8 * 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); + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + std::string data(1500, 'x'); // each record ≈ 1516 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); + }; + 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 erases the key from shard.map, decrementing + // the map's shared_ptr. If no reader holds a copy, the + // RefCountedAllocationHandle destructor calls freeAllocation + // and the extent returns to the allocator. + // 3. While a reader holds its shared_ptr copy (refcount >= 1), + // 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 = + 8 * 1024; // 8 KB — tight enough that 80 keys 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); + }; + + 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"; + } +} + } // namespace mooncake::test diff --git a/mooncake-store/tests/tenant_quota_test.cpp b/mooncake-store/tests/tenant_quota_test.cpp index 7faba55aff..c17f7257b3 100644 --- a/mooncake-store/tests/tenant_quota_test.cpp +++ b/mooncake-store/tests/tenant_quota_test.cpp @@ -2,11 +2,18 @@ #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 @@ -15,6 +22,19 @@ namespace mooncake { namespace { +#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 + TenantQuotaSnapshot Snapshot(const TenantQuotaTable& table, const std::string& tenant_id) { auto snapshot = table.GetTenantSnapshot(tenant_id); @@ -36,6 +56,46 @@ std::filesystem::path MakeTempPolicyPath(const std::string& suffix) { 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) { ASSERT_TRUE(table->UpsertTenantPolicy(tenant_id, bytes).has_value()); @@ -228,6 +288,77 @@ TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlFile) { std::filesystem::remove(path); } +TEST(TenantQuotaPolicyStoreTest, FileFactoryCreatesYamlStore) { + const auto path = MakeTempPolicyPath("factory-file"); + auto store = + CreateTenantQuotaPolicyStore("file", path.string(), "test_cluster"); + ASSERT_TRUE(store.has_value()) << store.error(); +} + +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); +} + +#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 + +#ifdef STORE_USE_ETCD +TEST(TenantQuotaPolicyStoreTest, EtcdMissingKeyLoadsEmptySnapshot) { + if (auto skip_reason = GetTenantQuotaEtcdSkipReason(); + skip_reason.has_value()) { + GTEST_SKIP() << skip_reason.value(); + } + + const std::string cluster_id = + "tenant_quota_missing_" + std::to_string(::getpid()); + CleanupTenantQuotaEtcdCluster(cluster_id); + + 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(TenantQuotaPolicyStoreTest, EtcdRoundTripsSnapshot) { + if (auto skip_reason = GetTenantQuotaEtcdSkipReason(); + skip_reason.has_value()) { + GTEST_SKIP() << skip_reason.value(); + } + + const std::string cluster_id = + "tenant_quota_roundtrip_" + std::to_string(::getpid()); + CleanupTenantQuotaEtcdCluster(cluster_id); + + auto store = CreateTenantQuotaPolicyStore( + "etcd", GetTenantQuotaEtcdEndpoints(), cluster_id); + ASSERT_TRUE(store.has_value()) << store.error(); + + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = {{"tenant-a", 1024}, {"tenant-b", 2048}}; + auto save = store.value()->Save(snapshot); + ASSERT_TRUE(save.has_value()) << save.error(); + + auto loaded = store.value()->Load(); + ASSERT_TRUE(loaded.has_value()) << loaded.error(); + EXPECT_EQ(loaded->tenant_quotas, snapshot.tenant_quotas); + + CleanupTenantQuotaEtcdCluster(cluster_id); +} +#endif + TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlSpecialScalarNames) { TenantQuotaPolicySnapshot snapshot; snapshot.tenant_quotas = {{"foo#bar", 1}, diff --git a/mooncake-transfer-engine/benchmark/CMakeLists.txt b/mooncake-transfer-engine/benchmark/CMakeLists.txt index b17a8a803d..54be43112b 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,22 @@ 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 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..cd0dc9958f 100644 --- a/mooncake-transfer-engine/benchmark/bench_runner.h +++ b/mooncake-transfer-engine/benchmark/bench_runner.h @@ -57,10 +57,10 @@ 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) = 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..1a55528b09 100644 --- a/mooncake-transfer-engine/benchmark/main.cpp +++ b/mooncake-transfer-engine/benchmark/main.cpp @@ -15,6 +15,7 @@ #include "utils.h" #include "bench_runner.h" +#include "qos_metrics_adapter.h" #include "te_backend.h" #ifdef USE_TENT #include "tent_backend.h" @@ -23,7 +24,8 @@ using namespace mooncake::tent; int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, - int num_threads) { + int num_threads, + const std::vector& qos_classes) { bool mixed_opcode = false; OpCode opcode = READ; if (XferBenchConfig::check_consistency || XferBenchConfig::op_type == "mix") @@ -38,6 +40,9 @@ 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; int rc = runner.runInitiatorTasks([&](int thread_id) -> int { runner.pinThread(thread_id); @@ -49,11 +54,24 @@ int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, local_gpu_offset + thread_id, max_block_size, max_batch_size); uint64_t target_addr = runner.getTargetBufferBase( target_gpu_offset + thread_id, max_block_size, max_batch_size); + const bool qos_enabled = !qos_classes.empty(); + const size_t qos_class = + qos_enabled ? qosClassForThread(qos_classes, thread_id) : 0; + const bool tight = XferBenchConfig::deadline_us > 0 && + thread_id < XferBenchConfig::deadline_tight_threads; + auto deadlineNs = [&]() -> uint64_t { + if (!tight) return 0; + const auto now = + std::chrono::steady_clock::now().time_since_epoch(); + return std::chrono::duration_cast(now) + .count() + + XferBenchConfig::deadline_us * 1000ull; + }; XferBenchTimer timer; while (timer.lap_us(false) < 1000000ull) { runner.runSingleTransfer(local_addr, target_addr, block_size, - batch_size, opcode); + batch_size, opcode, deadlineNs()); } timer.reset(); std::vector transfer_duration; @@ -64,12 +82,14 @@ int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, if (XferBenchConfig::check_consistency) pattern = fillData((void*)local_addr, block_size * batch_size); - auto val = runner.runSingleTransfer( - local_addr, target_addr, block_size, batch_size, WRITE); + auto val = runner.runSingleTransfer(local_addr, target_addr, + block_size, batch_size, + WRITE, deadlineNs()); 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); + block_size, batch_size, READ, + deadlineNs()); if (XferBenchConfig::check_consistency) verifyData((void*)local_addr, block_size * batch_size, pattern); @@ -78,21 +98,50 @@ int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, } else { while (timer.lap_us(false) < XferBenchConfig::duration * 1000000ull) { - auto val = runner.runSingleTransfer( - local_addr, target_addr, block_size, batch_size, opcode); + auto val = runner.runSingleTransfer(local_addr, target_addr, + block_size, batch_size, + opcode, deadlineNs()); 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 (!qos_classes.empty()) { + auto report = calculateQosMetricsFromBenchStats( + block_size, batch_size, num_threads, qos_classes, &qos_stats, + XferBenchConfig::qos_link_capacity_gbps); + 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) { + 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 +151,57 @@ int main(int argc, char* argv[]) { "Usage: ./tebench [options]"); gflags::ParseCommandLineFlags(&argc, &argv, true); XferBenchConfig::loadFromFlags(); + std::vector qos_classes; + 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::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 (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; + } std::unique_ptr runner; if (XferBenchConfig::backend == "classic") { runner = std::make_unique(); @@ -145,7 +245,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..05b1dc4cf1 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/qos_metrics_adapter.cpp @@ -0,0 +1,49 @@ +// 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) { + return calculateQosMetrics(block_size, batch_size, num_threads, classes, + makeQosClassSamples(classes, stats), + 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..1778600b37 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/qos_metrics_adapter.h @@ -0,0 +1,38 @@ +// 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); + +} // 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..11165efeff 100644 --- a/mooncake-transfer-engine/benchmark/te_backend.cpp +++ b/mooncake-transfer-engine/benchmark/te_backend.cpp @@ -349,7 +349,9 @@ 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) { + (void)deadline_ns; 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..ea2a5d3eab 100644 --- a/mooncake-transfer-engine/benchmark/te_backend.h +++ b/mooncake-transfer-engine/benchmark/te_backend.h @@ -69,7 +69,7 @@ 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); private: int allocateBuffers(); diff --git a/mooncake-transfer-engine/benchmark/tent_backend.cpp b/mooncake-transfer-engine/benchmark/tent_backend.cpp index 3aa24db0d4..3ca9156306 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,6 +94,26 @@ static TransportType getTransportType(const std::string& xport_type) { return UNSPEC; } +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; +} + int TENTBenchRunner::allocateBuffers() { const auto total_buffer_size = XferBenchConfig::total_buffer_size; const auto& seg_type = XferBenchConfig::seg_type; @@ -201,6 +223,7 @@ TENTBenchRunner::TENTBenchRunner() { engine_ = std::make_unique(loadConfig()); transport_hint_ = TransportSelector::parseTransportType( XferBenchConfig::tent_transport_hint); + intent_type_ = getIntentType(XferBenchConfig::tent_intent_type); allocateBuffers(); } @@ -337,7 +360,8 @@ 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) { auto batch_id = engine_->allocateBatch(batch_size); std::vector requests; for (uint64_t i = 0; i < batch_size; ++i) { @@ -348,6 +372,8 @@ 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_; 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..b648605c0f 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.h +++ b/mooncake-transfer-engine/benchmark/tent_backend.h @@ -72,7 +72,7 @@ class TENTBenchRunner : 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); private: int allocateBuffers(); @@ -87,6 +87,7 @@ class TENTBenchRunner : public BenchRunner { SegmentID handle_; SegmentInfo info_; TransportType transport_hint_{UNSPEC}; + IntentType intent_type_{IntentType::INTENT_UNSPEC}; std::vector> current_task_; std::vector threads_; @@ -99,4 +100,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..367bf13c04 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/tests/qos_metrics_test.cpp @@ -0,0 +1,162 @@ +// 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" + +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, 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, 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..ca69e87020 100644 --- a/mooncake-transfer-engine/benchmark/utils.cpp +++ b/mooncake-transfer-engine/benchmark/utils.cpp @@ -35,6 +35,25 @@ 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_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)."); +DEFINE_int32(deadline_tight_threads, 0, + "tent only: workers [0, N) that carry --deadline_us; remaining " + "workers have no deadline."); +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,6 +72,10 @@ 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 { @@ -70,6 +93,13 @@ 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; +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,6 +108,7 @@ 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; @@ -96,6 +127,13 @@ 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; + 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 +144,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 +207,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..ddb4e73b08 100644 --- a/mooncake-transfer-engine/benchmark/utils.h +++ b/mooncake-transfer-engine/benchmark/utils.h @@ -68,6 +68,13 @@ 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 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 +83,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 +115,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 +173,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/example/CMakeLists.txt b/mooncake-transfer-engine/example/CMakeLists.txt index 4adfcdba93..4aed903ea3 100644 --- a/mooncake-transfer-engine/example/CMakeLists.txt +++ b/mooncake-transfer-engine/example/CMakeLists.txt @@ -106,3 +106,7 @@ 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) 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/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..b6eb996173 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -90,10 +90,11 @@ DEFINE_string(operation, "read", "Operation type: read or write"); DEFINE_string(protocol, "rdma", "Transfer protocol: " - "rdma|barex|tcp|efa|nvlink|nvlink_intra|hip|sunrise_link"); + "rdma|barex|tcp|efa|ub|nvlink|nvlink_intra|hip|ubshmem|" + "sunrise_link"); DEFINE_string(device_name, "mlx5_2", - "Device name to use, valid if protocol=rdma"); + "Device name to use, valid if protocol=rdma|ub"); DEFINE_string(nic_priority_matrix, "", "Path to RDMA NIC priority matrix file (Advanced)"); @@ -673,6 +674,17 @@ std::shared_ptr createTentConfig() { config->set("local_segment_name", FLAGS_local_server_name); config->set("verbose", true); + // Propagate --device_name to the UB/RDMA transport so that + // UbTentTransport::install() can filter which device to use + // (important for bonded-device setups like bonding_dev_0). + if (!FLAGS_device_name.empty()) { + if (FLAGS_protocol == "rdma") { + config->set("transports/rdma/device_name", FLAGS_device_name); + } else if (FLAGS_protocol == "ub") { + config->set("transports/ub/device_name", FLAGS_device_name); + } + } + return config; } diff --git a/mooncake-transfer-engine/include/common.h b/mooncake-transfer-engine/include/common.h index 5ae9a66471..312795a284 100644 --- a/mooncake-transfer-engine/include/common.h +++ b/mooncake-transfer-engine/include/common.h @@ -528,6 +528,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 +536,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 +549,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 +578,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 +605,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 569279dd09..8a3fba87d9 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -60,6 +60,11 @@ struct GlobalConfig { // which is minutes. Override via MC_HANDSHAKE_CONNECT_TIMEOUT. int handshake_connect_timeout = 5; 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; @@ -71,6 +76,7 @@ struct GlobalConfig { 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; size_t eic_max_block_size = 64UL * 1024 * 1024; EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE; diff --git a/mooncake-transfer-engine/include/gpu_vendor/maca.h b/mooncake-transfer-engine/include/gpu_vendor/maca.h index 56384149b1..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 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/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/topology.h b/mooncake-transfer-engine/include/topology.h index e9d16605f1..2efbe0f4d5 100644 --- a/mooncake-transfer-engine/include/topology.h +++ b/mooncake-transfer-engine/include/topology.h @@ -40,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); @@ -84,6 +86,9 @@ 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); diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 9fa1d813de..afa9fce2fb 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -15,12 +15,15 @@ #ifndef MULTI_TRANSFER_ENGINE_H_ #define MULTI_TRANSFER_ENGINE_H_ +#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; @@ -72,9 +75,9 @@ class TransferEngine { TransferEngine(bool auto_discover, const std::vector& filter); - TransferEngine(TransferEngine&&) = default; + TransferEngine(TransferEngine&& other) noexcept; - TransferEngine& operator=(TransferEngine&&) = default; + TransferEngine& operator=(TransferEngine&& other) noexcept; ~TransferEngine(); @@ -198,9 +201,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..23c7d7f703 100644 --- a/mooncake-transfer-engine/include/transfer_engine_c.h +++ b/mooncake-transfer-engine/include/transfer_engine_c.h @@ -165,6 +165,10 @@ int freeBatchID(transfer_engine_t engine, batch_id_t batch_id); int syncSegmentCache(transfer_engine_t engine); +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 c5d919d232..a4048514d5 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -43,6 +43,8 @@ #endif namespace mooncake { +class TransferEngineImplTestPeer; + using TransferRequest = Transport::TransferRequest; using TransferStatus = Transport::TransferStatus; using TransferStatusEnum = Transport::TransferStatusEnum; @@ -56,6 +58,8 @@ using RegisteredBuffer = TransferEngine::RegisteredBuffer; #endif class TransferEngineImpl { + friend class TransferEngineImplTestPeer; + public: TransferEngineImpl(bool auto_discover = false) : metadata_(nullptr), @@ -342,6 +346,7 @@ class TransferEngineImpl { } Transport* getTransport(const std::string& proto) { + if (!multi_transports_) return nullptr; return multi_transports_->getTransport(proto); } @@ -400,12 +405,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); @@ -416,6 +425,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_; diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index 715cd7cabb..af917fed01 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,6 +49,8 @@ class TransferMetadata { uint16_t lid; std::string gid; std::string eid; // for ub + + bool operator==(const DeviceDesc &other) const = default; }; struct BufferDesc { @@ -70,12 +74,16 @@ class TransferMetadata { 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 { @@ -89,6 +97,8 @@ class TransferMetadata { uint64_t devicePort; uint64_t pid; std::vector endpoints; + + bool operator==(const RankInfoDesc &other) const = default; }; using SegmentID = uint64_t; @@ -111,6 +121,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 @@ -125,6 +139,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; }; @@ -144,11 +162,16 @@ class TransferMetadata { std::string peer_nic_path; #ifdef USE_UB std::vector jetty_num; // for ub/urma + std::string local_eid; // for ub/urma #endif #ifdef USE_BAREX 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) @@ -166,13 +189,13 @@ class TransferMetadata { public: TransferMetadata(const std::string &conn_string); - ~TransferMetadata(); + virtual ~TransferMetadata(); - std::shared_ptr getSegmentDescByName( + virtual std::shared_ptr getSegmentDescByName( const std::string &segment_name, bool force_update = false); - std::shared_ptr getSegmentDescByID(SegmentID segment_id, - bool force_update = false); + virtual std::shared_ptr getSegmentDescByID( + SegmentID segment_id, bool force_update = false); int updateLocalSegmentDesc(SegmentID segment_id = LOCAL_SEGMENT_ID); @@ -182,7 +205,7 @@ class TransferMetadata { std::shared_ptr getSegmentDesc( const std::string &segment_name); - SegmentID getSegmentID(const std::string &segment_name); + virtual SegmentID getSegmentID(const std::string &segment_name); int syncSegmentCache(const std::string &segment_name); @@ -212,12 +235,12 @@ class TransferMetadata { using OnReceiveHandShake = std::function; - int startHandshakeDaemon(OnReceiveHandShake on_receive_handshake, - uint16_t listen_port, int sockfd); + virtual int startHandshakeDaemon(OnReceiveHandShake on_receive_handshake, + uint16_t listen_port, int sockfd); - int sendHandshake(const std::string &peer_server_name, - const HandShakeDesc &local_desc, - HandShakeDesc &peer_desc); + virtual int sendHandshake(const std::string &peer_server_name, + const HandShakeDesc &local_desc, + HandShakeDesc &peer_desc); int sendNotify(const std::string &peer_server_name, const NotifyDesc &local_desc, NotifyDesc &peer_desc); @@ -238,6 +261,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_; @@ -258,6 +284,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/device/ibgda/mlx5gda.h b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h index 82030954d2..34584ff3ad 100644 --- a/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h +++ b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h @@ -87,6 +87,16 @@ struct mlx5gda_qp_devctx { uint16_t wq_tail; // last non-completed wqeid }; +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, 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/kunpeng_transport/ub_context.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h index 56767c9dee..0453350d63 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h @@ -174,6 +174,19 @@ class UbContext { public: virtual int registerMemoryRegion(uint64_t va, size_t length) = 0; + // Returns the most recently registered local segment handle on this + // context (the one produced by the last successful registerMemoryRegion), + // or nullptr if none. Used to share a single host-global URMA segment + // across all contexts (see UbTransport::registerLocalMemory). + virtual void* lastRegisteredSeg() = 0; + + // Adopts a segment that was registered on another context for the same + // host virtual address. URMA registers memory into a host-global ubva + // space, so a given host VA must be registered with the driver exactly + // once; the other contexts reference that single segment instead of + // re-registering it (which the driver rejects as a duplicate). + virtual int adoptLocalSeg(uint64_t va, size_t length, void* seg) = 0; + virtual int unregisterMemoryRegion(uint64_t va) = 0; virtual int doProcessContextEvents() = 0; diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h index 6b6fa7076a..2f781035e5 100644 --- a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include "common.h" #include "config.h" @@ -55,6 +56,8 @@ class UrmaContext : public UbContext { int max_endpoints); ~UrmaContext(); int registerMemoryRegion(uint64_t va, size_t length) override; + void* lastRegisteredSeg() override; + int adoptLocalSeg(uint64_t va, size_t length, void* seg) override; int unregisterMemoryRegion(uint64_t va) override; int doProcessContextEvents() override; void* retrieveRemoteSeg(const std::string& value) override; @@ -127,6 +130,10 @@ class UrmaContext : public UbContext { RWSpinlock seg_region_lock_; std::vector> seg_region_list_; std::vector local_tseg_list_; + // Local segments actually registered with the URMA driver by THIS context + // (as opposed to segments adopted from another context for the same + // host-global VA). Only owned segments are passed to urma_unregister_seg. + std::unordered_set owned_segs_; std::vector remote_seg_list_; std::vector imported_seg_list_; 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..6f63646c8b 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,16 @@ 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; }; class CUFileDescPool { + friend class CUFileDescPoolTestPeer; + public: explicit CUFileDescPool(size_t max_batch_size = 128); ~CUFileDescPool(); @@ -75,6 +83,9 @@ class CUFileDescPool { CUFileBatchDesc* getDesc(int idx); private: + static bool cachePolledEvent(std::vector& io_events, + const CUfileIOEvents_t& event); + static const size_t MAX_NR_DESC = 256; // Max number of descriptors size_t max_batch_size_; @@ -89,4 +100,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..aa979d8f49 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,11 @@ 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 startTransfer(Slice *slice); private: 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..655d04640c 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h @@ -59,10 +59,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,6 +87,19 @@ 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 { @@ -93,6 +121,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 +151,7 @@ class RdmaContext { private: int registerMemoryRegionInternal(void *addr, size_t length, int access, + const DmabufExport &exp, MemoryRegionMeta &mrMeta); using MemoryRegionMap = std::map; @@ -113,9 +162,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 @@ -170,6 +221,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 +238,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 +250,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 +274,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 +295,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_; @@ -254,7 +318,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 e98fa8e76a..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); @@ -75,23 +83,31 @@ 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); @@ -122,6 +138,8 @@ class RdmaEndPoint { // 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; @@ -151,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); @@ -166,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; @@ -189,15 +216,16 @@ 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 e7a171b8b1..334bba20e6 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h @@ -135,6 +135,10 @@ class RdmaTransport : public Transport { int &buffer_id, int &device_id, int retry_cnt = 0); + const std::vector> &getContextList() const { + return context_list_; + } + private: std::vector> context_list_; std::shared_ptr local_topology_; 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 e8ada2a7af..86d2c0245e 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h @@ -31,6 +31,11 @@ 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: void performPostSend(int thread_id); @@ -62,6 +67,8 @@ class WorkerPool { // and optionally deletes the endpoint void handlePathFailure(const std::string &peer_nic_path, RdmaEndPoint *endpoint = nullptr); + void refreshPublishedLocalTopology(); + GidRefreshResult refreshPublishedLocalGid(); // Context-level health tracking for catastrophic hardware failure. // When all rails through a local RNIC are unavailable, increment the @@ -90,6 +97,17 @@ class WorkerPool { std::atomic workers_running_; 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_; 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/tcp_transport/tcp_transport.h b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h index 3e537cfec2..be096e96df 100644 --- a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h +++ b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h @@ -142,6 +142,10 @@ class TcpTransport : public Transport { const std::string &host, uint16_t port); 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 532eebc641..01253de842 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -133,7 +133,7 @@ class Transport { 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 diff --git a/mooncake-transfer-engine/src/config.cpp b/mooncake-transfer-engine/src/config.cpp index 6adb935242..0e9399af7c 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -317,6 +317,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) { @@ -453,6 +474,14 @@ void loadGlobalConfig(GlobalConfig& config) { 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) { @@ -631,6 +660,8 @@ void dumpGlobalConfig() { LOG(INFO) << "parallel_reg_mr = " << config.parallel_reg_mr; 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; { std::ostringstream oss; for (size_t i = 0; i < config.mlx5_qp_udp_sports.size(); ++i) { @@ -645,6 +676,8 @@ void dumpGlobalConfig() { << (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..8fbf2f0cf0 --- /dev/null +++ b/mooncake-transfer-engine/src/graceful_shutdown.cpp @@ -0,0 +1,181 @@ +// 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 "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; + + try { + std::thread(signalWatcher).detach(); + } catch (...) { + 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/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 30475ede7d..81cb0174f3 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -19,6 +19,7 @@ #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" @@ -470,11 +471,20 @@ Status MultiTransport::selectTransport(const TransferRequest& entry, // 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) { @@ -486,6 +496,7 @@ Status MultiTransport::selectTransport(const TransferRequest& entry, : 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; @@ -503,6 +514,12 @@ Status MultiTransport::selectTransport(const TransferRequest& entry, 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(); } @@ -540,6 +557,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 38435f4328..bc6b15eccd 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -695,6 +695,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_) { diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 7bd2ced5d2..a24a66dbb8 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -14,10 +14,54 @@ #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 +70,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,6 +108,7 @@ int TransferEngine::init(const std::string& metadata_conn_string, } int TransferEngine::freeEngine() { + detachShutdownToken(shutdown_token_); if (impl_) { impl_->freeEngine(); impl_.reset(); @@ -221,6 +293,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" @@ -228,9 +313,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")) { @@ -251,6 +380,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( @@ -300,6 +457,7 @@ int TransferEngine::init(const std::string& metadata_conn_string, } int TransferEngine::freeEngine() { + detachShutdownToken(shutdown_token_); if (!use_tent_ && impl_) { impl_->freeEngine(); impl_.reset(); @@ -670,5 +828,20 @@ 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 diff --git a/mooncake-transfer-engine/src/transfer_engine_c.cpp b/mooncake-transfer-engine/src/transfer_engine_c.cpp index b643a15386..5ed7e99d0a 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,18 @@ 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 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 d1514bf210..98b55596a6 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -587,24 +587,30 @@ 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; + } + for (auto transport : multi_transports_->listTransports()) { int ret = transport->registerLocalMemory( addr, length, location, remote_accessible, update_metadata); - if (ret < 0) return ret; + if (ret < 0) { + releaseMemoryRegions(regions); + return ret; + } } - std::unique_lock lock(mutex_); - insertMemoryRegionLocked({addr, length, location, remote_accessible}); + commitMemoryRegions(regions); return 0; } @@ -742,22 +748,53 @@ 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; + regions.reserve(buffer_list.size()); + for (const auto& buffer : buffer_list) { + regions.push_back({buffer.addr, buffer.length, location, true}); + } + if (!tryReserveMemoryRegions(regions)) { + LOG(ERROR) + << "Transfer Engine does not support overlapped memory region"; + return ERR_ADDRESS_OVERLAPPED; + } + for (auto transport : multi_transports_->listTransports()) { int ret = transport->registerLocalMemoryBatch(buffer_list, location); - if (ret < 0) return ret; + if (ret < 0) { + 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; } @@ -775,51 +812,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)) { @@ -830,6 +843,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 30de576d86..aebe095050 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" @@ -76,6 +78,8 @@ 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 @@ -113,6 +117,12 @@ 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 @@ -161,6 +171,7 @@ TransferMetadata::TransferMetadata(const std::string &conn_string) { } if (conn_string == P2PHANDSHAKE) { p2p_handshake_mode_ = true; + startMetadataRefreshPollingIfNeeded(); return; } storage_plugin_ = MetadataStoragePlugin::Create(conn_string); @@ -169,9 +180,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 { @@ -269,7 +339,7 @@ static int encodeMultiProtocolSegmentDesc( bufferJSON["lkey"] = lkeyJSON; } else if (buffer.protocol == "tcp") { bufferJSON["addr"] = static_cast(buffer.addr); - } else if (buffer.protocol == "hip") { + } else if (buffer.protocol == "hip" || buffer.protocol == "maca") { bufferJSON["addr"] = static_cast(buffer.addr); bufferJSON["shm_name"] = buffer.shm_name; } @@ -278,6 +348,7 @@ static int encodeMultiProtocolSegmentDesc( 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; @@ -297,7 +368,7 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, is_multi_protocol = true; for (const auto &proto : protocols) { if (proto != "cxl" && proto != "tcp" && proto != "rdma" && - proto != "hip") { + proto != "hip" && proto != "maca") { is_multi_protocol = false; break; } @@ -305,7 +376,8 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, if (!is_multi_protocol) { LOG(ERROR) << "Unsupported multi-protocol combination: " << desc.protocol - << ". Only cxl, tcp, rdma and hip may be combined."; + << ". Only cxl, tcp, rdma, hip and maca may be " + "combined."; return ERR_INVALID_ARGUMENT; } } @@ -319,6 +391,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; @@ -515,6 +588,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")) @@ -608,7 +684,7 @@ decodeMultiProtocolSegmentDesc(Json::Value &segmentJSON, return nullptr; } desc->buffers.push_back(buffer); - } else if (buffer_protocol == "hip") { + } else if (buffer_protocol == "hip" || buffer_protocol == "maca") { TransferMetadata::BufferDesc buffer; buffer.name = bufferJSON["name"].asString(); buffer.addr = bufferJSON["addr"].asUInt64(); @@ -645,7 +721,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, for (const auto &protocolStr : segmentJSON["protocol"]) { std::string proto = protocolStr.asString(); if (proto != "cxl" && proto != "tcp" && proto != "rdma" && - proto != "hip") { + proto != "hip" && proto != "maca") { is_multi_protocol = false; break; } @@ -654,7 +730,7 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, LOG(ERROR) << "Unsupported multi-protocol combination in segment: " << segment_name - << ". Only cxl, tcp, rdma and hip may be combined."; + << ". Only cxl, tcp, rdma, hip and maca may be combined."; return nullptr; } } @@ -670,6 +746,9 @@ 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")) @@ -958,7 +1037,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; { @@ -971,25 +1064,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; } 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..51e97eeeb5 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 { 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..bc8e20dbd5 100644 --- a/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp +++ b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp @@ -25,11 +25,13 @@ #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" @@ -199,28 +201,7 @@ 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; + return allocateControlBuffer(false); } int createQueuePairs(void* stream_ptr) override { @@ -231,6 +212,8 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { ctrl_buf_heap_, pd_, 16384, 1, stream); if (!qp) { LOG(ERROR) << "[EP IBGDA] mlx5gda_create_rc_qp failed at " << i; + if (retryWithHostControlBuffer()) + return createQueuePairs(stream_ptr); return -1; } if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) { @@ -243,11 +226,11 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { .qpn = qp->qpn, .wqeid_mask = qp->num_wqebb - 1, .wq = reinterpret_cast( - static_cast(ctrl_buf_) + qp->wq_offset), + static_cast(ctrl_buf_dev_) + qp->wq_offset), .cq = reinterpret_cast( - static_cast(ctrl_buf_) + qp->send_cq->cq_offset), + static_cast(ctrl_buf_dev_) + qp->send_cq->cq_offset), .dbr = reinterpret_cast( - static_cast(ctrl_buf_) + qp->dbr_offset), + static_cast(ctrl_buf_dev_) + qp->dbr_offset), .bf = static_cast(qp->uar->reg_addr), }; cudaMemcpy( @@ -259,11 +242,7 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { } 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 +328,98 @@ 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; + } + + int allocateControlBuffer(bool host_backed) { + ctrl_buf_host_ = host_backed; + if (ctrl_buf_host_) { + 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; + ctrl_buf_host_ = false; + return -1; + } + LOG(INFO) << "[EP IBGDA] Using host-backed mapped control buffer"; + } 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] mlx5dv_devx_umem_reg failed (errno=" + << errno << ")"; + freeControlBuffer(); + 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"; + freeControlBuffer(); + return -1; + } + return 0; + } + + bool retryWithHostControlBuffer() { + auto failure = mlx5gda_last_create_qp_failure(); + if (ctrl_buf_host_ || !isCreateQpBadParam(failure)) return false; + + 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(true) == 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 +429,21 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { ctrl_buf_umem_ = nullptr; } if (ctrl_buf_) { - cudaFree(ctrl_buf_); + if (ctrl_buf_host_) { + cudaHostUnregister(ctrl_buf_); + free(ctrl_buf_); + } else { + cudaFree(ctrl_buf_); + } ctrl_buf_ = nullptr; + ctrl_buf_dev_ = nullptr; + ctrl_buf_host_ = false; } + } + + void teardown() { + destroyQueuePairs(); + freeControlBuffer(); if (mr_) { ibv_dereg_mr(mr_); mr_ = nullptr; @@ -407,6 +485,8 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { // Control buffer void* ctrl_buf_ = nullptr; // GPU VA + void* ctrl_buf_dev_ = nullptr; + bool ctrl_buf_host_ = false; mlx5dv_devx_umem* ctrl_buf_umem_ = nullptr; memheap* ctrl_buf_heap_ = nullptr; diff --git a/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp b/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp index ca601b2536..532e627027 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,47 @@ 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 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 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, 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). @@ -99,6 +141,7 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, struct ibv_context* ctx = pd->context; void* cq_context = NULL; + bool ctrl_host = is_host_control_buffer(ctrl_buf); if (cqe <= 0) { errno = EINVAL; @@ -120,11 +163,16 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, // 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; + if (ctrl_host) { + memset(static_cast(ctrl_buf) + cq_offset, -1, + num_cqe * sizeof(struct mlx5_cqe64)); + } else { + if (cudaMemsetAsync(static_cast(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) { @@ -162,9 +210,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, @@ -208,6 +258,8 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, struct memheap* ctrl_buf_heap, struct ibv_pd* pd, int wqe, uint8_t port_num, cudaStream_t stream) { + mlx5gda_reset_create_qp_failure(); + struct mlx5gda_qp* qp = NULL; struct mlx5gda_cq* send_cq = NULL; struct mlx5dv_devx_uar* uar = NULL; @@ -219,6 +271,7 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, void* qp_context = NULL; void* cap = NULL; uint32_t cqe_version = 0; + bool ctrl_host = is_host_control_buffer(ctrl_buf); if (wqe <= 0) { errno = EINVAL; @@ -291,10 +344,16 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, 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(ctrl_buf) + dbr_offset, 0, + sizeof(struct mlx5gda_wq_dbr)); + } else { + if (cudaMemsetAsync(static_cast(ctrl_buf) + 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); @@ -326,14 +385,19 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, // 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, + ctrl_buf_umem->umem_id, wq_offset, dbr_offset, errno); goto fail; } @@ -517,4 +581,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/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index 2dd3dc42ca..06be3f10ce 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -86,15 +86,49 @@ int UbTransport::registerLocalMemory(void* addr, size_t length, bool update_metadata) { (void)remote_accessible; BufferDesc buffer_desc; - for (auto& context : context_list_) { - int ret = context->registerMemoryRegion((uint64_t)addr, length); + if (context_list_.empty()) { + LOG(ERROR) << "UbTransport: no available context to register memory"; + return ERR_DEVICE_NOT_FOUND; + } + + // URMA registers host memory into a host-global ubva space, so a given host + // virtual address may be registered with the driver exactly once per host + // (re-registering the same VA on another context is rejected by the driver + // as a duplicate, e.g. "registered twice within 500ms"). Register the + // buffer on the primary context, then share the resulting segment with the + // remaining contexts so every device's data path still has a valid local + // segment handle for this buffer. + int ret = context_list_[0]->registerMemoryRegion((uint64_t)addr, length); + if (ret) { + LOG(ERROR) << "UbTransport: cannot register LocalMemory on primary " + "context"; + return ret; + } + void* shared_seg = context_list_[0]->lastRegisteredSeg(); + if (!shared_seg) { + LOG(ERROR) << "UbTransport: primary context returned null segment"; + context_list_[0]->unregisterMemoryRegion((uint64_t)addr); + return ERR_CONTEXT; + } + + for (size_t i = 1; i < context_list_.size(); ++i) { + ret = + context_list_[i]->adoptLocalSeg((uint64_t)addr, length, shared_seg); if (ret) { - LOG(ERROR) << "UbTransport: cannot register LocalMemory"; + LOG(ERROR) << "UbTransport: cannot share segment to context " << i; + for (size_t j = 0; j < i; ++j) + context_list_[j]->unregisterMemoryRegion((uint64_t)addr); return ret; } - ret = context->buildLocalBufferDesc((uint64_t)addr, buffer_desc); + } + + for (size_t i = 0; i < context_list_.size(); ++i) { + ret = + context_list_[i]->buildLocalBufferDesc((uint64_t)addr, buffer_desc); if (ret) { LOG(ERROR) << "UbTransport: build buffer description failed"; + for (size_t j = 0; j < context_list_.size(); ++j) + context_list_[j]->unregisterMemoryRegion((uint64_t)addr); return ret; } } @@ -203,19 +237,29 @@ Status UbTransport::submitTransfer( Status UbTransport::submitTransferTask( const std::vector& task_list) { - std::unordered_map, std::vector> - slices_to_post; auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); const size_t kBlockSize = globalConfig().slice_size; const int kMaxRetryCount = globalConfig().retry_cnt; const size_t kFragmentSize = globalConfig().fragment_limit; const size_t kSubmitWatermark = globalConfig().max_wr * globalConfig().num_qp_per_ep; - uint64_t nr_slices; - for (size_t index = 0; index < task_list.size(); ++index) { - assert(task_list[index]); - auto& task = *task_list[index]; - nr_slices = 0; + + struct PlannedSlice { + TransferTask* task; + std::shared_ptr context; + uint64_t offset; + size_t length; + int buffer_id; + int device_id; + }; + std::vector plan; + + // Validate and route the complete submission before mutating TransferTask + // state or posting any work. This guarantees that a non-OK return has no + // partial submission for the caller to roll back. + for (auto* task_ptr : task_list) { + assert(task_ptr); + auto& task = *task_ptr; assert(task.request); auto& request = *task.request; auto request_buffer_id = -1, request_device_id = -1; @@ -229,27 +273,11 @@ Status UbTransport::submitTransferTask( for (uint64_t offset = 0; offset < request.length; offset += kBlockSize) { - 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 = + auto* source_addr = static_cast(request.source) + offset; + size_t slice_length = merge_final_slice ? request.length - offset : kBlockSize; - slice->opcode = request.opcode; - // LOG(INFO) << "target_offset : " << request.target_offset << ", - // offset : " << offset; - slice->ub.dest_addr = request.target_offset + offset; - slice->ub.retry_cnt = 0; - slice->ub.max_retry_cnt = kMaxRetryCount; - slice->task = &task; - slice->target_id = request.target_id; - slice->ts = 0; - slice->status = Slice::PENDING; - task.slice_list.push_back(slice); int buffer_id = -1, device_id = -1, retry_cnt = request.advise_retry_cnt; @@ -261,8 +289,9 @@ Status UbTransport::submitTransferTask( } while (retry_cnt < kMaxRetryCount && !found_device) { if (selectDevice(local_segment_desc.get(), - (uint64_t)slice->source_addr, slice->length, - buffer_id, device_id, retry_cnt++)) + reinterpret_cast(source_addr), + slice_length, buffer_id, device_id, + retry_cnt++)) continue; assert(device_id >= 0 && static_cast(device_id) < context_list_.size()); @@ -277,10 +306,7 @@ Status UbTransport::submitTransferTask( found_device = true; break; } - if (device_id < 0) { - auto source_addr = slice->source_addr; - for (auto& entry : slices_to_post) - for (auto s : entry.second) getSliceCache().deallocate(s); + if (!found_device || device_id < 0) { LOG(ERROR) << "UbTransport: Address not registered by any device(s) " << source_addr; @@ -296,22 +322,54 @@ Status UbTransport::submitTransferTask( return Status::InvalidArgument( "Device " + std::to_string(device_id) + " is not active"); } - auto local_tseg_index = - local_segment_desc->buffers[buffer_id].l_seg_index[device_id]; - slice->ub.l_seg = context->localSegWithIndex(local_tseg_index); - slices_to_post[context].push_back(slice); - task.total_bytes += slice->length; - __sync_fetch_and_add(&task.slice_count, 1); - if (nr_slices >= kSubmitWatermark) { - for (auto& entry : slices_to_post) - entry.first->submitPostSend(entry.second); - slices_to_post.clear(); - nr_slices = 0; - } - if (merge_final_slice) { - break; - } + plan.push_back( + {&task, context, offset, slice_length, buffer_id, device_id}); + + if (merge_final_slice) break; + } + } + + std::unordered_map, std::vector> + slices_to_post; + TransferTask* current_task = nullptr; + uint64_t nr_slices = 0; + for (const auto& planned : plan) { + auto& task = *planned.task; + auto& request = *task.request; + if (current_task != &task) { + current_task = &task; + nr_slices = 0; + } + + Slice* slice = getSliceCache().allocate(); + assert(slice); + if (!slice->from_cache) ++nr_slices; + slice->source_addr = + static_cast(request.source) + planned.offset; + slice->length = planned.length; + slice->opcode = request.opcode; + slice->ub.dest_addr = request.target_offset + planned.offset; + slice->ub.retry_cnt = 0; + slice->ub.max_retry_cnt = kMaxRetryCount; + slice->task = &task; + slice->target_id = request.target_id; + slice->ts = 0; + slice->status = Slice::PENDING; + task.slice_list.push_back(slice); + + auto& context = planned.context; + auto local_tseg_index = local_segment_desc->buffers[planned.buffer_id] + .l_seg_index[planned.device_id]; + slice->ub.l_seg = context->localSegWithIndex(local_tseg_index); + slices_to_post[context].push_back(slice); + task.total_bytes += slice->length; + __sync_fetch_and_add(&task.slice_count, 1); + if (nr_slices >= kSubmitWatermark) { + for (auto& entry : slices_to_post) + entry.first->submitPostSend(entry.second); + slices_to_post.clear(); + nr_slices = 0; } } for (auto& entry : slices_to_post) 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 06c65d0385..e5c1bf32ed 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 @@ -35,13 +35,22 @@ UrmaContext::UrmaContext(UbTransport& engine, std::string device_name, next_jfr_list_index_(0) {} UrmaContext::~UrmaContext() { + // toString() calls getAsyncFd() which dereferences urma_context_. + // getAsyncFd() must tolerate nullptr before we call it here. auto thisString = toString(); + + // worker_pool_ is a shared_ptr; reset on null is safe. worker_pool_.reset(); LOG(INFO) << "destroy worker pool done."; + + // endpoint_store_ may be null if construction failed before doConstruct(). if (endpoint_store_) { endpoint_store_->destroy(); + LOG(INFO) << "destroy endpoint store done."; + } else { + LOG(INFO) << "endpoint store is null, skip destroy."; } - LOG(INFO) << "destroy endpoint store done."; + if (urma_context_) deconstruct(); LOG(WARNING) << "finished destroy context : " << thisString; } @@ -54,7 +63,10 @@ std::string UrmaContext::toString() { return ss.str(); } -int UrmaContext::getAsyncFd() { return urma_context_->async_fd; } +int UrmaContext::getAsyncFd() { + if (urma_context_ == nullptr) return -1; + return urma_context_->async_fd; +} int UrmaContext::submitPostSend( const std::vector& slice_list) { @@ -158,12 +170,29 @@ int UrmaContext::construct(GlobalConfig& config) { int UrmaContext::deconstruct() { for (auto& entry : seg_region_list_) { - int ret = urma_unregister_seg(entry.first); + urma_target_seg_t* seg_ptr = entry.first; + // Release the application-side reference in local_tseg_list_ BEFORE + // calling urma_unregister_seg. URMA uses reference counting: if the + // app still holds the pointer, the count stays > 0 and the VA range + // is not freed. A subsequent urma_register_seg for the same VA would + // then fail with "duplicate". + for (auto& tseg : local_tseg_list_) { + if (tseg == seg_ptr) { + tseg = nullptr; + break; + } + } + // Only unregister segments this context actually registered with the + // driver; adopted (shared) segments are owned by another context. + if (owned_segs_.find(seg_ptr) == owned_segs_.end()) continue; + int ret = urma_unregister_seg(seg_ptr); if (ret) { PLOG(ERROR) << "Failed to unregister segment"; } } seg_region_list_.clear(); + local_tseg_list_.clear(); + owned_segs_.clear(); for (auto& seg : imported_seg_list_) { int ret = urma_unimport_seg(seg); @@ -231,7 +260,11 @@ int UrmaContext::deconstruct() { urma_context_ = nullptr; } - urma_uninit(); + // urma_uninit() is a global teardown and must only be called once via + // the static UrmaContext::uninit() (invoked by UbTransport::uninit()). + // Calling it here per-context would tear down the URMA runtime while + // other UrmaContext objects are still alive, causing the second and + // subsequent contexts to fail on urma_delete_context(). return 0; } @@ -266,6 +299,21 @@ int UrmaContext::registerMemoryRegion(uint64_t va, size_t length) { << "shrink it to " << globalConfig().max_seg_size; length = (size_t)globalConfig().max_seg_size; } + + // urma_register_seg on Kunpeng hardware requires a page-aligned start + // address. Round down to the page boundary and extend the length to cover + // the original range, then round the length up to the next page boundary. + static const uint64_t kPageMask = ~(uint64_t)(4096 - 1); + uint64_t aligned_va = va & kPageMask; + if (aligned_va != va) { + LOG(WARNING) << "registerMemoryRegion: va " << va + << " is not page-aligned, rounding down to " << aligned_va; + length += (va - aligned_va); + va = aligned_va; + } + // Round length up to a multiple of the page size. + length = (length + 4095) & (size_t)kPageMask; + LOG(INFO) << "Register memory region " << va << " length " << length; urma_reg_seg_flag_t flag = {}; flag.bs.token_policy = URMA_TOKEN_NONE; @@ -290,12 +338,45 @@ int UrmaContext::registerMemoryRegion(uint64_t va, size_t length) { } LOG(INFO) << "Local seg token id : " << seg->seg.token_id; local_tseg_list_.push_back(seg); + owned_segs_.insert(seg); RWSpinlock::WriteGuard guard(seg_region_lock_); seg_region_list_.emplace_back(seg, length); return 0; } +void* UrmaContext::lastRegisteredSeg() { + for (auto iter = local_tseg_list_.rbegin(); iter != local_tseg_list_.rend(); + ++iter) { + if (*iter) return *iter; + } + return nullptr; +} + +int UrmaContext::adoptLocalSeg(uint64_t va, size_t length, void* seg) { + // URMA registers memory into a host-global ubva space: a given host VA can + // only be registered with the driver once (re-registering the same VA is + // rejected as a duplicate). When the transport spans multiple devices we + // register the buffer on the primary context and reference the resulting + // segment here instead of calling urma_register_seg again. This segment is + // NOT inserted into owned_segs_, so it will not be passed to + // urma_unregister_seg by this context. + (void)va; + (void)length; + auto* tseg = static_cast(seg); + if (!tseg) { + LOG(ERROR) << "adoptLocalSeg: null segment for va " << va; + return ERR_CONTEXT; + } + local_tseg_list_.push_back(tseg); + + RWSpinlock::WriteGuard guard(seg_region_lock_); + // Use the segment's registered (page-aligned) range so seg() lookups by the + // original, possibly unaligned, address still match. + seg_region_list_.emplace_back(tseg, tseg->seg.len); + return 0; +} + int UrmaContext::unregisterMemoryRegion(uint64_t addr) { RWSpinlock::WriteGuard guard(seg_region_lock_); bool has_removed; @@ -305,12 +386,32 @@ int UrmaContext::unregisterMemoryRegion(uint64_t addr) { iter != seg_region_list_.end(); ++iter) { if ((*iter).first->seg.ubva.va <= addr && addr < (*iter).first->seg.ubva.va + (*iter).second) { - if (urma_unregister_seg((*iter).first)) { - LOG(ERROR) << "Failed to unregister memory " - << (*iter).first->seg.ubva.va; - return ERR_CONTEXT; + urma_target_seg_t* seg_ptr = (*iter).first; + uint64_t seg_va = seg_ptr->seg.ubva.va; + + // l_seg_index is published in BufferDesc and remains valid for + // the lifetime of the context. Keep this slot as a tombstone; + // erasing it would shift indices for other registered buffers. + for (auto& tseg : local_tseg_list_) { + if (tseg == seg_ptr) { + tseg = nullptr; + break; + } } + seg_region_list_.erase(iter); + + // Only the context that actually registered the segment with + // the URMA driver may unregister it. Adopted (shared) segments + // belonging to another context are simply dereferenced here. + auto owned_it = owned_segs_.find(seg_ptr); + if (owned_it != owned_segs_.end()) { + owned_segs_.erase(owned_it); + if (urma_unregister_seg(seg_ptr)) { + LOG(ERROR) << "Failed to unregister memory " << seg_va; + return ERR_CONTEXT; + } + } has_removed = true; break; } @@ -747,6 +848,7 @@ int UrmaEndpoint::setupConnectionsByActive() { local_desc.local_nic_path = context_->nicPath(); local_desc.peer_nic_path = peer_nic_path_; local_desc.jetty_num = JettyNum(); + local_desc.local_eid = context_->getEid(); auto peer_server_name = getServerNameFromNicPath(peer_nic_path_); auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_); @@ -774,6 +876,10 @@ int UrmaEndpoint::setupConnectionsByActive() { return ERR_REJECT_HANDSHAKE; } + if (!peer_desc.local_eid.empty()) { + return doSetupConnection(peer_desc.local_eid, peer_desc.jetty_num); + } + auto segment_desc = context_->engine().meta()->getSegmentDescByName(peer_server_name); if (segment_desc) { @@ -846,6 +952,12 @@ int UrmaEndpoint::setupConnectionsByPassive(const HandShakeDesc& peer_desc, local_desc.local_nic_path = context_->nicPath(); local_desc.peer_nic_path = peer_nic_path_; local_desc.jetty_num = JettyNum(); + local_desc.local_eid = context_->getEid(); + + if (!peer_desc.local_eid.empty()) { + return doSetupConnection(peer_desc.local_eid, peer_desc.jetty_num, + &local_desc.reply_msg); + } auto segment_desc = context_->engine().meta()->getSegmentDescByName(peer_server_name); 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..b15fefea62 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); 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..79591f9ad7 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 @@ -110,7 +110,9 @@ 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_); descs_[idx] = desc; return idx; @@ -133,12 +135,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; } @@ -183,12 +191,36 @@ CUfileIOEvents_t CUFileDescPool::getTransferStatus(int idx, int slice_id) { } unsigned nr = desc->io_params.size(); + if (desc->polled_events.size() < nr) { + LOG(ERROR) << "Completion buffer is too small for descriptor " << idx; + CUfileIOEvents_t event; + event.status = CUFILE_FAILED; + event.ret = -1; + return event; + } CUFILE_CHECK(cuFileBatchIOGetStatus(desc->batch_handle->handle, 0, &nr, - desc->io_events.data(), nullptr)); + desc->polled_events.data(), nullptr)); + + 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 desc->io_events[slice_id]; } +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; +} + int CUFileDescPool::getSliceNum(int idx) { RWSpinlock::ReadGuard guard(mutex_); if (idx < 0 || idx >= (int)MAX_NR_DESC || descs_[idx] == nullptr) { @@ -236,4 +268,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..680c3d74eb 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,108 @@ 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"); + } + auto [slice_id, slice_num] = nvmeof_desc.task_to_slices[task_id]; + thread_local std::vector slice_statuses; + slice_statuses.clear(); + slice_statuses.reserve(slice_num); 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; - } + 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}); } - if (transfer_status.s == COMPLETED) { + + bool is_finished = false; + status = aggregateTransferStatus(slice_statuses, is_finished); + if (is_finished) { 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) { + result.s = has_pending ? PENDING : WAITING; + } + return result; } Status NVMeoFTransport::submitTransfer( 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 71fdcc9358..12bb7a1a75 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp @@ -181,6 +181,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) { @@ -202,6 +209,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 " @@ -343,34 +356,26 @@ int RdmaContext::deconstruct() { 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; - } +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) - // 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. + // 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); + out.method = DmabufExport::Method::kHostReg; #if defined(USE_CUDA) } 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) @@ -407,6 +412,8 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, } 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); @@ -421,17 +428,9 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, #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; - } + out.method = DmabufExport::Method::kDmabufReg; + out.fd = dmabuf_fd; + out.offset = (uintptr_t)addr - (uintptr_t)allocBase; #if defined(USE_CUDA) cuDevicePrimaryCtxRelease(cuDev); #endif @@ -443,8 +442,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 @@ -453,19 +451,33 @@ 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) { - // Pin to the owning device while exporting the dmabuf fd. - HipDeviceGuard dev_guard(hipAttr.device); - if (!dev_guard.set_ok()) { + // 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; + bool need_restore = false; + bool set_ok = false; + explicit HipDeviceGuard(int target_device) { + if (hipGetDevice(&prev_device) == hipSuccess) { + need_restore = (prev_device != target_device); + } + set_ok = (hipSetDevice(target_device) == hipSuccess); + } + ~HipDeviceGuard() { + if (need_restore) { + (void)hipSetDevice(prev_device); + } + } + } dev_guard(hipAttr.device); + if (!dev_guard.set_ok) { LOG(ERROR) << "Failed to set HIP device to " << hipAttr.device << " for dmabuf export of " << (uintptr_t)addr; return ERR_CONTEXT; @@ -498,23 +510,51 @@ 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) { + PLOG(WARNING) << "The buffer length exceeds device max_mr_size, " + << "shrink it to " << globalConfig().max_mr_size; + length = (size_t)globalConfig().max_mr_size; } -#else mrMeta.addr = addr; +#if defined(USE_MLU) || defined(USE_MACA) || defined(USE_CUDA) || \ + defined(USE_HIP_DMABUF) + 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) { @@ -524,9 +564,10 @@ 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) { MemoryRegionMeta mrMeta; - int ret = registerMemoryRegionInternal(addr, length, access, mrMeta); + int ret = registerMemoryRegionInternal(addr, length, access, exp, mrMeta); if (ret != 0) { return ret; } @@ -535,6 +576,20 @@ int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access) { return 0; } +int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access) { + // 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)); @@ -550,9 +605,17 @@ int RdmaContext::unregisterMemoryRegion(void *addr) { } int RdmaContext::preTouchMemory(void *addr, size_t length) { + 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; } @@ -601,7 +664,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; } @@ -745,9 +808,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; @@ -923,6 +988,120 @@ 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; + if (ibv_query_port(current_context, current_port, &port_attr)) { + PLOG(WARNING) << "Failed to refresh port attributes on " + << device_name_ << "/" + << static_cast(current_port); + 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; @@ -1134,6 +1313,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; @@ -1193,4 +1373,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 f5ada5d302..60491f8a55 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; } @@ -79,6 +81,7 @@ 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) {} @@ -102,19 +105,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; @@ -132,6 +135,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; } @@ -160,7 +164,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); @@ -177,15 +182,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); } } } @@ -221,9 +227,10 @@ 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. @@ -251,7 +258,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. @@ -269,13 +276,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"; @@ -320,19 +329,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); @@ -378,6 +390,7 @@ int RdmaEndPoint::setupConnectionsByActive() { } } RWSpinlock::ReadGuard guard(lock_); + if (readyToSend()) return 0; return connected() ? 0 : ERR_ENDPOINT; } @@ -433,127 +446,178 @@ int RdmaEndPoint::setupConnectionsByActive() { } 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; } } } @@ -561,12 +625,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) @@ -581,9 +678,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_) { @@ -617,9 +713,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; } @@ -679,6 +788,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; @@ -691,7 +801,8 @@ void RdmaEndPoint::disconnect() { int RdmaEndPoint::disconnectUnlocked() { auto curr_status = status_.load(std::memory_order_acquire); - if (curr_status != CONNECTED && curr_status != CONNECTING) return 0; + if (!isConnectedStatus(curr_status) && curr_status != CONNECTING) return 0; + ready_wait_start_ts_.store(0, std::memory_order_relaxed); if (!has_connected_) { // Pre-connected handshake retries are allowed to reuse this endpoint: @@ -699,7 +810,7 @@ int RdmaEndPoint::disconnectUnlocked() { // a QP that reached RTS cannot be reliably reset back to RTS. #ifdef CONFIG_ERDMA for (size_t i = 0; i < qp_list_.size(); ++i) { - CHECK_EQ(wr_depth_list_[i], 0) + CHECK_EQ(wr_depth_list_[i].load(std::memory_order_relaxed), 0) << "Pre-connected endpoint must not have outstanding WRs"; } return reconstruct(); @@ -714,7 +825,7 @@ int RdmaEndPoint::disconnectUnlocked() { PLOG(ERROR) << "Failed to modify pre-connected QP to RESET"; ret = ERR_ENDPOINT; } - CHECK_EQ(wr_depth_list_[i], 0) + CHECK_EQ(wr_depth_list_[i].load(std::memory_order_relaxed), 0) << "Pre-connected endpoint must not have outstanding WRs"; } peer_qp_num_list_.clear(); @@ -729,7 +840,8 @@ int RdmaEndPoint::disconnectUnlocked() { 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(); @@ -749,11 +861,42 @@ int RdmaEndPoint::resetConnection(const std::string &reason) { return ERR_ENDPOINT; } +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; + } + 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 { auto status = status_.load(std::memory_order_relaxed); 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)"; @@ -767,7 +910,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; @@ -776,7 +920,8 @@ 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_; + 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, @@ -792,7 +937,8 @@ int RdmaEndPoint::submitPostSend( 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; @@ -830,16 +976,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"; + 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; @@ -905,6 +1059,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()) { @@ -944,7 +1099,7 @@ int RdmaEndPoint::doSetupConnection(const std::string &peer_gid, peer_qp_num_list_ = std::move(peer_qp_num_list); has_connected_ = true; - status_.store(CONNECTED, std::memory_order_relaxed); + status_.store(connected_status, std::memory_order_relaxed); return 0; } 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 9b76dacb03..8472482b6d 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp @@ -250,8 +250,24 @@ int RdmaTransport::registerLocalMemoryInternal(void *addr, size_t length, } } + // Export a single dma_buf fd for the buffer and import it into every NIC's + // PD below, closing it once after all registrations. One dma_buf object + // shared across NICs lets the GPU driver reserve a single BAR1 window for + // the buffer instead of one per NIC. Host memory yields an empty export and + // takes the plain ibv_reg_mr path. The fd must stay open across all + // registrations (each MR takes its own reference); see exportDmabuf(). + 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; + } + } + auto reg_start = std::chrono::steady_clock::now(); + int reg_error = 0; if (use_parallel_reg) { std::vector reg_threads; reg_threads.reserve(context_list_.size()); @@ -259,10 +275,11 @@ int RdmaTransport::registerLocalMemoryInternal(void *addr, size_t length, 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, i, addr, length, ar]() { - ret_codes[i] = - context_list_[i]->registerMemoryRegion(addr, length, ar); - }); + reg_threads.emplace_back( + [this, &ret_codes, &dmabuf_exp, i, addr, length, ar]() { + ret_codes[i] = context_list_[i]->registerMemoryRegion( + addr, length, ar, dmabuf_exp); + }); } for (auto &thread : reg_threads) { @@ -273,21 +290,32 @@ int RdmaTransport::registerLocalMemoryInternal(void *addr, size_t length, if (ret_codes[i] != 0) { LOG(ERROR) << "Failed to register memory region with context " << i; - return ret_codes[i]; + reg_error = ret_codes[i]; + break; } } } else { for (size_t i = 0; i < context_list_.size(); ++i) { - int ret = context_list_[i]->registerMemoryRegion(addr, length, - access_rights); + int ret = context_list_[i]->registerMemoryRegion( + addr, length, access_rights, dmabuf_exp); if (ret) { LOG(ERROR) << "Failed to register memory region with context " << i; - return ret; + reg_error = ret; + break; } } } + // Close the single dma_buf fd now that all NIC registrations are done. + // Each successful MR holds its own reference, so the underlying dma_buf + // (and its BAR1 window) stays alive until those MRs are deregistered. + RdmaContext::closeDmabufExport(dmabuf_exp); + + if (reg_error != 0) { + return reg_error; + } + auto reg_end = std::chrono::steady_clock::now(); auto reg_duration_ms = std::chrono::duration_cast(reg_end - 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 ae309ce18c..651057ad3c 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp @@ -203,9 +203,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; @@ -271,6 +275,26 @@ int WorkerPool::submitPostSend( 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) { int post_tid = 0; int post_count = 0; @@ -344,6 +368,11 @@ void WorkerPool::performPostSend(int thread_id) { processed_slice_count_.fetch_add(entry.second.size()); entry.second.clear(); #else + if (!isRailAvailable(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()) @@ -365,6 +394,18 @@ void WorkerPool::performPostSend(int thread_id) { 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"; + handlePathFailure(entry.first, 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 for (auto &slice : entry.second) { slice->rdma.endpoint = endpoint.get(); @@ -390,9 +431,23 @@ void WorkerPool::performPostSend(int thread_id) { } 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)) { + } + } + int processed_slice_count = 0; const static size_t kPollCount = 64; - std::unordered_map qp_depth_set; + std::unordered_map *, int> qp_depth_set; SliceList failed_slice_list; // Unified: collect all slices for redispatch for (int cq_index = 0; cq_index < context_.cqCount(); cq_index++) { ibv_wc wc[kPollCount]; @@ -402,6 +457,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); @@ -452,12 +515,12 @@ void WorkerPool::performPollCq(int thread_id) { } } 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) { processed_slice_count_.fetch_add(processed_slice_count); @@ -503,6 +566,35 @@ 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->name, + 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) { + slice->markFailed(); + processed_slice_count_++; + continue; + } + } slice->peer_nic_path = peer_nic_path; if (globalConfig().log_rdma_slice_affinity) { VLOG(1) << "RDMA slice affinity: source_location=" @@ -521,6 +613,7 @@ void WorkerPool::redispatch(std::vector &slice_list, << ", 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); @@ -547,7 +640,9 @@ void WorkerPool::redispatch(std::vector &slice_list, 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) > 0) return true; + if (context_.cqOutstandingCount(cq_index)->load( + std::memory_order_relaxed) > 0) + return true; } return false; } @@ -633,6 +728,7 @@ int WorkerPool::doProcessContextEvents() { event.event_type == IBV_EVENT_PORT_ERR || event.event_type == IBV_EVENT_LID_CHANGE) { context_.set_active(false); + refreshPublishedLocalTopology(); /** * Similar deadlock might happen if we call @@ -651,9 +747,23 @@ int WorkerPool::doProcessContextEvents() { context_.disconnectAllEndpoints(); LOG(INFO) << "Worker: Context " << context_.deviceName() - << " is now inactive"; + << " is now inactive due to fatal event: " + << event.event_type; + } else if (event.event_type == IBV_EVENT_GID_CHANGE) { + auto gid_refresh_result = refreshPublishedLocalGid(); + ibv_ack_async_event(&event); + event_acked = true; + + if (gid_refresh_result != GidRefreshResult::UNCHANGED) { + context_.disconnectAllEndpoints(); + LOG(INFO) << "Worker: Context " << context_.deviceName() + << " GID refresh result=" + << static_cast(gid_refresh_result) + << ", disconnected all endpoints"; + } } else if (event.event_type == IBV_EVENT_PORT_ACTIVE) { context_.set_active(true); + refreshPublishedLocalTopology(); markContextSuccess(); // Reset failure counter on port recovery LOG(INFO) << "Worker: Context " << context_.deviceName() << " is now active"; @@ -666,11 +776,57 @@ int WorkerPool::doProcessContextEvents() { return 0; } +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()); if (current_ts - last_reset_ts > 1000000000ll) { // Drain endpoint_store_->waiting_list_ even when no new // insertions are happening. Without this, reclaim only runs @@ -679,6 +835,105 @@ void WorkerPool::monitorWorker() { context_.reclaimEndpoints(); 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) { 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..4cc7e93a38 100644 --- a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp @@ -15,7 +15,7 @@ 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 +33,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(); } @@ -395,7 +400,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 +417,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..28b6dab51f 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; 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 2ac73e1191..253d7ed32c 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" @@ -113,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) @@ -125,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( @@ -147,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_ && @@ -159,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(); + } }); } @@ -177,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; @@ -205,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 } @@ -230,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; @@ -246,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; } @@ -280,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 } @@ -305,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; @@ -319,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( @@ -353,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); }); } @@ -379,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; } @@ -413,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; } @@ -451,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; @@ -482,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; } @@ -516,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; } @@ -546,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; @@ -708,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; @@ -1026,6 +1255,28 @@ bool TcpTransport::validateAddress(uint64_t addr, uint64_t size) const { return false; } +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) { auto desc = metadata_->getSegmentDescByID(slice->target_id); if (!desc) { @@ -1045,6 +1296,15 @@ void TcpTransport::startTransfer(Slice* slice) { 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) { + slice->markSuccess(); + return; + } + // Get connection from pool auto socket = getConnection(meta_entry.ip_or_host_name, desc->tcp_data_port); @@ -1056,7 +1316,9 @@ void TcpTransport::startTransfer(Slice* slice) { } try { - auto session = std::make_shared(socket); + const bool use_v2 = + desc->tcp_proto_version >= 2 && !forceLegacyTcpProto(); + auto session = std::make_shared(socket, use_v2); session->on_finalize_ = [slice](TransferStatusEnum status) { if (status == TransferStatusEnum::COMPLETED) @@ -1065,15 +1327,21 @@ void TcpTransport::startTransfer(Slice* slice) { slice->markFailed(); }; - // Return connection to pool when transfer completes, or close if - // disabled + // 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 (enable_connection_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; @@ -1091,32 +1359,11 @@ 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); - } - } - } + // 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); slice->markFailed(); } } diff --git a/mooncake-transfer-engine/tent/docs/ub_tent_transport_guide.md b/mooncake-transfer-engine/tent/docs/ub_tent_transport_guide.md new file mode 100644 index 0000000000..10a7338a50 --- /dev/null +++ b/mooncake-transfer-engine/tent/docs/ub_tent_transport_guide.md @@ -0,0 +1,639 @@ +# TENT UB Transport Phase 3 Test Guide + +This guide documents how to validate UB transport support in TENT on the `UB_TENT` branch. + +This document focuses on the current Phase 3 work in this branch: enabling the existing Kunpeng UB/URMA data path to be reused by TENT through a TENT transport adapter, metadata bridge, and control-plane bootstrap path. + +## 1. Scope + +Compared with `main`, this branch adds TENT-side support for UB transport and adjusts the existing old Transfer Engine UB implementation so that it can be reused by TENT. + +| Area | `main` | `UB_TENT` | +| ------------------------ | ------------------------------------------ | ------------------------------------------------------------------ | +| TENT transport enum | No `UB` transport type | Adds `TransportType::UB` | +| Transport selector | Cannot parse or select `"ub"` | Supports `"ub"` as a transport policy entry | +| Transport loader | Cannot instantiate UB in TENT | Creates `UbTentTransport` when UB is enabled | +| TENT UB adapter | Not available | Adds `UbTentTransport` | +| Metadata bridge | Old UB depends on old `TransferMetadata` | Adds `UbTentMetadataBridge` to resolve TENT segments for old UB | +| UB bootstrap | Old UB handshake path only | Adds TENT `BootstrapUb` RPC path | +| Local segment publishing | No UB-specific TENT attrs | Publishes UB EID and tseg data through TENT transport attrs | +| Submit fallback | Submit failure may not be retried cleanly | Allows failed sub-batches to be retried through poll-time failover | +| URMA registration | Multiple contexts may register the same VA | Primary context registers once; other contexts adopt the segment | +| Tests | No TENT UB tests | Adds `tent_ub_transport_test` and `tent_ub_e2e_dual_node_test` | + +Main files to review: + +```text +mooncake-transfer-engine/tent/include/tent/transport/ub/* +mooncake-transfer-engine/tent/src/transport/ub/* +mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h +mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp +mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp +mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +mooncake-transfer-engine/tent/tests/ub_tent_transport_test.cpp +mooncake-transfer-engine/tent/tests/ub_e2e_dual_node_test.cpp +mooncake-transfer-engine/tent/tests/CMakeLists.txt +mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp +mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h +mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h +``` + +## 2. Build Configuration + +Use `BUILD_UNIT_TESTS`, not `BUILD_TESTS`. + +There is no `MOCK_URMA` CMake option in the current code. Mock URMA is selected by the existing UB CMake logic when the real URMA library is not found. + +Typical local build: + +```bash +cmake -S . -B build-ub-tent \ + -DUSE_TENT=ON \ + -DUSE_UB=ON \ + -DBUILD_UNIT_TESTS=ON \ + -DCMAKE_BUILD_TYPE=Debug + +cmake --build build-ub-tent \ + --target tent_xport_ub tent_ub_transport_test tent_ub_e2e_dual_node_test \ + --parallel +``` + +If URMA headers are installed in a non-default location: + +```bash +cmake -S . -B build-ub-tent \ + -DUSE_TENT=ON \ + -DUSE_UB=ON \ + -DBUILD_UNIT_TESTS=ON \ + -DURMA_INCLUDE_DIR=/path/to/urma/include \ + -DCMAKE_BUILD_TYPE=Debug +``` + +If the real URMA runtime is available on the test machine, make sure the runtime library and the selected UB device are usable before running hardware tests. + +## 3. Local Unit Test + +Unit test source: + +```text +mooncake-transfer-engine/tent/tests/ub_tent_transport_test.cpp +``` + +CTest command: + +```bash +ctest --test-dir build-ub-tent \ + -R '^tent_ub_transport_test$' \ + --output-on-failure \ + -V +``` + +Direct binary command: + +```bash +cd build-ub-tent + +GLOG_logtostderr=1 \ +./mooncake-transfer-engine/tent/tests/tent_ub_transport_test +``` + +Run one GTest case: + +```bash +cd build-ub-tent + +GLOG_logtostderr=1 \ +./mooncake-transfer-engine/tent/tests/tent_ub_transport_test \ + --gtest_filter="UbTentTransportTest.InstallWithMockUrma" +``` + +## 4. Unit Test Coverage + +| Test case | Coverage | +| ------------------------------------------------- | ---------------------------------------------------------------------------------------- | +| `UbSelectorTest.TypeNameRoundTrip` | Verifies that `TransportSelector::transportTypeName(UB)` returns `"ub"` | +| `UbSelectorTest.ParseUbString` | Verifies that `"ub"` parses to `TransportType::UB` | +| `UbSelectorTest.ParseUnknownStillReturnsUnspec` | Verifies that unknown transport strings still return `UNSPEC` | +| `UbSelectorTest.UbEnumValue` | Verifies that the UB enum value is inside the supported transport range | +| `UbSelectorTest.SelectorPicksUbWhenFirstInPolicy` | Verifies that a policy with `"ub"` first can select UB when UB is available | +| `UbTentTransportTest.InstallWithMockUrma` | Verifies basic install lifecycle with mock URMA and checks the transport name/capability | +| `UbTentTransportTest.AddAndRemoveMemoryBuffer` | Verifies page-aligned memory add/remove flow and updates to `desc.transports` | +| `UbTentTransportTest.AllocateAndFreeSubBatch` | Verifies sub-batch allocation and release through the old UB batch path | +| `UbTentTransportTest.SubmitAndPollMockTransfer` | Verifies that the mock submit/poll path does not crash or hang | +| `UbTentTransportTest.DoubleUninstallSafe` | Verifies that repeated uninstall is safe | + +The unit test mainly covers selector mapping, adapter lifecycle, memory buffer lifecycle, sub-batch lifecycle, and mock transfer smoke behavior. + +It does not prove that the real UB data plane completes on hardware. It also does not fully assert the serialized `transport_attrs[UB]` content. That deeper validation belongs to the dual-node integration test and hardware inspection. + +On a host where real `liburma.so` is present but no usable UB HCA exists, hardware-dependent setup may fail and the corresponding test path may skip. That is expected for local developer machines without Kunpeng UB hardware. + +## 5. Behavior To Validate + +### 5.1 Transport Selection + +The branch adds `UB` to the TENT transport enum and maps it to `"ub"`. + +Example policy: + +```json +{ + "policy": [ + { + "name": "kunpeng_ub_memory", + "segment_type": "memory", + "local_memory": "cpu", + "remote_memory": "cpu", + "same_machine": false, + "transports": ["ub", "rdma", "tcp"] + } + ] +} +``` + +Expected behavior: + +```text +The selector can parse "ub" and select TransportType::UB when UB is available. +``` + +### 5.2 UB Transport Loading + +When built with both `USE_UB=ON` and `USE_TENT=ON`, TENT can load UB if UB is enabled in config. + +Example: + +```json +{ + "transports": { + "ub": { + "enable": true + } + } +} +``` + +Disable UB explicitly: + +```json +{ + "transports": { + "ub": { + "enable": false + } + } +} +``` + +### 5.3 Device Selection + +`UbTentTransport::install()` should resolve UB devices in this order: + +```text +1. TENT config key: transports/ub/device_name +2. Environment variable: MC_UB_DEVICE_NAME +3. Auto-discovery of UB devices +``` + +For Kunpeng SuperNode environments, pin the expected logical or bonded UB device explicitly. + +Example config: + +```json +{ + "transports": { + "ub": { + "enable": true, + "device_name": "bonding_dev_0" + } + } +} +``` + +Equivalent environment override: + +```bash +export MC_UB_DEVICE_NAME=bonding_dev_0 +``` + +`device_name` and `MC_UB_DEVICE_NAME` may be comma-separated lists if multiple UB devices should be considered. + +### 5.4 Local Segment Publishing + +After UB transport installation, the local TENT segment should publish UB-specific device information. + +Expected behavior: + +```text +MemorySegmentDesc.devices[*].transport_attrs[UB] contains UB EID information. +MemorySegmentDesc.transport_attrs[UB] marks UB availability. +BufferDesc.transport_attrs[UB] contains the UB tseg handle after memory registration. +BufferDesc.transports contains TransportType::UB after successful registration. +``` + +This is needed because the remote node reconstructs the old-TE UB segment descriptor from TENT metadata. + +### 5.5 Metadata Bridge + +`UbTentMetadataBridge` allows old UB code to resolve TENT segments through the old `TransferMetadata` interface. + +Expected behavior: + +```text +Remote getSegmentDescByID() uses the TENT SegmentManager remote cache. +force_update=true invalidates both bridge-side and TENT-side remote cache. +getSegmentDescByName() opens the remote TENT segment and converts it. +getSegmentID(name) returns the TENT remote segment handle as the old-TE segment ID. +convertFromTent() extracts UB EID and tseg data from TENT transport attrs. +``` + +### 5.6 UB Bootstrap Through TENT RPC + +The old UB handshake daemon is not started by the bridge. Instead, UB endpoint bootstrap is routed through TENT control-plane RPC. + +Expected call path: + +```text +UbEndpoint + -> TransferMetadata::sendHandshake() + -> UbTentMetadataBridge::sendHandshake() + -> ControlClient::bootstrapUb() + -> remote ControlService::onBootstrapUb() + -> old-TE UB handshake callback +``` + +`UbBootstrapDesc` should carry: + +```text +local_nic_path +peer_nic_path +jetty_num +local_eid +reply_msg +``` + +This validates that UB connection setup can use the TENT control plane instead of the old standalone UB handshake daemon. + +## 6. Dual-Node Integration Test + +Integration test source: + +```text +mooncake-transfer-engine/tent/tests/ub_e2e_dual_node_test.cpp +``` + +Build target: + +```bash +cmake --build build-ub-tent \ + --target tent_ub_e2e_dual_node_test \ + --parallel +``` + +This executable is intentionally not registered with CTest because it requires: + +```text +1. Two Kunpeng UB-capable nodes +2. A usable URMA runtime +3. A shared TENT metadata backend +4. Network reachability between both TENT RPC servers +5. UB/URMA fabric connectivity between the two nodes +``` + +### 6.1 Hardware And Runtime Checks + +Run on both nodes: + +```bash +ls /dev/urma* || true +urma_cmd -q all +``` + +If needed, load the platform driver: + +```bash +modprobe urma_udrv +``` + +Also confirm that both nodes can reach: + +```text +The shared metadata backend, for example etcd +Each other's TENT RPC server address +The selected UB device, for example bonding_dev_0 +``` +### 6.2 Start A Shared etcd Metadata Backend + +The dual-node test requires both nodes to use the same TENT metadata backend. If `metadata_type` is set to `etcd`, start one etcd instance on either Node A or a third reachable node before running the test. + +Example on Node A: + +```bash +etcd \ + --name mooncake-ub-test \ + --data-dir /tmp/mooncake-ub-etcd \ + --listen-client-urls http://0.0.0.0:11451 \ + --advertise-client-urls http://NODE_A_IP:11451 \ + --listen-peer-urls http://127.0.0.1:11452 \ + --initial-advertise-peer-urls http://127.0.0.1:11452 \ + --initial-cluster mooncake-ub-test=http://127.0.0.1:11452 \ + --initial-cluster-state new +``` + +Both Node A and Node B should use the same metadata server address: + +```json +"metadata_type": "etcd", +"metadata_servers": "NODE_A_IP:11451" +``` + +Do not use `127.0.0.1:11451` in the config unless both the server and client run on the same host. In a two-node test, `127.0.0.1` on Node B points to Node B itself, not to the etcd instance on Node A. + +Before running the test, verify connectivity from both nodes: + +```bash +curl http://NODE_A_IP:11451/version +``` + +### 6.3 Recommended Config Files + +The test binary has a built-in UB-only config, but real two-node testing should use explicit config files. + +Server config example, `node_a_ub.json`: + +```json +{ + "metadata_type": "etcd", + "metadata_servers": "ETCD_IP:2379", + "local_segment_name": "node_a_seg", + "rpc_server_hostname": "NODE_A_IP", + "rpc_server_port": 0, + "transports": { + "tcp": { + "enable": false + }, + "rdma": { + "enable": false + }, + "ub": { + "enable": true, + "device_name": "bonding_dev_0" + } + }, + "policy": [ + { + "name": "ub_memory", + "segment_type": "memory", + "local_memory": "cpu", + "remote_memory": "cpu", + "same_machine": false, + "transports": ["ub"] + } + ] +} +``` + +Client config example, `node_b_ub.json`: + +```json +{ + "metadata_type": "etcd", + "metadata_servers": "ETCD_IP:2379", + "local_segment_name": "node_b_seg", + "rpc_server_hostname": "NODE_B_IP", + "rpc_server_port": 0, + "transports": { + "tcp": { + "enable": false + }, + "rdma": { + "enable": false + }, + "ub": { + "enable": true, + "device_name": "bonding_dev_0" + } + }, + "policy": [ + { + "name": "ub_memory", + "segment_type": "memory", + "local_memory": "cpu", + "remote_memory": "cpu", + "same_machine": false, + "transports": ["ub"] + } + ] +} +``` + +Notes: + +```text +1. Use the same metadata backend on both nodes. +2. Use different local segment names on the two nodes. +3. Keep the server's --segment_name consistent with local_segment_name to avoid confusion. +4. If metadata_type is p2p, TENT may replace the local segment name with the RPC address. +5. For stable cross-node tests, etcd is easier to reason about than p2p metadata. +``` + +### 6.4 Run The Test + +On Node A: + +```bash +cd build-ub-tent + +GLOG_logtostderr=1 GLOG_v=1 \ +./mooncake-transfer-engine/tent/tests/tent_ub_e2e_dual_node_test \ + --role=server \ + --segment_name=node_a_seg \ + --transport_config=/path/to/node_a_ub.json +``` + +On Node B: + +```bash +cd build-ub-tent + +GLOG_logtostderr=1 GLOG_v=1 \ +./mooncake-transfer-engine/tent/tests/tent_ub_e2e_dual_node_test \ + --role=client \ + --remote_segment=node_a_seg \ + --transport_config=/path/to/node_b_ub.json \ + --data_size=1048576 \ + --operation=write +``` + +Expected client-side result: + +```text +Client: WRITE 1048576 bytes ... COMPLETED +Client: READ 1048576 bytes ... COMPLETED +Client: data integrity VERIFIED +Client: test PASSED +``` + +`--operation=write` writes a known pattern to the remote buffer, reads it back, and verifies the returned bytes. + +`--operation=read` only executes the read path and does not verify a known pattern. + +### 6.5 What The Integration Test Covers + +The dual-node integration test validates: + +```text +1. Server-side TENT segment publication +2. UB EID publication through TENT segment device attrs +3. UB tseg publication through TENT buffer attrs +4. Client-side remote segment open +5. Metadata conversion through UbTentMetadataBridge +6. TENT request conversion through UbTentTransport +7. UB endpoint bootstrap through BootstrapUb RPC +8. Old-TE UB data path reuse from TENT +9. Remote WRITE +10. Remote READ +11. Data integrity after write + read-back +``` + +## 7. Debugging Integration Failures + +| Symptom | Checks | +| ---------------------------------- | ------------------------------------------------------------------------------------------------------------------------------- | +| `openSegment("node_a_seg") failed` | Confirm both nodes use the same metadata backend, the server is still running, and `local_segment_name` is `node_a_seg` | +| UB transport not installed | Check `USE_UB=ON`, `USE_TENT=ON`, `transports.ub.enable=true`, URMA headers/runtime, and `device_name` | +| Wrong UB device selected | Set `transports.ub.device_name` or `MC_UB_DEVICE_NAME` explicitly | +| Remote segment has no buffers | Confirm server-side memory registration succeeded and segment synchronization reached metadata | +| Remote segment has no UB attrs | Confirm UB transport installation happened before segment synchronization | +| `BootstrapUb` RPC failed | Check `rpc_server_hostname`, firewall, remote RPC reachability, and callback registration | +| URMA duplicate registration error | Confirm the branch includes the primary-register plus adopt-segment lifecycle change | +| Transfer hangs | Check UB fabric connectivity, selected UB device, endpoint creation logs, and poll/fallback logs | +| Data mismatch | Use `--operation=write`, confirm both nodes use the same `data_size`, and make sure no unintended fallback transport is enabled | + +Useful etcd inspection: + +```bash +etcdctl get mooncake/tent/ --prefix + +etcdctl get mooncake/tent/node_a_seg --print-value-only | python3 -m json.tool +``` + +Look for: + +```text +detail.devices[*].transport_attrs containing UB EID information +detail.buffers[*].transport_attrs containing UB tseg information +rpc_server_addr pointing to Node A's reachable TENT RPC address +``` + +## 8. Fallback And Regression Tests + +This branch also changes submit failure handling so that a failed sub-batch submit can be retried through poll-time failover. + +Run existing TENT regression tests: + +```bash +ctest --test-dir build-ub-tent \ + -R '^tent_engine_failover_e2e_test$' \ + --output-on-failure \ + -V + +ctest --test-dir build-ub-tent \ + -R '^tent_transport_hint_test$' \ + --output-on-failure \ + -V + +ctest --test-dir build-ub-tent \ + -R '^tent_transport_selector_test$' \ + --output-on-failure \ + -V +``` + +Old Transfer Engine UB regression target: + +```bash +cmake --build build-ub-tent \ + --target ub_transport_test \ + --parallel +``` + +Example run: + +```bash +GLOG_logtostderr=1 \ +build-ub-tent/mooncake-transfer-engine/tests/ub_transport_test \ + --device_name=mock_urma_device +``` + +`ub_transport_test` may allocate a large NUMA buffer, so it may not be suitable for every developer machine. + +## 9. Test Matrix + +| Test | Hardware | CTest | Coverage | +| ------------------------------- | ---------------------------------------: | ----: | ------------------------------------------------------------------------------------------------------------------------------- | +| `tent_ub_transport_test` | No real UB hardware if mock URMA is used | Yes | UB enum/string mapping, selector, install lifecycle, memory registration lifecycle, sub-batch lifecycle, mock submit/poll smoke | +| `tent_ub_e2e_dual_node_test` | Yes, two Kunpeng UB-capable nodes | No | TENT segment publishing, metadata bridge, BootstrapUb RPC, old-TE UB data path through TENT, data integrity | +| `tent_engine_failover_e2e_test` | No | Yes | Submit-failure poll-time resubmit behavior | +| `tent_transport_hint_test` | No | Yes | Transport hint behavior after adding UB | +| `tent_transport_selector_test` | No | Yes | Selector regression around transport policy handling | +| `ub_transport_test` | Mock or real UB, depending on build | No | Old Transfer Engine UB registration and transfer regression | + +## 10. Common Pitfalls + +Do not use: + +```bash +-DMOCK_URMA=ON +``` + +The current code does not define this CMake option. + +Do not use: + +```bash +-DBUILD_TESTS=ON +``` + +Use: + +```bash +-DBUILD_UNIT_TESTS=ON +``` + +Use the real TENT UB targets: + +```text +tent_ub_transport_test +tent_ub_e2e_dual_node_test +``` + +Do not use old or planned names such as: + +```text +tent_ub_transfer_test +tent_ub_e2e_test +``` + +The dual-node executable is not registered as a CTest test. + +For dual-node tests, prefer explicit config files and shared metadata. Keep the server `--segment_name` aligned with the config `local_segment_name`. + +UB memory registration should use page-aligned buffers. The tests use page-aligned allocation for this reason. + +## 11. Expected Validation Summary + +A complete validation should include: + +```text +1. Build with USE_TENT=ON and USE_UB=ON. +2. Run tent_ub_transport_test locally. +3. Run selector, hint, and failover regression tests. +4. Run old Transfer Engine ub_transport_test when the environment allows it. +5. Run tent_ub_e2e_dual_node_test on two Kunpeng UB-capable nodes. +6. Confirm metadata contains UB EID and tseg attrs. +7. Confirm BootstrapUb RPC is exercised during connection setup. +8. Confirm write + read-back data integrity in the dual-node test. +``` 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..a323a5ffed --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/common/qos_metrics.h @@ -0,0 +1,101 @@ +// 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; +}; + +struct QosClassSample { + size_t operations = 0; + double total_duration_us = 0.0; + double p99_us = 0.0; + std::optional slo_attainment; +}; + +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 27a7e98d32..a4ceb1c870 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/types.h +++ b/mooncake-transfer-engine/tent/include/tent/common/types.h @@ -39,9 +39,10 @@ struct Notification { std::string msg; }; -#ifndef LOCAL_SEGMENT_ID -#define LOCAL_SEGMENT_ID (0ull) -#endif +// Local segment handle. Use constexpr (not #define) so this does not clash +// with mooncake::LOCAL_SEGMENT_ID in common.h when USE_UB pulls in old-TE +// headers in the same translation unit. +static constexpr SegmentID LOCAL_SEGMENT_ID = 0; enum TransportType : int { UNSPEC = 0, @@ -54,6 +55,8 @@ enum TransportType : int { TCP, AscendDirect, SUNRISE_LINK, + UB, // Kunpeng UB / URMA transport + TPU, // Sentinel: must remain the last enumerator. kNumTransportTypes, }; @@ -65,6 +68,16 @@ inline TransportType c_to_transport_hint(int v) { return static_cast(v); } +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; @@ -85,6 +98,7 @@ struct Request { // (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 { 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 4abc12faf7..79e4072dfb 100644 --- a/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h +++ b/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h @@ -87,6 +87,15 @@ class TentMetrics { // transfer met its deadline; mlu >= 1 means it missed. Observability only. void recordDeadlineMLU(double mlu); + enum class Stage { + QueueWait, + Dispatch, + Transport, + }; + + // Causal chain: record per-stage latency breakdown (microseconds). + void recordStageLatency(Stage stage, double latency_us); + // Get metrics for HTTP server std::string getPrometheusMetrics(); std::string getJsonMetrics(); @@ -178,6 +187,21 @@ class TentMetrics { "Deadline feasibility ratio (MLU x 1000) distribution", kMluPerMilleBuckets}; + // 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::histogram_t stage_queue_wait_{ + "tent_stage_queue_wait_us", + "Causal chain: queue wait latency in microseconds", kStageBuckets}; + ylt::metric::histogram_t stage_dispatch_{ + "tent_stage_dispatch_us", + "Causal chain: dispatch latency in microseconds", kStageBuckets}; + ylt::metric::histogram_t stage_transport_{ + "tent_stage_transport_us", + "Causal chain: transport execution latency in microseconds", + kStageBuckets}; + // Helper to register all metrics to the vectors void registerMetrics(); #endif // TENT_METRICS_ENABLED @@ -283,6 +307,14 @@ class ScopedLatencyRecorder { ::mooncake::tent::ScopedLatencyRecorder _tent_latency_recorder_( \ ::mooncake::tent::ScopedLatencyRecorder::OperationType::Write, bytes) +#define TENT_RECORD_STAGE_LATENCY(stage, latency_us) \ + do { \ + if (::mooncake::tent::TentMetrics::isEnabled()) { \ + ::mooncake::tent::TentMetrics::instance().recordStageLatency( \ + stage, latency_us); \ + } \ + } while (0) + #else // !TENT_METRICS_ENABLED // No-op stub class for ScopedLatencyRecorder when metrics are disabled @@ -301,6 +333,7 @@ class ScopedLatencyRecorder { #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_STAGE_LATENCY(stage, 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/rpc/rpc.h b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h index 3105c6fdaf..7eb0aab5d9 100644 --- a/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h +++ b/mooncake-transfer-engine/tent/include/tent/rpc/rpc.h @@ -47,6 +47,7 @@ enum RpcFuncID { Unpin, SubscribeSegmentUpdate, NotifySegmentUpdated, + BootstrapUb, }; class ClientPool; 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 e52f70602a..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, @@ -103,6 +167,7 @@ class LocalTransferAdmissionQueue { uint64_t batch_token{0}; Request request{}; QueueOwnerKind kind{QueueOwnerKind::User}; + bool degradation_eligible{false}; QueueState state{QueueState::Queued}; TransferStatusEnum terminal_status{TransferStatusEnum::PENDING}; }; @@ -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/control_plane.h b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h index a4739f27aa..fed076ff59 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/control_plane.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -54,6 +55,23 @@ struct BootstrapDesc { notify_qp_num); }; +// UB/URMA-specific bootstrap descriptor for peer-to-peer jetty exchange. +struct UbBootstrapDesc { + std::string local_nic_path; + std::string peer_nic_path; + std::vector jetty_num; + std::string local_eid; + std::string reply_msg; + + public: + NLOHMANN_DEFINE_TYPE_INTRUSIVE(UbBootstrapDesc, local_nic_path, + peer_nic_path, jetty_num, local_eid, + reply_msg); +}; + +using OnReceiveUbBootstrap = + std::function; + struct XferDataDesc { uint64_t peer_mem_addr; size_t length; @@ -78,6 +96,10 @@ class ControlClient { const BootstrapDesc& request, BootstrapDesc& response); + static Status bootstrapUb(const std::string& server_addr, + const UbBootstrapDesc& request, + UbBootstrapDesc& response); + static Status sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, size_t length); @@ -125,6 +147,11 @@ class ControlService { bootstrap_callback_ = callback; } + void setBootstrapUbCallback(const OnReceiveUbBootstrap& callback) { + std::lock_guard lock(ub_bootstrap_callback_mutex_); + ub_bootstrap_callback_ = callback; + } + void setNotifyCallback(const OnNotify& callback) { notify_callback_ = callback; } @@ -138,6 +165,8 @@ class ControlService { void onBootstrapRdma(const std::string_view& request, std::string& response); + void onBootstrapUb(const std::string_view& request, std::string& response); + void onSendData(const std::string_view& request, std::string& response); void onRecvData(const std::string_view& request, std::string& response); @@ -165,6 +194,8 @@ class ControlService { std::shared_ptr rpc_server_; OnReceiveBootstrap bootstrap_callback_; + std::mutex ub_bootstrap_callback_mutex_; + OnReceiveUbBootstrap ub_bootstrap_callback_; OnNotify notify_callback_; TransferEngineImpl* impl_; }; 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/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/transfer_engine_impl.h b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h index 41aca93ad3..5ae33785bc 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 @@ -55,11 +55,15 @@ 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{}; // Submit time + std::chrono::steady_clock::time_point dispatch_time{}; // Dispatch entry + std::chrono::steady_clock::time_point post_time{}; // Transport post }; class TransferEngineImpl { @@ -139,6 +143,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); @@ -245,6 +251,8 @@ class TransferEngineImpl { 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, diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index b28cf7cf2a..e8433b2fe2 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -55,18 +55,22 @@ class Transport { } 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) { @@ -77,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); } @@ -134,17 +149,19 @@ 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; } 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..99d0929f54 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; }; /** diff --git a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h index ff93e0ed02..b21f1efc7e 100644 --- a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h +++ b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h @@ -25,9 +25,9 @@ extern "C" { #define tent_batch_id_t uint64_t #define tent_segment_id_t uint64_t -#ifndef LOCAL_SEGMENT_ID -#define LOCAL_SEGMENT_ID (0ull) -#endif +// C API local-segment sentinel. Do not alias to LOCAL_SEGMENT_ID here: that +// name is used by mooncake::tent (constexpr) and mooncake (old TE const). +#define TENT_LOCAL_SEGMENT_ID 0ull #define OPCODE_READ (0) #define OPCODE_WRITE (1) @@ -174,6 +174,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); @@ -299,6 +302,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); @@ -328,4 +336,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 7507453b00..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 @@ -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. 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 c619faec6c..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 @@ -48,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. 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 245715871c..287b624a5f 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]; }; @@ -160,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; } @@ -172,6 +173,10 @@ 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); @@ -193,11 +198,15 @@ class RdmaEndPoint : public std::enable_shared_from_this { 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_; 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 f67906c235..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; 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..8cdc95b8dd 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h @@ -149,6 +149,8 @@ class DeviceSelector { void printTrafficStats(); + double getAggregateEwmaBandwidth() const; + void fillDevicePriorities(); int getDevicePriority(int dev_id) const; 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 3dc6deba4b..ca946fd30b 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 @@ -79,6 +79,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 +95,8 @@ class RdmaTransport : public Transport { virtual const char* getName() const { return "rdma"; } + double getEstimatedBandwidth() const override; + virtual bool supportNotification() const override { return true; } virtual Status sendNotification(SegmentID target_id, 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..8332a7c34a 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}; @@ -87,6 +96,9 @@ struct RdmaSlice { int qp_index = 0; int retry_count = 0; 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 +118,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 +139,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 b8ba52d78d..1db9ed9da3 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h @@ -43,7 +43,7 @@ class Workers { using BoundedSliceQueue = BoundedMPSCQueue; public: - Workers(RdmaTransport *transport); + Workers(RdmaTransport* transport); ~Workers(); @@ -51,16 +51,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,38 +69,45 @@ 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); + int handleContextEvents(std::shared_ptr& context); - Status generatePostPath(RdmaSlice *slice); + 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); - 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_; @@ -110,7 +118,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; @@ -118,7 +126,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); @@ -128,10 +136,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); } @@ -188,7 +196,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; @@ -207,17 +215,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..38926d6cd1 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_; @@ -112,4 +115,4 @@ class ShmTransport : public Transport { } // 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/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/ub_tent_metadata_bridge.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_tent_metadata_bridge.h new file mode 100644 index 0000000000..74b5a13713 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_tent_metadata_bridge.h @@ -0,0 +1,108 @@ +// 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 UB_TENT_METADATA_BRIDGE_H +#define UB_TENT_METADATA_BRIDGE_H + +#include +#include +#include +#include + +// Old-TE headers must come first so that common.h declares +// LOCAL_SEGMENT_ID as a static const BEFORE tent/common/types.h +// redefines it as a macro. +#include "transfer_metadata.h" + +// TENT headers (tent/common/types.h defines #define LOCAL_SEGMENT_ID) +#include "tent/common/types.h" +#include "tent/runtime/control_plane.h" +#include "tent/runtime/segment.h" + +namespace mooncake { +namespace tent { + +// UbTentMetadataBridge bridges the legacy mooncake::TransferMetadata interface +// (used by UbTransport / UbWorkerPool internally) to TENT's ControlService and +// SegmentManager. +// +// Key responsibilities: +// - Remote segment lookups (getSegmentDescByID, getSegmentDescByName, +// getSegmentID) are redirected to the TENT SegmentManager so that the +// UbWorkerPool obtains the correct tseg / eid fields from TENT segment +// descriptors. +// - Local memory operations (addLocalMemoryBuffer, addLocalSegment, etc.) +// flow through to the base-class in-memory cache so that +// UbTransport::selectDevice() and UbContext::buildLocalBufferDesc() work +// correctly. +// - startHandshakeDaemon becomes a no-op: the UB handshake is handled by +// TENT's BootstrapUb RPC (registered via ControlService). +// - sendHandshake is redirected to ControlClient::bootstrapUb() so that +// the URMA connection setup uses TENT's RPC channel. +class UbTentMetadataBridge : public mooncake::TransferMetadata { + public: + // conn_string passed to old-TE TransferMetadata base. Use P2PHANDSHAKE for + // local-only mode (no external metadata store). + explicit UbTentMetadataBridge( + std::shared_ptr control_service, + const std::string& conn_string = P2PHANDSHAKE); + + ~UbTentMetadataBridge() override = default; + + // Remote segment lookup: queries TENT SegmentManager and converts the + // TENT SegmentDesc to the old-TE format (with tseg / eid fields). + // LOCAL_SEGMENT_ID falls through to the base-class in-memory cache. + std::shared_ptr getSegmentDescByID( + SegmentID segment_id, bool force_update = false) override; + + std::shared_ptr getSegmentDescByName( + const std::string& segment_name, bool force_update = false) override; + + // Maps a segment name to the corresponding TENT SegmentID (used as the + // old-TE SegmentID for remote segments). + SegmentID getSegmentID(const std::string& segment_name) override; + + // No-op: the UB handshake daemon is replaced by TENT's BootstrapUb RPC. + // Stores the callback so that onBootstrapUb can dispatch to it. + int startHandshakeDaemon(OnReceiveHandShake callback, uint16_t listen_port, + int sockfd) override; + + // Routes the UB handshake through TENT's ControlClient::bootstrapUb RPC + // instead of the old-TE TCP handshake plugin. + int sendHandshake(const std::string& peer_server_name, + const HandShakeDesc& local_desc, + HandShakeDesc& peer_desc) override; + + // Accessor for the stored handshake callback (invoked by onBootstrapUb). + OnReceiveHandShake handshakeCallback() const { return ub_handshake_cb_; } + + private: + // Converts a TENT SegmentDesc to the old-TE SegmentDesc format, extracting + // device EIDs and buffer tseg handles from transport_attrs[UB]. + std::shared_ptr convertFromTent( + const tent::SegmentDesc* tent_seg) const; + + std::shared_ptr control_service_; + OnReceiveHandShake ub_handshake_cb_; + + // Cache: old-TE SegmentID (= TENT SegmentID for remotes) → converted desc. + mutable std::mutex cache_mutex_; + mutable std::unordered_map> + remote_desc_cache_; +}; + +} // namespace tent +} // namespace mooncake + +#endif // UB_TENT_METADATA_BRIDGE_H diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_tent_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_tent_transport.h new file mode 100644 index 0000000000..d9b86c3538 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/ub_tent_transport.h @@ -0,0 +1,129 @@ +// 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 UB_TENT_TRANSPORT_H +#define UB_TENT_TRANSPORT_H + +// Old-TE headers first so that common.h declares LOCAL_SEGMENT_ID as a +// static const before tent/common/types.h redefines it as a macro. +#include "topology.h" +#include "transport/transport.h" +#include "transport/kunpeng_transport/ub_transport.h" + +// TENT headers (tent/common/types.h defines #define LOCAL_SEGMENT_ID) +#include "tent/runtime/transport.h" +#include "tent/runtime/control_plane.h" +#include "tent/common/types.h" +#include "tent/common/config.h" +#include "tent/common/status.h" +#include "tent/transport/ub/ub_tent_metadata_bridge.h" + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace tent { + +// UbTentTransport adapts the legacy UbTransport (old TE) to the TENT +// Transport interface. It reuses the proven URMA data-plane in UbTransport +// and bridges the semantic gap between TENT's request/segment model and the +// old TE's BatchID/TransferTask model. +// +// Threading: install/uninstall are single-threaded; all other methods may be +// called concurrently from multiple TENT worker threads. +class UbTentTransport : public Transport { + public: + // SubBatch holds the old-TE BatchID and the converted TransferRequest + // objects whose pointers are stored inside the old-TE TransferTask list. + // The requests must remain alive until freeSubBatch() is called. + struct UbSubBatch : public Transport::SubBatch { + // Old TE BatchID (pointer-as-integer to BatchDesc). + mooncake::Transport::BatchID ub_batch_id{0}; + + // Lifetime-managed copies of converted old-TE requests. + // UbTransport::submitTransferTask() stores bare pointers into this + // vector, so the vector must not be reallocated after submission. + std::vector te_requests; + + size_t size() const override { return task_count_; } + + private: + friend class UbTentTransport; + size_t task_count_{0}; + }; + + UbTentTransport() = default; + ~UbTentTransport() override; + + // TENT Transport interface + Status install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf = nullptr) override; + + Status uninstall() override; + + Status allocateSubBatch(SubBatchRef& batch, size_t max_size) override; + Status freeSubBatch(SubBatchRef& batch) override; + + Status submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) override; + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override; + + Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) override; + + Status removeMemoryBuffer(BufferDesc& desc) override; + + const char* getName() const override { return "ub"; } + + private: + // Translate a TENT SegmentID to the corresponding old-TE SegmentID by + // looking up the segment name via the TENT SegmentManager and querying + // the bridge for the matching ID. + mooncake::Transport::SegmentID getTESegmentID(SegmentID tent_id); + + // Publish UB device EIDs and buffer tseg handles to the TENT local segment + // so that remote nodes can read them via the SegmentManager. + Status setupUbLocalSegment(); + + private: + // Old-TE UbTransport instance and the topology used during install. + std::unique_ptr ub_transport_; + std::shared_ptr te_topology_; + + // Bridge: replaces the standalone te_metadata_; routes remote lookups to + // the TENT SegmentManager and makes the handshake daemon a no-op. + std::shared_ptr te_metadata_bridge_; + + // TENT control plane (for segment synchronization and RPC). + std::shared_ptr control_service_; + + std::string local_segment_name_; + + // Cache: TENT SegmentID → old-TE SegmentID. + std::mutex seg_id_cache_mutex_; + std::unordered_map + tent_to_te_seg_id_; +}; + +} // namespace tent +} // namespace mooncake + +#endif // UB_TENT_TRANSPORT_H diff --git a/mooncake-transfer-engine/tent/src/CMakeLists.txt b/mooncake-transfer-engine/tent/src/CMakeLists.txt index 73766ce120..597380c479 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() @@ -80,12 +81,19 @@ 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() +if(USE_UB) + target_compile_definitions(tent_interface INTERFACE USE_UB) + message(STATUS "UB (Kunpeng URMA): Enabled") +else() + message(STATUS "UB (Kunpeng URMA): Disabled") +endif() + set(YALANTING_TARGET yalantinglibs::yalantinglibs) if(NOT TARGET ${YALANTING_TARGET}) if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) @@ -131,6 +139,7 @@ foreach( platform_rocm platform_ascend platform_sunrise + platform_tpu tent_xport_gds tent_xport_uring tent_xport_bufio @@ -141,6 +150,8 @@ foreach( tent_xport_tcp tent_xport_ascend_direct tent_xport_sunrise_link + tent_xport_ub + tent_xport_tpu tent_metrics) if(TARGET ${tgt}) target_link_libraries(tent_link_group INTERFACE ${tgt}) diff --git a/mooncake-transfer-engine/tent/src/common/config.cpp b/mooncake-transfer-engine/tent/src/common/config.cpp index 2a0bf8f751..a7cc9dfe8e 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(); @@ -119,6 +137,14 @@ Status ConfigHelper::loadFromEnv(Config& config) { "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; } 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..6bbbf3519b --- /dev/null +++ b/mooncake-transfer-engine/tent/src/common/qos_metrics.cpp @@ -0,0 +1,374 @@ +// 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.p99_us = sample.p99_us; + + const double duration_s = sample.total_duration_us / 1e6; + const double bytes = + static_cast(block_size) * batch_size * metrics.operations; + if (duration_s > 0.0) + metrics.throughput_gbps = 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 + << " 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}, + {"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/metrics/tent_metrics.cpp b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp index 4eb9f28e64..94184e4538 100644 --- a/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp +++ b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp @@ -173,8 +173,8 @@ void TentMetrics::shutdown() { void TentMetrics::registerMetrics() { // Pre-allocate vectors to avoid reallocation counters_.reserve(7); - histograms_.reserve(5); - histogram_boundaries_.reserve(5); + histograms_.reserve(8); + histogram_boundaries_.reserve(8); // Register all counters - add new counters here counters_ = { @@ -186,12 +186,12 @@ void TentMetrics::registerMetrics() { // Register all histograms - add new histograms here // Note: histogram_boundaries_ must match the order of histograms_ histograms_ = { - &read_latency_, &write_latency_, &read_size_, - &write_size_, &deadline_mlu_, + &read_latency_, &write_latency_, &read_size_, &write_size_, + &deadline_mlu_, &stage_queue_wait_, &stage_dispatch_, &stage_transport_, }; histogram_boundaries_ = { - kLatencyBuckets, kLatencyBuckets, kSizeBuckets, - kSizeBuckets, kMluPerMilleBuckets, + kLatencyBuckets, kLatencyBuckets, kSizeBuckets, kSizeBuckets, + kMluPerMilleBuckets, kStageBuckets, kStageBuckets, kStageBuckets, }; } @@ -234,6 +234,24 @@ void TentMetrics::recordDeadlineMLU(double mlu) { deadline_mlu_.observe(static_cast(mlu * 1000.0)); } +void TentMetrics::recordStageLatency(Stage stage, double latency_us) { + if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) + return; + if (latency_us < 0.0) return; + int64_t val = static_cast(latency_us); + switch (stage) { + case Stage::QueueWait: + stage_queue_wait_.observe(val); + break; + case Stage::Dispatch: + stage_dispatch_.observe(val); + break; + case Stage::Transport: + stage_transport_.observe(val); + break; + } +} + void TentMetrics::recordReadFailed(size_t bytes) { // Fast path: check runtime switch first if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) @@ -390,6 +408,7 @@ void TentMetrics::recordReadFailed(size_t) {} void TentMetrics::recordWriteFailed(size_t) {} void TentMetrics::recordTransportFailover() {} void TentMetrics::recordDeadlineMLU(double) {} +void TentMetrics::recordStageLatency(Stage, 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/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/pybind.cpp b/mooncake-transfer-engine/tent/src/python/pybind.cpp index bc1d3a23b7..c8d6fb7410 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,17 @@ PYBIND11_MODULE(tent, m) { .value("TCP", TransportType::TCP) .value("AscendDirect", TransportType::AscendDirect) .value("SUNRISE_LINK", TransportType::SUNRISE_LINK) + .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 +333,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 +344,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 +365,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 +704,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/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index 3d30747a77..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,6 +26,13 @@ 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 || @@ -167,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}] = @@ -174,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); } @@ -185,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) { @@ -204,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; @@ -246,6 +369,37 @@ Status LocalTransferAdmissionQueue::complete( 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) { + --outstanding_user_owners_; + outstanding_user_bytes_ -= owner.request.length; + } + return Status::OK(); +} + Status LocalTransferAdmissionQueue::retireBatch(uint64_t batch_token) { if (batch_token == 0) { return Status::InvalidArgument("invalid batch token" LOC_MARK); diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index d82fbb5410..124180d6f3 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -45,6 +45,18 @@ Status ControlClient::bootstrap(const std::string& server_addr, return Status::OK(); } +Status ControlClient::bootstrapUb(const std::string& server_addr, + const UbBootstrapDesc& request, + UbBootstrapDesc& response) { + std::string request_raw, response_raw; + json j = request; + request_raw = j.dump(); + CHECK_STATUS( + tl_rpc_agent.call(server_addr, BootstrapUb, request_raw, response_raw)); + response = json::parse(response_raw).get(); + return Status::OK(); +} + Status ControlClient::sendData(const std::string& server_addr, uint64_t peer_mem_addr, void* local_mem_addr, size_t length) { @@ -175,6 +187,11 @@ ControlService::ControlService(const std::string& type, [this](const std::string_view& request, std::string& response) { onBootstrapRdma(request, response); }); + rpc_server_->registerFunction( + BootstrapUb, + [this](const std::string_view& request, std::string& response) { + onBootstrapUb(request, response); + }); rpc_server_->registerFunction( SendData, [this](const std::string_view& request, std::string& response) { @@ -226,7 +243,7 @@ Status ControlService::start(uint16_t& port, bool ipv6_) { void ControlService::onGetSegmentDesc(const std::string_view& request, std::string& response) { - // Re-use the cached dump shared across concurrent peer fetches. + // Reuse the cached dump shared across concurrent peer fetches. auto cached = manager_->getLocalDumpedJson(); response = *cached; } @@ -242,6 +259,31 @@ void ControlService::onBootstrapRdma(const std::string_view& request, response = j.dump(); } +void ControlService::onBootstrapUb(const std::string_view& request, + std::string& response) { + UbBootstrapDesc request_desc = + json::parse(std::string(request)).get(); + UbBootstrapDesc response_desc; + int ret = 0; + { + // Serialize callback replacement with invocation so uninstall waits + // for an in-flight bootstrap before destroying the UB transport. + std::lock_guard lock(ub_bootstrap_callback_mutex_); + if (ub_bootstrap_callback_) { + ret = ub_bootstrap_callback_(request_desc, response_desc); + } else { + ret = -1; + response_desc.reply_msg = "BootstrapUb callback is not registered"; + } + } + if (ret != 0 && response_desc.reply_msg.empty()) { + response_desc.reply_msg = + "BootstrapUb callback failed, ret=" + std::to_string(ret); + } + json j = response_desc; + response = j.dump(); +} + void ControlService::onSendData(const std::string_view& request, std::string& response) { if (request.size() < sizeof(XferDataDesc)) { @@ -249,7 +291,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 +315,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/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/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index 2e3f345540..214b78fc98 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -265,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(); } @@ -299,6 +301,12 @@ Status TransferEngineImpl::construct() { 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 = @@ -366,6 +374,30 @@ 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, runtime_queue_config_.enabled @@ -430,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(); }); @@ -514,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)); } @@ -641,10 +677,12 @@ std::vector TransferEngineImpl::getSupportedTransports( if (transport_list_[NVLINK]) result.push_back(NVLINK); if (transport_list_[RDMA]) result.push_back(RDMA); if (transport_list_[SUNRISE_LINK]) result.push_back(SUNRISE_LINK); + if (transport_list_[UB]) result.push_back(UB); if (transport_list_[AscendDirect]) result.push_back(AscendDirect); 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; } @@ -877,7 +915,12 @@ class TransferEngineImpl::BatchRef { }; 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, @@ -905,6 +948,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; } @@ -929,9 +973,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); @@ -964,7 +1009,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)) { @@ -992,6 +1041,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 @@ -1042,6 +1092,12 @@ static const char* transportTypeName(TransportType type) { return "AscendDirect"; case SUNRISE_LINK: return "SUNRISE_LINK"; + case UB: + return "UB"; + default: + break; + case TPU: + return "TPU"; } return "UNKNOWN"; } @@ -1212,7 +1268,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) { @@ -1272,8 +1329,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) @@ -1313,7 +1372,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) { @@ -1328,6 +1387,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, @@ -1432,8 +1519,10 @@ Status TransferEngineImpl::commitPreparedSubmit( task.staging = false; task.start_time = 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); @@ -1444,6 +1533,7 @@ Status TransferEngineImpl::commitPreparedSubmit( if (owner.staging) { task.staging = true; staging_proxy_->submit(&task, (BatchID)batch, owner.staging_params); + task.post_time = std::chrono::steady_clock::now(); continue; } @@ -1483,6 +1573,8 @@ Status TransferEngineImpl::commitPreparedSubmit( // 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 status = transport->submitTransferTasks( @@ -1490,8 +1582,17 @@ Status TransferEngineImpl::commitPreparedSubmit( if (!status.ok()) { // LOG(WARNING) << "Failed to submit SubBatch " << type << ":" // << status.ToString(); + for (auto& task_id : task_id_list[type]) { + // Mark as UNSPEC so pollTaskStatus returns FAILED and + // updateTaskStatusAfterPoll can try the next transport. + batch->task_list[task_id].type = UNSPEC; + } for (auto& task_id : task_id_list[type]) batch->task_list[task_id].type = UNSPEC; + } else { + auto now = std::chrono::steady_clock::now(); + for (auto& task_id : task_id_list[type]) + batch->task_list[task_id].post_time = now; } } @@ -1524,6 +1625,8 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, 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)); } @@ -1545,6 +1648,7 @@ Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, 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; } @@ -1590,6 +1694,25 @@ Status TransferEngineImpl::finishQueuedOwner( 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); @@ -1620,9 +1743,11 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { 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); } @@ -1635,6 +1760,7 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { 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); } } @@ -1651,13 +1777,17 @@ Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { 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; + 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(); auto status = transport->submitTransferTasks(sub_batch, {task.request}); if (!status.ok()) { task.type = UNSPEC; return finishQueuedOwner(owner_id, FAILED); } + task.post_time = std::chrono::steady_clock::now(); return markQueuedOwnerSubmitted(owner_id); } @@ -1834,6 +1964,75 @@ Status TransferEngineImpl::submitTransfer( QueueOwnerKind::User); } +Status TransferEngineImpl::cancelTransfer(BatchID batch_id, size_t task_id) { + 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); + } + 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(); +} + Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { auto& task = batch->task_list[task_id]; auto prev_type = task.type; @@ -1903,9 +2102,13 @@ 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_status.s != FAILED) return; + if (!allow_failover || task.cancel_requested || task_status.s != FAILED || + task.type == UNSPEC) return; + // Allow resubmission for UNSPEC tasks: these occur when submitTransferTasks + // failed (e.g., UB transport submission error). if (resubmitTransferTask(batch, task_id).ok()) { task_status.s = PENDING; task.status = PENDING; @@ -2214,6 +2417,31 @@ void TransferEngineImpl::recordTaskCompletionMetrics( TentMetrics::instance().recordWriteCompleted( task.request.length, latency_seconds); } +#if TENT_METRICS_ENABLED + // Causal chain stage decomposition + 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, + 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, + dispatch_us); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::Transport, + transport_us); + } + } +#endif // Observability only (RFC #2519): if this transfer carried a // deadline, emit the post-hoc feasibility ratio MLU = // actual_transfer_time / available_window, where the window is diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp index 7ecd51276e..a7eee07a9f 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp @@ -41,6 +41,14 @@ #include "tent/transport/sunrise_link/sunrise_link_transport.h" #endif +#ifdef USE_UB +#include "tent/transport/ub/ub_tent_transport.h" +#endif + +#ifdef USE_TPU +#include "tent/transport/tpu/tpu_transport.h" +#endif + namespace mooncake { namespace tent { @@ -94,6 +102,16 @@ Status TransferEngineImpl::loadTransports() { } #endif +#ifdef USE_UB + if (conf_->get("transports/ub/enable", true)) + transport_list_[UB] = std::make_shared(); +#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..01fea6e45c 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp @@ -17,11 +17,10 @@ #include "tent/runtime/platform.h" #include "tent/thirdparty/nlohmann/json.h" -#include -#include - #include #include +#include +#include namespace mooncake { namespace tent { @@ -29,20 +28,34 @@ 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}, + {"unspec", UNSPEC}, + {"rdma", RDMA}, + {"mnnvl", MNNVL}, + {"shm", SHM}, + {"nvlink", NVLINK}, + {"gds", GDS}, + {"io_uring", IOURING}, + {"tcp", TCP}, + {"ascend", AscendDirect}, + {"sunrise_link", SUNRISE_LINK}, + {"tpu", TPU}, + {"ub", UB}, }; 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"}, + {UNSPEC, "unspec"}, + {RDMA, "rdma"}, + {MNNVL, "mnnvl"}, + {SHM, "shm"}, + {NVLINK, "nvlink"}, + {GDS, "gds"}, + {IOURING, "io_uring"}, + {TCP, "tcp"}, + {AscendDirect, "ascend"}, + {SUNRISE_LINK, "sunrise_link"}, + {TPU, "tpu"}, + {UB, "ub"}, }; // Memory type name mapping for pattern matching @@ -51,6 +64,46 @@ static const std::string kMemoryTypeCuda = "cuda"; static const std::string kMemoryTypeNpu = "npu"; static const std::string kMemoryTypeWildcard = "*"; +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; + } + + if (value.is_number_unsigned()) { + const auto raw = value.get(); + if (raw <= static_cast(IntentType::STAGING_INTERNAL)) { + return static_cast(raw); + } + return std::nullopt; + } + + 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::string TransportSelector::transportTypeName(TransportType type) { auto it = kTransportTypeNames.find(type); if (it != kTransportTypeNames.end()) { @@ -74,14 +127,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 +146,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 +245,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"]) { @@ -206,6 +280,42 @@ void TransportSelector::loadPolicies() { } } + // 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 +346,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 +413,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 +439,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 +510,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..8d495763f6 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); diff --git a/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp b/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp index 150df8b8e6..f9c88258f0 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; diff --git a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt index e73f7cb063..8eede69cf2 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(ub) +add_subdirectory(tpu) add_library(tent_transport_all INTERFACE) foreach( @@ -21,7 +23,9 @@ foreach( tent_xport_shm tent_xport_tcp tent_xport_ascend_direct - tent_xport_sunrise_link) + tent_xport_sunrise_link + tent_xport_ub + 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 d50fed7fef..99f7be72c0 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."); 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/mnnvl/mnnvl_transport.cpp b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp index 046b253be4..2e869018ea 100644 --- a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp @@ -174,7 +174,8 @@ Status MnnvlTransport::submitTransferTasks( // Get local segment for buffer lookup auto &segment_manager = metadata_->segmentManager(); - SegmentDesc *local_segment = segment_manager.getLocal().get(); + // 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); @@ -237,6 +238,7 @@ Status MnnvlTransport::submitTransferTasks( stream_device)); CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->async_stream, stream_device)); + mnnvl_batch->stream_device_id = stream_device; } startTransfer(new_tasks, mnnvl_batch); @@ -311,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; } @@ -324,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; } @@ -553,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 70707beea8..f3de90d7f3 100644 --- a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp @@ -134,7 +134,8 @@ Status NVLinkTransport::submitTransferTasks( // Get local segment for buffer lookup auto& segment_manager = metadata_->segmentManager(); - SegmentDesc* local_segment = segment_manager.getLocal().get(); + // 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); @@ -193,6 +194,7 @@ Status NVLinkTransport::submitTransferTasks( stream_device)); CHECK_STATUS(platform_->getStreamFromPool(shm_batch->async_stream, stream_device)); + shm_batch->stream_device_id = stream_device; } startTransfer(new_tasks, shm_batch); @@ -272,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; } @@ -285,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; } @@ -459,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/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 d1f5109f7d..64760d25cd 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp @@ -82,14 +82,30 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params, context_ = context; params_ = params; endpoint_name_ = endpoint_name; - inflight_slices_ = 0; - qp_list_.resize(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[params_->qp_mul_factor](); + 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(); @@ -255,9 +271,10 @@ int RdmaEndPoint::deconstructUnlocked() { bool all_qps_destroyed = true; for (size_t i = 0; i < qp_list_.size(); ++i) { - if (wr_depth_list_ && wr_depth_list_[i].value != 0) { - const int outstanding = wr_depth_list_[i].value; - cancelQuota(i, outstanding); + 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])) { @@ -351,7 +368,7 @@ bool RdmaEndPoint::finishDestroy() { // 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 != 0) { + if (wr_depth_list_[i].value.load(std::memory_order_relaxed) != 0) { has_outstanding = true; break; } @@ -648,6 +665,14 @@ int RdmaEndPoint::resetConnection(const std::string& reason) { 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) { @@ -694,14 +719,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; @@ -817,15 +853,17 @@ std::vector RdmaEndPoint::qpNum() { 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; @@ -835,8 +873,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); } @@ -847,6 +886,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)); @@ -886,9 +938,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; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp index 53a8bb4296..c4fd6bea9e 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp @@ -292,7 +292,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(); } @@ -356,5 +358,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 741181441a..a3878ea0e1 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,6 +193,41 @@ 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); @@ -401,9 +437,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; @@ -455,6 +496,7 @@ Status RdmaTransport::submitTransferTasks( slice->length = length; slice->task = task; slice->retry_count = 0; + slice->quota_charged = false; slice->ep_weak_ptr.reset(); slice->word = PENDING; slice->next = nullptr; @@ -462,8 +504,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]; @@ -499,6 +543,19 @@ 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); +} + bool RdmaTransport::warmupMemory(void* addr, size_t length) { if (length < kMrWarmupMinBytes) return false; unsigned hwc = std::thread::hardware_concurrency(); @@ -543,22 +600,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(); } @@ -606,13 +664,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); @@ -624,6 +680,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( @@ -642,8 +699,7 @@ std::shared_ptr RdmaTransport::getEndpoint(SegmentID target_id, 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; @@ -775,5 +831,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 500cc9e3e4..4058dc340b 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -16,11 +16,14 @@ #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" @@ -32,6 +35,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. @@ -46,7 +55,10 @@ 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_; @@ -142,6 +154,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) // ============================================================ @@ -248,8 +270,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) { @@ -260,8 +316,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( @@ -347,11 +403,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, @@ -367,29 +435,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) { @@ -400,14 +528,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); } } @@ -420,40 +548,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; } - for (auto& slice_list : promoted) { - worker.queues[PRIO_LOW].push(slice_list); + + std::vector promote(drained.size(), false); + for (size_t idx : decision.promote_indices) { + if (idx < drained.size()) promote[idx] = true; } - } + 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() { @@ -500,10 +649,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); @@ -531,7 +677,12 @@ 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); @@ -670,7 +821,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) @@ -727,6 +878,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) 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 ccd737eeb6..cb2ba911c8 100644 --- a/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp @@ -60,11 +60,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(); @@ -241,31 +241,48 @@ 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) { + 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( @@ -273,8 +290,17 @@ 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( @@ -299,20 +325,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/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..4a68d72bfc --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/CMakeLists.txt @@ -0,0 +1,32 @@ +if(USE_UB) + file(GLOB UB_TENT_SOURCES "*.cpp") + + # Keep this library limited to the TENT UB adapter sources. The old-TE + # ub_transport objects are already provided by the transfer_engine target when + # USE_UB is enabled; embedding them here as well makes final executables link + # the same object files twice. + add_library(tent_xport_ub STATIC ${UB_TENT_SOURCES}) + + # Old-TE public headers (transfer_metadata.h, topology.h, transport/*.h). + # These are already on the global include path when building inside the + # mooncake-transfer-engine tree (set by include_directories(include) in + # mooncake-transfer-engine/CMakeLists.txt). The explicit entry below handles + # standalone TENT builds where that global path is absent. + target_include_directories( + tent_xport_ub + PUBLIC + $ + $ + ${urma_INCLUDE_DIR}) + + target_link_libraries( + tent_xport_ub + PUBLIC tent_common + PRIVATE JsonCpp::JsonCpp glog::glog pthread) + + if(URMA_LIBRARY) + target_link_libraries(tent_xport_ub PUBLIC ${URMA_LIBRARY}) + endif() + + message(STATUS "tent_xport_ub: built (USE_UB=ON)") +endif() diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_tent_metadata_bridge.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_tent_metadata_bridge.cpp new file mode 100644 index 0000000000..9e53a37006 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_tent_metadata_bridge.cpp @@ -0,0 +1,289 @@ +// 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/ub/ub_tent_metadata_bridge.h" + +#include + +#include "tent/runtime/segment_manager.h" +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { + +using json = nlohmann::json; + +// --------------------------------------------------------------------------- +// Construction +// --------------------------------------------------------------------------- + +UbTentMetadataBridge::UbTentMetadataBridge( + std::shared_ptr control_service, + const std::string& conn_string) + : mooncake::TransferMetadata(conn_string), + control_service_(std::move(control_service)) {} + +// --------------------------------------------------------------------------- +// Remote segment conversion +// --------------------------------------------------------------------------- + +std::shared_ptr +UbTentMetadataBridge::convertFromTent(const tent::SegmentDesc* tent_seg) const { + if (!tent_seg) return nullptr; + + if (tent_seg->type != tent::SegmentType::Memory) { + LOG(WARNING) + << "UbTentMetadataBridge: segment '" << tent_seg->name + << "' is not a memory segment; cannot extract UB attributes"; + return nullptr; + } + + auto desc = std::make_shared(); + desc->name = tent_seg->name; + desc->protocol = "ub"; + + const auto& mem = std::get(tent_seg->detail); + + // Collect device names for topology construction. + std::vector ub_device_names; + for (const auto& dev : mem.devices) { + TransferMetadata::DeviceDesc d; + d.name = dev.name; + auto it = dev.transport_attrs.find(TransportType::UB); + if (it != dev.transport_attrs.end()) { + d.eid = it->second; + } + desc->devices.push_back(d); + ub_device_names.push_back(dev.name); + } + + // Extract per-buffer tseg handles stored in transport_attrs[UB] as a + // JSON array of strings. + for (const auto& buf : mem.buffers) { + TransferMetadata::BufferDesc b; + // The old-TE selectDevice() uses buffer.name as the topology storage + // key (e.g. "cpu:0" or "*"). Map TENT buf.location to this field. + b.name = buf.location.empty() ? kWildcardLocation : buf.location; + b.addr = buf.addr; + b.length = buf.length; + auto it = buf.transport_attrs.find(TransportType::UB); + if (it != buf.transport_attrs.end()) { + try { + auto j = json::parse(it->second); + for (auto& t : j) b.tseg.push_back(t.get()); + } catch (const std::exception& e) { + LOG(WARNING) << "UbTentMetadataBridge: failed to parse tseg " + "for buffer " + << buf.addr << ": " << e.what(); + } + } + desc->buffers.push_back(std::move(b)); + } + + // Build a minimal old-TE topology from the TENT device list so that + // selectDevice() can resolve storage_type → device index. The old-TE + // topology is a JSON object of the form: + // { "storage_type": [ ["preferred_hca"], ["avail_hca"] ] } + // We construct a wildcard "*" entry covering all UB devices, plus + // entries for each distinct buffer.location found in the remote segment. + if (!ub_device_names.empty()) { + json topo_j(json::object()); + auto build_entry = [&ub_device_names](const std::string& key) { + json entry(json::array()); + json preferred(json::array()); + for (const auto& name : ub_device_names) preferred.push_back(name); + entry.push_back(std::move(preferred)); + entry.push_back(json::array()); // empty avail_hca list + return entry; + }; + topo_j["*"] = build_entry("*"); + // Also register each buffer location so selectDevice can match it. + for (const auto& buf : mem.buffers) { + if (!buf.location.empty() && buf.location != kWildcardLocation) { + topo_j[buf.location] = build_entry(buf.location); + } + } + desc->topology.parse(topo_j.dump()); + } + + LOG(INFO) << "UbTentMetadataBridge::convertFromTent segment=" << desc->name + << " devices=" << desc->devices.size() + << " buffers=" << desc->buffers.size() + << " topology_entries=" << desc->topology.getMatrix().size(); + + return desc; +} + +// --------------------------------------------------------------------------- +// Remote segment lookup +// --------------------------------------------------------------------------- + +std::shared_ptr +UbTentMetadataBridge::getSegmentDescByID(SegmentID segment_id, + bool force_update) { + // Local segment: served by the base-class in-memory cache. + if (segment_id == LOCAL_SEGMENT_ID) { + return TransferMetadata::getSegmentDescByID(segment_id, force_update); + } + + if (!control_service_) { + LOG(ERROR) << "UbTentMetadataBridge: no ControlService for ID lookup " + << segment_id; + return nullptr; + } + + if (!force_update) { + std::lock_guard lock(cache_mutex_); + auto it = remote_desc_cache_.find(segment_id); + if (it != remote_desc_cache_.end()) return it->second; + } else { + // force_update: invalidate both the bridge cache and the TENT + // thread-local cache so that getRemoteCached() re-fetches from + // the metadata registry (etcd/p2p) instead of returning a stale + // copy (TENT's default TTL is ~1 hour). + { + std::lock_guard lock(cache_mutex_); + remote_desc_cache_.erase(segment_id); + } + control_service_->segmentManager().invalidateRemote(segment_id); + } + + tent::SegmentDesc* tent_desc = nullptr; + auto status = control_service_->segmentManager().getRemoteCached( + tent_desc, segment_id); + if (!status.ok() || !tent_desc) { + LOG(WARNING) << "UbTentMetadataBridge: TENT segment ID " << segment_id + << " not found: " << status.message(); + return nullptr; + } + + auto converted = convertFromTent(tent_desc); + if (converted) { + std::lock_guard lock(cache_mutex_); + remote_desc_cache_[segment_id] = converted; + } + return converted; +} + +std::shared_ptr +UbTentMetadataBridge::getSegmentDescByName(const std::string& segment_name, + bool force_update) { + if (!control_service_) { + return TransferMetadata::getSegmentDescByName(segment_name, + force_update); + } + + tent::SegmentDescRef tent_desc_ref; + auto status = control_service_->segmentManager().getRemote(tent_desc_ref, + segment_name); + if (!status.ok() || !tent_desc_ref) { + LOG(WARNING) << "UbTentMetadataBridge: segment '" << segment_name + << "' not found in TENT: " << status.message(); + return nullptr; + } + + return convertFromTent(tent_desc_ref.get()); +} + +// --------------------------------------------------------------------------- +// Segment ID lookup +// --------------------------------------------------------------------------- + +TransferMetadata::SegmentID UbTentMetadataBridge::getSegmentID( + const std::string& segment_name) { + if (!control_service_) { + return TransferMetadata::getSegmentID(segment_name); + } + + SegmentID handle = 0; + auto status = + control_service_->segmentManager().openRemote(handle, segment_name); + if (!status.ok()) { + LOG(ERROR) << "UbTentMetadataBridge: cannot open remote segment '" + << segment_name << "': " << status.message(); + return static_cast(-1); + } + return static_cast(handle); +} + +// --------------------------------------------------------------------------- +// Handshake daemon: no-op (replaced by TENT BootstrapUb RPC) +// --------------------------------------------------------------------------- + +int UbTentMetadataBridge::startHandshakeDaemon(OnReceiveHandShake callback, + uint16_t /*listen_port*/, + int /*sockfd*/) { + // Store the callback so that the ControlService onBootstrapUb handler can + // dispatch incoming UB bootstrap requests to it. + ub_handshake_cb_ = std::move(callback); + return 0; +} + +// --------------------------------------------------------------------------- +// sendHandshake: routes through TENT BootstrapUb RPC +// --------------------------------------------------------------------------- + +int UbTentMetadataBridge::sendHandshake(const std::string& peer_server_name, + const HandShakeDesc& local_desc, + HandShakeDesc& peer_desc) { + if (!control_service_) { + return TransferMetadata::sendHandshake(peer_server_name, local_desc, + peer_desc); + } + + // Resolve the TENT RPC address for this peer. The peer_server_name + // matches the TENT segment name (= the remote local_segment_name_). + std::string rpc_addr; + { + tent::SegmentDescRef tent_desc_ref; + auto status = control_service_->segmentManager().getRemote( + tent_desc_ref, peer_server_name); + if (!status.ok() || !tent_desc_ref) { + LOG(ERROR) << "UbTentMetadataBridge: cannot resolve peer '" + << peer_server_name + << "' to a TENT segment: " << status.message(); + return -1; + } + rpc_addr = tent_desc_ref->rpc_server_addr; + } + + UbBootstrapDesc request; + request.local_nic_path = local_desc.local_nic_path; + request.peer_nic_path = local_desc.peer_nic_path; +#ifdef USE_UB + request.jetty_num = local_desc.jetty_num; + request.local_eid = local_desc.local_eid; +#endif + + UbBootstrapDesc response; + auto status = ControlClient::bootstrapUb(rpc_addr, request, response); + if (!status.ok()) { + LOG(ERROR) << "UbTentMetadataBridge: BootstrapUb RPC to " << rpc_addr + << " failed: " << status.message(); + return -1; + } + + peer_desc.local_nic_path = response.local_nic_path; + peer_desc.peer_nic_path = response.peer_nic_path; + peer_desc.reply_msg = response.reply_msg; +#ifdef USE_UB + peer_desc.jetty_num = response.jetty_num; + peer_desc.local_eid = response.local_eid; +#endif + + return 0; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/ub/ub_tent_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ub/ub_tent_transport.cpp new file mode 100644 index 0000000000..6f0c75b9d0 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/ub_tent_transport.cpp @@ -0,0 +1,524 @@ +// 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/ub/ub_tent_transport.h" + +#include + +#include +#include +#include +#include + +#include "tent/runtime/segment.h" +#include "tent/runtime/segment_manager.h" +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { + +using json = nlohmann::json; + +namespace { +// Splits a comma-separated device list ("a,b ,c") into trimmed, non-empty +// names. +std::vector parseDeviceFilter(const std::string& csv) { + std::vector out; + std::stringstream ss(csv); + std::string item; + while (std::getline(ss, item, ',')) { + size_t b = item.find_first_not_of(" \t"); + if (b == std::string::npos) continue; + size_t e = item.find_last_not_of(" \t"); + out.push_back(item.substr(b, e - b + 1)); + } + return out; +} +} // namespace + +// --------------------------------------------------------------------------- +// Lifecycle +// --------------------------------------------------------------------------- + +UbTentTransport::~UbTentTransport() { uninstall(); } + +Status UbTentTransport::install(std::string& local_segment_name, + std::shared_ptr metadata, + std::shared_ptr /*local_topology*/, + std::shared_ptr conf) { + local_segment_name_ = local_segment_name; + control_service_ = metadata; + + // Resolve which UB device(s) to use. Production pins a specific (often + // bonded) device, e.g. device_name=bonding_dev_0, so that only ONE logical + // UB context is created -- matching the legacy old-TE UB behavior. Without + // a filter we auto-discover every UB NIC on the host (used by unit tests). + // Resolution order: + // 1. TENT config key transports/ub/device_name (comma-separated) + // 2. Env var MC_UB_DEVICE_NAME (comma-separated) + // 3. (none) discover all UB devices + std::string device_csv; + if (conf) + device_csv = + conf->get("transports/ub/device_name", std::string()); + if (device_csv.empty()) { + if (const char* env = std::getenv("MC_UB_DEVICE_NAME")) + device_csv = env; + } + std::vector device_filter = parseDeviceFilter(device_csv); + + // Discover UB HCAs on this host. Falls back to mock_urma_device when no + // real HCAs are present (handled inside + // UbTransport::initializeUbResources). + te_topology_ = std::make_shared(); + if (device_filter.empty()) { + LOG(INFO) << "UbTentTransport: no device filter set, discovering all " + "UB devices"; + te_topology_->discover(); + } else { + LOG(INFO) << "UbTentTransport: restricting to UB device(s): " + << device_csv; + te_topology_->discover(device_filter); + } + + // Build the bridge that replaces the standalone te_metadata_. + // Use P2PHANDSHAKE so old-TE TransferMetadata stays in local-only mode + // (no etcd/http storage plugin). Remote lookups go through TENT + // ControlService instead. + te_metadata_bridge_ = + std::make_shared(control_service_, P2PHANDSHAKE); + + // Instantiate and install UbTransport, using the bridge as metadata. + ub_transport_ = std::make_unique(URMA_ENDPOINT); + int rc = ub_transport_->install(local_segment_name_, te_metadata_bridge_, + te_topology_); + if (rc != 0) { + LOG(ERROR) << "UbTentTransport: UbTransport::install() failed, rc=" + << rc; + ub_transport_.reset(); + te_metadata_bridge_.reset(); + te_topology_.reset(); + return Status::InternalError( + "UbTentTransport: UbTransport install failed, rc=" + + std::to_string(rc) + LOC_MARK); + } + + // Register the UB bootstrap callback so that incoming BootstrapUb RPCs + // (from remote nodes initiating URMA connections) are dispatched to + // UbTransport::onSetupConnections. + if (control_service_) { + auto cb = te_metadata_bridge_->handshakeCallback(); + if (cb) { + control_service_->setBootstrapUbCallback( + [cb](const UbBootstrapDesc& peer, + UbBootstrapDesc& local) -> int { + mooncake::TransferMetadata::HandShakeDesc peer_hs, local_hs; + peer_hs.local_nic_path = peer.local_nic_path; + peer_hs.peer_nic_path = peer.peer_nic_path; +#ifdef USE_UB + peer_hs.jetty_num = peer.jetty_num; + peer_hs.local_eid = peer.local_eid; +#endif + int ret = cb(peer_hs, local_hs); + local.local_nic_path = local_hs.local_nic_path; + local.peer_nic_path = local_hs.peer_nic_path; + local.reply_msg = local_hs.reply_msg; +#ifdef USE_UB + local.jetty_num = local_hs.jetty_num; + local.local_eid = local_hs.local_eid; +#endif + return ret; + }); + } + } + + caps.dram_to_dram = true; + + // Publish UB device EIDs to the TENT local segment. + Status s = setupUbLocalSegment(); + if (!s.ok()) { + LOG(WARNING) << "UbTentTransport: setupUbLocalSegment() failed: " + << s.message() + << " (non-fatal; remote nodes may lack UB device info)"; + } + + LOG(INFO) << "UbTentTransport: installed on segment '" + << local_segment_name_ << "'"; + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// setupUbLocalSegment +// +// Writes UB device EIDs into the TENT local MemorySegmentDesc so that remote +// nodes can extract them via the bridge's convertFromTent(). +// --------------------------------------------------------------------------- + +Status UbTentTransport::setupUbLocalSegment() { + if (!control_service_ || !ub_transport_) { + return Status::InternalError( + "UbTentTransport: setupUbLocalSegment: not ready" LOC_MARK); + } + + auto& manager = control_service_->segmentManager(); + auto segment = manager.getLocal(); + if (!segment) { + return Status::InternalError( + "UbTentTransport: setupUbLocalSegment: local segment not " + "found" LOC_MARK); + } + + if (segment->type != tent::SegmentType::Memory) { + return Status::InvalidArgument( + "UbTentTransport: local segment is not a memory segment" LOC_MARK); + } + + auto& detail = std::get(segment->detail); + + // The old-TE local SegmentDesc (built by allocateLocalSegmentID) already + // has devices with EIDs. Mirror them into the TENT MemorySegmentDesc. + auto local_te_desc = + te_metadata_bridge_->getSegmentDescByID(LOCAL_SEGMENT_ID); + if (local_te_desc) { + for (const auto& te_dev : local_te_desc->devices) { + // Check if this device is already present in the TENT descriptor. + bool found = false; + for (auto& tent_dev : detail.devices) { + if (tent_dev.name == te_dev.name) { + tent_dev.transport_attrs[TransportType::UB] = te_dev.eid; + found = true; + break; + } + } + if (!found) { + tent::DeviceDesc d; + d.name = te_dev.name; + d.transport_attrs[TransportType::UB] = te_dev.eid; + detail.devices.push_back(std::move(d)); + } + } + } + + // Tag the segment so remote nodes know UB is available. + detail.transport_attrs[static_cast(TransportType::UB)] = "ub"; + + return manager.synchronizeLocal(); +} + +Status UbTentTransport::uninstall() { + if (control_service_) { + control_service_->setBootstrapUbCallback(nullptr); + } + ub_transport_.reset(); + te_metadata_bridge_.reset(); + te_topology_.reset(); + control_service_.reset(); + { + std::lock_guard lock(seg_id_cache_mutex_); + tent_to_te_seg_id_.clear(); + } + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// Memory registration +// --------------------------------------------------------------------------- + +Status UbTentTransport::addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) { + if (!ub_transport_) { + return Status::InternalError("UbTentTransport: not installed" LOC_MARK); + } + + void* addr = reinterpret_cast(desc.addr); + size_t length = static_cast(desc.length); + const std::string& location = + options.location.empty() ? kWildcardLocation : options.location; + bool remote_accessible = (options.perm != kLocalReadWrite); + + int rc = ub_transport_->registerLocalMemory(addr, length, location, + remote_accessible); + if (rc != 0) { + LOG(ERROR) << "UbTentTransport: registerLocalMemory failed, rc=" << rc + << " addr=" << addr << " length=" << length; + return Status::InternalError( + "UbTentTransport: registerLocalMemory failed, rc=" + + std::to_string(rc) + LOC_MARK); + } + + // After registration, the old-TE local SegmentDesc (accessible via bridge) + // contains the tseg handles for this buffer. Serialize them into the TENT + // BufferDesc so that remote nodes can extract them via convertFromTent(). + auto local_te_desc = + te_metadata_bridge_->getSegmentDescByID(LOCAL_SEGMENT_ID); + if (local_te_desc) { + for (const auto& te_buf : local_te_desc->buffers) { + if (te_buf.addr == desc.addr && !te_buf.tseg.empty()) { + try { + json j(te_buf.tseg); + desc.transport_attrs[TransportType::UB] = j.dump(); + } catch (const std::exception& e) { + LOG(WARNING) + << "UbTentTransport: failed to serialize tseg: " + << e.what(); + } + break; + } + } + } + + desc.transports.push_back(TransportType::UB); + return Status::OK(); +} + +Status UbTentTransport::removeMemoryBuffer(BufferDesc& desc) { + if (!ub_transport_) { + return Status::InternalError("UbTentTransport: not installed" LOC_MARK); + } + + void* addr = reinterpret_cast(desc.addr); + int rc = ub_transport_->unregisterLocalMemory(addr); + if (rc != 0) { + LOG(WARNING) << "UbTentTransport: unregisterLocalMemory failed, rc=" + << rc << " addr=" << addr; + } + + // Remove UB from the transports list in the descriptor. + auto& ts = desc.transports; + ts.erase(std::remove(ts.begin(), ts.end(), TransportType::UB), ts.end()); + + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// SubBatch management +// --------------------------------------------------------------------------- + +Status UbTentTransport::allocateSubBatch(SubBatchRef& batch, size_t max_size) { + if (!ub_transport_) { + return Status::InternalError("UbTentTransport: not installed" LOC_MARK); + } + + auto* ub_batch = new UbSubBatch(); + ub_batch->ub_batch_id = ub_transport_->allocateBatchID(max_size); + if (ub_batch->ub_batch_id == 0) { + delete ub_batch; + return Status::InternalError( + "UbTentTransport: allocateBatchID failed" LOC_MARK); + } + // Pre-reserve request storage to avoid reallocation after submit. + ub_batch->te_requests.reserve(max_size); + + batch = ub_batch; + return Status::OK(); +} + +Status UbTentTransport::freeSubBatch(SubBatchRef& batch) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch) { + return Status::InvalidArgument( + "UbTentTransport: invalid sub-batch" LOC_MARK); + } + + if (ub_transport_ && ub_batch->ub_batch_id != 0) { + auto s = ub_transport_->freeBatchID(ub_batch->ub_batch_id); + if (!s.ok()) { + LOG(WARNING) << "UbTentTransport: freeBatchID failed: " + << s.message(); + } + } + + delete ub_batch; + batch = nullptr; + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// Transfer submission and status +// --------------------------------------------------------------------------- + +Status UbTentTransport::submitTransferTasks( + SubBatchRef batch, const std::vector& request_list) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch) { + return Status::InvalidArgument( + "UbTentTransport: invalid sub-batch" LOC_MARK); + } + if (!ub_transport_) { + return Status::InternalError("UbTentTransport: not installed" LOC_MARK); + } + + auto& batch_desc = mooncake::Transport::toBatchDesc(ub_batch->ub_batch_id); + if (batch_desc.task_list.size() + request_list.size() > + batch_desc.batch_size) { + return Status::TooManyRequests( + "UbTentTransport: exceed batch capacity" LOC_MARK); + } + + // Convert TENT Requests → old-TE TransferRequests. + // The converted requests are stored inside ub_batch->te_requests so that + // the raw pointers assigned to TransferTask::request remain valid until + // freeSubBatch() is called. + std::vector converted_requests; + converted_requests.reserve(request_list.size()); + for (const auto& req : request_list) { + mooncake::Transport::TransferRequest te_req{}; + te_req.opcode = (req.opcode == Request::READ) + ? mooncake::Transport::TransferRequest::READ + : mooncake::Transport::TransferRequest::WRITE; + te_req.source = req.source; + te_req.target_offset = req.target_offset; + te_req.length = req.length; + + if (req.target_id == LOCAL_SEGMENT_ID) { + te_req.target_id = + static_cast(LOCAL_SEGMENT_ID); + } else { + auto te_id = getTESegmentID(req.target_id); + if (te_id == static_cast(-1)) { + return Status::InvalidArgument( + "UbTentTransport: cannot resolve TENT segment ID " + + std::to_string(req.target_id) + LOC_MARK); + } + te_req.target_id = te_id; + } + converted_requests.push_back(te_req); + } + + size_t first_new = ub_batch->te_requests.size(); + ub_batch->te_requests.insert(ub_batch->te_requests.end(), + converted_requests.begin(), + converted_requests.end()); + + // Set up old-TE task list inside the existing BatchDesc. + size_t first_task = batch_desc.task_list.size(); + batch_desc.task_list.resize(first_task + request_list.size()); + + std::vector task_ptrs; + task_ptrs.reserve(request_list.size()); + for (size_t i = 0; i < request_list.size(); ++i) { + auto& task = batch_desc.task_list[first_task + i]; + task.batch_id = ub_batch->ub_batch_id; + task.request = &ub_batch->te_requests[first_new + i]; + task_ptrs.push_back(&task); + } + + ub_batch->task_count_ += request_list.size(); + + auto old_s = ub_transport_->submitTransferTask(task_ptrs); + if (!old_s.ok()) { + batch_desc.task_list.resize(first_task); + ub_batch->te_requests.resize(first_new); + ub_batch->task_count_ -= request_list.size(); + return Status::InternalError( + "UbTentTransport: submitTransferTask failed: " + + std::string(old_s.message()) + LOC_MARK); + } + return Status::OK(); +} + +Status UbTentTransport::getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) { + auto* ub_batch = dynamic_cast(batch); + if (!ub_batch) { + return Status::InvalidArgument( + "UbTentTransport: invalid sub-batch" LOC_MARK); + } + if (!ub_transport_) { + return Status::InternalError("UbTentTransport: not installed" LOC_MARK); + } + + mooncake::Transport::TransferStatus te_status{}; + auto old_s = ub_transport_->getTransferStatus( + ub_batch->ub_batch_id, static_cast(task_id), te_status); + if (!old_s.ok()) { + return Status::InternalError( + "UbTentTransport: getTransferStatus failed: " + + std::string(old_s.message()) + LOC_MARK); + } + + status.transferred_bytes = te_status.transferred_bytes; + + // Map old-TE status enum → TENT status enum. + switch (te_status.s) { + case mooncake::Transport::WAITING: + case mooncake::Transport::PENDING: + status.s = TransferStatusEnum::PENDING; + break; + case mooncake::Transport::COMPLETED: + status.s = TransferStatusEnum::COMPLETED; + break; + case mooncake::Transport::FAILED: + status.s = TransferStatusEnum::FAILED; + break; + case mooncake::Transport::TIMEOUT: + status.s = TransferStatusEnum::TIMEOUT; + break; + case mooncake::Transport::CANCELED: + status.s = TransferStatusEnum::CANCELED; + break; + case mooncake::Transport::INVALID: + status.s = TransferStatusEnum::INVALID; + break; + default: + status.s = TransferStatusEnum::INVALID; + break; + } + return Status::OK(); +} + +// --------------------------------------------------------------------------- +// Segment ID translation +// --------------------------------------------------------------------------- + +mooncake::Transport::SegmentID UbTentTransport::getTESegmentID( + SegmentID tent_id) { + // Check the per-transport cache first. + { + std::lock_guard lock(seg_id_cache_mutex_); + auto it = tent_to_te_seg_id_.find(tent_id); + if (it != tent_to_te_seg_id_.end()) return it->second; + } + + if (!control_service_ || !te_metadata_bridge_) { + LOG(ERROR) << "UbTentTransport: cannot resolve TENT segment ID " + << tent_id << " (not installed)"; + return static_cast(-1); + } + + // Look up the TENT segment name from the segment manager. + tent::SegmentDesc* tent_desc = nullptr; + auto s = + control_service_->segmentManager().getRemoteCached(tent_desc, tent_id); + if (!s.ok() || !tent_desc) { + LOG(ERROR) << "UbTentTransport: cannot resolve TENT segment ID " + << tent_id << ": " << s.message(); + return static_cast(-1); + } + + // Obtain the old-TE ID via the bridge (which maps name → TENT SegmentID). + auto old_te_id = te_metadata_bridge_->getSegmentID(tent_desc->name); + if (old_te_id == static_cast(-1)) { + LOG(ERROR) << "UbTentTransport: bridge could not resolve segment '" + << tent_desc->name << "'"; + return static_cast(-1); + } + + std::lock_guard lock(seg_id_cache_mutex_); + tent_to_te_seg_id_[tent_id] = old_te_id; + return old_te_id; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 4bd5553035..17db279568 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -1,5 +1,14 @@ # TENT Tests and Examples +# tent_xport_ub embeds old-TE ub_transport objects. Executables that link +# tent_link_group must also link transfer_engine for TransferMetadata and +# related symbols. Keep this at the test executable level to avoid a CMake +# target cycle (transfer_engine already links tent_link_group when USE_TENT). +set(TENT_GTEST_LINK_LIBS gtest gtest_main tent_link_group) +if(USE_UB AND TARGET transfer_engine) + list(APPEND TENT_GTEST_LINK_LIBS transfer_engine) +endif() + # TENT Metrics Example add_executable(tent_metrics_example tent_metrics_example.cpp) # asio_shared must be linked explicitly for executables: @@ -27,12 +36,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,9 +102,18 @@ 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) +target_link_libraries(request_merge_test PRIVATE ${TENT_GTEST_LINK_LIBS}) target_include_directories(request_merge_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME request_merge_test COMMAND request_merge_test) @@ -50,7 +121,7 @@ add_test(NAME request_merge_test COMMAND request_merge_test) add_executable(transfer_engine_config_override_test transfer_engine_config_override_test.cpp) target_link_libraries(transfer_engine_config_override_test - PRIVATE gtest gtest_main tent_link_group) + PRIVATE ${TENT_GTEST_LINK_LIBS}) if(TARGET asio_shared) target_link_libraries(transfer_engine_config_override_test PRIVATE asio_shared) @@ -61,36 +132,46 @@ add_test(NAME transfer_engine_config_override_test COMMAND transfer_engine_config_override_test) add_executable(tent_tcp_transport_test tcp_transport_test.cpp) -target_link_libraries(tent_tcp_transport_test PRIVATE gtest gtest_main - tent_link_group) +target_link_libraries(tent_tcp_transport_test PRIVATE ${TENT_GTEST_LINK_LIBS}) 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) +target_link_libraries(tent_failover_test PRIVATE ${TENT_GTEST_LINK_LIBS}) target_include_directories(tent_failover_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) 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 tent_link_group) + PRIVATE ${TENT_GTEST_LINK_LIBS}) +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_link_libraries(tent_endpoint_store_test PRIVATE ${TENT_GTEST_LINK_LIBS}) +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_link_libraries(tent_rdma_transport_test PRIVATE ${TENT_GTEST_LINK_LIBS} + ibverbs) +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) @@ -98,9 +179,10 @@ 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 + target_link_libraries(tent_rocm_platform_test PRIVATE ${TENT_GTEST_LINK_LIBS} 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) @@ -109,9 +191,9 @@ 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 ${TENT_GTEST_LINK_LIBS}) + PRIVATE gtest gtest_main tent_link_group) target_include_directories( tent_sunrise_link_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include @@ -120,23 +202,27 @@ if(USE_SUNRISE) COMMAND tent_sunrise_link_transport_test) endif() add_executable(tent_fault_proxy_test fault_proxy_test.cpp) +target_link_libraries(tent_fault_proxy_test PRIVATE ${TENT_GTEST_LINK_LIBS}) 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 ${TENT_GTEST_LINK_LIBS}) 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 add_executable(tent_transport_selector_test transport_selector_test.cpp) +target_link_libraries(tent_transport_selector_test + PRIVATE ${TENT_GTEST_LINK_LIBS}) 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) @@ -145,7 +231,9 @@ add_test(NAME tent_transport_selector_test COMMAND tent_transport_selector_test) # 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) + PRIVATE ${TENT_GTEST_LINK_LIBS}) +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() @@ -157,8 +245,9 @@ 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 ${TENT_GTEST_LINK_LIBS}) +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() @@ -166,23 +255,31 @@ 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 ${TENT_GTEST_LINK_LIBS}) +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) +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) + PRIVATE ${TENT_GTEST_LINK_LIBS}) +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() @@ -190,3 +287,80 @@ 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) + +# UB TENT transport unit test (mock URMA; no real Kunpeng hardware required). +if(USE_UB) + add_executable(tent_ub_transport_test ub_tent_transport_test.cpp) + target_link_libraries(tent_ub_transport_test PRIVATE ${TENT_GTEST_LINK_LIBS}) + if(TARGET asio_shared) + target_link_libraries(tent_ub_transport_test PRIVATE asio_shared) + endif() + target_include_directories( + tent_ub_transport_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../include + PRIVATE ${CMAKE_SOURCE_DIR}/mooncake-common/include) + add_test(NAME tent_ub_transport_test COMMAND tent_ub_transport_test) + + # UB + TENT dual-node e2e integration test (requires real URMA hardware and + # two machines sharing a TENT metadata backend). + add_executable(tent_ub_e2e_dual_node_test ub_e2e_dual_node_test.cpp) + target_link_libraries(tent_ub_e2e_dual_node_test + PRIVATE ${TENT_GTEST_LINK_LIBS} gflags glog) + if(TARGET asio_shared) + target_link_libraries(tent_ub_e2e_dual_node_test PRIVATE asio_shared) + endif() + target_include_directories( + tent_ub_e2e_dual_node_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../include + PRIVATE ${CMAKE_SOURCE_DIR}/mooncake-common/include) +endif() +# 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 b26ff889f9..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; @@ -388,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..b4b927700d --- /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, 100.0); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::Dispatch, 50.0); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::Transport, 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/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/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/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 5963ed008d..fc05c7a411 100644 --- a/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp +++ b/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp @@ -104,6 +104,69 @@ TEST(RailMonitorConfigTest, CustomJsonOverridesAutomaticPeerMapping) { 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 index bc4c4971ed..197cca1246 100644 --- a/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp +++ b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -81,6 +82,8 @@ std::shared_ptr makeRdmaConfig() { 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; @@ -122,6 +125,9 @@ 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); @@ -135,7 +141,7 @@ TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { TransferEngine server(makeRdmaConfig()); if (!server.available()) _exit(2); - std::vector buffer(kDataLength); + std::vector buffer(kBufferLength); if (!server.registerLocalMemory(buffer.data(), buffer.size()).ok()) _exit(3); @@ -172,7 +178,7 @@ TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { TransferEngine client(makeRdmaConfig()); ASSERT_TRUE(client.available()); - std::vector buffer(kDataLength * 2); + std::vector buffer(kBufferLength); for (size_t i = 0; i < kDataLength; ++i) { buffer[i] = static_cast((i * 31) & 0xff); } @@ -214,6 +220,45 @@ TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { 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()); 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 index 2a94c9e5ff..b164b642a3 100644 --- a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -60,6 +60,8 @@ class FakeTransport : public Transport { 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, @@ -110,6 +112,23 @@ class FakeTransport : public Transport { 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(); @@ -358,6 +377,113 @@ TEST(RuntimeQueueDispatch, KeepsDispatchWindowUntilOwnerIsTerminal) { 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); 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..47f8c0b557 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/shm_transport_test.cpp @@ -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. + +#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(); + } +}; + +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}; +}; + +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(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); + BufferDesc buffer; + buffer.addr = kRemoteAddress; + buffer.length = page_size; + buffer.location = "cpu:0"; + buffer.shm_path = shm_file.name(); + memory.buffers.push_back(std::move(buffer)); + return Status::OK(); + }) + .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()); +} + +} // namespace +} // namespace tent +} // namespace mooncake 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 223e2cf1ea..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 @@ -293,6 +293,49 @@ TEST(TransferEngineConfigOverrideTest, 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..277e0567d7 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 @@ -641,6 +643,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_e2e_dual_node_test.cpp b/mooncake-transfer-engine/tent/tests/ub_e2e_dual_node_test.cpp new file mode 100644 index 0000000000..d0885b83a8 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/ub_e2e_dual_node_test.cpp @@ -0,0 +1,335 @@ +// 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. +// +// UB + TENT dual-node end-to-end test (requires real Kunpeng URMA hardware). +// +// Usage (two machines, shared etcd): +// +// # Machine A — server (holds and exposes data) +// ./tent_ub_e2e_dual_node_test --role=server --segment_name=node_a \ +// --data_size=1048576 +// +// # Machine B — client (reads / writes remote data) +// ./tent_ub_e2e_dual_node_test --role=client --remote_segment=node_a \ +// --data_size=1048576 --operation=write # or --operation=read +// +// The test writes a known pattern, reads it back, and verifies data integrity. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#include "tent/transfer_engine.h" + +DEFINE_string(role, "", + "Role: 'server' (holds data and waits) or 'client' (initiates " + "transfer and verifies)"); +DEFINE_string(segment_name, "", "Local segment name (used by the server)"); +DEFINE_string(remote_segment, "", "Remote segment name to open as client"); +DEFINE_string(transport_config, "", + "Path to a TENT JSON config file (optional). When omitted a " + "default config with UB enabled is built automatically."); +DEFINE_int32(data_size, 1 * 1024 * 1024, + "Data size in bytes for the test buffer (default 1 MiB)"); +DEFINE_string(operation, "write", + "Operation for the client: 'write' or 'read'"); + +namespace { + +std::atomic running{true}; + +void signalHandler(int /*signum*/) { running.store(false); } + +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); +} + +std::shared_ptr buildDefaultUbConfig() { + auto conf = std::make_shared(); + std::string cfg_json = R"({ + "transports": { + "tcp": { "enable": false }, + "rdma": { "enable": false }, + "ub": { "enable": true } + } + })"; + auto status = conf->load(cfg_json); + LOG_ASSERT(status.ok()) + << "Failed to load built-in UB config: " << status.ToString(); + return conf; +} + +int runServer() { + setupSignalHandler(); + + std::shared_ptr config; + if (!FLAGS_transport_config.empty()) { + config = std::make_shared(); + auto s = config->loadFile(FLAGS_transport_config); + LOG_ASSERT(s.ok()) << "Failed to load config: " << s.ToString(); + } else { + config = buildDefaultUbConfig(); + } + + auto engine = std::make_unique(config); + LOG_ASSERT(engine->available()) + << "TENT TransferEngine not available (check config / transports)"; + + std::string seg_name = FLAGS_segment_name; + LOG_ASSERT(!seg_name.empty()) + << "Server requires --segment_name (e.g. --segment_name=node_a)"; + + // Allocate page-aligned buffer for UB URMA registration. + const size_t buf_size = static_cast(FLAGS_data_size); + void* buf = nullptr; + LOG_ASSERT(posix_memalign(&buf, 4096, buf_size) == 0) + << "Failed to allocate page-aligned buffer (" << buf_size << " bytes)"; + std::memset(buf, 0xAB, buf_size); + LOG(INFO) << "Server: allocated " << buf_size << " bytes at " << buf; + + auto reg_s = engine->registerLocalMemory(buf, buf_size, + mooncake::tent::kGlobalReadWrite); + LOG_ASSERT(reg_s.ok()) << "Server: registerLocalMemory failed: " + << reg_s.ToString(); + LOG(INFO) << "Server: registered local memory"; + + LOG(INFO) << "Server: segment '" << seg_name + << "' is ready. Press Ctrl-C to stop."; + while (running) { + std::this_thread::sleep_for(std::chrono::seconds(1)); + } + + engine->unregisterLocalMemory(buf, buf_size); + free(buf); + LOG(INFO) << "Server: shutdown complete"; + return 0; +} + +int runClient() { + std::shared_ptr config; + if (!FLAGS_transport_config.empty()) { + config = std::make_shared(); + auto s = config->loadFile(FLAGS_transport_config); + LOG_ASSERT(s.ok()) << "Failed to load config: " << s.ToString(); + } else { + config = buildDefaultUbConfig(); + } + + auto engine = std::make_unique(config); + LOG_ASSERT(engine->available()) + << "TENT TransferEngine not available (check config / transports)"; + + std::string remote = FLAGS_remote_segment; + LOG_ASSERT(!remote.empty()) + << "Client requires --remote_segment (e.g. --remote_segment=node_a)"; + + // Allocate page-aligned buffer. + const size_t buf_size = static_cast(FLAGS_data_size); + void* buf = nullptr; + LOG_ASSERT(posix_memalign(&buf, 4096, buf_size) == 0) + << "Failed to allocate page-aligned buffer (" << buf_size << " bytes)"; + std::memset(buf, 0x00, buf_size); + LOG(INFO) << "Client: allocated " << buf_size << " bytes at " << buf; + + auto reg_s = engine->registerLocalMemory(buf, buf_size, + mooncake::tent::kGlobalReadWrite); + LOG_ASSERT(reg_s.ok()) << "Client: registerLocalMemory failed: " + << reg_s.ToString(); + LOG(INFO) << "Client: registered local memory"; + + // Open the remote segment. + mooncake::tent::SegmentID seg_id = 0; + auto open_s = engine->openSegment(seg_id, remote); + LOG_ASSERT(open_s.ok()) << "Client: openSegment('" << remote + << "') failed: " << open_s.ToString(); + LOG(INFO) << "Client: opened remote segment '" << remote + << "' → segment_id " << seg_id; + + // Fetch segment info to get remote base address (buffer at index 0). + mooncake::tent::SegmentInfo info; + auto info_s = engine->getSegmentInfo(seg_id, info); + LOG_ASSERT(info_s.ok()) + << "Client: getSegmentInfo failed: " << info_s.ToString(); + LOG(INFO) << "Client: remote segment has " << info.buffers.size() + << " buffer(s)"; + if (info.buffers.empty()) { + LOG(ERROR) << "Client: remote segment has no buffers"; + engine->unregisterLocalMemory(buf, buf_size); + free(buf); + return 1; + } + uint64_t remote_addr = info.buffers[0].base; + LOG(INFO) << "Client: remote buffer base = " << remote_addr; + + // --- WRITE test --- + std::string op = FLAGS_operation; + if (op == "write") { + // Fill local buffer with a known pattern. + for (size_t i = 0; i < buf_size; ++i) + static_cast(buf)[i] = static_cast(i & 0xFF); + + auto batch_id = engine->allocateBatch(1); + LOG_ASSERT(batch_id != 0) << "Client: allocateBatch failed"; + + mooncake::tent::Request req; + req.opcode = mooncake::tent::Request::WRITE; + req.source = buf; + req.target_id = seg_id; + req.target_offset = remote_addr; + req.length = buf_size; + + auto sub_s = engine->submitTransfer(batch_id, {req}); + LOG_ASSERT(sub_s.ok()) + << "Client: submitTransfer(WRITE) failed: " << sub_s.ToString(); + + // Poll until completed. + mooncake::tent::TransferStatus ts; + bool done = false; + for (int i = 0; i < 1000 && !done; ++i) { + auto gs = engine->getTransferStatus(batch_id, ts); + LOG_ASSERT(gs.ok()) + << "Client: getTransferStatus failed: " << gs.ToString(); + if (ts.s == mooncake::tent::COMPLETED) { + done = true; + } else if (ts.s == mooncake::tent::FAILED || + ts.s == mooncake::tent::TIMEOUT) { + LOG(ERROR) << "Client: WRITE transfer failed, status=" + << static_cast(ts.s); + engine->freeBatch(batch_id); + engine->unregisterLocalMemory(buf, buf_size); + free(buf); + return 1; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + LOG_ASSERT(done) << "Client: WRITE did not complete within 10s"; + + auto free_s = engine->freeBatch(batch_id); + LOG_ASSERT(free_s.ok()) + << "Client: freeBatch failed: " << free_s.ToString(); + LOG(INFO) << "Client: WRITE " << buf_size << " bytes → COMPLETED ✓"; + } + + // --- READ-back-and-verify test --- + { + std::memset(buf, 0x00, buf_size); + auto batch_id = engine->allocateBatch(1); + LOG_ASSERT(batch_id != 0) << "Client: allocateBatch(READ) failed"; + + mooncake::tent::Request req; + req.opcode = mooncake::tent::Request::READ; + req.source = buf; + req.target_id = seg_id; + req.target_offset = remote_addr; + req.length = buf_size; + + auto sub_s = engine->submitTransfer(batch_id, {req}); + LOG_ASSERT(sub_s.ok()) + << "Client: submitTransfer(READ) failed: " << sub_s.ToString(); + + mooncake::tent::TransferStatus ts; + bool done = false; + for (int i = 0; i < 1000 && !done; ++i) { + auto gs = engine->getTransferStatus(batch_id, ts); + LOG_ASSERT(gs.ok()) + << "Client: getTransferStatus(READ) failed: " << gs.ToString(); + if (ts.s == mooncake::tent::COMPLETED) + done = true; + else if (ts.s == mooncake::tent::FAILED || + ts.s == mooncake::tent::TIMEOUT) { + LOG(ERROR) << "Client: READ transfer failed, status=" + << static_cast(ts.s); + engine->freeBatch(batch_id); + engine->unregisterLocalMemory(buf, buf_size); + free(buf); + return 1; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + LOG_ASSERT(done) << "Client: READ did not complete within 10s"; + + engine->freeBatch(batch_id); + LOG(INFO) << "Client: READ " << buf_size << " bytes → COMPLETED ✓"; + } + + // Data integrity: if we wrote first, verify the content. + if (op == "write") { + bool ok = true; + size_t first_mismatch = 0; + uint8_t expected = 0, got = 0; + for (size_t i = 0; i < buf_size && ok; ++i) { + expected = static_cast(i & 0xFF); + got = static_cast(buf)[i]; + if (got != expected) { + ok = false; + first_mismatch = i; + } + } + if (!ok) { + LOG(ERROR) << "Client: data MISMATCH at offset " << first_mismatch + << " (expected 0x" << std::hex + << static_cast(expected) << ", got 0x" + << static_cast(got) << std::dec << ")"; + engine->unregisterLocalMemory(buf, buf_size); + free(buf); + return 1; + } + LOG(INFO) << "Client: data integrity VERIFIED ✓ (all " << buf_size + << " bytes match)"; + } + + engine->unregisterLocalMemory(buf, buf_size); + free(buf); + LOG(INFO) << "Client: test PASSED"; + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + gflags::SetUsageMessage( + "UB + TENT dual-node e2e test.\n\n" + "Server: --role=server --segment_name=\n" + "Client: --role=client --remote_segment=\n\n" + "Both nodes must share a TENT metadata backend (etcd or built-in " + "RPC). Use --transport_config to point at a JSON file, or omit it " + "to use the built-in UB-only config."); + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = 1; + + if (FLAGS_role == "server" || FLAGS_role.empty()) { + return runServer(); + } else if (FLAGS_role == "client") { + return runClient(); + } + + LOG(ERROR) << "Unknown --role '" << FLAGS_role + << "'. Use 'server' or 'client'."; + return 1; +} diff --git a/mooncake-transfer-engine/tent/tests/ub_tent_transport_test.cpp b/mooncake-transfer-engine/tent/tests/ub_tent_transport_test.cpp new file mode 100644 index 0000000000..a0f464d323 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/ub_tent_transport_test.cpp @@ -0,0 +1,306 @@ +// 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 UbTentTransport and the "ub" selector string mapping. +// +// These tests run without real Kunpeng hardware. The UbTransport falls back +// to mock_urma_device automatically when liburma.so is absent or no real HCA +// is found (see UbTransport::initializeUbResources()). +// +// To build and run: +// cmake -DUSE_UB=ON -DUSE_TENT=ON ... +// ctest -R tent_ub_transport_test -V + +#include +#include + +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#include "tent/runtime/platform.h" +#include "tent/runtime/transport_selector.h" +#include "tent/runtime/segment.h" +#include "tent/transport/ub/ub_tent_transport.h" + +namespace mooncake { +namespace tent { +namespace { + +// --------------------------------------------------------------------------- +// 1. Selector / string-mapping tests (no hardware required) +// --------------------------------------------------------------------------- + +TEST(UbSelectorTest, TypeNameRoundTrip) { + EXPECT_EQ(TransportSelector::transportTypeName(UB), "ub"); +} + +TEST(UbSelectorTest, ParseUbString) { + EXPECT_EQ(TransportSelector::parseTransportType("ub"), UB); +} + +TEST(UbSelectorTest, ParseUnknownStillReturnsUnspec) { + EXPECT_EQ(TransportSelector::parseTransportType("ub_typo"), UNSPEC); +} + +TEST(UbSelectorTest, UbEnumValue) { + // UB must be between SUNRISE_LINK and kNumTransportTypes. + EXPECT_GT(static_cast(UB), static_cast(SUNRISE_LINK)); + EXPECT_LT(static_cast(UB), static_cast(kNumTransportTypes)); +} + +// A policy JSON with "ub" first in the transports array must cause the +// selector to pick UB when a UB transport is available. +TEST(UbSelectorTest, SelectorPicksUbWhenFirstInPolicy) { + auto conf = std::make_shared(); + const std::string policy_json = R"({ + "policy": [ + { + "name": "kunpeng_ub_memory", + "segment_type": "memory", + "local_memory": "cpu", + "remote_memory": "cpu", + "same_machine": false, + "transports": ["ub", "rdma", "tcp"] + } + ] + })"; + ASSERT_TRUE(conf->load(policy_json).ok()); + + TransportSelector selector(conf); + + // Register a fake UB and TCP transport. + std::array, kSupportedTransportTypes> + transports{}; + + struct MinimalFake : public Transport { + void setDram() { caps.dram_to_dram = true; } + Status allocateSubBatch(SubBatchRef&, size_t) override { + return Status::OK(); + } + Status freeSubBatch(SubBatchRef&) override { return Status::OK(); } + Status submitTransferTasks(SubBatchRef, + const std::vector&) override { + return Status::OK(); + } + Status getTransferStatus(SubBatchRef, int, TransferStatus&) override { + return Status::OK(); + } + const char* getName() const override { return "fake"; } + }; + + auto ub_fake = std::make_shared(); + ub_fake->setDram(); + auto tcp_fake = std::make_shared(); + tcp_fake->setDram(); + transports[UB] = ub_fake; + transports[TCP] = tcp_fake; + + 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 = 0; + ctx.buffer_transports = nullptr; + + auto result = selector.select(ctx, transports); + EXPECT_EQ(result.transport, UB) + << "Selector should pick UB first per policy"; +} + +// --------------------------------------------------------------------------- +// 2. UbTentTransport control-flow tests (mock URMA, no network) +// --------------------------------------------------------------------------- + +// install() with "p2p" metadata (no etcd) and a null Config should succeed +// when mock URMA is active (no real liburma.so). +TEST(UbTentTransportTest, InstallWithMockUrma) { + UbTentTransport transport; + + std::string seg_name = "test_segment"; + // Null ControlService and Topology are accepted; UbTentTransport creates + // its own old-TE Topology via discover() and a p2p TransferMetadata. + auto status = transport.install(seg_name, nullptr, nullptr, nullptr); + + // Success is expected on machines where mock_urma_device is available. + // If install fails (e.g. liburma.so found but no real HCA), we still + // accept that gracefully and skip dependent sub-tests. + if (!status.ok()) { + GTEST_SKIP() << "UbTentTransport::install() failed (likely no mock " + "URMA): " + << status.message(); + } + EXPECT_STREQ(transport.getName(), "ub"); + EXPECT_TRUE(transport.capabilities().dram_to_dram); + EXPECT_FALSE(transport.capabilities().dram_to_gpu); +} + +TEST(UbTentTransportTest, AddAndRemoveMemoryBuffer) { + UbTentTransport transport; + std::string seg_name = "test_segment"; + auto status = transport.install(seg_name, nullptr, nullptr, nullptr); + if (!status.ok()) { + GTEST_SKIP() << "install failed: " << status.message(); + } + + // Allocate a page-aligned CPU buffer. urma_register_seg on Kunpeng + // hardware requires the start address to be page-aligned (4 KiB). + // std::vector uses malloc which only guarantees 16-byte alignment, so we + // use posix_memalign here instead. + const size_t kBufLen = 4096; + void* addr = nullptr; + ASSERT_EQ(posix_memalign(&addr, 4096, kBufLen), 0); + std::memset(addr, 0, kBufLen); + + BufferDesc desc; + desc.addr = reinterpret_cast(addr); + desc.length = kBufLen; + desc.location = "*"; + + MemoryOptions opts; + opts.perm = kGlobalReadWrite; + + auto add_s = transport.addMemoryBuffer(desc, opts); + ASSERT_TRUE(add_s.ok()) << add_s.message(); + + // UB must appear in the transport list after registration. + auto it = std::find(desc.transports.begin(), desc.transports.end(), UB); + EXPECT_NE(it, desc.transports.end()) + << "UB not found in desc.transports after addMemoryBuffer()"; + + auto rm_s = transport.removeMemoryBuffer(desc); + EXPECT_TRUE(rm_s.ok()) << rm_s.message(); + + // UB should be removed from the transport list. + it = std::find(desc.transports.begin(), desc.transports.end(), UB); + EXPECT_EQ(it, desc.transports.end()) + << "UB still present in desc.transports after removeMemoryBuffer()"; + + free(addr); +} + +TEST(UbTentTransportTest, AllocateAndFreeSubBatch) { + UbTentTransport transport; + std::string seg_name = "test_segment"; + auto status = transport.install(seg_name, nullptr, nullptr, nullptr); + if (!status.ok()) { + GTEST_SKIP() << "install failed: " << status.message(); + } + + Transport::SubBatchRef batch = nullptr; + auto alloc_s = transport.allocateSubBatch(batch, 8); + ASSERT_TRUE(alloc_s.ok()) << alloc_s.message(); + ASSERT_NE(batch, nullptr); + + auto free_s = transport.freeSubBatch(batch); + EXPECT_TRUE(free_s.ok()) << free_s.message(); + EXPECT_EQ(batch, nullptr); +} + +// Submit a local-to-local mock transfer and verify that getTransferStatus() +// returns a terminal state (COMPLETED or FAILED) rather than hanging. +// This exercises the BatchDesc / TransferTask lifetime path without network. +TEST(UbTentTransportTest, SubmitAndPollMockTransfer) { + UbTentTransport transport; + std::string seg_name = "test_segment"; + auto status = transport.install(seg_name, nullptr, nullptr, nullptr); + if (!status.ok()) { + GTEST_SKIP() << "install failed: " << status.message(); + } + + // Allocate page-aligned buffers (urma_register_seg requires 4 KiB + // alignment). + const size_t kBufLen = 4096; + void* src_raw = nullptr; + void* dst_raw = nullptr; + ASSERT_EQ(posix_memalign(&src_raw, 4096, kBufLen), 0); + ASSERT_EQ(posix_memalign(&dst_raw, 4096, kBufLen), 0); + std::memset(src_raw, 0xAB, kBufLen); + std::memset(dst_raw, 0x00, kBufLen); + + BufferDesc src_desc; + src_desc.addr = reinterpret_cast(src_raw); + src_desc.length = kBufLen; + src_desc.location = "*"; + + MemoryOptions opts; + opts.perm = kGlobalReadWrite; + + auto add_s = transport.addMemoryBuffer(src_desc, opts); + if (!add_s.ok()) { + GTEST_SKIP() << "addMemoryBuffer failed: " << add_s.message(); + } + + // Allocate a sub-batch. + Transport::SubBatchRef batch = nullptr; + ASSERT_TRUE(transport.allocateSubBatch(batch, 4).ok()); + ASSERT_NE(batch, nullptr); + + // Build a local WRITE request (LOCAL_SEGMENT_ID). + Request req{}; + req.opcode = Request::WRITE; + req.source = src_raw; + req.target_id = LOCAL_SEGMENT_ID; + req.target_offset = reinterpret_cast(dst_raw); + req.length = kBufLen; + + auto sub_s = transport.submitTransferTasks(batch, {req}); + // submitTransferTasks may fail if mock URMA rejects the send (e.g. no + // connected peer). Accept either outcome; we mainly verify no crash. + if (!sub_s.ok()) { + LOG(WARNING) + << "submitTransferTasks returned non-OK (expected in mock): " + << sub_s.message(); + } + + // Poll for up to ~1s. + TransferStatus ts{}; + bool terminal = false; + for (int i = 0; i < 100; ++i) { + auto gs = transport.getTransferStatus(batch, 0, ts); + if (!gs.ok()) break; + if (ts.s == COMPLETED || ts.s == FAILED || ts.s == TIMEOUT) { + terminal = true; + break; + } + usleep(10000); // 10 ms + } + // In mock mode, completion may not happen; just don't crash/hang. + (void)terminal; + + transport.removeMemoryBuffer(src_desc); + transport.freeSubBatch(batch); + free(src_raw); + free(dst_raw); +} + +// Calling uninstall() twice must not crash. +TEST(UbTentTransportTest, DoubleUninstallSafe) { + UbTentTransport transport; + std::string seg_name = "test_segment"; + auto s = transport.install(seg_name, nullptr, nullptr, nullptr); + if (!s.ok()) GTEST_SKIP() << "install failed: " << s.message(); + EXPECT_TRUE(transport.uninstall().ok()); + EXPECT_TRUE(transport.uninstall().ok()); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 1cca40a7da..2b2eb3d4b7 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -24,6 +24,12 @@ target_link_libraries(endpoint_store_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME endpoint_store_test COMMAND endpoint_store_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 @@ -66,6 +72,11 @@ if(USE_CXL) 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) @@ -77,6 +88,12 @@ if(USE_TCP) 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 @@ -84,11 +101,17 @@ add_executable(tcp_address_validation_test 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) +# 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() if(USE_HIP) @@ -181,20 +204,22 @@ if(USE_SUNRISE) ${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) + 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) + 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) @@ -202,31 +227,32 @@ if(USE_SUNRISE) ${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) + 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) + 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) + 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) + 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) + 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() @@ -244,6 +270,13 @@ 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(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 @@ -285,3 +318,12 @@ if(ENABLE_MULTI_PROTOCOL) 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) diff --git a/mooncake-transfer-engine/tests/common_test.cpp b/mooncake-transfer-engine/tests/common_test.cpp index b07971be36..a68aee92b4 100644 --- a/mooncake-transfer-engine/tests/common_test.cpp +++ b/mooncake-transfer-engine/tests/common_test.cpp @@ -1,16 +1,81 @@ #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 //------------------------------------------------------------------------------ diff --git a/mooncake-transfer-engine/tests/config_test.cpp b/mooncake-transfer-engine/tests/config_test.cpp index 85c28652cf..44f6f98d2f 100644 --- a/mooncake-transfer-engine/tests/config_test.cpp +++ b/mooncake-transfer-engine/tests/config_test.cpp @@ -27,6 +27,7 @@ class PkeyIndexEnvTest : public ::testing::Test { ::unsetenv("MC_PKEY_INDEX"); ::unsetenv("MC_AUTO_GID_MAX_RETRIES"); ::unsetenv("MC_IB_SL"); + ::unsetenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS"); } }; @@ -158,5 +159,43 @@ TEST_F(PkeyIndexEnvTest, IbSlNonNumericKeepsDefault) { 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); +} + } // namespace } // 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/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/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..c5a7e8f63d --- /dev/null +++ b/mooncake-transfer-engine/tests/nvmeof_status_test.cpp @@ -0,0 +1,114 @@ +// 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, WaitsForEverySliceBeforeReportingTerminalFailure) { + bool is_finished = true; + auto status = NVMeoFTransportTestPeer::aggregate( + {{Transport::FAILED, 0}, {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_context_reprobe_test.cpp b/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp index 319d57b242..35f8706d74 100644 --- a/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp @@ -19,6 +19,10 @@ #include #include +#ifdef __linux__ +#include +#endif + #include "common.h" #include "error.h" #include "transfer_metadata.h" @@ -39,6 +43,105 @@ using 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 FakeVerbsDeviceScope { + public: + 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; } +}; + +} // 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; } + +} // extern "C" +#endif // __linux__ + namespace mooncake { class RdmaTransportTestPeer { @@ -53,6 +156,10 @@ class RdmaTransportTestPeer { 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) { @@ -73,6 +180,53 @@ class RdmaContextTestPeer { 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_; +}; + +TEST_F(RdmaContextConstructionTest, RejectsZeroCompletionQueuesBeforeSetup) { + EXPECT_EQ(context_->construct(/*num_cq_list=*/0, + /*num_comp_channels=*/1), + ERR_INVALID_ARGUMENT); + EXPECT_FALSE(RdmaContextTestPeer::hasEndpointStore(*context_)); +} + +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 +} + ibv_gid makeGid(const std::array &bytes) { ibv_gid gid = {}; std::memcpy(gid.raw, bytes.data(), bytes.size()); 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/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/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/transfer_metadata_test.cpp b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp index d8fd0d889c..c1e6141604 100644 --- a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp +++ b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp @@ -19,9 +19,13 @@ #include #include +#include #include +#include +#include "config.h" #include "transport/transport.h" +#include "transfer_metadata_plugin.h" using namespace mooncake; @@ -123,9 +127,115 @@ 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"; +} + } // 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..3c5b10319e 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -17,18 +17,105 @@ #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)); + } +}; + +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 { @@ -171,9 +258,117 @@ 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); +} } // 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-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..b37cfb4006 100644 --- a/mooncake-wheel/mooncake/http_metadata_server.py +++ b/mooncake-wheel/mooncake/http_metadata_server.py @@ -61,7 +61,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 e311296865..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): @@ -103,9 +159,9 @@ def _parse_bool(value) -> bool: s = str(value).strip().lower() if not s: return False - if s in ("true", "1", "yes", "on"): + if s in ("true", "1", "yes", "on", "enable"): return True - if s in ("false", "0", "no", "off"): + if s in ("false", "0", "no", "off", "disable"): return False raise ValueError(f"Invalid boolean value: {value!r}") @@ -119,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: { @@ -140,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", @@ -153,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 @@ -165,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) @@ -179,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"), @@ -192,26 +319,39 @@ def from_file(file_path: str) -> 'MooncakeConfig': device_name=config.get("device_name", ""), master_server_address=config.get("master_server_address"), enable_ssd_offload=_parse_bool(config.get("enable_ssd_offload", False)), - ssd_offload_path=str(config.get("ssd_offload_path", "")), + 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) @@ -219,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=_parse_bool(os.getenv("MOONCAKE_OFFLOAD_ENABLED", "false")), + 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_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_store_service.py b/mooncake-wheel/mooncake/mooncake_store_service.py index a91584070d..796b1e6c66 100644 --- a/mooncake-wheel/mooncake/mooncake_store_service.py +++ b/mooncake-wheel/mooncake/mooncake_store_service.py @@ -20,6 +20,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 @@ -63,10 +64,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,7 +91,7 @@ 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): @@ -122,22 +125,30 @@ 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 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 +157,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}. " @@ -157,24 +170,40 @@ async def start_store_service(self, max_wait_time: float = 60): if actual_sleep_time > 0: 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,24 +224,34 @@ 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") + logging.info( + "Reconfigure decode: unmounting previous segments before remount" + ) ret = self.store.unmount_segment(self.mounted_segment_ids) if ret != 0: 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"Unmount of previous segments failed, ret={ret}" + } + ), + content_type="application/json", ) self.mounted_segment_ids.clear() - result = self.store.mount_segment(path, size, offset, protocol, location) + result = self.store.mount_segment( + path, size, offset, protocol, location + ) if result["ret"] != 0: self.current_mode = "prefill" self.mounted_segment_ids.clear() @@ -228,24 +267,29 @@ async def handle_reconfigure(self, request): "mode": self.current_mode, } ), - content_type="application/json" + content_type="application/json", ) self.mounted_segment_ids = list(result["segment_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 +299,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 +312,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 +345,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 +353,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 +371,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 +422,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 +435,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 +467,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 +484,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 +502,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,124 +523,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'] + key = request.match_info["key"] 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' + 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' + 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' + 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' + 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' + 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' + text=json.dumps({"error": "Exist check failed"}), + content_type="application/json", ) return web.Response( status=200, - text=json.dumps({'exists': exists > 0}), - 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): @@ -602,21 +648,21 @@ 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): @@ -624,32 +670,44 @@ async def stop(self): self.store.close() 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}") 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 26ead11351..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,42 @@ 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. @@ -107,6 +150,20 @@ def test_enable_ssd_offload_string_values(self): 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""" for field in ["local_hostname", "metadata_server", "master_server_address"]: @@ -117,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): @@ -193,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) @@ -233,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_store_service_api.py b/mooncake-wheel/tests/test_mooncake_store_service_api.py index 62d28e0e42..e416f63ae2 100644 --- a/mooncake-wheel/tests/test_mooncake_store_service_api.py +++ b/mooncake-wheel/tests/test_mooncake_store_service_api.py @@ -2,10 +2,12 @@ import asyncio import json import sys +import tempfile import types import unittest from pathlib import Path from types import SimpleNamespace +from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -50,6 +52,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)) @@ -107,6 +114,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( diff --git a/mooncake-wheel/tests/test_put_get_tensor.py b/mooncake-wheel/tests/test_put_get_tensor.py index 0c84f65a70..99554d76f8 100644 --- a/mooncake-wheel/tests/test_put_get_tensor.py +++ b/mooncake-wheel/tests/test_put_get_tensor.py @@ -5,6 +5,12 @@ 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 +19,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 +128,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 diff --git a/scripts/tone_tests/scripts/common.sh b/scripts/tone_tests/scripts/common.sh index 1c422bf2e3..143d3f78c7 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,81 @@ 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 "nvidia-smi not available, skipping GPU drain wait" + return 0 + 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) + [ -z "$max_used" ] && { echo "nvidia-smi query failed, skipping GPU drain wait"; return 0; } + 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 +} + +# Force-kill every host process still holding GPU memory. This catches leftovers +# that an in-container pkill cannot reach: processes reparented to the host +# (orphans) or in a different PID namespace. Only GPU-holding PIDs are targeted. +force_kill_gpu_procs() { + command -v nvidia-smi >/dev/null 2>&1 || return 0 + local pids + 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 "$pids" ] && return 0 + echo "Force-killing GPU-holding PIDs: $(echo $pids | tr '\n' ' ')" + for pid in $pids; do + kill -9 "$pid" 2>/dev/null || true + done + sleep 3 +} + +# Fully clear GPU memory on the current node: graceful in-container kill first, +# then host-level force-kill of any process still holding GPU memory. +# 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. +# force_kill_gpu_procs stays as a fallback for leftovers restart cannot reclaim. +drain_gpu_local() { + echo "Restarting container ${CONTAINER_NAME} to reset GPU/ERDMA state..." + docker restart ${CONTAINER_NAME} >/dev/null 2>&1 || true + if ! wait_gpu_idle 60; then + force_kill_gpu_procs + wait_gpu_idle 45 || echo "WARNING: GPU still occupied after restart+force-kill (possible stuck/zombie process or GPU fault)" + fi +} + +# 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 =====" + drain_gpu_local + + if [ -n "$REMOTE_IP" ]; then + echo "Resetting environment on remote node $REMOTE_IP..." + ${SSH_CMD} "$REMOTE_IP" " + source ${REMOTE_TEST_DIR}/run/.shrc && \ + source ${REMOTE_TEST_DIR}/scripts/common.sh && \ + drain_gpu_local + " 2>/dev/null || true + fi +} + setup_node_env() { local registry_addr=$1 echo "===== Setting up docker environment =====" @@ -416,6 +491,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 +916,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 +944,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 +995,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..a768a0b1cd --- 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 @@ -282,6 +285,10 @@ run_all_tests(){ fi [ $exit_code -ne 0 ] && all_passed=false + + # Container is shared across cases in run-all; drain leftover GPU + # memory before the next case so it does not OOM / get SIGKILLed. + drain_gpu_between_tests done cleanup_test_env "double" diff --git a/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh b/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh index fd1ff7e047..018b44772c 100644 --- a/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh +++ b/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh @@ -16,11 +16,20 @@ 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) only when both models are already downloaded. + local off_mla=$(hf_offline_prefix "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct") + local off_llama=$(hf_offline_prefix "meta-llama/Llama-3.2-3B-Instruct") + local offline_prefix="" + if [ -n "$off_mla" ] && [ -n "$off_llama" ]; then + offline_prefix="$off_mla" + fi + ${docker_exec} "\ - cd /sgl-workspace/sglang/test/registered/distributed && \ + 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' && \ - python3 -m pytest test_disaggregation_different_tp.py -v -s --tb=long" | tee "$log_file" + ${offline_prefix}python3 -m pytest test_disaggregation_different_tp.py -v -s --tb=long" | tee "$log_file" return ${PIPESTATUS[0]} } 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)"