diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 420420fee3..85679b75c6 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -369,12 +369,16 @@ jobs: - configuration: simulation features: simulation native: false + # native compiles the decrypt kernels with nvcc; native-and-simulation + # keeps covering the build without them. - configuration: native features: native native: true + kernels: true - configuration: native-and-simulation features: native,simulation native: true + kernels: false env: CARGO_TARGET_DIR: ${{ github.workspace }}/bin/racer-gpu-cargo CARGO_BUILD_JOBS: "2" @@ -387,11 +391,26 @@ jobs: - name: Set up pinned Rust run: timeout --signal=TERM --kill-after=10s 300s rustup toolchain install 1.96.0 --profile minimal --component rustfmt,clippy - # No CUDA toolkit on hosted runners: the shim builds without decrypt - # kernels and loads without a driver. Kernel math is covered by the - # C-core tests below and the kernel emulation step. + # Hosted runners have no GPU. nvcc (same package as the dataplane image) + # compiles the decrypt kernels to PTX so they at least build; running + # them needs .github/workflows/racer-gpu.yaml on a GPU runner. Kernel + # math is covered by the C-core tests below and the kernel emulation step. + - name: Install nvcc + if: matrix.kernels + run: | + timeout --signal=TERM --kill-after=10s 120s curl -fsSL -o "$RUNNER_TEMP/cuda-keyring.deb" \ + https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb + sudo dpkg -i "$RUNNER_TEMP/cuda-keyring.deb" + timeout --signal=TERM --kill-after=10s 300s sudo apt-get update + timeout --signal=TERM --kill-after=10s 600s sudo apt-get install -y --no-install-recommends cuda-nvcc-12-9 + echo /usr/local/cuda/bin >> "$GITHUB_PATH" + + - name: Compile CUDA shim with kernels (not hardware validation) + if: matrix.kernels + run: RACER_GPU_REQUIRE_PTX=1 timeout --signal=TERM --kill-after=10s 300s sh images/racer-dataplane/build-native-cuda.sh cmd/racer-dataplane/gpu/native bin/racer-gpu-native/libracer_gpu_cuda.so.1 + - name: Compile CUDA shim without kernels (not hardware validation) - if: matrix.native + if: matrix.native && !matrix.kernels run: timeout --signal=TERM --kill-after=10s 300s sh images/racer-dataplane/build-native-cuda.sh cmd/racer-dataplane/gpu/native bin/racer-gpu-native/libracer_gpu_cuda.so.1 - name: Require CUDA shim ABI and symbols (no hardware) diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml index 2e532bbf94..d6c9b67552 100644 --- a/.github/workflows/nightly.yaml +++ b/.github/workflows/nightly.yaml @@ -268,6 +268,7 @@ jobs: # - orca : deployed separately by deploy-orca # - racer-object : injected sidecars and per-ClusterCache # origin DaemonSets from the racer component + # - racer-vllm : vLLM with the Racer loader; not deployed # # machine-ops-controller is omitted on purpose: the operator synthesizes # metalman per Site instead, so machine-ops is never rolled out here. @@ -297,6 +298,9 @@ jobs: file: images/racer-dataplane/Containerfile - name: racer-object file: images/racer-object/Containerfile + - name: racer-vllm + file: images/racer-vllm/Containerfile + large: true - name: orca file: images/orca/Containerfile - name: inventory-aggregator @@ -316,6 +320,14 @@ jobs: - name: Normalize container registry to lowercase run: echo "REGISTRY=${REGISTRY,,}" >> "$GITHUB_ENV" + - name: Free disk space for large images + # The vLLM base image alone is larger than a hosted runner's free disk. + if: matrix.component.large + run: | + set -euo pipefail + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL + df -h / + - uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4.4.1 - name: Log in to ghcr.io diff --git a/.github/workflows/racer-gpu.yaml b/.github/workflows/racer-gpu.yaml new file mode 100644 index 0000000000..52b34eb9c8 --- /dev/null +++ b/.github/workflows/racer-gpu.yaml @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +name: racer gpu + +# Runs `make racer-gpu-test` on a GPU runner: the CUDA shim with its decrypt +# kernels, device discovery, client allocation import, GPU decrypt against the +# CPU result, and cuFile reads. Hosted runners have no GPU, so this runs only +# when the repository variable RACER_GPU_RUNNER holds the runner labels as a +# JSON array, for example ["self-hosted", "linux", "gpu"]. Without it the job +# is skipped instead of queuing forever (forks, repos without GPU runners). +# +# The runner needs an NVIDIA driver, gcc, curl and sudo apt access (for nvcc +# if it is not installed). Set RACER_GPU_CUFILE=require to fail when cuFile +# cannot open, on runners with GPUDirect Storage. +on: + workflow_dispatch: + schedule: + - cron: "17 6 * * *" + +permissions: + contents: read + +concurrency: + group: racer-gpu-${{ github.ref }} + cancel-in-progress: true + +jobs: + gpu: + name: Racer GPU hardware tests + if: vars.RACER_GPU_RUNNER != '' + runs-on: ${{ fromJSON(vars.RACER_GPU_RUNNER || '["ubuntu-24.04"]') }} + timeout-minutes: 45 + env: + CARGO_TARGET_DIR: ${{ github.workspace }}/bin/racer-gpu-cargo + RACER_GPU_TESTS_CUFILE: ${{ vars.RACER_GPU_CUFILE }} + RACER_GPU_TEST_DIR: ${{ github.workspace }}/tmp/racer-gpu + steps: + - name: Checkout + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Show GPUs + run: timeout --signal=TERM --kill-after=10s 60s nvidia-smi + + - name: Set up pinned Rust + run: | + if ! command -v rustup >/dev/null; then + timeout --signal=TERM --kill-after=10s 300s sh -c 'curl -fsSL https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain none' + echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" + export PATH="$HOME/.cargo/bin:$PATH" + fi + timeout --signal=TERM --kill-after=10s 300s rustup toolchain install 1.96.0 --profile minimal + rustup override set 1.96.0 + + - name: Install nvcc if missing + run: | + if command -v nvcc >/dev/null || [ -x /usr/local/cuda/bin/nvcc ]; then + echo /usr/local/cuda/bin >> "$GITHUB_PATH" + exit 0 + fi + . /etc/os-release + repo="${ID}${VERSION_ID//./}/$(uname -m | sed 's/aarch64/sbsa/')" + timeout --signal=TERM --kill-after=10s 120s curl -fsSL -o "$RUNNER_TEMP/cuda-keyring.deb" \ + "https://developer.download.nvidia.com/compute/cuda/repos/$repo/cuda-keyring_1.1-1_all.deb" + sudo dpkg -i "$RUNNER_TEMP/cuda-keyring.deb" + timeout --signal=TERM --kill-after=10s 300s sudo apt-get update + timeout --signal=TERM --kill-after=10s 600s sudo apt-get install -y --no-install-recommends cuda-nvcc-12-9 + echo /usr/local/cuda/bin >> "$GITHUB_PATH" + + - name: Run GPU hardware tests + run: | + mkdir -p "$RACER_GPU_TEST_DIR" + timeout --signal=TERM --kill-after=30s 2400s make racer-gpu-test RACER_GPU_TESTS=require diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index c709d6f398..c391fcdee9 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -26,12 +26,14 @@ name: unbounded release # -> metalman-image (stamped with the matching default netboot image) # -> component-images (matrix: machina, machine-ops-controller, unbounded-operator, # token-refresher, orca, gantry, racer-dataplane, -# racer-object, host-ubuntu2404, +# racer-object, racer-vllm, +# host-ubuntu2404, # host-azurelinux3, host-azurelinux3-nvidia, # inventory-aggregator, inventory-inspector, # inventory-viewer) # -> gantry-chart (OCI chart + signed downloadable archive) # -> manifests-tarball (make release-manifests + cosign sign) +# -> python-packages (racer-client and racer-vllm sdists and wheels) # -> finalize (attach extras + krew manifest) # # All artifacts land on a single GitHub release that is left as a draft. @@ -867,6 +869,8 @@ jobs: file: images/gantry/Containerfile platforms: linux/amd64,linux/arm64 trivy: true + # Includes GPU reads (CUDA shim and cuFile). The dataplane touches no + # GPU until a client sends the first gpu-read. - name: racer-dataplane file: images/racer-dataplane/Containerfile platforms: linux/amd64,linux/arm64 @@ -875,6 +879,13 @@ jobs: file: images/racer-object/Containerfile platforms: linux/amd64,linux/arm64 trivy: true + # vLLM with racer-client and racer-vllm. The upstream vLLM base is + # not ours to patch, so it is not gated on Trivy. + - name: racer-vllm + file: images/racer-vllm/Containerfile + platforms: linux/amd64 + trivy: false + large: true - name: host-ubuntu2404 file: images/host-ubuntu2404/Containerfile platforms: linux/amd64,linux/arm64 @@ -905,6 +916,14 @@ jobs: - name: Normalize container registry to lowercase run: echo "REGISTRY=${REGISTRY,,}" >> "$GITHUB_ENV" + - name: Free disk space for large images + # The vLLM base image alone is larger than a hosted runner's free disk. + if: matrix.component.large + run: | + set -euo pipefail + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL + df -h / + - name: Install cosign uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 @@ -1143,6 +1162,38 @@ jobs: # Past the window the only recovery is a full re-run. retention-days: 7 + # --------------------------------------------------------------------------- + # Build sdists and wheels for the Racer Python packages (racer-client and + # racer-vllm) and attach them to the draft release. The release is the only + # place they are published; nothing is uploaded to PyPI. + # --------------------------------------------------------------------------- + python-packages: + needs: version + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Build sdists and wheels + run: | + set -euo pipefail + python3 -m venv "$RUNNER_TEMP/build-venv" + "$RUNNER_TEMP/build-venv/bin/pip" install --quiet build==1.3.0 + for package in racer racer-vllm; do + timeout --signal=TERM --kill-after=10s 300s \ + "$RUNNER_TEMP/build-venv/bin/python" -m build --sdist --wheel \ + --outdir dist/python "clients/python/$package" + done + ls -l dist/python + + - name: Attach Python packages to the release + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GH_REPO: ${{ github.repository }} + TAG: ${{ needs.version.outputs.tag }} + run: | + set -euo pipefail + gh release upload "$TAG" dist/python/* --clobber + # --------------------------------------------------------------------------- # Attach the manifests tarball + krew manifest to the draft release that # GoReleaser created and set release notes. The release stays a draft here; @@ -1160,6 +1211,7 @@ jobs: - component-images - gantry-chart - manifests-tarball + - python-packages runs-on: ubuntu-latest steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/Makefile b/Makefile index 538adfb8b1..3ea72d1397 100644 --- a/Makefile +++ b/Makefile @@ -366,6 +366,7 @@ help: ## Show this help @echo "" @echo "Racer:" @echo " racer-test Test the Rust Racer workspace and Go Racer API and wire contracts" + @echo " racer-gpu-test Build the CUDA shim with kernels and run GPU hardware tests" @echo " racer-envtest Run Racer API-server tests with KUBEBUILDER_ASSETS" @echo " racer-envtest-ci Provision pinned assets and run Racer API-server tests" @echo " racer-generate Generate Racer deepcopy and CRD artifacts" @@ -575,6 +576,23 @@ racer-rust-test: ## Test the Rust Racer workspace, including the embedded cluste racer-test: racer-go-test racer-rust-test ## Check Racer +RACER_GPU_NATIVE ?= $(CURDIR)/bin/racer-gpu-native +# 1 skips when no GPU is found; require fails instead (dedicated GPU runners). +RACER_GPU_TESTS ?= require + +.PHONY: racer-gpu-test +racer-gpu-test: ## Build the CUDA shim with kernels (needs nvcc) and run GPU hardware tests (needs a GPU) + RACER_GPU_REQUIRE_PTX=1 timeout --signal=TERM --kill-after=10s 600s sh images/racer-dataplane/build-native-cuda.sh cmd/racer-dataplane/gpu/native "$(RACER_GPU_NATIVE)/libracer_gpu_cuda.so.1" + LD_LIBRARY_PATH="$(RACER_GPU_NATIVE)$${LD_LIBRARY_PATH:+:$$LD_LIBRARY_PATH}" timeout --signal=TERM --kill-after=10s 300s cargo test --locked --manifest-path cmd/racer-dataplane/Cargo.toml -p racer-gpu --no-default-features --features native --lib ffi::tests::required_native_shim_loads_abi_and_symbols -- --exact --ignored + LD_LIBRARY_PATH="$(RACER_GPU_NATIVE)$${LD_LIBRARY_PATH:+:$$LD_LIBRARY_PATH}" RACER_GPU_TESTS=$(RACER_GPU_TESTS) timeout --signal=TERM --kill-after=10s 900s cargo test --locked --manifest-path cmd/racer-dataplane/Cargo.toml -p racer-gpu --no-default-features --features native --test hardware -- --nocapture --test-threads=1 + @if [ -n "$$RACER_GPU_E2E_SOCKET" ]; then \ + test -x "$(RACER_PY_VENV)/bin/pytest" || { timeout --signal=TERM --kill-after=10s 120s python3 -m venv "$(RACER_PY_VENV)" && \ + timeout --signal=TERM --kill-after=10s 300s "$(RACER_PY_VENV)/bin/pip" install -q pytest; }; \ + cd clients/python/racer && PYTHONPATH=src timeout --signal=TERM --kill-after=10s 300s "$(RACER_PY_VENV)/bin/pytest" -q -p no:cacheprovider -m gpu -rs; \ + elif [ "$$RACER_GPU_E2E" = require ]; then \ + echo "RACER_GPU_E2E=require but RACER_GPU_E2E_SOCKET is unset" >&2; exit 1; \ + else echo "RACER_GPU_E2E_SOCKET unset: skipping the Python gpu-read test against a live dataplane"; fi + RACER_PY_VENV ?= $(CURDIR)/tmp/racer-python-venv .PHONY: racer-python-test @@ -584,6 +602,14 @@ racer-python-test: ## Run the Racer Python client and vLLM plugin tests (no GPU cd clients/python/racer && PYTHONPATH=src timeout --signal=TERM --kill-after=10s 300s "$(RACER_PY_VENV)/bin/pytest" -q -p no:cacheprovider cd clients/python/racer-vllm && PYTHONPATH=src:../racer/src timeout --signal=TERM --kill-after=10s 300s "$(RACER_PY_VENV)/bin/pytest" -q -p no:cacheprovider +.PHONY: racer-python-dist +racer-python-dist: ## Build sdists and wheels for racer-client and racer-vllm into tmp/racer-python-dist + @test -x "$(RACER_PY_VENV)/bin/pyproject-build" || { timeout --signal=TERM --kill-after=10s 120s python3 -m venv "$(RACER_PY_VENV)" && \ + timeout --signal=TERM --kill-after=10s 300s "$(RACER_PY_VENV)/bin/pip" install -q build; } + rm -rf tmp/racer-python-dist + timeout --signal=TERM --kill-after=10s 300s "$(RACER_PY_VENV)/bin/python" -m build -o tmp/racer-python-dist clients/python/racer + timeout --signal=TERM --kill-after=10s 300s "$(RACER_PY_VENV)/bin/python" -m build -o tmp/racer-python-dist clients/python/racer-vllm + $(SETUP_ENVTEST): @mkdir -p bin tmp/envtest-tools TMPDIR="$(CURDIR)/tmp/envtest-tools" GOBIN="$(CURDIR)/bin" timeout --signal=TERM --kill-after=10s 300s $(GOCMD) install sigs.k8s.io/controller-runtime/tools/setup-envtest@$(SETUP_ENVTEST_VERSION) diff --git a/clients/python/racer-vllm/LICENSE b/clients/python/racer-vllm/LICENSE new file mode 100644 index 0000000000..261eeb9e9f --- /dev/null +++ b/clients/python/racer-vllm/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/clients/python/racer-vllm/README.md b/clients/python/racer-vllm/README.md index 501919aea9..e4bf5d741a 100644 --- a/clients/python/racer-vllm/README.md +++ b/clients/python/racer-vllm/README.md @@ -5,10 +5,27 @@ format `racer`. It uses the `racer-client` package. ## Install +The packages are not on PyPI. Each GitHub release has their wheels and sdists +as assets. To install from a release: + +``` +pip install --no-index \ + --find-links https://github.com/Azure/unbounded/releases/expanded_assets/ \ + racer-vllm +``` + +From a checkout: + ``` pip install ./clients/python/racer ./clients/python/racer-vllm ``` +To build sdists and wheels, run `make racer-python-dist` from +the repository root (output in `tmp/racer-python-dist`), or +`python -m build clients/python/racer` and +`python -m build clients/python/racer-vllm`. Both packages share one version; +`racer-vllm` requires the matching `racer-client` release. + vLLM finds the plugin through the `vllm.general_plugins` entry point. It was written for vLLM v0.25.x. @@ -23,7 +40,9 @@ names the S3 object. vLLM reads the model config and tokenizer before it loads weights, so these small files must be on local disk first. Fetch them, for example in an init -container: +container. The admission webhook mounts the client socket into init +containers as well as regular containers (limited by +`racer.unbounded-cloud.io/client-containers` when set): ``` racer-vllm-fetch racer://models/my-bucket/llama-3-8b /models/llama-3-8b @@ -31,7 +50,14 @@ racer-vllm-fetch racer://models/my-bucket/llama-3-8b /models/llama-3-8b This copies `config.json`, tokenizer files, `model.safetensors.index.json` and similar files, plus small files named in the index metadata. It also -writes `racer-model.json`, which tells the loader where the weights are. +writes `racer-model.json`, which tells the loader where the weights are and +records the ETag of every fetched file and every weight file. The loader reads +the weights at those ETags, so they match the config and tokenizer. If the +model at the origin changes after the fetch, the load fails with +`ModelChangedError`; run `racer-vllm-fetch` again. The loader also refuses +any weight file that `racer-model.json` does not pin, so a layout change +(such as an index uploaded after a single-file fetch) fails the same way. +Without `racer-model.json`, nothing is pinned. Then start vLLM on that directory: @@ -50,6 +76,9 @@ Extra config keys, all optional: | `clone` | true | Copy each tensor out of staging before handing it to vLLM. | | `socket` | none | Client socket path. | | `use_vmm` | true | Allow exportable `cuMemCreate` staging, needed for `gpu-read`. | +| `mode_retry_seconds` | 300 | Retry a delivery mode the dataplane refused after this many seconds. `null` retries only after a dataplane restart. | +| `sharded_state` | false | Read a vLLM `sharded_state` checkpoint; each rank reads only its files. | +| `pattern` | as vLLM | File name pattern for `sharded_state`, such as `model-rank-{rank}-part-{part}.safetensors`. | On fast local disks, more reads in flight help only when the dataplane also fetches further ahead. Set `concurrency` to 8 together with a larger @@ -60,18 +89,61 @@ the GPU reads section of the Racer reference. 1. Find the safetensors files from `model.safetensors.index.json`, or use `model.safetensors`. -2. For each file, read the 8-byte header length and the JSON header. -3. Group tensors into ranges that fit a staging buffer. Each range ends on a +2. Stat every file before reading any, at the ETag from `racer-model.json` + when there is one. Every later read sends `If-Match` with that ETag, so the + whole model is read at one version. A file that changes during the load + fails it with `ModelChangedError`. +3. For each file, read the 8-byte header length and the JSON header. All + headers are read before any tensor data. +4. Group tensors into ranges that fit a staging buffer. Each range ends on a Racer page boundary (16 MiB), so no page is read twice. A tensor that crosses the boundary is split across two ranges and joined again before vLLM sees it. Read each range into GPU memory with `gpu-read`. If the dataplane cannot do that, use `host-read` into pinned memory and copy, else the normal stream. -4. Give each tensor (right dtype and shape) to `model.load_weights`. While +5. Give each tensor (right dtype and shape) to `model.load_weights`. While vLLM uses one buffer, the next range is read into another. A tensor bigger than a staging buffer gets its own GPU allocation. +After each load the loader logs one line (logger `vllm.racer_vllm.loader`, +so it uses vLLM's log format) such as: + +``` +racer: loaded 4 files, 14.96 GiB in 9.12s (1.64 GiB/s) from racer://models/llama/Llama-3.1-8B-Instruct; reads by delivery mode: gpu=60 (14.96 GiB) host=0 (0.00 GiB) stream=0 (0.00 GiB) +``` + +The counts cover tensor data reads. Safetensors header reads use the stream +and are not counted. + +When the dataplane refuses a mode, the client logs one warning (logger +`racer.client`) with the status, for example +`racer: gpu-read is not available from /run/racer/models/client/socket +(status 501: ...); falling back to host mode ...`. It warns again only after the +mode is retried and refused again. + +## Tensor parallelism + +With a regular checkpoint, every tensor-parallel rank reads every weight file, +and vLLM's `weight_loader` keeps the rank's slice. vLLM decides the slice +inside each model's loader, so the plugin cannot read only the needed bytes in +general. Ranks on one node share the node's Racer dataplane, so only the first +read of each page goes to a peer or the origin; repeat reads by the other ranks +come from the dataplane's memory or disk cache. + +To have each rank read only its own data, save a `sharded_state` checkpoint +with vLLM (`examples/features/sharded_state/save_sharded_state_offline.py`, +which uses `ShardedStateLoader.save_model`), upload it, and load with +`{"sharded_state": true}`. The loader probes this rank's parts +(`part-0`, `part-1`, ... until one is missing) and copies each tensor straight +into the model's state dict, as vLLM's `sharded_state` loader does. The +checkpoint must be saved with the same tensor and pipeline parallel sizes. +`racer-vllm-fetch` finds no safetensors weights for such a checkpoint, so it +probes vLLM's default patterns (or `--pattern`) and records every rank's +parts. Each rank then loads only the parts pinned there; a new, missing, or +replaced part fails with `ModelChangedError`. Without `racer-model.json`, each +rank pins its own files at the ETags seen when its load starts. + With `clone` set to false, tensors are views into staging memory. They stay valid only until vLLM asks for the next tensors after the buffer's last one. Only use this when the model copies each weight at once. diff --git a/clients/python/racer-vllm/pyproject.toml b/clients/python/racer-vllm/pyproject.toml index 55540c18d5..b12d579a83 100644 --- a/clients/python/racer-vllm/pyproject.toml +++ b/clients/python/racer-vllm/pyproject.toml @@ -1,15 +1,35 @@ [build-system] -requires = ["setuptools>=64"] +requires = ["setuptools>=77"] build-backend = "setuptools.build_meta" [project] name = "racer-vllm" -version = "0.1.0" +version = "0.2.0" description = "vLLM model loader that reads weights from a Racer cache" readme = "README.md" -license = { text = "Apache-2.0" } +license = "Apache-2.0" +license-files = ["LICENSE"] requires-python = ">=3.10" -dependencies = ["racer-client"] +authors = [{ name = "Microsoft Corporation" }] +keywords = ["racer", "unbounded", "kubernetes", "cache", "gpu"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] +dependencies = ["racer-client>=0.2.0,<0.3"] + +[project.urls] +Homepage = "https://github.com/Azure/unbounded" +Documentation = "https://github.com/Azure/unbounded/tree/main/docs/content/guides/racer-vllm.md" +Source = "https://github.com/Azure/unbounded/tree/main/clients/python/racer-vllm" +Issues = "https://github.com/Azure/unbounded/issues" [project.optional-dependencies] test = ["pytest"] @@ -18,10 +38,11 @@ test = ["pytest"] racer = "racer_vllm:register" [project.scripts] -racer-vllm-fetch = "racer_vllm.model:main" +racer-vllm-fetch = "racer_vllm:main" -[tool.setuptools.packages.find] -where = ["src"] +[tool.setuptools] +py-modules = ["racer_vllm"] +package-dir = {"" = "src"} [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/clients/python/racer-vllm/src/racer_vllm.py b/clients/python/racer-vllm/src/racer_vllm.py new file mode 100644 index 0000000000..be2ddcc398 --- /dev/null +++ b/clients/python/racer-vllm/src/racer_vllm.py @@ -0,0 +1,773 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 +"""vLLM plugin that loads model weights from a Racer cache. + +Adds the load format "racer" (vLLM v0.25.x) and the racer-vllm-fetch +command. See README.md for the options. +""" + +from __future__ import annotations + +import collections +import json +import logging +import math +import os +import string +import struct +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import racer + +# Under "vllm." so messages use vLLM's log handler and format. +log = logging.getLogger("vllm.racer_vllm.loader") + +LOAD_FORMAT = "racer" +SCHEME = "racer://" +MARKER = "racer-model.json" +INDEX = "model.safetensors.index.json" +SINGLE = "model.safetensors" +DEFAULT_STAGING_BYTES = 256 * 1024 * 1024 +DEFAULT_CONCURRENCY = 4 +MAX_HEADER_BYTES = 100 * 1024 * 1024 +MAX_PARTS = 100_000 +TP_ONLY_PATTERN = "model-rank-{rank}-part-{part}.safetensors" +PP_AND_TP_PATTERN = "model-pp-{pp_rank}-tp-{tp_rank}-part-{part}.safetensors" + +# Small files vLLM and transformers read to build the model and tokenizer. +CONFIG_FILES = ( + "config.json", + "generation_config.json", + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "tokenizer.model", + "vocab.json", + "vocab.txt", + "merges.txt", + "added_tokens.json", + "chat_template.jinja", + "chat_template.json", + "preprocessor_config.json", + "processor_config.json", + "params.json", + "tekken.json", + INDEX, +) +SMALL_SUFFIXES = (".json", ".txt", ".model", ".jinja", ".py", ".tiktoken") + +# safetensors dtype: (bytes per element, torch dtype name). +DTYPES = { + "BOOL": (1, "bool"), + "U8": (1, "uint8"), + "I8": (1, "int8"), + "F8_E4M3": (1, "float8_e4m3fn"), + "F8_E5M2": (1, "float8_e5m2"), + "F8_E8M0": (1, "float8_e8m0fnu"), + "I16": (2, "int16"), + "U16": (2, "uint16"), + "F16": (2, "float16"), + "BF16": (2, "bfloat16"), + "I32": (4, "int32"), + "U32": (4, "uint32"), + "F32": (4, "float32"), + "I64": (8, "int64"), + "U64": (8, "uint64"), + "F64": (8, "float64"), +} + +_registered = False + + +class ModelChangedError(RuntimeError): + """A model file no longer has the version the load started with.""" + + +@dataclass(frozen=True) +class ModelURI: + """A parsed racer:////.""" + + cache: str + bucket: str + prefix: str # no leading or trailing slash; may be empty + + def key(self, name: str) -> str: + return f"{self.prefix}/{name}" if self.prefix else name + + def request(self, name: str) -> racer.Request: + return racer.object_request(self.bucket, self.key(name)) + + def __str__(self) -> str: + return f"{SCHEME}{self.cache}/{self.bucket}/{self.prefix}".rstrip("/") + + +def parse_uri(uri: str) -> ModelURI: + """Parse racer:////.""" + if not uri.startswith(SCHEME): + raise ValueError(f"not a racer:// URI: {uri!r}") + parts = uri[len(SCHEME) :].split("/", 2) + if len(parts) < 2 or not parts[0] or not parts[1]: + raise ValueError(f"racer URI needs a cache and a bucket: {uri!r}") + prefix = parts[2].strip("/") if len(parts) == 3 else "" + return ModelURI(parts[0], parts[1], prefix) + + +def _unsafe(name: str) -> bool: + return name.startswith("/") or ".." in name.split("/") + + +def _changed(uri: ModelURI, name: str, etag: str) -> ModelChangedError: + return ModelChangedError( + f"racer: {uri.key(name)} in {uri} changed during the load (expected ETag {etag}); " + "restart the load, and rerun racer-vllm-fetch if the model was replaced" + ) + + +def _unpinned(uri: ModelURI, name: str) -> ModelChangedError: + return ModelChangedError( + f"racer: {uri.key(name)} in {uri} is not pinned in {MARKER}; the checkpoint layout changed " + "since racer-vllm-fetch ran, so rerun racer-vllm-fetch" + ) + + +def check_pinned(uri: ModelURI, files: list, pins: dict) -> None: + """With pins, every weight file must have one, so the load uses the layout fetch saw.""" + for name in files: + if pins and name not in pins: + raise _unpinned(uri, name) + + +def stat_pinned(client, uri: ModelURI, name: str, etag: str | None = None): + """Stat a file, at version etag when given.""" + try: + return client.stat(uri.request(name), if_match=etag) + except (racer.VersionUnavailable, racer.NotFound): + if etag is None: + raise + raise _changed(uri, name, etag) from None + + +def read_pinned(client, uri: ModelURI, name: str, etag: str | None = None): + """Read a whole small file as (bytes, etag), or None when it does not exist. + + The bytes are read at the ETag of the stat, so they match. With etag + set, any other version raises ModelChangedError. + """ + try: + meta = stat_pinned(client, uri, name, etag) + except racer.NotFound: + return None + try: + return client.read(uri.request(name), if_match=meta.etag), meta.etag + except (racer.VersionUnavailable, racer.NotFound): + raise _changed(uri, name, meta.etag) from None + + +def weight_files(client, uri: ModelURI, pins: dict | None = None) -> list: + """Safetensors files named by the index, else model.safetensors.""" + pins = pins or {} + got = read_pinned(client, uri, INDEX, pins.get(INDEX)) + if got is not None and pins and INDEX not in pins: + raise _unpinned(uri, INDEX) # fetch saw no index; one was uploaded since + if got is None: + stat_pinned(client, uri, SINGLE, pins.get(SINGLE)) # raises NotFound when absent + return [SINGLE] + files = sorted(set((json.loads(got[0]).get("weight_map") or {}).values())) + if not files: + raise ValueError(f"{INDEX} has an empty weight_map") + for f in files: + if _unsafe(f): + raise ValueError(f"unsafe file name in {INDEX}: {f!r}") + return files + + +def fetch(uri, dest: str, client=None, pattern=None) -> list: + """Copy the small model files to dest and write racer-model.json. + + racer-model.json records the URI and the ETag of every fetched file and + every weight file, so the loader reads the same model version. Without + safetensors weights, it records the files of a sharded_state checkpoint + (pattern, else vLLM's default patterns) for every rank. + """ + u = parse_uri(uri) if isinstance(uri, str) else uri + client = client or racer.Client(cache=u.cache) + os.makedirs(dest, exist_ok=True) + names = list(CONFIG_FILES) + fetched = [] + etags = {} + i = 0 + while i < len(names): + name = names[i] + i += 1 + if name in etags: + continue + got = read_pinned(client, u, name) + if got is None: + if name == "config.json": + raise FileNotFoundError(f"{u.key(name)} not found in {u}") + continue + data, etags[name] = got + _write(os.path.join(dest, name), data) + fetched.append(name) + if name == INDEX: + # Small files named in the index metadata, such as a chat template. + for v in (json.loads(data).get("metadata") or {}).values(): + if isinstance(v, str) and v.endswith(SMALL_SUFFIXES) and not v.startswith("/") and ".." not in v: + names.append(v) + try: + weights = weight_files(client, u, etags) + except racer.NotFound: + weights = sharded_files(client, u, pattern) + for name in weights: + etags[name] = stat_pinned(client, u, name).etag + _write(os.path.join(dest, MARKER), json.dumps({"uri": str(u), "etags": etags}, sort_keys=True).encode()) + return fetched + + +def _write(path: str, data: bytes) -> None: + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path + ".tmp", "wb") as f: + f.write(data) + os.replace(path + ".tmp", path) + + +def _read_marker(model_path: str): + path = os.path.join(model_path, MARKER) + if not os.path.exists(path): + return None + with open(path) as f: + return json.load(f) + + +def resolve_uri(model_path: str, extra: dict) -> ModelURI: + """The model URI from the extra config, the model name, or racer-model.json.""" + if extra.get("uri"): + return parse_uri(extra["uri"]) + if model_path.startswith(SCHEME): + return parse_uri(model_path) + doc = _read_marker(model_path) + if doc is not None: + return parse_uri(doc["uri"]) + raise ValueError( + "racer loader: no model URI; set model_loader_extra_config.uri or run racer-vllm-fetch into the model directory" + ) + + +def resolve_pins(model_path: str, uri: ModelURI) -> dict: + """ETags that racer-vllm-fetch recorded for uri, or {}.""" + doc = _read_marker(model_path) + if not doc or doc.get("uri") is None or parse_uri(doc["uri"]) != uri: + return {} + pins = doc.get("etags") or {} + if not isinstance(pins, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in pins.items()): + raise ValueError(f"bad etags in {os.path.join(model_path, MARKER)}") + return pins + + +# safetensors: an 8-byte little-endian header length N, N bytes of JSON, then +# the data. Each entry has dtype, shape and data_offsets relative to the data. + + +@dataclass(frozen=True) +class Tensor: + name: str + dtype: str + shape: tuple + start: int # absolute file offset + end: int + + +def header_length(prefix: bytes) -> int: + if len(prefix) != 8: + raise ValueError("short safetensors prefix") + (n,) = struct.unpack(" MAX_HEADER_BYTES: + raise ValueError(f"bad safetensors header length {n}") + return n + + +def _uints(v) -> bool: + # JSON integers only; bool is an int subclass, so check the exact type. + return isinstance(v, list) and all(type(x) is int and x >= 0 for x in v) + + +def parse_header(raw: bytes, file_size: int | None = None) -> list: + """Parse the header JSON (after the 8-byte length) into Tensors sorted by offset.""" + try: + doc = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, ValueError) as exc: + raise ValueError("bad safetensors header JSON") from exc + if not isinstance(doc, dict): + raise ValueError("safetensors header is not an object") + base = 8 + len(raw) + doc.pop("__metadata__", None) + tensors = [] + for name, spec in doc.items(): + if not isinstance(spec, dict) or spec.get("dtype") not in DTYPES: + raise ValueError(f"bad entry for {name!r}") + shape, offsets = spec.get("shape"), spec.get("data_offsets") + if not _uints(shape) or not _uints(offsets) or len(offsets) != 2 or offsets[1] < offsets[0]: + raise ValueError(f"bad shape or offsets for {name!r}") + if offsets[1] - offsets[0] != math.prod(shape) * DTYPES[spec["dtype"]][0]: + raise ValueError(f"size mismatch for {name!r}") + tensors.append(Tensor(name, spec["dtype"], tuple(shape), base + offsets[0], base + offsets[1])) + tensors.sort(key=lambda t: (t.start, t.end)) + prev = base + for t in tensors: + if t.start < prev: + raise ValueError(f"overlapping tensor {t.name!r}") + prev = t.end + if file_size is not None and prev > file_size: + raise ValueError("tensor data past end of file") + return tensors + + +@dataclass +class Chunk: + """A file range read with one request. + + tensors lie wholly inside it. parts are (tensor, offset, length) pieces of + tensors split across chunks. A dedicated chunk holds one tensor bigger + than a staging buffer. + """ + + file: str + etag: str + start: int + end: int + tensors: list + parts: list + dedicated: bool = False + + +def plan_chunks(file: str, etag: str, tensors: list, staging_bytes: int, page_bytes: int = 0) -> list: + """Group sorted tensors into chunks of at most staging_bytes. + + With page_bytes no larger than staging_bytes, a chunk ends on a page + boundary, and a tensor across the boundary is split into parts, so no + page is read twice. + """ + paged = 0 < page_bytes <= staging_bytes + chunks = [] + cur = None + limit = 0 + for t in tensors: + size = t.end - t.start + if size > staging_bytes: + if cur is not None: + chunks.append(cur) + cur = None + chunks.append(Chunk(file, etag, t.start, t.end, [t], [], dedicated=True)) + continue + if cur is not None and t.end > limit and (t.start >= limit or not paged): + chunks.append(cur) + cur = None + if cur is None: + cur = Chunk(file, etag, t.start, t.start, [], []) + limit = _limit(t.start, staging_bytes, page_bytes if paged else 0) + if t.end <= limit: + cur.tensors.append(t) + cur.end = max(cur.end, t.end) + continue + pos = t.start + while t.end > limit: + cur.parts.append((t, pos - t.start, limit - pos)) + cur.end = pos = limit + chunks.append(cur) + cur = Chunk(file, etag, pos, pos, [], []) + limit = _limit(pos, staging_bytes, page_bytes) + cur.parts.append((t, pos - t.start, t.end - pos)) + cur.end = t.end + if cur is not None: + chunks.append(cur) + return chunks + + +def _limit(start: int, staging_bytes: int, page_bytes: int) -> int: + end = start + staging_bytes + return end - end % page_bytes if page_bytes else end + + +def read_chunks(client, uri: ModelURI, files: list, staging_bytes: int) -> list: + """Read every file's header and plan its chunks. files holds (name, Metadata).""" + chunks = [] + for name, meta in files: + req = uri.request(name) + try: + n = header_length(client.read(req, 0, 8, if_match=meta.etag)) + tensors = parse_header(client.read(req, 8, n, if_match=meta.etag), meta.size) + except (racer.VersionUnavailable, racer.NotFound): + raise _changed(uri, name, meta.etag) from None + chunks += plan_chunks(name, meta.etag, tensors, staging_bytes, racer.PAGE_SIZE) + return chunks + + +def iter_weights( + client, uri: ModelURI, chunks: list, staging_bytes: int, concurrency: int, clone=True, use_vmm=True, stats=None +): + """Yield (name, CUDA tensor) for every tensor in chunks. + + Up to concurrency chunks are read ahead, each into its own staging + buffer. A buffer is reused only after the consumer has moved past its + tensors and the GPU work queued on it has finished (a CUDA event). Every + buffer is fenced before it is dropped, also on early close or error. + stats, a dict, gets counts and bytes of reads per delivery mode. + """ + import torch + import racer.cuda as rcuda + import racer.tensors as rtensors + + if concurrency < 1: + raise ValueError("concurrency must be at least 1") + device = torch.cuda.current_device() + cuda_device = f"cuda:{device}" + + def new_buffer(size): + # Exportable cuMemCreate memory allows gpu-read; torch memory needs host-read. + nonlocal use_vmm + if use_vmm: + try: + alloc = rcuda.GpuAllocation(size, device) + return {"target": alloc, "flat": rtensors.as_tensor(alloc), "event": None} + except racer.RacerError as exc: + log.info("racer: exportable GPU memory unavailable (%s); using torch memory", exc) + use_vmm = False + flat = torch.empty(size, dtype=torch.uint8, device=cuda_device) + return {"target": rtensors.TensorTarget(flat), "flat": flat, "event": None} + + def done(buf): + # Record on the target device's stream; this thread's device may differ. + with torch.cuda.device(device): + buf["event"] = torch.cuda.Event() + buf["event"].record(torch.cuda.current_stream(device)) + + def wait(buf): + if buf["event"] is not None: + buf["event"].synchronize() + buf["event"] = None + + def read(chunk, buf): + wait(buf) + if chunk.end == chunk.start: + return None + try: + return client.read_into_gpu( + uri.request(chunk.file), + buf["target"], + 0, + offset=chunk.start, + length=chunk.end - chunk.start, + if_match=chunk.etag, + ) + except (racer.VersionUnavailable, racer.NotFound): + raise _changed(uri, chunk.file, chunk.etag) from None + + def tensor(t, raw): + dtype = getattr(torch, DTYPES[t.dtype][1]) + return raw.view(dtype).view(t.shape) + + free = [new_buffer(staging_bytes) for _ in range(concurrency)] + todo = collections.deque(chunks) + inflight = collections.deque() + joined = {} + active = None + pool = ThreadPoolExecutor(max_workers=concurrency, thread_name_prefix="racer-read") + try: + while todo or inflight: + while todo and len(inflight) < concurrency: + chunk = todo.popleft() + buf = new_buffer(chunk.end - chunk.start) if chunk.dedicated else free.pop() + inflight.append((pool.submit(read, chunk, buf), chunk, buf)) + fut, chunk, buf = inflight.popleft() + result = fut.result() + if result is not None and stats is not None: + stats["reads"][result.delivery] += 1 + stats["bytes"][result.delivery] += result.length + active = buf + flat = buf["flat"] + for t in chunk.tensors: + size = t.end - t.start + if size == 0: + yield t.name, torch.empty(t.shape, dtype=getattr(torch, DTYPES[t.dtype][1]), device=cuda_device) + continue + rel = t.start - chunk.start + raw = flat[rel : rel + size] + # Realign odd offsets; copy out of staging unless clone is off. + if rel % DTYPES[t.dtype][0] or (clone and not chunk.dedicated): + raw = raw.clone() + yield t.name, tensor(t, raw) + for t, offset, length in chunk.parts: + if offset == 0: + joined[t.name] = torch.empty(t.end - t.start, dtype=torch.uint8, device=cuda_device) + rel = t.start + offset - chunk.start + joined[t.name][offset : offset + length].copy_(flat[rel : rel + length]) + if offset + length == t.end - t.start: + yield t.name, tensor(t, joined.pop(t.name)) + active = None + done(buf) + if chunk.dedicated: + wait(buf) # then drop it; tensors handed out keep their own reference + else: + free.append(buf) + finally: + for fut, _, _ in inflight: + fut.cancel() + pool.shutdown(wait=True) + if active is not None: + done(active) + # Reads have stopped; wait for queued GPU work before buffers are dropped. + for buf in free + [b for _, _, b in inflight] + ([active] if active else []): + wait(buf) + + +def summary(stats: dict, seconds: float, uri) -> str: + """The one-line load summary.""" + total = sum(stats["bytes"].values()) / (1 << 30) + rate = total / seconds if seconds > 0 else 0.0 + modes = " ".join( + f"{m}={stats['reads'][m]} ({stats['bytes'][m] / (1 << 30):.2f} GiB)" + for m in (racer.GPU, racer.HOST, racer.STREAM) + ) + return ( + f"racer: loaded {stats['files']} files, {total:.2f} GiB in {seconds:.2f}s ({rate:.2f} GiB/s) " + f"from {uri}; reads by delivery mode: {modes}" + ) + + +# vLLM sharded_state checkpoints, as written by ShardedStateLoader.save_model: +# each tensor/pipeline parallel rank has its own files, already split. + + +def pattern_fields(pattern) -> set: + """Placeholder names in a pattern, ignoring format specs such as {rank:02d}.""" + return {f for _, f, _, _ in string.Formatter().parse(pattern) if f is not None} + + +def render_pattern(pattern, part: int, tp_rank=0, tp_size=1, pp_rank=0, pp_size=1) -> str: + """File name of one part, with vLLM's pattern checks.""" + if pattern is None: + pattern = PP_AND_TP_PATTERN if pp_size > 1 else TP_ONLY_PATTERN + fields = pattern_fields(pattern) + unsupported = fields - {"pp_rank", "tp_rank", "rank", "part"} + if unsupported: + raise ValueError(f"pattern {pattern!r} uses unsupported placeholders {sorted(unsupported)}") + if "part" not in fields: + raise ValueError(f"pattern {pattern!r} needs a {{part}} placeholder") + if pp_size > 1 and "pp_rank" not in fields: + raise ValueError(f"pattern {pattern!r} does not embed the pipeline parallel rank") + if pp_size == 1 and "pp_rank" in fields: + raise ValueError(f"pattern {pattern!r} requires pipeline parallelism, but the pipeline parallel size is 1") + if tp_size > 1 and not fields & {"tp_rank", "rank"}: + raise ValueError(f"pattern {pattern!r} does not embed the tensor parallel rank") + name = pattern.format(pp_rank=pp_rank, tp_rank=tp_rank, rank=tp_rank, part=part) + if _unsafe(name): + raise ValueError(f"unsafe file name from pattern: {name!r}") + return name + + +def rank_files(client, uri: ModelURI, pattern, ranks: dict) -> list: + """This rank's files: part 0, 1, ... until one is missing (Racer has no listing).""" + files = [] + for part in range(MAX_PARTS): + name = render_pattern(pattern, part, **ranks) + try: + client.stat(uri.request(name)) + except racer.NotFound: + break + files.append(name) + if not files: + first = render_pattern(pattern, 0, **ranks) + raise ValueError(f"racer loader: no sharded_state files for this rank in {uri} (looked for {uri.key(first)})") + return files + + +def sharded_files(client, uri: ModelURI, pattern=None) -> list: + """Every rank's sharded_state files, probing rank 0, 1, ... until one has no part 0.""" + for pat in [pattern] if pattern else [TP_ONLY_PATTERN, PP_AND_TP_PATTERN]: + fields = pattern_fields(pat) + pp = "pp_rank" in fields + tp = bool(fields & {"tp_rank", "rank"}) + files = [] + for pp_rank in range(MAX_PARTS if pp else 1): + n = len(files) + for tp_rank in range(MAX_PARTS if tp else 1): + ranks = {"tp_rank": tp_rank, "tp_size": 2 if tp else 1, "pp_rank": pp_rank, "pp_size": 2 if pp else 1} + try: + files += rank_files(client, uri, pat, ranks) + except ValueError: + break + if len(files) == n: + break + if files: + return files + return [] + + +def vllm_ranks() -> dict: + from vllm.distributed import ( + get_pipeline_model_parallel_rank, + get_pp_group, + get_tensor_model_parallel_rank, + get_tp_group, + ) + + return { + "tp_rank": get_tensor_model_parallel_rank(), + "tp_size": get_tp_group().world_size, + "pp_rank": get_pipeline_model_parallel_rank(), + "pp_size": get_pp_group().world_size, + } + + +def load_state(state_dict: dict, weights) -> None: + """Copy (name, tensor) pairs into the parameters, as vLLM's ShardedStateLoader does. + + A tensor smaller than its parameter (for example LoRA padding) goes into + a narrowed view. Every key must be loaded exactly once. + """ + remaining = dict(state_dict) + for key, tensor in weights: + param = remaining.pop(key, None) + if param is None: + what = "appears twice in" if key in state_dict else "is unexpected in" + raise ValueError(f"racer loader: tensor {key!r} {what} the sharded checkpoint") + if param.dtype != tensor.dtype: + log.warning( + "racer: loading weight %s with dtype %s into parameter with dtype %s", key, tensor.dtype, param.dtype + ) + shape = tuple(param.shape) + if len(tensor.shape) != len(shape): + raise ValueError(f"racer loader: tensor {key!r} has shape {tuple(tensor.shape)}, parameter has {shape}") + data = param.data + for dim, size in enumerate(tensor.shape): + if size < shape[dim]: + data = data.narrow(dim, 0, size) + if tuple(tensor.shape) != shape: + log.warning( + "racer: loading tensor of shape %s into parameter %r of shape %s", tuple(tensor.shape), key, shape + ) + data.copy_(tensor) + if remaining: + raise ValueError(f"racer loader: missing keys {tuple(sorted(remaining))[:10]} in the sharded checkpoint") + + +def settings(extra: dict) -> dict: + """Loader options from model_loader_extra_config, with defaults.""" + retry = extra.get("mode_retry_seconds", racer.DEFAULT_MODE_RETRY_SECONDS) + return { + "uri": extra.get("uri"), + "staging_bytes": int(extra.get("staging_bytes", DEFAULT_STAGING_BYTES)), + "concurrency": int(extra.get("concurrency", DEFAULT_CONCURRENCY)), + "clone": bool(extra.get("clone", True)), + "socket": extra.get("socket"), + "use_vmm": bool(extra.get("use_vmm", True)), + "mode_retry_seconds": None if retry is None else float(retry), + "sharded_state": bool(extra.get("sharded_state", False)), + "pattern": extra.get("pattern"), + } + + +def register() -> None: + """vllm.general_plugins entry point: add --load-format racer. Safe to call twice.""" + global _registered + if _registered: + return + from vllm.model_executor.model_loader import register_model_loader + from vllm.model_executor.model_loader.base_loader import BaseModelLoader + + class RacerModelLoader(BaseModelLoader): + """Loads safetensors weights from a Racer cache.""" + + def __init__(self, load_config): + super().__init__(load_config) + self.opts = settings(dict(getattr(load_config, "model_loader_extra_config", None) or {})) + + def _open(self, model_config): + """Return (uri, client, file names, pins).""" + o = self.opts + uri = resolve_uri(str(model_config.model), o) + pins = resolve_pins(str(model_config.model), uri) + client = racer.Client( + cache=uri.cache, + socket_path=o["socket"], + staging_bytes=o["staging_bytes"], + mode_retry_seconds=o["mode_retry_seconds"], + ) + try: + if o["sharded_state"]: + ranks = vllm_ranks() + files = rank_files(client, uri, o["pattern"], ranks) + # A pinned part past the last one found was deleted at the origin. + after = render_pattern(o["pattern"], len(files), **ranks) + if after in pins: + raise _changed(uri, after, pins[after]) + else: + files = weight_files(client, uri, pins) + check_pinned(uri, files, pins) + except BaseException: + client.close() + raise + return uri, client, files, pins + + def download_model(self, model_config) -> None: + # Nothing to download; check that the weights are there. + _, client, _, _ = self._open(model_config) + client.close() + + def load_weights(self, model, model_config) -> None: + o = self.opts + started = time.monotonic() + uri, client, files, pins = self._open(model_config) + stats = {"files": len(files), "reads": dict.fromkeys((racer.GPU, racer.HOST, racer.STREAM), 0)} + stats["bytes"] = dict(stats["reads"]) + with client: + # Stat every file before any data read, so the model is read at one version. + pinned = [(name, stat_pinned(client, uri, name, pins.get(name))) for name in files] + chunks = read_chunks(client, uri, pinned, o["staging_bytes"]) + weights = iter_weights( + client, uri, chunks, o["staging_bytes"], o["concurrency"], o["clone"], o["use_vmm"], stats + ) + try: + if o["sharded_state"]: + load_state(_filter_subtensors(model.state_dict()), weights) + loaded = None + else: + loaded = model.load_weights(weights) + finally: + # Fence and stop reads before the client closes. + weights.close() + log.info("%s", summary(stats, time.monotonic() - started, uri)) + if loaded is not None and getattr(model_config, "quantization", None) is None: + missing = {n for n, _ in model.named_parameters()} - set(loaded) + if missing: + raise ValueError(f"racer loader: weights not loaded: {sorted(missing)[:10]}") + + register_model_loader(LOAD_FORMAT)(RacerModelLoader) + _registered = True + + +def _filter_subtensors(state_dict: dict) -> dict: + """vLLM's filter for tensors that share memory.""" + from vllm.model_executor.model_loader.sharded_state_loader import ShardedStateLoader + + return ShardedStateLoader._filter_subtensors(state_dict) + + +def main(argv=None) -> int: + """racer-vllm-fetch: copy a model's config and tokenizer files from Racer.""" + import argparse + + p = argparse.ArgumentParser(prog="racer-vllm-fetch", description="Fetch model config files from Racer.") + p.add_argument("uri", help="racer:////") + p.add_argument("dest", help="local directory") + p.add_argument("--pattern", help="sharded_state file name pattern, as the loader's pattern option") + a = p.parse_args(argv) + for name in fetch(a.uri, a.dest, pattern=a.pattern): + print(name) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/clients/python/racer-vllm/src/racer_vllm/__init__.py b/clients/python/racer-vllm/src/racer_vllm/__init__.py deleted file mode 100644 index f1801da289..0000000000 --- a/clients/python/racer-vllm/src/racer_vllm/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 -"""vLLM plugin that loads model weights from a Racer cache.""" - -from .loader import LOAD_FORMAT, register -from .model import ModelURI, fetch, parse_uri - -__all__ = ["LOAD_FORMAT", "ModelURI", "fetch", "parse_uri", "register"] diff --git a/clients/python/racer-vllm/src/racer_vllm/loader.py b/clients/python/racer-vllm/src/racer_vllm/loader.py deleted file mode 100644 index e578482be2..0000000000 --- a/clients/python/racer-vllm/src/racer_vllm/loader.py +++ /dev/null @@ -1,97 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 -"""vLLM model loader for --load-format racer. - -Targets vLLM v0.25.x. The loader interface used is -vllm/model_executor/model_loader/base_loader.py (BaseModelLoader with -download_model(model_config) and load_weights(model, model_config)), and -registration uses register_model_loader from -vllm/model_executor/model_loader/__init__.py. Recent main adds optional hooks -(get_external_weight_memory, create_model) with defaults, so this loader -works there too. -""" - -from __future__ import annotations - -import logging - -from .model import client_for, resolve_uri, weight_files -from .weights import DEFAULT_DEPTH, DEFAULT_STAGING_BYTES, TorchBackend, iter_weights - -log = logging.getLogger(__name__) - -LOAD_FORMAT = "racer" -_registered = False - - -def _settings(extra: dict) -> dict: - return { - "staging_bytes": int(extra.get("staging_bytes", DEFAULT_STAGING_BYTES)), - "concurrency": int(extra.get("concurrency", DEFAULT_DEPTH)), - "clone": bool(extra.get("clone", True)), - "socket": extra.get("socket"), - "use_vmm": bool(extra.get("use_vmm", True)), - } - - -def make_loader_class(base): - class RacerModelLoader(base): - """Loads safetensors weights from a Racer cache. - - model_loader_extra_config keys: - uri racer:////; optional when the - model directory has racer-model.json (racer-vllm-fetch) - staging_bytes size of each GPU staging buffer (default 256 MiB) - concurrency number of staging buffers and reads in flight (default 4) - clone copy each tensor out of staging (default true) - socket client socket path (else RACER_SOCKET, else the cache) - use_vmm allow exportable cuMemCreate staging (default true) - """ - - def __init__(self, load_config): - super().__init__(load_config) - self.extra = dict(getattr(load_config, "model_loader_extra_config", None) or {}) - self.settings = _settings(self.extra) - - def _client(self, model_config): - uri = resolve_uri(str(model_config.model), self.extra) - client = client_for(uri, self.settings["socket"], staging_bytes=self.settings["staging_bytes"]) - return uri, client - - def download_model(self, model_config) -> None: - # Nothing to download; check the weights are reachable. - uri, client = self._client(model_config) - with client: - weight_files(client, uri) - - def load_weights(self, model, model_config) -> None: - uri, client = self._client(model_config) - s = self.settings - with client: - backend = TorchBackend(clone=s["clone"], use_vmm=s["use_vmm"]) - weights = iter_weights(client, uri, backend, s["staging_bytes"], s["concurrency"]) - try: - loaded = model.load_weights(weights) - finally: - # Run the generator's cleanup (executor shutdown and GPU - # fences) before the client closes, even if load_weights - # raised or returned without exhausting the iterator. - weights.close() - if loaded is not None and getattr(model_config, "quantization", None) is None: - missing = {n for n, _ in model.named_parameters()} - set(loaded) - if missing: - raise ValueError(f"racer loader: weights not loaded: {sorted(missing)[:10]}") - - return RacerModelLoader - - -def register() -> None: - """vllm.general_plugins entry point. Safe to call more than once.""" - global _registered - if _registered: - return - from vllm.model_executor.model_loader import register_model_loader - from vllm.model_executor.model_loader.base_loader import BaseModelLoader - - register_model_loader(LOAD_FORMAT)(make_loader_class(BaseModelLoader)) - _registered = True diff --git a/clients/python/racer-vllm/src/racer_vllm/model.py b/clients/python/racer-vllm/src/racer_vllm/model.py deleted file mode 100644 index 8eafc2dfba..0000000000 --- a/clients/python/racer-vllm/src/racer_vllm/model.py +++ /dev/null @@ -1,168 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 -"""racer:// model locations and small-file fetching.""" - -from __future__ import annotations - -import json -import os -from dataclasses import dataclass - -import racer - -SCHEME = "racer://" -MARKER = "racer-model.json" -INDEX = "model.safetensors.index.json" -SINGLE = "model.safetensors" - -# Small files vLLM and transformers read to build the model and tokenizer. -CONFIG_FILES = ( - "config.json", - "generation_config.json", - "tokenizer.json", - "tokenizer_config.json", - "special_tokens_map.json", - "tokenizer.model", - "vocab.json", - "vocab.txt", - "merges.txt", - "added_tokens.json", - "chat_template.jinja", - "chat_template.json", - "preprocessor_config.json", - "processor_config.json", - "params.json", - "tekken.json", - INDEX, -) -REQUIRED_FILES = ("config.json",) -_SMALL_SUFFIXES = (".json", ".txt", ".model", ".jinja", ".py", ".tiktoken") - - -@dataclass(frozen=True) -class ModelURI: - cache: str - bucket: str - prefix: str # no leading or trailing slash; may be empty - - def key(self, name: str) -> str: - return f"{self.prefix}/{name}" if self.prefix else name - - def request(self, name: str) -> racer.Request: - return racer.object_request(self.bucket, self.key(name)) - - def __str__(self) -> str: - return f"{SCHEME}{self.cache}/{self.bucket}/{self.prefix}".rstrip("/") - - -def parse_uri(uri: str) -> ModelURI: - """Parse racer:////.""" - if not uri.startswith(SCHEME): - raise ValueError(f"not a racer:// URI: {uri!r}") - parts = uri[len(SCHEME) :].split("/", 2) - if len(parts) < 2 or not parts[0] or not parts[1]: - raise ValueError(f"racer URI needs a cache and a bucket: {uri!r}") - prefix = parts[2].strip("/") if len(parts) == 3 else "" - return ModelURI(parts[0], parts[1], prefix) - - -def client_for(uri: ModelURI, socket_path: str | None = None, **kw) -> racer.Client: - """Client for the URI's cache. RACER_SOCKET or socket_path override the path.""" - return racer.Client(cache=uri.cache, socket_path=socket_path, **kw) - - -def read_small(client: racer.Client, uri: ModelURI, name: str) -> bytes | None: - """Read a whole small object, or None when it does not exist.""" - try: - return client.read(uri.request(name)) - except racer.NotFound: - return None - - -def index_extra_files(index: dict) -> list: - """Small file names listed as string values in the index metadata.""" - out = [] - for v in (index.get("metadata") or {}).values(): - if isinstance(v, str) and v.endswith(_SMALL_SUFFIXES) and not v.startswith("/") and ".." not in v: - out.append(v) - return out - - -def weight_files(client: racer.Client, uri: ModelURI) -> list: - """List safetensors files from the index, else the single model.safetensors.""" - raw = read_small(client, uri, INDEX) - if raw is None: - client.stat(uri.request(SINGLE)) # raises NotFound with a clear key - return [SINGLE] - index = json.loads(raw) - files = sorted(set((index.get("weight_map") or {}).values())) - if not files: - raise ValueError(f"{INDEX} has an empty weight_map") - for f in files: - if f.startswith("/") or ".." in f.split("/"): - raise ValueError(f"unsafe file name in {INDEX}: {f!r}") - return files - - -def fetch(uri: str | ModelURI, dest: str, client: racer.Client | None = None, extra=()) -> list: - """Copy the small model files to dest and write the racer-model.json marker. - - Point vLLM at dest (--model dest, and --tokenizer dest if needed) with - --load-format racer. The loader reads the marker to find the weights. - """ - u = parse_uri(uri) if isinstance(uri, str) else uri - client = client or client_for(u) - os.makedirs(dest, exist_ok=True) - names = list(CONFIG_FILES) + list(extra) - fetched = [] - seen = set() - while names: - name = names.pop(0) - if name in seen: - continue - seen.add(name) - data = read_small(client, u, name) - if data is None: - if name in REQUIRED_FILES: - raise FileNotFoundError(f"{u.key(name)} not found in {u}") - continue - path = os.path.join(dest, name) - os.makedirs(os.path.dirname(path), exist_ok=True) - tmp = path + ".tmp" - with open(tmp, "wb") as f: - f.write(data) - os.replace(tmp, path) - fetched.append(name) - if name == INDEX: - names.extend(index_extra_files(json.loads(data))) - with open(os.path.join(dest, MARKER), "w") as f: - json.dump({"uri": str(u)}, f) - return fetched - - -def resolve_uri(model_path: str, extra: dict) -> ModelURI: - """Find the racer:// URI from extra config, the model name, or the marker.""" - if extra.get("uri"): - return parse_uri(extra["uri"]) - if model_path.startswith(SCHEME): - return parse_uri(model_path) - marker = os.path.join(model_path, MARKER) - if os.path.exists(marker): - with open(marker) as f: - return parse_uri(json.load(f)["uri"]) - raise ValueError( - "racer loader: no model URI; set model_loader_extra_config.uri or run racer-vllm-fetch into the model directory" - ) - - -def main(argv=None) -> int: - import argparse - - p = argparse.ArgumentParser(prog="racer-vllm-fetch", description="Fetch model config files from Racer.") - p.add_argument("uri", help="racer:////") - p.add_argument("dest", help="local directory") - p.add_argument("--file", action="append", default=[], help="extra small file to fetch") - a = p.parse_args(argv) - for name in fetch(a.uri, a.dest, extra=a.file): - print(name) - return 0 diff --git a/clients/python/racer-vllm/src/racer_vllm/safetensors.py b/clients/python/racer-vllm/src/racer_vllm/safetensors.py deleted file mode 100644 index 2574e35a53..0000000000 --- a/clients/python/racer-vllm/src/racer_vllm/safetensors.py +++ /dev/null @@ -1,205 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 -"""Torch-free safetensors header parsing and read planning. - -Format: an 8-byte little-endian header length N, N bytes of JSON, then the -tensor data. Each JSON entry has dtype, shape and data_offsets [begin, end) -relative to the end of the header. "__metadata__" holds string metadata. -""" - -from __future__ import annotations - -import json -import math -import struct -from dataclasses import dataclass, field - -MAX_HEADER_BYTES = 100 * 1024 * 1024 - -# Bytes per element for each safetensors dtype. -DTYPE_SIZES = { - "BOOL": 1, - "U8": 1, - "I8": 1, - "F8_E4M3": 1, - "F8_E5M2": 1, - "F8_E8M0": 1, - "I16": 2, - "U16": 2, - "F16": 2, - "BF16": 2, - "I32": 4, - "U32": 4, - "F32": 4, - "I64": 8, - "U64": 8, - "F64": 8, -} - - -class SafetensorsError(ValueError): - pass - - -@dataclass(frozen=True) -class TensorInfo: - name: str - dtype: str - shape: tuple - start: int # absolute file offset - end: int - - @property - def nbytes(self) -> int: - return self.end - self.start - - @property - def itemsize(self) -> int: - return DTYPE_SIZES[self.dtype] - - -@dataclass -class Header: - data_start: int - tensors: list - metadata: dict = field(default_factory=dict) - - -def header_length(prefix: bytes) -> int: - if len(prefix) != 8: - raise SafetensorsError("short safetensors prefix") - (n,) = struct.unpack(" MAX_HEADER_BYTES: - raise SafetensorsError(f"bad safetensors header length {n}") - return n - - -def _is_uint_list(v) -> bool: - # JSON integers only: reject floats, numeric strings, and bools (bool is an int subclass). - return isinstance(v, list) and all(type(x) is int and x >= 0 for x in v) - - -def parse_header(raw: bytes, file_size: int | None = None) -> Header: - """Parse header JSON (without the 8-byte prefix) into sorted TensorInfo.""" - try: - doc = json.loads(raw.decode("utf-8").rstrip(" ")) - except (UnicodeDecodeError, ValueError) as exc: - raise SafetensorsError("bad safetensors header JSON") from exc - if not isinstance(doc, dict): - raise SafetensorsError("safetensors header is not an object") - data_start = 8 + len(raw) - metadata = doc.pop("__metadata__", None) or {} - tensors = [] - for name, spec in doc.items(): - try: - dtype = spec["dtype"] - raw_shape = spec["shape"] - raw_offsets = spec["data_offsets"] - except (KeyError, TypeError) as exc: - raise SafetensorsError(f"bad entry for {name!r}") from exc - if not _is_uint_list(raw_shape) or not _is_uint_list(raw_offsets) or len(raw_offsets) != 2: - raise SafetensorsError(f"bad shape or offsets for {name!r}") - shape = tuple(raw_shape) - begin, end = raw_offsets - if dtype not in DTYPE_SIZES: - raise SafetensorsError(f"unsupported dtype {dtype} for {name!r}") - if any(d < 0 for d in shape) or begin < 0 or end < begin: - raise SafetensorsError(f"bad shape or offsets for {name!r}") - if end - begin != math.prod(shape) * DTYPE_SIZES[dtype]: - raise SafetensorsError(f"size mismatch for {name!r}") - tensors.append(TensorInfo(name, dtype, shape, data_start + begin, data_start + end)) - tensors.sort(key=lambda t: (t.start, t.end)) - prev = data_start - for t in tensors: - if t.start < prev: - raise SafetensorsError(f"overlapping tensor {t.name!r}") - prev = t.end - if file_size is not None and prev > file_size: - raise SafetensorsError("tensor data past end of file") - return Header(data_start, tensors, metadata) - - -@dataclass(frozen=True) -class Part: - """Bytes [offset, offset + length) of a tensor that spans several chunks.""" - - tensor: TensorInfo - offset: int - length: int - - @property - def last(self) -> bool: - return self.offset + self.length == self.tensor.nbytes - - -@dataclass -class Chunk: - """A contiguous file range read with one request.""" - - start: int - end: int - tensors: list # tensors that lie wholly inside the chunk - dedicated: bool = False # larger than the staging buffer - parts: list = field(default_factory=list) # Parts of tensors split across chunks - - @property - def length(self) -> int: - return self.end - self.start - - -def plan_chunks(tensors: list, staging_bytes: int, page_bytes: int = 0) -> list: - """Group sorted tensors into ranges of at most staging_bytes. - - A tensor bigger than staging_bytes gets its own dedicated chunk. Zero-size - tensors join the current chunk and need no bytes. - - With page_bytes no larger than staging_bytes, a chunk that would end in - the middle of a page ends on the page boundary instead, and the tensor that - crosses the boundary is split into Parts read by consecutive chunks. Without - this, both chunks read the boundary page. - """ - paged = 0 < page_bytes <= staging_bytes - chunks: list = [] - cur: Chunk | None = None - limit = 0 - - def open_chunk(start: int) -> Chunk: - nonlocal limit - limit = start + staging_bytes - if paged: - limit -= limit % page_bytes - return Chunk(start, start, []) - - for t in tensors: - if t.nbytes > staging_bytes: - if cur is not None: - chunks.append(cur) - cur = None - chunks.append(Chunk(t.start, t.end, [t], dedicated=True)) - continue - if cur is not None and t.end > limit and (t.start >= limit or not paged): - chunks.append(cur) - cur = None - if cur is None: - cur = open_chunk(t.start) - if t.end <= limit: - cur.tensors.append(t) - cur.end = max(cur.end, t.end) - continue - pos = t.start - while t.end > limit: - cur.parts.append(Part(t, pos - t.start, limit - pos)) - cur.end = pos = limit - chunks.append(cur) - cur = open_chunk(pos) - cur.parts.append(Part(t, pos - t.start, t.end - pos)) - cur.end = t.end - if cur is not None: - chunks.append(cur) - return chunks - - -def tensor_bytes(buf, chunk: Chunk, t: TensorInfo) -> memoryview: - """Return the bytes of t inside a buffer holding chunk (torch-free).""" - rel = t.start - chunk.start - return memoryview(buf)[rel : rel + t.nbytes] diff --git a/clients/python/racer-vllm/src/racer_vllm/weights.py b/clients/python/racer-vllm/src/racer_vllm/weights.py deleted file mode 100644 index 5637aeec18..0000000000 --- a/clients/python/racer-vllm/src/racer_vllm/weights.py +++ /dev/null @@ -1,293 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 -"""Stream safetensors weights from Racer into GPU staging buffers. - -The pipeline is torch-free; a Backend supplies memory and tensor views. -""" - -from __future__ import annotations - -import collections -import logging -from concurrent.futures import ThreadPoolExecutor -from typing import Iterator - -import racer - -from .model import ModelURI, weight_files -from .safetensors import Chunk, TensorInfo, header_length, parse_header, plan_chunks - -log = logging.getLogger(__name__) - -DEFAULT_STAGING_BYTES = 256 * 1024 * 1024 -DEFAULT_DEPTH = 4 - - -class Backend: - """Memory and tensor views for the pipeline.""" - - def new_staging(self, size: int): - raise NotImplementedError - - def new_dedicated(self, size: int): - raise NotImplementedError - - def tensor(self, target, rel: int, info: TensorInfo, staged: bool): - raise NotImplementedError - - def new_joined(self, info: TensorInfo): - """Memory for a tensor that is split across chunks.""" - raise NotImplementedError - - def copy_part(self, joined, offset: int, target, rel: int, length: int) -> None: - """Copy length bytes at rel in target to offset in joined.""" - raise NotImplementedError - - def joined_tensor(self, joined, info: TensorInfo): - """Return the tensor once every part has been copied into joined.""" - raise NotImplementedError - - def done(self, target) -> None: - """The consumer is finished with every tensor from target.""" - - def wait(self, target) -> None: - """Block until work queued by the consumer on target has finished.""" - - -def _file_chunks(client: racer.Client, uri: ModelURI, staging_bytes: int, page_bytes: int): - for name in weight_files(client, uri): - req = uri.request(name) - meta = client.stat(req) - prefix = client.read(req, 0, 8, if_match=meta.etag) - n = header_length(prefix) - header = parse_header(client.read(req, 8, n, if_match=meta.etag), meta.size) - log.debug("racer: %s has %d tensors", name, len(header.tensors)) - for chunk in plan_chunks(header.tensors, staging_bytes, page_bytes): - yield req, meta.etag, chunk - - -def iter_weights( - client: racer.Client, - uri: ModelURI, - backend: Backend, - staging_bytes: int = DEFAULT_STAGING_BYTES, - depth: int = DEFAULT_DEPTH, - page_bytes: int = racer.PAGE_SIZE, -) -> Iterator[tuple]: - """Yield (name, tensor) for every tensor in the model. - - Up to ``depth`` chunks are read ahead, each into its own staging buffer, - so reads overlap with the consumer. A staging buffer is reused only after - the consumer has moved past all of its tensors and backend.wait() says - queued GPU work that reads it has finished. Every staging and dedicated - target is fenced with done() and wait() before the pipeline drops its - reference, including on early close or error, so external allocations - are never released while queued GPU work may still read them. Dedicated - targets are fenced and dropped as soon as the consumer resumes past - them, so oversized tensors do not accumulate in GPU memory. - - Chunks end on Racer page boundaries (page_bytes) where they can, so - neighboring chunks do not both read the boundary page. A tensor that - crosses a boundary is copied out of each chunk into memory from - backend.new_joined() and yielded with the chunk that holds its last byte. - """ - if depth < 1: - raise ValueError("depth must be at least 1") - free = [backend.new_staging(staging_bytes) for _ in range(depth)] - chunks = _file_chunks(client, uri, staging_bytes, page_bytes) - joined: dict = {} - inflight: collections.deque = collections.deque() - pool = ThreadPoolExecutor(max_workers=depth, thread_name_prefix="racer-read") - - def read(req, etag, chunk: Chunk, target): - backend.wait(target) - if chunk.length: - client.read_into_gpu(req, target, 0, offset=chunk.start, length=chunk.length, if_match=etag) - return target - - def fill() -> bool: - while len(inflight) < depth: - nxt = next(chunks, None) - if nxt is None: - return bool(inflight) - req, etag, chunk = nxt - if chunk.dedicated: - target = backend.new_dedicated(chunk.length) - else: - target = free.pop() - inflight.append((pool.submit(read, req, etag, chunk, target), chunk)) - return True - - active = None - try: - while fill(): - fut, chunk = inflight.popleft() - target = fut.result() - active = target - for t in chunk.tensors: - yield t.name, backend.tensor(target, t.start - chunk.start, t, not chunk.dedicated) - for part in chunk.parts: - t = part.tensor - if part.offset == 0: - joined[t.name] = backend.new_joined(t) - rel = t.start + part.offset - chunk.start - backend.copy_part(joined[t.name], part.offset, target, rel, part.length) - if part.last: - yield t.name, backend.joined_tensor(joined.pop(t.name), t) - active = None - backend.done(target) - if chunk.dedicated: - # Fence and drop now so oversized targets do not accumulate; - # tensors the consumer kept hold their own storage reference. - backend.wait(target) - del target - else: - free.append(target) - finally: - for fut, _ in inflight: - fut.cancel() - pool.shutdown(wait=True) - if active is not None: - backend.done(active) - backend.wait(active) - # Reads have stopped; fence consumer GPU work before release. Targets - # still in flight were fenced by read() and the consumer never saw them. - for target in free: - backend.wait(target) - - -class MemoryBackend(Backend): - """Host memory backend: plain buffers, bytes views. Used for tests and CPU.""" - - class Target: - def __init__(self, size: int) -> None: - self.size = size - self.buf = bytearray(size) - - def copy_from_host(self, offset: int, data) -> None: - mv = memoryview(data).cast("B") - self.buf[offset : offset + len(mv)] = mv - - def new_staging(self, size: int): - return self.Target(size) - - def new_dedicated(self, size: int): - return self.Target(size) - - def tensor(self, target, rel, info, staged): - return bytes(target.buf[rel : rel + info.nbytes]) - - def new_joined(self, info): - return bytearray(info.nbytes) - - def copy_part(self, joined, offset, target, rel, length): - joined[offset : offset + length] = target.buf[rel : rel + length] - - def joined_tensor(self, joined, info): - return bytes(joined) - - -TORCH_DTYPES = { - "BOOL": "bool", - "U8": "uint8", - "I8": "int8", - "F8_E4M3": "float8_e4m3fn", - "F8_E5M2": "float8_e5m2", - "F8_E8M0": "float8_e8m0fnu", - "I16": "int16", - "U16": "uint16", - "F16": "float16", - "BF16": "bfloat16", - "I32": "int32", - "U32": "uint32", - "F32": "float32", - "I64": "int64", - "U64": "uint64", - "F64": "float64", -} - - -class TorchBackend(Backend): - """CUDA backend. - - Staging buffers are exportable cuMemCreate allocations when the CUDA - driver allows it (so gpu-read works), else plain torch tensors (host or - stream mode plus a copy). Tensors from staging buffers are cloned unless - clone=False; then they are views that are valid only until the consumer - asks for the next chunk's tensors. - """ - - def __init__(self, device: int | None = None, clone: bool = True, use_vmm: bool = True) -> None: - import torch - - self.torch = torch - self.device = torch.cuda.current_device() if device is None else device - self.clone = clone - self.use_vmm = use_vmm - self._flat: dict = {} - self._events: dict = {} - - def _new(self, size: int): - import racer.cuda as cuda - import racer.tensors as tensors - - if self.use_vmm: - try: - alloc = cuda.GpuAllocation(size, self.device) - self._flat[id(alloc)] = tensors.as_tensor(alloc) - return alloc - except racer.RacerError as exc: - log.info("racer: exportable GPU memory unavailable (%s); using torch memory", exc) - self.use_vmm = False - t = self.torch.empty(size, dtype=self.torch.uint8, device=f"cuda:{self.device}") - target = tensors.TensorTarget(t) - self._flat[id(target)] = t - return target - - def new_staging(self, size): - return self._new(size) - - def new_dedicated(self, size): - return self._new(size) - - def tensor(self, target, rel, info, staged): - torch = self.torch - dtype = getattr(torch, TORCH_DTYPES[info.dtype]) - if info.nbytes == 0: - return torch.empty(info.shape, dtype=dtype, device=f"cuda:{self.device}") - flat = self._flat[id(target)] - if not staged: - # Dedicated memory: the tensor keeps it alive; drop our reference. - self._flat.pop(id(target), None) - raw = flat[rel : rel + info.nbytes] - if rel % info.itemsize: - raw = raw.clone() # realign; the clone owns its memory - elif staged and self.clone: - raw = raw.clone() - return raw.view(dtype).view(info.shape) - - def new_joined(self, info): - return self.torch.empty(info.nbytes, dtype=self.torch.uint8, device=f"cuda:{self.device}") - - def copy_part(self, joined, offset, target, rel, length): - # Queued before done(target), so the staging buffer is not reused - # until the copy has finished. - joined[offset : offset + length].copy_(self._flat[id(target)][rel : rel + length]) - - def joined_tensor(self, joined, info): - dtype = getattr(self.torch, TORCH_DTYPES[info.dtype]) - return joined.view(dtype).view(info.shape) - - def done(self, target): - # Record on the target device's stream; the calling thread's current - # device may differ, and its stream would not cover our copies. - cuda = self.torch.cuda - with cuda.device(self.device): - ev = cuda.Event() - ev.record(cuda.current_stream(self.device)) - self._events[id(target)] = ev - - def wait(self, target): - ev = self._events.pop(id(target), None) - if ev is not None: - ev.synchronize() diff --git a/clients/python/racer-vllm/tests/conftest.py b/clients/python/racer-vllm/tests/conftest.py index ea8395a24c..eef26bb1ca 100644 --- a/clients/python/racer-vllm/tests/conftest.py +++ b/clients/python/racer-vllm/tests/conftest.py @@ -1,15 +1,18 @@ # Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 +import importlib.util import json +import math import os import struct import sys +import types + +import pytest HERE = os.path.dirname(__file__) sys.path.insert(0, os.path.join(HERE, "..", "..", "racer", "tests")) -import importlib.util # noqa: E402 - # Reuse the racer client's fake server fixtures without a module name clash. _spec = importlib.util.spec_from_file_location( "racer_test_fixtures", os.path.join(HERE, "..", "..", "racer", "tests", "conftest.py") @@ -31,3 +34,103 @@ def make_safetensors(tensors, metadata=None, pad=0): header["__metadata__"] = metadata h = json.dumps(header).encode() + b" " * pad return struct.pack("= 2 - - -class _Param: - pass -class _Param: - pass - - -def install_fake_vllm(monkeypatch): - registry = {} - vllm = types.ModuleType("vllm") - me = types.ModuleType("vllm.model_executor") - ml = types.ModuleType("vllm.model_executor.model_loader") - bl = types.ModuleType("vllm.model_executor.model_loader.base_loader") - - class BaseModelLoader: - def __init__(self, load_config): - self.load_config = load_config - - def register_model_loader(name): - def wrap(cls): - assert issubclass(cls, BaseModelLoader) - registry[name] = cls - return cls - - return wrap - - ml.register_model_loader = register_model_loader - bl.BaseModelLoader = BaseModelLoader - for name, mod in { - "vllm": vllm, - "vllm.model_executor": me, - "vllm.model_executor.model_loader": ml, - "vllm.model_executor.model_loader.base_loader": bl, - }.items(): - monkeypatch.setitem(sys.modules, name, mod) - monkeypatch.setattr(loader, "_registered", False) - return registry - - -def test_register_is_idempotent(monkeypatch): - install_fake_vllm(monkeypatch) - ml = sys.modules["vllm.model_executor.model_loader"] - inner = ml.register_model_loader - calls = [] - - def counting(name): - calls.append(name) - return inner(name) - - ml.register_model_loader = counting - loader.register() - loader.register() - assert calls == ["racer"] - assert loader._registered is True - - -def test_register_and_load(server, uri, monkeypatch, tmp_path): - registry = install_fake_vllm(monkeypatch) - loader.register() - loader.register() - cls = registry["racer"] - setup_model(server, uri) - model.fetch(uri, str(tmp_path / "m"), client=racer.Client(socket_path=server.path)) - monkeypatch.setenv("RACER_SOCKET", server.path) - cfg = types.SimpleNamespace(model_loader_extra_config={"staging_bytes": 400, "concurrency": 2}) - inst = cls(cfg) - mc = types.SimpleNamespace(model=str(tmp_path / "m"), quantization=None) - inst.download_model(mc) - - # Swap the torch backend for the memory one, as there is no GPU here. - monkeypatch.setattr(loader, "TorchBackend", lambda **kw: MemoryBackend()) - - class Model: - def __init__(self): - self.got = {} - - def load_weights(self, weights): - for name, t in weights: - self.got[name] = t - return set(self.got) - - def named_parameters(self): - return [(n, _Param()) for n in ("w1", "w2", "w3", "w4")] - - m = Model() - inst.load_weights(m, mc) - assert m.got["w3"] == b"\3" * 900 - - class Missing(Model): - def named_parameters(self): - return super().named_parameters() + [("w5", _Param())] - - with pytest.raises(ValueError, match="w5"): - inst.load_weights(Missing(), mc) - - -def test_entry_point_declared(): - here = os.path.dirname(__file__) - with open(os.path.join(here, "..", "pyproject.toml"), encoding="utf-8") as f: - text = f.read() - assert '[project.entry-points."vllm.general_plugins"]' in text - assert 'racer = "racer_vllm:register"' in text - - -def test_torch_backend_records_event_on_target_device_stream(monkeypatch): - calls = [] - - class Event: - def record(self, stream=None): - calls.append(("record", stream)) - - def synchronize(self): - calls.append(("sync",)) - - class DeviceCtx: - def __init__(self, dev): - self.dev = dev - - def __enter__(self): - calls.append(("enter", self.dev)) - - def __exit__(self, *a): - calls.append(("exit", self.dev)) - - cuda = types.SimpleNamespace( - current_device=lambda: 0, - device=DeviceCtx, - Event=lambda: (calls.append(("event",)), Event())[1], - current_stream=lambda dev=None: ("stream", dev), - ) - monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(cuda=cuda)) - from racer_vllm.weights import TorchBackend - - b = TorchBackend(device=3) - target = object() - b.done(target) - b.wait(target) - assert calls == [("enter", 3), ("event",), ("record", ("stream", 3)), ("exit", 3), ("sync",)] - - -class _FencingBackend(MemoryBackend): - """Records done/wait so tests can check every target is fenced.""" - - def __init__(self): - self.targets = [] - self.done_ids = set() - self.waited = set() - - def _track(self, t): - self.targets.append(t) - return t - - def new_staging(self, size): - return self._track(self.Target(size)) - - def new_dedicated(self, size): - return self._track(self.Target(size)) - - def done(self, target): - self.done_ids.add(id(target)) - self.waited.discard(id(target)) - - def wait(self, target): - if id(target) in self.done_ids: - self.waited.add(id(target)) - - def used(self): - return [t for t in self.targets if id(t) in self.done_ids] - - def assert_fenced(self): - assert self.used() - for t in self.used(): - assert id(t) in self.waited - - -@pytest.mark.parametrize("depth", [1, 2]) -def test_iter_weights_fences_on_completion(server, uri, depth): - setup_model(server, uri) - c = racer.Client(socket_path=server.path, staging_bytes=4096) - b = _FencingBackend() - names = [n for n, _ in iter_weights(c, uri, b, staging_bytes=400, depth=depth)] - assert "w3" in names - b.assert_fenced() - # The dedicated target (w3) was fenced too. - assert any(t.size == 900 for t in b.used()) - - -@pytest.mark.parametrize("stop", ["w1", "w3"]) -def test_iter_weights_fences_on_early_close(server, uri, stop): - setup_model(server, uri) - c = racer.Client(socket_path=server.path, staging_bytes=4096) - b = _FencingBackend() - gen = iter_weights(c, uri, b, staging_bytes=400, depth=2) - for name, _ in gen: - if name == stop: - break - gen.close() - b.assert_fenced() - - -def test_iter_weights_fences_on_exception(server, uri): - setup_model(server, uri) - c = racer.Client(socket_path=server.path, staging_bytes=4096) - b = _FencingBackend() - gen = iter_weights(c, uri, b, staging_bytes=400, depth=2) - next(gen) - with pytest.raises(RuntimeError): - gen.throw(RuntimeError("consumer failed")) - b.assert_fenced() - - -@pytest.mark.parametrize("mode", ["raise", "early"]) -def test_load_weights_closes_iterator_before_client(server, uri, monkeypatch, tmp_path, mode): - registry = install_fake_vllm(monkeypatch) - loader.register() - setup_model(server, uri) - model.fetch(uri, str(tmp_path / "m"), client=racer.Client(socket_path=server.path)) - monkeypatch.setenv("RACER_SOCKET", server.path) - monkeypatch.setattr(loader, "TorchBackend", lambda **kw: MemoryBackend()) - events = [] - real_iter = loader.iter_weights - - def tracking_iter(*a, **kw): - try: - yield from real_iter(*a, **kw) - finally: - events.append("weights closed") - - monkeypatch.setattr(loader, "iter_weights", tracking_iter) - real_client_for = loader.client_for - - def tracking_client_for(*a, **kw): - c = real_client_for(*a, **kw) - inner_exit = c.__exit__ - - class Wrap: - def __getattr__(self, n): - return getattr(c, n) - - def __enter__(self): - c.__enter__() - return self - - def __exit__(self, *exc): - events.append("client closed") - return inner_exit(*exc) - - return Wrap() - - monkeypatch.setattr(loader, "client_for", tracking_client_for) - inst = registry["racer"](types.SimpleNamespace(model_loader_extra_config={"staging_bytes": 400})) - mc = types.SimpleNamespace(model=str(tmp_path / "m"), quantization=None) - - class Model: - def load_weights(self, weights): - next(weights) - if mode == "raise": - raise RuntimeError("boom") - return None - - if mode == "raise": - with pytest.raises(RuntimeError, match="boom"): - inst.load_weights(Model(), mc) - else: - inst.load_weights(Model(), mc) - assert events == ["weights closed", "client closed"] - - -def test_iter_weights_releases_dedicated_when_consumed(server, uri): - import gc - import weakref - - setup_model(server, uri) - c = racer.Client(socket_path=server.path, staging_bytes=4096) - refs = [] - - class _WeakBackend(_FencingBackend): - def new_dedicated(self, size): - t = self.Target(size) - refs.append(weakref.ref(t)) - return t - - b = _WeakBackend() - gen = iter_weights(c, uri, b, staging_bytes=400, depth=2) - for name, _ in gen: - if name == "w3": - break - assert len(refs) == 1 and refs[0]() is not None - target_id = id(refs[0]()) - # Resuming past w3 fences the dedicated target and drops it. - next(gen, None) - assert target_id in b.waited - gc.collect() - assert refs[0]() is None - gen.close() diff --git a/clients/python/racer-vllm/tests/test_racer_vllm.py b/clients/python/racer-vllm/tests/test_racer_vllm.py new file mode 100644 index 0000000000..8fa691da00 --- /dev/null +++ b/clients/python/racer-vllm/tests/test_racer_vllm.py @@ -0,0 +1,787 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 +import json +import logging +import os +import struct +import sys +import types + +import pytest +from conftest import FakeTensor, make_safetensors +from fakeserver import Obj + +import racer +import racer_vllm as rv + +ALL = {"w1": b"\1" * 300, "w2": b"\2" * 50, "w3": b"\3" * 900, "w4": b"\0" * 8} +HERE = os.path.dirname(__file__) + + +@pytest.fixture +def uri(): + return rv.parse_uri("racer://cache1/models/org/llama") + + +def put(server, uri, name, data, etag='"v1"'): + server.objects[uri.request(name).hex] = Obj(data, etag) + + +def replace(server, uri, name, etag='"v2"'): + """Swap in a new version of an object, as the origin would.""" + put(server, uri, name, server.objects[uri.request(name).hex].data, etag) + + +def setup_model(server, uri, index=True): + s1 = make_safetensors([("w1", "U8", (300,), ALL["w1"]), ("w2", "U8", (50,), ALL["w2"])]) + s2 = make_safetensors([("w3", "U8", (900,), ALL["w3"]), ("w4", "F32", (2,), ALL["w4"])]) + put(server, uri, "config.json", b'{"model_type":"llama"}') + put(server, uri, "tokenizer.json", b"{}") + if not index: + put(server, uri, rv.SINGLE, s1) + return + idx = { + "metadata": {"total_size": 1258, "extra": "chat_template.jinja", "bad": "../x.json"}, + "weight_map": {"w1": "a.safetensors", "w2": "a.safetensors", "w3": "b.safetensors", "w4": "b.safetensors"}, + } + put(server, uri, rv.INDEX, json.dumps(idx).encode()) + put(server, uri, "chat_template.jinja", b"{{x}}") + put(server, uri, "a.safetensors", s1) + put(server, uri, "b.safetensors", s2) + + +def client(server): + return racer.Client(socket_path=server.path, staging_bytes=4096) + + +def chunks_for(c, uri, staging=400, pins=None): + files = rv.weight_files(c, uri, pins) + pinned = [(n, rv.stat_pinned(c, uri, n, (pins or {}).get(n))) for n in files] + return rv.read_chunks(c, uri, pinned, staging) + + +def weights(c, uri, depth=2, **kw): + return rv.iter_weights(c, uri, chunks_for(c, uri), 400, depth, use_vmm=False, **kw) + + +def range_reads(server): + out = [] + for s in server.seen: + if s.method == "POST" and "range" in s.headers: + first, last = (int(x) for x in s.headers["range"][len("bytes=") :].split("-")) + out.append((s, first, last)) + return out + + +# URIs and racer-model.json + + +def test_parse_uri(): + u = rv.parse_uri("racer://c/b/a/b/c/") + assert (u.cache, u.bucket, u.prefix) == ("c", "b", "a/b/c") + assert u.key("config.json") == "a/b/c/config.json" + assert str(u) == "racer://c/b/a/b/c" + assert rv.parse_uri("racer://c/b").key("x") == "x" + for bad in ["s3://c/b", "racer://c", "racer:///b", "racer://c/"]: + with pytest.raises(ValueError): + rv.parse_uri(bad) + + +def test_fetch_and_marker(server, uri, tmp_path): + setup_model(server, uri) + replace(server, uri, "b.safetensors", '"vb"') + dest = tmp_path / "m" + got = rv.fetch(uri, str(dest), client=client(server)) + assert set(got) == {"config.json", "tokenizer.json", rv.INDEX, "chat_template.jinja"} + assert (dest / "chat_template.jinja").read_bytes() == b"{{x}}" + assert not (tmp_path / "x.json").exists() + doc = json.loads((dest / rv.MARKER).read_text()) + assert doc["uri"] == str(uri) + assert doc["etags"]["config.json"] == '"v1"' and doc["etags"]["a.safetensors"] == '"v1"' + assert doc["etags"]["b.safetensors"] == '"vb"' + assert rv.resolve_uri(str(dest), {}) == uri + assert rv.resolve_pins(str(dest), uri) == doc["etags"] + # Pins belong to the fetched URI only. + assert rv.resolve_pins(str(dest), rv.parse_uri("racer://x/y/z")) == {} + assert rv.resolve_pins(str(tmp_path / "nowhere"), uri) == {} + + +def test_resolve_uri_order(tmp_path): + assert rv.resolve_uri("/nowhere", {"uri": "racer://x/y/z"}) == rv.parse_uri("racer://x/y/z") + assert rv.resolve_uri("racer://a/b/c", {}) == rv.parse_uri("racer://a/b/c") + with pytest.raises(ValueError, match="no model URI"): + rv.resolve_uri(str(tmp_path), {}) + + +def test_fetch_single_file(server, uri, tmp_path): + setup_model(server, uri, index=False) + rv.fetch(uri, str(tmp_path), client=client(server)) + assert json.loads((tmp_path / rv.MARKER).read_text())["etags"][rv.SINGLE] == '"v1"' + + +def test_fetch_without_weights_still_writes_marker(server, uri, tmp_path): + put(server, uri, "config.json", b"{}") + rv.fetch(uri, str(tmp_path), client=client(server)) + assert json.loads((tmp_path / rv.MARKER).read_text())["etags"] == {"config.json": '"v1"'} + + +def test_fetch_requires_config(server, uri, tmp_path): + with pytest.raises(FileNotFoundError): + rv.fetch(uri, str(tmp_path), client=client(server)) + + +def test_main(server, uri, tmp_path, monkeypatch, capsys): + setup_model(server, uri) + monkeypatch.setenv("RACER_SOCKET", server.path) + assert rv.main([str(uri), str(tmp_path)]) == 0 + assert "config.json" in capsys.readouterr().out.split() + assert (tmp_path / rv.MARKER).exists() + with pytest.raises(SystemExit): + rv.main([str(uri)]) + + +def test_bad_pins_rejected(uri, tmp_path): + (tmp_path / rv.MARKER).write_text(json.dumps({"uri": str(uri), "etags": {"a": 1}})) + with pytest.raises(ValueError, match="bad etags"): + rv.resolve_pins(str(tmp_path), uri) + + +def test_read_pinned(server, uri): + setup_model(server, uri) + c = client(server) + assert rv.read_pinned(c, uri, "config.json") == (b'{"model_type":"llama"}', '"v1"') + assert rv.read_pinned(c, uri, "config.json", '"v1"')[1] == '"v1"' + assert rv.read_pinned(c, uri, "missing.json") is None + with pytest.raises(rv.ModelChangedError, match="config.json"): + rv.read_pinned(c, uri, "config.json", '"v0"') + # A pinned file that was deleted also counts as a change. + with pytest.raises(rv.ModelChangedError): + rv.read_pinned(c, uri, "missing.json", '"v1"') + + +def test_weight_files(server, uri): + c = client(server) + with pytest.raises(racer.NotFound): + rv.weight_files(c, uri) + setup_model(server, uri, index=False) + assert rv.weight_files(c, uri) == [rv.SINGLE] + setup_model(server, uri) + assert rv.weight_files(c, uri) == ["a.safetensors", "b.safetensors"] + replace(server, uri, rv.INDEX) + with pytest.raises(rv.ModelChangedError, match="index"): + rv.weight_files(c, uri, {rv.INDEX: '"v1"'}) + + +@pytest.mark.parametrize("weight_map", [{}, {"w": "../x.safetensors"}, {"w": "/x.safetensors"}]) +def test_weight_files_rejects_bad_index(server, uri, weight_map): + put(server, uri, rv.INDEX, json.dumps({"weight_map": weight_map}).encode()) + with pytest.raises(ValueError): + rv.weight_files(client(server), uri) + + +# safetensors headers and chunk planning + + +def test_parse_header(): + f = make_safetensors( + [ + ("a", "F32", (2, 3), struct.pack("<6f", *range(6))), + ("b", "BF16", (4,), b"\1\2" * 4), + ("c", "U8", (0,), b""), + ("d", "I64", (), struct.pack(" c.start else set() + assert not pages & seen + seen |= pages + # Each tensor is whole in one chunk, or covered exactly once by its parts. + covered = {t.name: [] for t in tensors} + for c in chunks: + for t in c.tensors: + assert c.start <= t.start and t.end <= c.end + covered[t.name].append((0, t.end - t.start)) + for t, off, n in c.parts: + assert c.start <= t.start + off and t.start + off + n <= c.end + covered[t.name].append((off, off + n)) + for t in tensors: + spans = sorted(covered[t.name]) + assert spans[0][0] == 0 and spans[-1][1] == t.end - t.start + assert all(a[1] == b[0] for a, b in zip(spans, spans[1:])) + assert any(c.parts for c in chunks) + + +def test_plan_chunks_splits_across_several_chunks(): + _, chunks = plan([100, 127], 127, 64) + parts = [(off, n) for c in chunks for t, off, n in c.parts if t.name == "t1"] + assert len(parts) >= 3 and parts[0][0] == 0 + assert all(a[0] + a[1] == b[0] for a, b in zip(parts, parts[1:])) + assert sum(n for _, n in parts) == 127 + assert all(c.end - c.start <= 127 for c in chunks) + + +def test_plan_chunks_ignores_pages_larger_than_staging(): + _, paged = plan([100] * 5, 250, 1024) + _, plain = plan([100] * 5, 250) + assert paged == plain + + +# The weight pipeline, with a fake torch + + +@pytest.mark.parametrize("depth", [1, 2, 3, 4]) +def test_iter_weights(server, uri, torch, depth): + setup_model(server, uri) + c = client(server) + got = {n: t for n, t in weights(c, uri, depth)} + assert {n: t.bytes() for n, t in got.items()} == ALL + assert got["w4"].dtype == "float32" and got["w4"].shape == (2,) + # w3 is larger than staging, so it was read on its own. + assert 900 in [last + 1 - first for _, first, last in range_reads(server)] + # Every recorded event was waited on before its buffer was reused or dropped. + assert torch.events and all(e.synced for e in torch.events) + + +def test_iter_weights_reads_each_page_once(server, uri, torch, monkeypatch): + monkeypatch.setattr(racer, "PAGE_SIZE", 64) + setup_model(server, uri) + c = client(server) + assert {n: t.bytes() for n, t in weights(c, uri)} == ALL + pages = {} + reads = 0 + for s, first, last in range_reads(server): + if first in (0, 8) or last + 1 - first == 900: + continue # header reads and the dedicated w3 read + reads += 1 + mine = set(range(first // 64, last // 64 + 1)) + assert not mine & pages.setdefault(s.target, set()) + pages[s.target] |= mine + assert reads >= 2 + + +def test_iter_weights_clone_off_gives_staging_views(server, uri, torch): + setup_model(server, uri) + c = client(server) + got = dict(rv.iter_weights(c, uri, chunks_for(c, uri), 400, 1, clone=False, use_vmm=False)) + # w1 and w2 share one staging buffer; with clone=True they would not. + assert got["w1"].buf is got["w2"].buf + + +def test_iter_weights_falls_back_from_vmm(server, uri, torch, monkeypatch, caplog): + def no_vmm(size, device): + raise racer.RacerError("no VMM") + + monkeypatch.setattr("racer.cuda.GpuAllocation", no_vmm) + setup_model(server, uri) + c = client(server) + with caplog.at_level(logging.INFO, logger="vllm.racer_vllm.loader"): + got = {n: t.bytes() for n, t in rv.iter_weights(c, uri, chunks_for(c, uri), 400, 2)} + assert got == ALL + assert [r for r in caplog.records if "exportable GPU memory unavailable" in r.getMessage()] + + +def test_iter_weights_counts_reads(server, uri, torch): + setup_model(server, uri) + stats = {"reads": dict.fromkeys((racer.GPU, racer.HOST, racer.STREAM), 0)} + stats["bytes"] = dict(stats["reads"]) + assert {n: t.bytes() for n, t in weights(client(server), uri, stats=stats)} == ALL + assert sum(stats["bytes"].values()) == sum(len(v) for v in ALL.values()) + # Fake torch memory cannot be exported, so host-read is used. + assert stats["reads"][racer.HOST] >= 2 and stats["reads"][racer.GPU] == 0 + + +@pytest.mark.parametrize("stop", ["w1", "w3", None]) +def test_iter_weights_fences_on_early_close_and_error(server, uri, torch, stop): + setup_model(server, uri) + gen = weights(client(server), uri) + if stop is None: + next(gen) + with pytest.raises(RuntimeError): + gen.throw(RuntimeError("consumer failed")) + else: + for name, _ in gen: + if name == stop: + break + gen.close() + assert torch.events and all(e.synced for e in torch.events) + + +def test_iter_weights_rejects_zero_concurrency(server, uri, torch): + setup_model(server, uri) + with pytest.raises(ValueError): + list(weights(client(server), uri, depth=0)) + + +@pytest.mark.parametrize("stop", ["w1", "w3"]) +def test_model_changes_mid_load(server, uri, torch, stop): + setup_model(server, uri) + gen = weights(client(server), uri, depth=1) + for name, _ in gen: + if name == stop: + break + replace(server, uri, "b.safetensors") + with pytest.raises(rv.ModelChangedError, match="b.safetensors"): + for _ in gen: + pass + + +def test_summary(): + stats = {"files": 2, "reads": {racer.GPU: 1, racer.HOST: 0, racer.STREAM: 1}} + stats["bytes"] = {racer.GPU: 1 << 30, racer.HOST: 0, racer.STREAM: 1 << 29} + text = rv.summary(stats, 2.0, "racer://c/b/p") + assert "loaded 2 files, 1.50 GiB in 2.00s (0.75 GiB/s) from racer://c/b/p" in text + assert "gpu=1 (1.00 GiB) host=0 (0.00 GiB) stream=1 (0.50 GiB)" in text + assert "(0.00 GiB/s)" in rv.summary(stats, 0.0, "x") + + +# sharded_state + + +def test_render_pattern(): + assert rv.render_pattern(None, 0) == "model-rank-0-part-0.safetensors" + assert rv.render_pattern(None, 3, tp_rank=1, tp_size=2) == "model-rank-1-part-3.safetensors" + assert rv.render_pattern(None, 0, 1, 2, 1, 2) == "model-pp-1-tp-1-part-0.safetensors" + assert rv.render_pattern("w/{tp_rank}-{part}.st", 0, 1, 2) == "w/1-0.st" + + +@pytest.mark.parametrize( + "pattern,ranks,match", + [ + ("m-{x}-{part}", {}, "unsupported"), + ("m-{rank}", {}, "part"), + ("m-{rank}-{part}", {"pp_size": 2}, "pipeline parallel rank"), + ("m-{pp_rank}-{rank}-{part}", {}, "pipeline parallel size is 1"), + ("m-{part}", {"tp_size": 2}, "tensor parallel rank"), + ("../m-{rank}-{part}", {}, "unsafe"), + ("/m-{rank}-{part}", {}, "unsafe"), + ], +) +def test_render_pattern_rejects(pattern, ranks, match): + with pytest.raises(ValueError, match=match): + rv.render_pattern(pattern, 0, **ranks) + + +def test_rank_files(server, uri): + for name in ( + "model-rank-0-part-0.safetensors", + "model-rank-0-part-1.safetensors", + "model-rank-1-part-0.safetensors", + ): + put(server, uri, name, b"x") + c = client(server) + assert rv.rank_files(c, uri, None, {"tp_rank": 0, "tp_size": 2}) == [ + "model-rank-0-part-0.safetensors", + "model-rank-0-part-1.safetensors", + ] + assert rv.rank_files(c, uri, None, {"tp_rank": 1, "tp_size": 2}) == ["model-rank-1-part-0.safetensors"] + with pytest.raises(ValueError, match="model-rank-2-part-0"): + rv.rank_files(c, uri, None, {"tp_rank": 2, "tp_size": 4}) + + +def ft(data, dtype="uint8"): + return FakeTensor(bytearray(data), dtype=dtype) + + +def test_load_state_copies_and_narrows(caplog): + params = {"a": ft(b"\0" * 3), "b": ft(b"\0" * 4)} + with caplog.at_level(logging.WARNING): + rv.load_state(params, [("a", ft(b"abc")), ("b", ft(b"xyz"))]) + assert params["a"].bytes() == b"abc" + # b is padded (as with LoRA): only the first 3 bytes are written. + assert params["b"].bytes() == b"xyz\0" + assert any("shape" in r.getMessage() for r in caplog.records) + + +def test_load_state_errors(caplog): + with pytest.raises(ValueError, match="missing keys"): + rv.load_state({"a": ft(b"1"), "b": ft(b"1")}, [("a", ft(b"1"))]) + with pytest.raises(ValueError, match="unexpected"): + rv.load_state({"a": ft(b"1")}, [("z", ft(b"1"))]) + with pytest.raises(ValueError, match="twice"): + rv.load_state({"a": ft(b"1")}, [("a", ft(b"1")), ("a", ft(b"1"))]) + with pytest.raises(ValueError, match="shape"): + rv.load_state({"a": ft(b"1")}, [("a", ft(b"1").view((1, 1)))]) + with caplog.at_level(logging.WARNING): + rv.load_state({"a": ft(b"1", "bfloat16")}, [("a", ft(b"2", "float32"))]) + assert any("dtype" in r.getMessage() for r in caplog.records) + + +# The vLLM loader + + +def install_fake_vllm(monkeypatch, tp_rank=0, tp_size=1): + registry = {} + mods = { + n: types.ModuleType(n) + for n in ( + "vllm", + "vllm.distributed", + "vllm.model_executor", + "vllm.model_executor.model_loader", + "vllm.model_executor.model_loader.base_loader", + "vllm.model_executor.model_loader.sharded_state_loader", + ) + } + + class BaseModelLoader: + def __init__(self, load_config): + self.load_config = load_config + + def register_model_loader(name): + def wrap(cls): + assert issubclass(cls, BaseModelLoader) + registry.setdefault("calls", []).append(name) + registry[name] = cls + return cls + + return wrap + + mods["vllm.model_executor.model_loader"].register_model_loader = register_model_loader + mods["vllm.model_executor.model_loader.base_loader"].BaseModelLoader = BaseModelLoader + mods["vllm.model_executor.model_loader.sharded_state_loader"].ShardedStateLoader = types.SimpleNamespace( + _filter_subtensors=lambda sd: sd + ) + dist = mods["vllm.distributed"] + dist.get_tensor_model_parallel_rank = lambda: tp_rank + dist.get_pipeline_model_parallel_rank = lambda: 0 + dist.get_tp_group = lambda: types.SimpleNamespace(world_size=tp_size) + dist.get_pp_group = lambda: types.SimpleNamespace(world_size=1) + for name, mod in mods.items(): + monkeypatch.setitem(sys.modules, name, mod) + monkeypatch.setattr(rv, "_registered", False) + return registry + + +class Model: + def __init__(self, names=tuple(ALL)): + self.names = names + self.got = {} + + def load_weights(self, weights): + for name, t in weights: + self.got[name] = t.bytes() + return set(self.got) + + def named_parameters(self): + return [(n, None) for n in self.names] + + +def make_loader(monkeypatch, server, **extra): + registry = install_fake_vllm(monkeypatch) + rv.register() + rv.register() + assert registry["calls"] == ["racer"] + monkeypatch.setenv("RACER_SOCKET", server.path) + extra = {"staging_bytes": 400, "concurrency": 2, "use_vmm": False, **extra} + return registry["racer"](types.SimpleNamespace(model_loader_extra_config=extra)) + + +def test_loader_loads_logs_and_uses_pins(server, uri, torch, monkeypatch, tmp_path, caplog): + setup_model(server, uri) + rv.fetch(uri, str(tmp_path), client=client(server)) + inst = make_loader(monkeypatch, server) + mc = types.SimpleNamespace(model=str(tmp_path), quantization=None) + inst.download_model(mc) + m = Model() + with caplog.at_level(logging.INFO, logger="vllm.racer_vllm.loader"): + inst.load_weights(m, mc) + assert m.got == ALL + msgs = [r.getMessage() for r in caplog.records if r.name == "vllm.racer_vllm.loader"] + assert len(msgs) == 1 and "loaded 2 files" in msgs[0] and "host=" in msgs[0] and str(uri) in msgs[0] + + with pytest.raises(ValueError, match="w5"): + inst.load_weights(Model(tuple(ALL) + ("w5",)), mc) + # Quantized models skip the check. + inst.load_weights(Model(tuple(ALL) + ("w5",)), types.SimpleNamespace(model=str(tmp_path), quantization="fp8")) + + # The model was replaced after racer-vllm-fetch: the load fails before reading data. + replace(server, uri, "a.safetensors") + before = len(range_reads(server)) + with pytest.raises(rv.ModelChangedError, match="racer-vllm-fetch"): + inst.load_weights(Model(), mc) + assert len(range_reads(server)) == before + + +def test_loader_stats_every_file_before_data_reads(server, uri, torch, monkeypatch): + setup_model(server, uri) + inst = make_loader(monkeypatch, server, uri=str(uri)) + inst.load_weights(Model(), types.SimpleNamespace(model="/unused", quantization=None)) + a, b = uri.request("a.safetensors").hex, uri.request("b.safetensors").hex + heads = [i for i, s in enumerate(server.seen) if s.method == "HEAD" and s.target.endswith((a, b))] + posts = [i for i, s in enumerate(server.seen) if s.method == "POST" and (a in s.target or b in s.target)] + assert len(heads) == 2 and max(heads) < min(posts) + assert all(server.seen[i].headers.get("if-match") == '"v1"' for i in posts) + + +@pytest.mark.parametrize("fail", [True, False]) +def test_loader_closes_weights_before_client(server, uri, torch, monkeypatch, fail): + setup_model(server, uri) + inst = make_loader(monkeypatch, server, uri=str(uri)) + events = [] + real_iter = rv.iter_weights + + def tracking_iter(*a, **kw): + try: + yield from real_iter(*a, **kw) + finally: + events.append("weights closed") + + real_close = racer.Client.close + monkeypatch.setattr(rv, "iter_weights", tracking_iter) + monkeypatch.setattr(racer.Client, "close", lambda self: (events.append("client closed"), real_close(self))) + + class Partial: + def load_weights(self, weights): + next(weights) + if fail: + raise RuntimeError("boom") + + mc = types.SimpleNamespace(model="/unused", quantization=None) + if fail: + with pytest.raises(RuntimeError, match="boom"): + inst.load_weights(Partial(), mc) + else: + inst.load_weights(Partial(), mc) + assert events[:2] == ["weights closed", "client closed"] + + +def test_loader_without_uri_fails(server, monkeypatch, tmp_path): + inst = make_loader(monkeypatch, server) + with pytest.raises(ValueError, match="no model URI"): + inst.download_model(types.SimpleNamespace(model=str(tmp_path))) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_loader_sharded_state_reads_only_own_rank(server, uri, torch, monkeypatch, rank): + for r in (0, 1): + put(server, uri, f"model-rank-{r}-part-0.safetensors", make_safetensors([("w", "U8", (4,), bytes([r]) * 4)])) + put( + server, + uri, + f"model-rank-{r}-part-1.safetensors", + make_safetensors([("v", "U8", (2,), bytes([r + 10]) * 2)]), + ) + inst = make_loader(monkeypatch, server, uri=str(uri), sharded_state=True) + install_fake_vllm(monkeypatch, tp_rank=rank, tp_size=2) + params = {"w": ft(b"\0" * 4), "v": ft(b"\0" * 2)} + + class Sharded: + def state_dict(self): + return params + + mc = types.SimpleNamespace(model="/unused", quantization=None) + inst.download_model(mc) + inst.load_weights(Sharded(), mc) + assert params["w"].bytes() == bytes([rank]) * 4 + assert params["v"].bytes() == bytes([rank + 10]) * 2 + other = {uri.request(f"model-rank-{1 - rank}-part-{p}.safetensors").hex for p in (0, 1)} + assert not [s for s in server.seen if s.method == "POST" and s.target.split("/")[-1] in other] + + +def test_settings(): + s = rv.settings({}) + assert (s["staging_bytes"], s["concurrency"], s["clone"], s["use_vmm"], s["sharded_state"]) == ( + 256 << 20, + 4, + True, + True, + False, + ) + assert s["mode_retry_seconds"] == racer.DEFAULT_MODE_RETRY_SECONDS + assert rv.settings({"mode_retry_seconds": None})["mode_retry_seconds"] is None + assert rv.settings({"mode_retry_seconds": "30"})["mode_retry_seconds"] == 30.0 + + +# Packaging + + +def test_pyproject(): + tomllib = pytest.importorskip("tomllib") + with open(os.path.join(HERE, "..", "pyproject.toml"), "rb") as f: + proj = tomllib.load(f) + with open(os.path.join(HERE, "..", "..", "racer", "pyproject.toml"), "rb") as f: + client_proj = tomllib.load(f)["project"] + vllm_proj = proj["project"] + assert vllm_proj["entry-points"]["vllm.general_plugins"] == {"racer": "racer_vllm:register"} + assert vllm_proj["scripts"] == {"racer-vllm-fetch": "racer_vllm:main"} + assert proj["tool"]["setuptools"]["py-modules"] == ["racer_vllm"] + # racer-vllm needs the racer-client release built from the same tree. + assert [d for d in vllm_proj["dependencies"] if d.startswith("racer-client")] == [ + f"racer-client>={client_proj['version']},<0.3" + ] + for p in (vllm_proj, client_proj): + assert p["license"] == "Apache-2.0" and p["license-files"] == ["LICENSE"] + assert p["urls"]["Source"].startswith("https://github.com/Azure/unbounded") + for d in ("racer", "racer-vllm"): + assert os.path.exists(os.path.join(HERE, "..", "..", d, "LICENSE")) + + +def test_loader_rejects_index_added_after_single_file_fetch(server, uri, torch, monkeypatch, tmp_path): + setup_model(server, uri, index=False) + rv.fetch(uri, str(tmp_path), client=client(server)) + setup_model(server, uri) # an index and shards are uploaded later + inst = make_loader(monkeypatch, server) + mc = types.SimpleNamespace(model=str(tmp_path), quantization=None) + with pytest.raises(rv.ModelChangedError, match=f"{rv.INDEX}.*not pinned.*racer-vllm-fetch"): + inst.download_model(mc) + assert not range_reads(server) + + +def test_loader_rejects_unpinned_weight_file(server, uri, torch, monkeypatch, tmp_path): + # Fetch found no weights, so the marker pins none; a single file appears later. + put(server, uri, "config.json", b"{}") + rv.fetch(uri, str(tmp_path), client=client(server)) + setup_model(server, uri, index=False) + inst = make_loader(monkeypatch, server) + with pytest.raises(rv.ModelChangedError, match=f"{rv.SINGLE}.*not pinned"): + inst.download_model(types.SimpleNamespace(model=str(tmp_path), quantization=None)) + + +def test_loader_without_marker_is_unpinned(server, uri, torch, monkeypatch, tmp_path): + setup_model(server, uri, index=False) + inst = make_loader(monkeypatch, server, uri=str(uri)) + mc = types.SimpleNamespace(model=str(tmp_path), quantization=None) + setup_model(server, uri) # layout changes, but nothing was pinned + inst.load_weights(Model(), mc) + # A marker for another URI is ignored too. + (tmp_path / rv.MARKER).write_text(json.dumps({"uri": "racer://cache1/models/other", "etags": {}})) + inst.load_weights(Model(), mc) + + +def put_sharded(server, uri, ranks=2, parts=2, etag='"v1"'): + for r in range(ranks): + for p in range(parts): + data = make_safetensors([(f"w{p}", "U8", (2,), bytes([r * 10 + p]) * 2)]) + put(server, uri, f"model-rank-{r}-part-{p}.safetensors", data, etag) + + +def test_fetch_pins_sharded_state_files(server, uri, tmp_path): + put(server, uri, "config.json", b"{}") + put_sharded(server, uri) + rv.fetch(uri, str(tmp_path), client=client(server)) + etags = json.loads((tmp_path / rv.MARKER).read_text())["etags"] + shards = [f"model-rank-{r}-part-{p}.safetensors" for r in (0, 1) for p in (0, 1)] + assert sorted(etags) == sorted(["config.json", *shards]) + + pp = tmp_path / "pp" + put(server, uri, "model-pp-0-tp-0-part-0.safetensors", b"x") + put(server, uri, "model-pp-1-tp-0-part-0.safetensors", b"x") + pattern = "model-pp-{pp_rank}-tp-{tp_rank}-part-{part}.safetensors" + rv.fetch(uri, str(pp), client=client(server), pattern=pattern) + etags = json.loads((pp / rv.MARKER).read_text())["etags"] + assert sorted(etags) == ["config.json", "model-pp-0-tp-0-part-0.safetensors", "model-pp-1-tp-0-part-0.safetensors"] + + +@pytest.mark.parametrize( + "pattern", + [ + "model-rank-{rank:02d}-part-{part:02d}.safetensors", + "model-pp-{pp_rank:02d}-tp-{tp_rank:02d}-part-{part}.safetensors", + ], +) +def test_fetch_pins_sharded_state_files_with_format_specs(server, uri, tmp_path, pattern): + put(server, uri, "config.json", b"{}") + pp = "pp_rank" in pattern + names = [] + for pp_rank in (0, 1) if pp else (0,): + for tp_rank in (0, 1): + ranks = {"pp_rank": pp_rank, "tp_rank": tp_rank, "rank": tp_rank} + name = pattern.format(part=0, **ranks) + put(server, uri, name, b"x") + names.append(name) + rv.fetch(uri, str(tmp_path), client=client(server), pattern=pattern) + etags = json.loads((tmp_path / rv.MARKER).read_text())["etags"] + assert sorted(etags) == sorted(["config.json", *names]) + + +@pytest.mark.parametrize("change", ["none", "new part", "deleted part", "replaced"]) +def test_loader_sharded_state_uses_pins(server, uri, torch, monkeypatch, tmp_path, change): + put(server, uri, "config.json", b"{}") + put_sharded(server, uri) + rv.fetch(uri, str(tmp_path), client=client(server)) + if change == "new part": + put(server, uri, "model-rank-1-part-2.safetensors", make_safetensors([("w2", "U8", (2,), b"zz")])) + elif change == "deleted part": + del server.objects[uri.request("model-rank-1-part-1.safetensors").hex] + elif change == "replaced": + replace(server, uri, "model-rank-1-part-0.safetensors") + inst = make_loader(monkeypatch, server, sharded_state=True) + install_fake_vllm(monkeypatch, tp_rank=1, tp_size=2) + params = {"w0": ft(b"\0" * 2), "w1": ft(b"\0" * 2)} + + class Sharded: + def state_dict(self): + return params + + mc = types.SimpleNamespace(model=str(tmp_path), quantization=None) + if change == "none": + inst.load_weights(Sharded(), mc) + assert params["w0"].bytes() == b"\x0a\x0a" and params["w1"].bytes() == b"\x0b\x0b" + return + match = {"new part": "part-2.*not pinned", "deleted part": "part-1", "replaced": "part-0"}[change] + with pytest.raises(rv.ModelChangedError, match=match): + inst.load_weights(Sharded(), mc) diff --git a/clients/python/racer-vllm/tests/test_safetensors.py b/clients/python/racer-vllm/tests/test_safetensors.py deleted file mode 100644 index e57988740b..0000000000 --- a/clients/python/racer-vllm/tests/test_safetensors.py +++ /dev/null @@ -1,158 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# SPDX-License-Identifier: Apache-2.0 -import struct - -import pytest -from conftest import make_safetensors - -from racer_vllm.safetensors import ( - SafetensorsError, - header_length, - parse_header, - plan_chunks, - tensor_bytes, -) - - -def sample(): - return make_safetensors( - [ - ("a", "F32", (2, 3), struct.pack("<6f", *range(6))), - ("b", "BF16", (4,), b"\1\2" * 4), - ("c", "U8", (0,), b""), - ("d", "I64", (), struct.pack(" start else set() - - -def test_plan_chunks_on_page_boundaries(): - f = make_safetensors([(f"t{i}", "U8", (90,), bytes([i]) * 90) for i in range(8)] + [("z", "U8", (0,), b"")]) - n = header_length(f[:8]) - h = parse_header(f[8 : 8 + n]) - chunks = plan_chunks(h.tensors, 250, 64) - # Every chunk but the last ends on a page boundary, and no two chunks share a page. - assert all(c.end % 64 == 0 for c in chunks[:-1]) - assert all(c.length <= 250 for c in chunks) - seen = set() - for c in chunks: - pages = _pages(c.start, c.end, 64) - assert not pages & seen - seen |= pages - # Each tensor is either whole in one chunk or covered exactly once by its parts. - covered = {t.name: [] for t in h.tensors} - for c in chunks: - for t in c.tensors: - assert c.start <= t.start and t.end <= c.end - covered[t.name].append((0, t.nbytes)) - for p in c.parts: - assert c.start <= p.tensor.start + p.offset and p.tensor.start + p.offset + p.length <= c.end - covered[p.tensor.name].append((p.offset, p.offset + p.length)) - for t in h.tensors: - spans = sorted(covered[t.name]) - assert spans[0][0] == 0 and spans[-1][1] == t.nbytes, t.name - assert all(a[1] == b[0] for a, b in zip(spans, spans[1:])), t.name - assert any(c.parts for c in chunks) - assert [p.last for c in chunks for p in c.parts].count(True) == len({p.tensor.name for c in chunks for p in c.parts}) - - -def test_plan_chunks_splits_across_several_chunks(): - # With staging of just under two pages, each chunk after the first holds - # one page, so a tensor that starts mid-chunk spans three chunks. - f = make_safetensors([("a", "U8", (100,), b"a" * 100), ("b", "U8", (127,), b"b" * 127)]) - n = header_length(f[:8]) - h = parse_header(f[8 : 8 + n]) - chunks = plan_chunks(h.tensors, 127, 64) - parts = [p for c in chunks for p in c.parts if p.tensor.name == "b"] - assert len(parts) >= 3 - assert [(p.offset, p.length) for p in parts][0][0] == 0 - assert all(a.offset + a.length == b.offset for a, b in zip(parts, parts[1:])) - assert sum(p.length for p in parts) == 127 - assert [p.last for p in parts] == [False] * (len(parts) - 1) + [True] - assert all(c.length <= 127 for c in chunks) - - -def test_plan_chunks_ignores_pages_larger_than_staging(): - f = make_safetensors([(f"t{i}", "U8", (100,), bytes([i]) * 100) for i in range(5)]) - n = header_length(f[:8]) - h = parse_header(f[8 : 8 + n]) - assert [[t.name for t in c.tensors] for c in plan_chunks(h.tensors, 250, 1024)] == [ - [t.name for t in c.tensors] for c in plan_chunks(h.tensors, 250) - ] - assert not any(c.parts for c in plan_chunks(h.tensors, 250, 1024)) - - -@pytest.mark.parametrize( - "header", - [ - b"[]", - b"{not json", - b'{"a":{"dtype":"F4","shape":[1],"data_offsets":[0,1]}}', - b'{"a":{"dtype":"F32","shape":[2],"data_offsets":[0,4]}}', - b'{"a":{"dtype":"U8","shape":[2],"data_offsets":[0,2]},"b":{"dtype":"U8","shape":[2],"data_offsets":[1,3]}}', - b'{"a":{"dtype":"U8","shape":[-1],"data_offsets":[0,0]}}', - b'{"a":{"dtype":"U8"}}', - b'{"a":{"dtype":"U8","shape":[1.9],"data_offsets":[0.9,1.9]}}', - b'{"a":{"dtype":"U8","shape":[1],"data_offsets":[0.0,1.0]}}', - b'{"a":{"dtype":"U8","shape":[-0.5],"data_offsets":[0,0]}}', - b'{"a":{"dtype":"U8","shape":["1"],"data_offsets":[0,1]}}', - b'{"a":{"dtype":"U8","shape":[1],"data_offsets":["0","1"]}}', - b'{"a":{"dtype":"U8","shape":[true],"data_offsets":[0,1]}}', - b'{"a":{"dtype":"U8","shape":[1],"data_offsets":[false,true]}}', - b'{"a":{"dtype":"U8","shape":"1","data_offsets":[0,1]}}', - b'{"a":{"dtype":"U8","shape":[1],"data_offsets":[0,1,2]}}', - b'{"a":{"dtype":"U8","shape":[1],"data_offsets":{"0":0}}}', - b'{"a":[]}', - ], -) -def test_bad_headers(header): - with pytest.raises(SafetensorsError): - parse_header(header) - - -def test_bad_length(): - with pytest.raises(SafetensorsError): - header_length(struct.pack("=64"] +requires = ["setuptools>=77"] build-backend = "setuptools.build_meta" [project] name = "racer-client" -version = "0.1.0" +version = "0.2.0" description = "Python client for the Racer node-local object cache" readme = "README.md" -license = { text = "Apache-2.0" } +license = "Apache-2.0" +license-files = ["LICENSE"] requires-python = ">=3.10" +authors = [{ name = "Microsoft Corporation" }] +keywords = ["racer", "unbounded", "kubernetes", "cache", "gpu"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] dependencies = [] +[project.urls] +Homepage = "https://github.com/Azure/unbounded" +Documentation = "https://github.com/Azure/unbounded/tree/main/docs/content/guides/racer-vllm.md" +Source = "https://github.com/Azure/unbounded/tree/main/clients/python/racer" +Issues = "https://github.com/Azure/unbounded/issues" + [project.optional-dependencies] torch = ["torch"] test = ["pytest"] @@ -20,3 +40,4 @@ where = ["src"] [tool.pytest.ini_options] testpaths = ["tests"] +markers = ["gpu: needs a GPU and a live dataplane (see tests/test_gpu_e2e.py)"] diff --git a/clients/python/racer/src/racer/__init__.py b/clients/python/racer/src/racer/__init__.py index 2a338da74e..2ae063f064 100644 --- a/clients/python/racer/src/racer/__init__.py +++ b/clients/python/racer/src/racer/__init__.py @@ -6,6 +6,7 @@ GPU, HOST, STREAM, + DEFAULT_MODE_RETRY_SECONDS, Client, Metadata, ReadResult, @@ -26,6 +27,7 @@ from .keys import Request, canonical_json, object_key, object_request __all__ = [ + "DEFAULT_MODE_RETRY_SECONDS", "GPU", "HOST", "PAGE_SIZE", diff --git a/clients/python/racer/src/racer/client.py b/clients/python/racer/src/racer/client.py index c2146395b5..d99c8df716 100644 --- a/clients/python/racer/src/racer/client.py +++ b/clients/python/racer/src/racer/client.py @@ -11,9 +11,12 @@ import array import contextlib +import logging import os +import re import socket import threading +import time from dataclasses import dataclass from typing import Callable, Optional @@ -33,12 +36,35 @@ DEFAULT_STAGING_BYTES = 4 * wire.PAGE_SIZE +# Seconds before a mode remembered as unsupported is tried again. +DEFAULT_MODE_RETRY_SECONDS = 300.0 + _SUPPORTED = "supported" _UNSUPPORTED = "unsupported" -# Mode support per socket path, kept until the process exits. +log = logging.getLogger(__name__) + + +@dataclass +class _ModeEntry: + state: str + since: float # clock() when the state was recorded + socket_id: tuple | None # (st_dev, st_ino) of the socket file then + + +# Mode support per socket path. An entry is dropped when the socket file is +# replaced (the dataplane restarted), and an unsupported entry also expires +# after the client's mode_retry_seconds. _mode_lock = threading.Lock() -_mode_state: dict[str, dict[str, str]] = {} +_mode_state: dict[str, dict[str, _ModeEntry]] = {} + + +def _socket_id(path: str) -> tuple | None: + try: + st = os.stat(path) + except OSError: + return None + return (st.st_dev, st.st_ino) def _dest_device(dest) -> int | None: @@ -50,6 +76,30 @@ def _dest_device(dest) -> int | None: return index if isinstance(index, int) else None +_UUID = re.compile(r"GPU-[0-9a-fA-F]{8}(-[0-9a-fA-F]{4}){3}-[0-9a-fA-F]{12}") + + +def _dest_uuid(dest) -> str | None: + """Return the GPU UUID of a gpu-read destination, or None if unknown. + + The dataplane uses it to pick its context for the client's GPU, because + device ordinals differ between containers. Without it the dataplane tries + each GPU, so failing to learn the UUID is never an error. + """ + get = getattr(dest, "gpu_uuid", None) + if get is None: + return None + try: + uuid = get() if callable(get) else get + except Exception: + return None + if not isinstance(uuid, str): + return None + if not uuid.upper().startswith("GPU-"): + uuid = "GPU-" + uuid + return uuid if _UUID.fullmatch(uuid) else None + + def socket_path_for(cache: str | None = None, socket_path: str | None = None) -> str: """Resolve the client socket path. @@ -186,6 +236,8 @@ def __init__( timeout: float | None = 60.0, staging_bytes: int = DEFAULT_STAGING_BYTES, pin_staging: bool = True, + mode_retry_seconds: float | None = DEFAULT_MODE_RETRY_SECONDS, + clock: Callable[[], float] = time.monotonic, ) -> None: self.socket_path = socket_path_for(cache, socket_path) self.timeout = timeout @@ -193,18 +245,62 @@ def __init__( raise RacerError("staging_bytes must be positive") self.staging_bytes = staging_bytes self.pin_staging = pin_staging + if mode_retry_seconds is not None and mode_retry_seconds < 0: + raise RacerError("mode_retry_seconds must not be negative") + # None keeps an unsupported mode until the socket file changes. + self.mode_retry_seconds = mode_retry_seconds + self._clock = clock self._pool: list = [] self._pool_lock = threading.Lock() # Mode memory ----------------------------------------------------------- def mode_state(self, mode: str) -> str | None: - with _mode_lock: - return _mode_state.get(self.socket_path, {}).get(mode) + """Return "supported", "unsupported", or None when unknown. - def _set_mode(self, mode: str, state: str) -> None: + Memory is dropped when the socket file was replaced since it was + recorded, and "unsupported" expires after mode_retry_seconds. + """ + sid = _socket_id(self.socket_path) + with _mode_lock: + modes = _mode_state.get(self.socket_path) + if not modes: + return None + if any(e.socket_id != sid for e in modes.values()): + # The dataplane restarted (or the socket is gone): forget all. + _mode_state.pop(self.socket_path, None) + return None + e = modes.get(mode) + if e is None: + return None + if ( + e.state == _UNSUPPORTED + and self.mode_retry_seconds is not None + and self._clock() - e.since >= self.mode_retry_seconds + ): + del modes[mode] + return None + return e.state + + def _set_mode(self, mode: str, state: str, reason: str | None = None, next_mode: str | None = None) -> None: + sid = _socket_id(self.socket_path) with _mode_lock: - _mode_state.setdefault(self.socket_path, {})[mode] = state + modes = _mode_state.setdefault(self.socket_path, {}) + prev = modes.get(mode) + modes[mode] = _ModeEntry(state, self._clock(), sid) + if state == _UNSUPPORTED and (prev is None or prev.state != _UNSUPPORTED): + log.warning( + "racer: %s-read is not available from %s (%s); falling back to %s until %s", + mode, + self.socket_path, + reason or "unsupported", + f"{next_mode} mode" if next_mode else "the next mode", + "the dataplane restarts" + if self.mode_retry_seconds is None + else f"a retry in {self.mode_retry_seconds:g}s or a dataplane restart", + ) + elif state == _SUPPORTED and prev is not None and prev.state == _UNSUPPORTED: + log.info("racer: %s-read is available again from %s", mode, self.socket_path) # Staging pool ---------------------------------------------------------- @@ -421,6 +517,7 @@ def _rpc( if_match: str | None, on_frame: Callable[[int, int], None] | None, on_head: Callable[[], None] | None = None, + gpu_uuid: str | None = None, ) -> ReadResult: if dest_offset < 0 or dest_length <= 0: raise RacerError("invalid destination window") @@ -432,6 +529,10 @@ def _rpc( head += f"Racer-Destination-Offset: {dest_offset}\r\nRacer-Destination-Length: {dest_length}\r\n" if allocation_size is not None: head += f"Racer-Allocation-Size: {allocation_size}\r\n" + if gpu_uuid is not None: + if not _UUID.fullmatch(gpu_uuid): + raise RacerError("invalid GPU UUID") + head += f"Racer-Gpu-Uuid: {gpu_uuid}\r\n" head += self._context_headers(request) + "\r\n" conn = self._connect() try: @@ -489,11 +590,29 @@ def gpu_read( length: int | None = None, if_match: str | None = None, on_page: OnPage = None, + gpu_uuid: str | None = None, ) -> ReadResult: - """One gpu-read RPC with no fallback. fd comes from cuMemExportToShareableHandle.""" + """One gpu-read RPC with no fallback. fd comes from cuMemExportToShareableHandle. + + gpu_uuid (``GPU-xxxxxxxx-...``, see racer.cuda.device_uuid) names the + GPU that holds the allocation. Send it when the process sees more than + one GPU or its device numbers may differ from the host's. + """ if dest_offset + dest_length > allocation_size: raise RacerError("destination window is outside the allocation") - return self._rpc(GPU, request, fd, dest_offset, dest_length, allocation_size, offset, length, if_match, on_page) + return self._rpc( + GPU, + request, + fd, + dest_offset, + dest_length, + allocation_size, + offset, + length, + if_match, + on_page, + gpu_uuid=gpu_uuid, + ) # Fallback driver ------------------------------------------------------- @@ -505,15 +624,17 @@ def _run_modes(self, attempts: list[tuple[str, Callable[[Callable[[], None]], Re unsupported only once a later mode gets a 200 head, so a missing object (404 everywhere) does not disable a mode for good. """ - pending: list[str] = [] + pending: list[tuple[str, str]] = [] last: StatusError | None = None + current = [""] def understood() -> None: - for m in pending: - self._set_mode(m, _UNSUPPORTED) + for m, why in pending: + self._set_mode(m, _UNSUPPORTED, why, current[0]) pending.clear() - for mode, run in attempts: + for i, (mode, run) in enumerate(attempts): + current[0] = mode state = self.mode_state(mode) if mode != STREAM else _SUPPORTED if state == _UNSUPPORTED: continue @@ -524,13 +645,15 @@ def understood() -> None: except StatusError as exc: if mode == STREAM or state is not None or exc.status not in FALLBACK_STATUSES: raise + why = f"status {exc.status}: {exc}" if exc.status == 501: - self._set_mode(mode, _UNSUPPORTED) + nxt = attempts[i + 1][0] if i + 1 < len(attempts) else None + self._set_mode(mode, _UNSUPPORTED, why, nxt) else: - pending.append(mode) + pending.append((mode, why)) last = exc continue - if mode != STREAM: + if mode != STREAM and state != _SUPPORTED: self._set_mode(mode, _SUPPORTED) return result if last is not None: # pragma: no cover - stream always runs last @@ -603,9 +726,9 @@ def read_into_gpu( dest is a racer.cuda.GpuAllocation, a racer.tensors.TensorTarget, or any object with ``size``, ``copy_from_host(offset, data)`` and - optionally ``export_fd()``. Modes: gpu-read, then host-read into a - pinned staging memfd plus a host-to-device copy, then stream into the - staging buffer plus a copy. on_page(object_offset, length) runs once a + optionally ``export_fd()`` and ``gpu_uuid()``. Modes: gpu-read, then + host-read into a pinned staging memfd plus a host-to-device copy, then + stream into the staging buffer plus a copy. on_page(object_offset, length) runs once a page is in GPU memory. """ window = dest.size - dest_offset @@ -624,7 +747,18 @@ def gpu(ok): raise _Skip() from None try: return self._rpc( - GPU, request, fd, dest_offset, window, dest.size, offset, length, if_match, on_page, ok + GPU, + request, + fd, + dest_offset, + window, + dest.size, + offset, + length, + if_match, + on_page, + ok, + gpu_uuid=_dest_uuid(dest), ) finally: os.close(fd) diff --git a/clients/python/racer/src/racer/cuda.py b/clients/python/racer/src/racer/cuda.py index 0ab7e3e390..ca5dcc0d91 100644 --- a/clients/python/racer/src/racer/cuda.py +++ b/clients/python/racer/src/racer/cuda.py @@ -177,6 +177,35 @@ def available() -> bool: return True +class CUuuid(ctypes.Structure): + _fields_ = [("bytes", ctypes.c_ubyte * 16)] + + +def format_uuid(raw: bytes) -> str: + """Format 16 UUID bytes as ``GPU-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx``, like nvidia-smi.""" + if len(raw) != 16: + raise ValueError("a GPU UUID is 16 bytes") + h = raw.hex() + return f"GPU-{h[:8]}-{h[8:12]}-{h[12:16]}-{h[16:20]}-{h[20:]}" + + +def device_uuid(device: int = 0) -> str: + """Return the UUID of CUDA device ``device`` (an ordinal in this process). + + Ordinals depend on CUDA_VISIBLE_DEVICES and the container's GPU set, so + they differ between this process and the dataplane. The UUID does not. + """ + lib = _load() + fn = getattr(lib, "cuDeviceGetUuid_v2", None) or getattr(lib, "cuDeviceGetUuid", None) + if fn is None: + raise CudaUnavailable("libcuda has no cuDeviceGetUuid") + dev = ctypes.c_int(0) + _check("cuDeviceGet", lib.cuDeviceGet(ctypes.byref(dev), device)) + uuid = CUuuid() + _check("cuDeviceGetUuid", fn(ctypes.byref(uuid), dev.value)) + return format_uuid(bytes(uuid.bytes)) + + def ensure_context(device: int = 0) -> None: """Make sure a context is current; retain the primary context if not. @@ -310,6 +339,10 @@ def __init__(self, size: int, device: int = 0) -> None: self.close() raise + def gpu_uuid(self) -> str: + """The UUID of the GPU that holds the allocation, sent with gpu-read.""" + return device_uuid(self.device) + def export_fd(self) -> int: """Return a new POSIX file descriptor for the allocation. The caller closes it.""" lib = _load() diff --git a/clients/python/racer/tests/conftest.py b/clients/python/racer/tests/conftest.py index 4478c36e02..5eec4739ab 100644 --- a/clients/python/racer/tests/conftest.py +++ b/clients/python/racer/tests/conftest.py @@ -33,7 +33,7 @@ def server(tmp_path, small_pages): class FakeGpu: """A GPU destination backed by a memfd, so the fake server can write it.""" - def __init__(self, size, exportable=True): + def __init__(self, size, exportable=True, uuid=None): self.size = size self.device = 0 self.fd = os.memfd_create("fake-gpu") @@ -41,6 +41,8 @@ def __init__(self, size, exportable=True): self.copies = 0 if not exportable: self.export_fd = self._no_export + if uuid is not None: + self.gpu_uuid = uuid if callable(uuid) else (lambda: uuid) def _no_export(self): raise NotImplementedError diff --git a/clients/python/racer/tests/test_client.py b/clients/python/racer/tests/test_client.py index 071daee8f6..ed757f0c3e 100644 --- a/clients/python/racer/tests/test_client.py +++ b/clients/python/racer/tests/test_client.py @@ -203,6 +203,52 @@ def test_gpu_read_success(server, gpu_writer): g.close() +UUID = "GPU-00112233-4455-6677-8899-aabbccddeeff" + + +def _boom(): + raise RuntimeError("no uuid") + + +@pytest.mark.parametrize( + "uuid, sent", + [ + (None, None), # older destinations: no header, the dataplane tries each GPU + (UUID, UUID), + (UUID[4:], UUID), # torch prints the UUID without the GPU- prefix + ("MIG-00112233-4455-6677-8899-aabbccddeeff", None), + ("junk", None), + (_boom, None), # a failed lookup never fails the read + ], +) +def test_gpu_read_sends_gpu_uuid(server, gpu_writer, uuid, sent): + req = add(server) + server.gpu_status = 200 + server.gpu_write = gpu_writer + g = FakeGpu(len(DATA), uuid=uuid) + try: + r = client(server).read_into_gpu(req, g) + assert r.delivery == racer.GPU + assert server.seen[-1].headers.get("racer-gpu-uuid") == sent + finally: + g.close() + + +def test_gpu_read_unknown_uuid_falls_back_to_host(server): + # The dataplane answers 501 for a GPU it cannot use. + req = add(server) + server.gpu_status = 501 + g = FakeGpu(len(DATA), uuid=UUID) + try: + c = client(server) + r = c.read_into_gpu(req, g) + assert r.delivery == racer.HOST + assert g.read(0, len(DATA)) == DATA + assert c.mode_state(racer.GPU) == "unsupported" + finally: + g.close() + + def test_gpu_falls_back_to_host_chunks(server): req = add(server) g = FakeGpu(len(DATA)) @@ -253,7 +299,7 @@ def test_modes_remembered_per_socket(server, tmp_path): server.old = True with racer.HostBuffer(len(DATA)) as buf: client(server).read_into_host(req, buf) - assert racer.client._mode_state[server.path][racer.HOST] == "unsupported" + assert racer.client._mode_state[server.path][racer.HOST].state == "unsupported" other = racer.Client(socket_path=str(tmp_path / "other")) assert other.mode_state(racer.HOST) is None @@ -320,3 +366,100 @@ def test_response_head_parsing(): wire.check_error_status(wire.parse_response_head(b"HTTP/1.1 404 N\r\nContent-Length: 3\r\n\r\n")) with pytest.raises(racer.ProtocolError): wire.check_error_status(wire.parse_response_head(b"HTTP/1.1 416 N\r\nContent-Length: 0\r\n\r\n")) + + +class FakeClock: + def __init__(self): + self.now = 1000.0 + + def __call__(self): + return self.now + + +def test_unsupported_mode_expires_after_retry_interval(server): + req = add(server) + server.old = True + clock = FakeClock() + c = client(server, mode_retry_seconds=60, clock=clock) + with racer.HostBuffer(len(DATA)) as buf: + assert c.read_into_host(req, buf).delivery == racer.STREAM + assert c.mode_state(racer.HOST) == "unsupported" + clock.now += 59 + assert c.mode_state(racer.HOST) == "unsupported" + # The dataplane is upgraded; after the interval host-read is retried. + server.old = False + clock.now += 1 + assert c.mode_state(racer.HOST) is None + assert c.read_into_host(req, buf).delivery == racer.HOST + assert c.mode_state(racer.HOST) == "supported" + + +def test_unsupported_mode_without_expiry(server): + req = add(server) + server.old = True + clock = FakeClock() + c = client(server, mode_retry_seconds=None, clock=clock) + with racer.HostBuffer(len(DATA)) as buf: + c.read_into_host(req, buf) + clock.now += 1e9 + assert c.mode_state(racer.HOST) == "unsupported" + + +def test_negative_retry_rejected(server): + with pytest.raises(racer.RacerError): + client(server, mode_retry_seconds=-1) + + +def test_mode_memory_reset_when_socket_replaced(server, tmp_path): + req = add(server) + server.old = True + c = client(server, mode_retry_seconds=None) + with racer.HostBuffer(len(DATA)) as buf: + c.read_into_host(req, buf) + assert c.mode_state(racer.HOST) == "unsupported" + # A restarted dataplane binds a new socket file at the same path. + os.rename(server.path, str(tmp_path / "moved")) + with open(server.path, "w"): + pass + try: + assert c.mode_state(racer.HOST) is None + assert server.path not in racer.client._mode_state + finally: + os.unlink(server.path) + os.rename(str(tmp_path / "moved"), server.path) + + +def test_fallback_warns_once_per_transition(server, caplog): + req = add(server) + g = FakeGpu(len(DATA)) + clock = FakeClock() + try: + c = client(server, mode_retry_seconds=10, clock=clock) + with caplog.at_level("WARNING", logger="racer.client"): + for _ in range(3): + assert c.read_into_gpu(req, g).delivery == racer.HOST + warns = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warns) == 1 + msg = warns[0].getMessage() + assert "gpu-read" in msg and "status 501" in msg and "host mode" in msg + # After expiry the mode is retried; a second refusal warns again. + clock.now += 10 + caplog.clear() + with caplog.at_level("WARNING", logger="racer.client"): + c.read_into_gpu(req, g) + c.read_into_gpu(req, g) + assert len([r for r in caplog.records if r.levelname == "WARNING"]) == 1 + finally: + g.close() + + +def test_host_to_stream_fallback_warns_with_reason(server, caplog): + req = add(server) + server.old = True + c = client(server) + with racer.HostBuffer(len(DATA)) as buf, caplog.at_level("WARNING", logger="racer.client"): + c.read_into_host(req, buf) + c.read_into_host(req, buf) + warns = [r.getMessage() for r in caplog.records if r.levelname == "WARNING"] + assert len(warns) == 1 + assert "host-read" in warns[0] and "status 400" in warns[0] and "stream mode" in warns[0] diff --git a/clients/python/racer/tests/test_cuda.py b/clients/python/racer/tests/test_cuda.py index 406803abe5..1c8a1ae89b 100644 --- a/clients/python/racer/tests/test_cuda.py +++ b/clients/python/racer/tests/test_cuda.py @@ -66,6 +66,11 @@ def cuDeviceGet(self, out, ordinal): out._obj.value = ordinal return 0 + def cuDeviceGetUuid_v2(self, out, dev): + self.calls.append(("uuid", dev)) + out._obj.bytes[:] = list(range(dev, dev + 16)) + return self._rc("cuDeviceGetUuid_v2") + def cuDevicePrimaryCtxRetain(self, out, dev): out._obj.value = 100 + dev self.calls.append(("retain", dev)) @@ -258,3 +263,22 @@ def test_tensor_target_staging_on_tensor_device(fake, monkeypatch): assert ("register", 100) not in lib.calls finally: buf.close() + + +def test_format_uuid(): + assert cuda.format_uuid(bytes(range(16))) == "GPU-00010203-0405-0607-0809-0a0b0c0d0e0f" + with pytest.raises(ValueError): + cuda.format_uuid(b"short") + + +def test_device_uuid_uses_the_requested_ordinal(fake): + lib = fake() + assert cuda.device_uuid(3) == "GPU-03040506-0708-090a-0b0c-0d0e0f101112" + assert ("uuid", 3) in lib.calls + assert _alloc(2).gpu_uuid() == cuda.format_uuid(bytes(range(2, 18))) + + +def test_device_uuid_error(fake): + fake(fail="cuDeviceGetUuid_v2") + with pytest.raises(cuda.CudaError): + cuda.device_uuid(0) diff --git a/clients/python/racer/tests/test_gpu_e2e.py b/clients/python/racer/tests/test_gpu_e2e.py new file mode 100644 index 0000000000..a6c65b5b55 --- /dev/null +++ b/clients/python/racer/tests/test_gpu_e2e.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 +"""gpu-read against a live dataplane on a real GPU. + +Runs only when RACER_GPU_E2E_SOCKET names a running dataplane's client +socket, RACER_GPU_E2E_KEY is the 64-hex-digit key of an object it can serve +(RACER_GPU_E2E_METADATA and RACER_GPU_E2E_AUTHORIZATION are passed through), +and libcuda can allocate GPU memory. `make racer-gpu-test` runs this. Set RACER_GPU_E2E=require to fail instead of +skipping. +""" +import os + +import pytest + +from racer import Client, Request, cuda + +SOCKET = os.environ.get("RACER_GPU_E2E_SOCKET") +KEY = os.environ.get("RACER_GPU_E2E_KEY") +REQUIRED = os.environ.get("RACER_GPU_E2E") == "require" + + +def _missing() -> str | None: + if not SOCKET or not KEY: + return "set RACER_GPU_E2E_SOCKET and RACER_GPU_E2E_KEY to run" + if not cuda.available(): + return "libcuda is unavailable" + return None + + +@pytest.fixture +def live(): + why = _missing() + if why: + if REQUIRED: + pytest.fail(f"RACER_GPU_E2E=require but {why}") + pytest.skip(why) + request = Request.from_hex( + KEY, + os.environ.get("RACER_GPU_E2E_METADATA"), + os.environ.get("RACER_GPU_E2E_AUTHORIZATION"), + ) + return Client(socket_path=SOCKET), request + + +@pytest.mark.gpu +def test_gpu_read_matches_stream_read(live): + client, request = live + expected = client.read(request) + assert expected, "the object is empty" + offset = 4096 + with cuda.GpuAllocation(offset + len(expected)) as g: + fd = g.export_fd() + try: + result = client.gpu_read(request, fd, g.size, offset, len(expected)) + finally: + os.close(fd) + assert result.length == len(expected) + out = bytearray(len(expected)) + g.copy_to_host(offset, out) + assert out == expected, "GPU bytes differ from the stream read" + + +@pytest.mark.gpu +def test_read_into_gpu_partial_range(live): + client, request = live + expected = client.read(request) + start = len(expected) // 3 + length = max(1, len(expected) // 2) + with cuda.GpuAllocation(length) as g: + client.read_into_gpu(request, g, 0, offset=start, length=length) + out = bytearray(length) + g.copy_to_host(0, out) + assert out == expected[start : start + length] diff --git a/clients/python/racer/tests/test_wire_golden.py b/clients/python/racer/tests/test_wire_golden.py index a8d5a22466..61ab779474 100644 --- a/clients/python/racer/tests/test_wire_golden.py +++ b/clients/python/racer/tests/test_wire_golden.py @@ -96,6 +96,25 @@ def test_gpu_read_request_bytes(): assert len(got["flags"]) == 1 +def test_gpu_read_request_bytes_with_gpu_uuid(): + req = racer.Request(bytes.fromhex(KEY)) + uuid = "GPU-00112233-4455-6677-8899-aabbccddeeff" + with racer.HostBuffer(1) as buf: + got = _capture("gpu", lambda c: c.gpu_read(req, buf.fileno(), 4196, 4096, 100, gpu_uuid=uuid)) + assert got["head"] == ( + f"POST /v1/objects/{KEY}/gpu-read HTTP/1.1\r\nHost: racer\r\nContent-Length: 0\r\n" + "Racer-Destination-Offset: 4096\r\nRacer-Destination-Length: 100\r\nRacer-Allocation-Size: 4196\r\n" + f"Racer-Gpu-Uuid: {uuid}\r\n\r\n" + ).encode() + + +def test_gpu_read_rejects_invalid_gpu_uuid(): + req = racer.Request(bytes.fromhex(KEY)) + with racer.HostBuffer(1) as buf: + with pytest.raises(racer.RacerError, match="UUID"): + racer.Client(socket_path="/nonexistent").gpu_read(req, buf.fileno(), 4196, 4096, 100, gpu_uuid="x\r\nEvil: 1") + + def test_direct_head_from_rust_is_accepted(): # direct_head for a 3-byte range of an object of 2 pages + 5 bytes: # the Get head with Content-Length 21 * (pages + 1) and Racer-Delivery diff --git a/cmd/racer-dataplane/admission/src/lib.rs b/cmd/racer-dataplane/admission/src/lib.rs index db37eb7f75..e9de8b7a14 100644 --- a/cmd/racer-dataplane/admission/src/lib.rs +++ b/cmd/racer-dataplane/admission/src/lib.rs @@ -1,6 +1,6 @@ //! Mutating admission logic for Racer pods. [`ENABLED_LABEL`] injects the //! racer-object sidecar as a native sidecar (a restartable init container). -//! [`CLIENT_SOCKET_LABEL`] mounts the Racer client socket into app containers. +//! [`CLIENT_SOCKET_LABEL`] mounts the Racer client socket into init and app containers. //! The two opt-ins are independent and both apply when a pod sets both labels. //! //! It reads an `admission.k8s.io/v1` `AdmissionReview` and returns the @@ -50,8 +50,8 @@ pub const INJECTED_ANNOTATION: &str = "racer.unbounded-cloud.io/object-sidecar-i pub const CLIENT_SOCKET_LABEL: &str = "racer.unbounded-cloud.io/client-socket"; /// Names the ClusterCache whose client socket is mounted. pub const CLIENT_CACHE_ANNOTATION: &str = "racer.unbounded-cloud.io/client-cache"; -/// Optionally lists, comma-separated, the app containers that get the mount. -/// The default is every app container. +/// Optionally lists, comma-separated, the init and app containers that get the +/// mount. The default is every init and app container except racer-object. pub const CLIENT_CONTAINERS_ANNOTATION: &str = "racer.unbounded-cloud.io/client-containers"; /// Records the mounted cache. pub const CLIENT_SOCKET_INJECTED_ANNOTATION: &str = diff --git a/cmd/racer-dataplane/admission/src/mutate.rs b/cmd/racer-dataplane/admission/src/mutate.rs index 34c720125b..184682ee66 100644 --- a/cmd/racer-dataplane/admission/src/mutate.rs +++ b/cmd/racer-dataplane/admission/src/mutate.rs @@ -132,7 +132,7 @@ fn mutate_sidecar(pod: &mut Value, image: &str) -> Result<(), Reject> { } /// Mounts `/run/racer//client` read-only at the same path in the -/// selected app containers and sets [`SOCKET_ENV`]. A read-only mount still +/// selected init and app containers and sets [`SOCKET_ENV`]. A read-only mount still /// permits connect(2) on the socket: write permission on the socket inode is /// checked, not the mount's read-only flag. Pods without the label, or already /// marked with [`CLIENT_SOCKET_INJECTED_ANNOTATION`], are left unchanged. @@ -161,8 +161,9 @@ fn mutate_client_socket(pod: &mut Value) -> Result<(), Reject> { } let cache = cache.to_owned(); + let init_containers = p.init_containers.map(Vec::as_slice).unwrap_or_default(); let containers = p.containers.map(Vec::as_slice).unwrap_or_default(); - let targets = client_targets(p.annotations, containers)?; + let targets = client_targets(p.annotations, init_containers, containers)?; check_linux(&p, "Racer client socket")?; let client_dir = format!("/run/racer/{cache}/client"); @@ -182,10 +183,15 @@ fn mutate_client_socket(pod: &mut Value) -> Result<(), Reject> { ))); } - for &i in &targets { - let c = containers[i].as_object().expect("checked by Pod::read"); + for &(field, i) in &targets { + let list = if field == "initContainers" { + init_containers + } else { + containers + }; + let c = list[i].as_object().expect("checked by Pod::read"); let name = get_str(Some(c), "name").unwrap_or(""); - let base = format!("spec.containers[{i}]"); + let base = format!("spec.{field}[{i}]"); let mounts = array(c, "volumeMounts", &format!("{base}.volumeMounts"))?; for m in mounts.into_iter().flatten() { let m = m @@ -219,9 +225,10 @@ fn mutate_client_socket(pod: &mut Value) -> Result<(), Reject> { CLIENT_VOLUME_NAME.to_owned() } }; - let containers = child_array(s, "containers"); - for &i in &targets { - let c = containers[i].as_object_mut().expect("checked by Pod::read"); + for &(field, i) in &targets { + let c = child_array(s, field)[i] + .as_object_mut() + .expect("checked by Pod::read"); child_array(c, "volumeMounts") .push(json!({"name": volume, "mountPath": client_dir, "readOnly": true})); child_array(c, "env").push(json!({"name": SOCKET_ENV, "value": socket})); @@ -231,24 +238,47 @@ fn mutate_client_socket(pod: &mut Value) -> Result<(), Reject> { Ok(()) } -/// Returns indexes into `spec.containers`. Ports Go `clientTargets`. +/// A selected container: its pod spec field (`initContainers` or +/// `containers`) and its index there. +type Target = (&'static str, usize); + +/// Returns the containers to mount the client socket into. Without +/// [`CLIENT_CONTAINERS_ANNOTATION`] that is every init and app container, +/// except the injected racer-object sidecar, so an init container can fetch +/// through the socket before the app starts. With it, each listed name must +/// match one init or app container. Ports Go `clientTargets`. fn client_targets( annotations: Option<&Map>, + init_containers: &[Value], containers: &[Value], -) -> Result, Reject> { +) -> Result, Reject> { + let name = |c: &Value| c.get("name").and_then(Value::as_str).map(str::to_owned); + let all = init_containers + .iter() + .enumerate() + .filter(|(_, c)| name(c).as_deref() != Some(CONTAINER_NAME)) + .map(|(i, c)| (("initContainers", i), name(c))) + .chain( + containers + .iter() + .enumerate() + .map(|(i, c)| (("containers", i), name(c))), + ) + .collect::>(); let Some(value) = get_str(annotations, CLIENT_CONTAINERS_ANNOTATION) else { - return Ok((0..containers.len()).collect()); + return Ok(all.into_iter().map(|(t, _)| t).collect()); }; let mut targets = Vec::new(); - for name in value.split(',').map(str::trim) { - let i = containers + for wanted in value.split(',').map(str::trim) { + let found = all .iter() - .position(|c| c.get("name").and_then(Value::as_str) == Some(name)); - match i { - Some(i) if !targets.contains(&i) => targets.push(i), + .find(|(_, n)| n.as_deref() == Some(wanted)) + .map(|(t, _)| *t); + match found { + Some(t) if !targets.contains(&t) => targets.push(t), _ => { return Err(deny(format!( - "annotation {CLIENT_CONTAINERS_ANNOTATION} must list distinct app container names, separated by commas" + "annotation {CLIENT_CONTAINERS_ANNOTATION} must list distinct init or app container names, separated by commas" ))); } } diff --git a/cmd/racer-dataplane/admission/src/tests.rs b/cmd/racer-dataplane/admission/src/tests.rs index 7fa4a5f633..d26f5c1b92 100644 --- a/cmd/racer-dataplane/admission/src/tests.rs +++ b/cmd/racer-dataplane/admission/src/tests.rs @@ -865,11 +865,9 @@ fn client_socket_all_containers() { for c in out["spec"]["containers"].as_array().unwrap() { assert!(socket_mounted(c, CLIENT_VOLUME_NAME), "{c}"); } - assert_eq!( - out["spec"]["initContainers"], - json!([{"name": "setup", "image": "app"}]), - "init containers must be untouched and no sidecar added" - ); + let init = out["spec"]["initContainers"].as_array().unwrap(); + assert_eq!(init.len(), 1, "no sidecar added: {init:?}"); + assert!(socket_mounted(&init[0], CLIENT_VOLUME_NAME), "{init:?}"); assert_eq!( out["spec"]["volumes"], json!([{"name": CLIENT_VOLUME_NAME, "hostPath": {"path": CLIENT_DIR, "type": "Directory"}}]) @@ -898,6 +896,49 @@ fn client_socket_selected_containers() { assert!(socket_mounted(&c[1], CLIENT_VOLUME_NAME)); } +#[test] +fn client_socket_selected_init_container() { + let out = run(&client_pod( + json!({CLIENT_CONTAINERS_ANNOTATION: "setup, helper"}), + )) + .unwrap(); + assert!(socket_mounted( + &out["spec"]["initContainers"][0], + CLIENT_VOLUME_NAME + )); + let c = out["spec"]["containers"].as_array().unwrap(); + assert!(!socket_mounted(&c[0], CLIENT_VOLUME_NAME)); + assert!(socket_mounted(&c[1], CLIENT_VOLUME_NAME)); +} + +#[test] +fn client_socket_native_sidecar_init_container() { + let mut p = client_pod(json!({})); + p["spec"]["initContainers"] = json!([ + {"name": "setup", "image": "app"}, + {"name": "agent", "image": "app", "restartPolicy": "Always"}, + ]); + let out = run(&p).unwrap(); + for c in out["spec"]["initContainers"].as_array().unwrap() { + assert!(socket_mounted(c, CLIENT_VOLUME_NAME), "{c}"); + assert_eq!(c["image"], "app"); + } +} + +#[test] +fn client_socket_skips_object_sidecar() { + let mut p = pod(json!({CLIENT_CACHE_ANNOTATION: CACHE})); + p["metadata"]["labels"][CLIENT_SOCKET_LABEL] = json!(ENABLED_VALUE); + let out = run(&p).unwrap(); + let sidecar = &out["spec"]["initContainers"][0]; + assert_eq!(sidecar["name"], CONTAINER_NAME); + let env = sidecar["env"].as_array().cloned().unwrap_or_default(); + assert!( + !env.iter().any(|e| e["name"] == SOCKET_ENV), + "sidecar must not get {SOCKET_ENV}: {sidecar}" + ); +} + #[test] fn client_socket_with_sidecar_shares_volume() { let mut p = pod(json!({CLIENT_CACHE_ANNOTATION: CACHE})); @@ -964,9 +1005,22 @@ fn client_socket_denials() { CLIENT_CONTAINERS_ANNOTATION, ), ( - with(json!({CLIENT_CONTAINERS_ANNOTATION: "setup"}), &none), + with(json!({CLIENT_CONTAINERS_ANNOTATION: "setup,setup"}), &none), CLIENT_CONTAINERS_ANNOTATION, ), + ( + with( + json!({CLIENT_CONTAINERS_ANNOTATION: CONTAINER_NAME}), + &|p| p["spec"]["initContainers"] = json!([{"name": CONTAINER_NAME, "image": "x"}]), + ), + CLIENT_CONTAINERS_ANNOTATION, + ), + ( + with(json!({}), &|p| { + p["spec"]["initContainers"][0]["env"] = json!([{"name": SOCKET_ENV, "value": "x"}]) + }), + SOCKET_ENV, + ), ( with(json!({CLIENT_CONTAINERS_ANNOTATION: "app,app"}), &none), CLIENT_CONTAINERS_ANNOTATION, diff --git a/cmd/racer-dataplane/gpu/native/cuda_private.h b/cmd/racer-dataplane/gpu/native/cuda_private.h index cf85676cd4..34432075ec 100644 --- a/cmd/racer-dataplane/gpu/native/cuda_private.h +++ b/cmd/racer-dataplane/gpu/native/cuda_private.h @@ -170,6 +170,11 @@ typedef struct { } CUfileDescr_t; #define CU_FILE_SUCCESS 0 +/* CUfileOpError values that mean cuFile cannot use one file, its + * filesystem, or its device. Other files may still work. */ +#define CU_FILE_IO_NOT_SUPPORTED 5008 +#define CU_FILE_DEVICE_NOT_SUPPORTED 5009 +#define CU_FILE_INVALID_FILE_TYPE 5018 #define CU_FILE_HANDLE_TYPE_OPAQUE_FD 1 _Static_assert(sizeof(CUfileError_t) == 8, "CUfileError_t layout"); @@ -178,6 +183,7 @@ _Static_assert(offsetof(CUfileDescr_t, handle) == 8, "CUfileDescr_t layout"); /* cuFileSetParameterBool (cuFile 1.10+); resolved optionally. */ #define CUFILE_PARAM_PROPERTIES_ALLOW_COMPAT_MODE 1 +#define CUFILE_PARAM_USE_PCIP2PDMA 7 typedef CUfileError_t (*cuFileSetParameterBool_fn)(int param, bool value); #define RACER_CUFILE_FUNCTIONS(X) \ diff --git a/cmd/racer-dataplane/gpu/native/gpu_cuda.c b/cmd/racer-dataplane/gpu/native/gpu_cuda.c index fcbe1066a6..a7212d2874 100644 --- a/cmd/racer-dataplane/gpu/native/gpu_cuda.c +++ b/cmd/racer-dataplane/gpu/native/gpu_cuda.c @@ -20,6 +20,7 @@ #include #include #include +#include #include #ifdef RACER_GPU_PTX @@ -372,6 +373,18 @@ static int import_release(struct gpu_import *a) { return RACER_GPU_OK; } +/* Whether device `ordinal` in this process is the context's GPU. */ +static int same_device(struct gpu_context *c, int ordinal) { + CUdevice dev; + CUuuid mine, theirs; + if (ordinal < 0 || cu.cuDeviceGet(&dev, ordinal) != CUDA_SUCCESS) return 0; + if (dev == c->dev) return 1; + if (cu.cuDeviceGetUuid(&theirs, dev) != CUDA_SUCCESS || + cu.cuDeviceGetUuid(&mine, c->dev) != CUDA_SUCCESS) + return 0; + return !memcmp(mine.bytes, theirs.bytes, sizeof(mine.bytes)); +} + int racer_gpu_import(void *context, int fd, uint64_t size, void **out) { struct gpu_context *c = context; if (!c || fd < 0 || !size || size > SIZE_MAX || !out) return RACER_GPU_EINVAL; @@ -395,7 +408,15 @@ int racer_gpu_import(void *context, int fd, uint64_t size, void **out) { a->has_handle = 1; rc = drv(cu.cuMemGetAllocationPropertiesFromHandle(&prop, a->handle)); if (rc) goto fail; - if (prop.location.type != CU_MEM_LOCATION_TYPE_DEVICE || prop.location.id != c->ordinal) { + if (prop.location.type != CU_MEM_LOCATION_TYPE_DEVICE) { + rc = RACER_GPU_EDEVICE; + goto fail; + } + /* The driver reports the location in this process's device numbering, + * which can differ from the exporter's (CUDA_VISIBLE_DEVICES). Compare + * physical GPUs by UUID, not ordinals, so a stale or foreign number can + * never pass. A number this process cannot resolve is another GPU. */ + if (!same_device(c, prop.location.id)) { rc = RACER_GPU_EDEVICE; goto fail; } @@ -635,7 +656,7 @@ static int file_error(CUfileError_t e) { * for P2PDMA, which can take most of a second, so a failure is not retried. */ static int cf_failed; -int racer_gpu_file_open(void) { +int racer_gpu_file_open(uint32_t flags) { int rc = RACER_GPU_OK; pthread_mutex_lock(&load_lock); if (cf_failed) { @@ -653,13 +674,38 @@ int racer_gpu_file_open(void) { /* Without nvidia-fs or P2PDMA, cuFile's compatibility mode does a * POSIX read into a bounce buffer and a copy on this helper thread, * so every GPU job waits behind the disk. Racer's own host path does - * that work on I/O threads. Turn it off unless the operator chose a - * value with CUFILE_ALLOW_COMPAT_MODE; cuFileDriverOpen then fails - * without GPUDirect Storage and Racer uses its host path. */ + * that work on I/O threads. Turn it off unless CUFILE_ALLOW_COMPAT_MODE + * is true, 1 or yes; then turn it on, since cufile.json turns it off. + * With it off, cuFileDriverOpen fails without GPUDirect Storage and + * Racer uses its host path. */ cuFileSetParameterBool_fn set_bool = (cuFileSetParameterBool_fn)dlsym(table.lib, "cuFileSetParameterBool"); - if (set_bool && !getenv("CUFILE_ALLOW_COMPAT_MODE")) - set_bool(CUFILE_PARAM_PROPERTIES_ALLOW_COMPAT_MODE, false); + if (set_bool) { + const char *compat = getenv("CUFILE_ALLOW_COMPAT_MODE"); + bool allow = compat && (!strcasecmp(compat, "true") || + !strcmp(compat, "1") || + !strcasecmp(compat, "yes")); + set_bool(CUFILE_PARAM_PROPERTIES_ALLOW_COMPAT_MODE, allow); + } + /* PCI P2PDMA (cuFile 1.13+) reads NVMe into GPU memory without + * nvidia-fs. Only turn it on: leave cufile.json's value otherwise. + * Without cuFileSetParameterBool, or if cuFile rejects the value, + * fail the open so Racer reports cuFile off and uses its host path + * rather than run without the path the operator asked for. */ + if (flags & RACER_GPU_FILE_PCI_P2PDMA) { + if (!set_bool) { + dlclose(table.lib); + cf_failed = 1; + rc = RACER_GPU_EUNAVAILABLE; + goto out; + } + if (file_error(set_bool(CUFILE_PARAM_USE_PCIP2PDMA, true))) { + dlclose(table.lib); + cf_failed = 1; + rc = RACER_GPU_EUNAVAILABLE; + goto out; + } + } if (file_error(table.cuFileDriverOpen())) { dlclose(table.lib); cf_failed = 1; @@ -690,7 +736,11 @@ int racer_gpu_file(int fd, void **out) { memset(&descr, 0, sizeof(descr)); descr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; descr.handle.fd = fd; - int rc = file_error(cf.cuFileHandleRegister(&f->handle, &descr)); + CUfileError_t e = cf.cuFileHandleRegister(&f->handle, &descr); + int rc = file_error(e); + if (rc && (e.err == CU_FILE_IO_NOT_SUPPORTED || e.err == CU_FILE_DEVICE_NOT_SUPPORTED || + e.err == CU_FILE_INVALID_FILE_TYPE)) + rc = RACER_GPU_EUNSUPPORTED; if (rc) { free(f); return rc; diff --git a/cmd/racer-dataplane/gpu/native/gpu_cuda.h b/cmd/racer-dataplane/gpu/native/gpu_cuda.h index 220d6a1bb0..4bc805926b 100644 --- a/cmd/racer-dataplane/gpu/native/gpu_cuda.h +++ b/cmd/racer-dataplane/gpu/native/gpu_cuda.h @@ -16,7 +16,7 @@ #include #include -#define RACER_GPU_ABI_VERSION 1u +#define RACER_GPU_ABI_VERSION 2u #define RACER_GPU_OK 0 #define RACER_GPU_EINVAL (-1) /* bad argument or handle type */ @@ -29,6 +29,7 @@ #define RACER_GPU_EIO (-8) /* cuFile read failed or was short */ #define RACER_GPU_EBUSY (-9) /* wrong state, for example no verdict to read */ #define RACER_GPU_ELEAKED (-10) /* failed, and cleanup leaked an internal resource */ +#define RACER_GPU_EUNSUPPORTED (-11) /* cuFile does not support this file or its device */ /* Device capability bits in racer_gpu_device.flags. */ #define RACER_GPU_FLAG_VMM 1u /* virtual memory management */ @@ -137,8 +138,12 @@ int racer_gpu_decrypt(void *stream, const struct racer_gpu_decrypt *params, void uint64_t src_offset, void *dst, uint64_t dst_offset); int racer_gpu_verdict(void *stream, struct racer_gpu_verdict *out); -/* Open the cuFile driver (libcufile.so.0). Safe to call more than once. */ -int racer_gpu_file_open(void); +/* racer_gpu_file_open flags. */ +#define RACER_GPU_FILE_PCI_P2PDMA 1u /* set cuFile use_pci_p2pdma before the driver opens */ + +/* Open the cuFile driver (libcufile.so.0). Safe to call more than once; + * `flags` take effect only on the call that opens the driver. */ +int racer_gpu_file_open(uint32_t flags); /* Register an O_DIRECT fd with cuFile. The fd stays owned by the caller. */ int racer_gpu_file(int fd, void **out); int racer_gpu_file_free(void *file); diff --git a/cmd/racer-dataplane/gpu/src/ffi.rs b/cmd/racer-dataplane/gpu/src/ffi.rs index 901ec043af..59559ca044 100644 --- a/cmd/racer-dataplane/gpu/src/ffi.rs +++ b/cmd/racer-dataplane/gpu/src/ffi.rs @@ -15,7 +15,7 @@ pub(crate) use crate::gcm::RawVerdict; /// `RACER_GPU_ABI_VERSION`. #[cfg(all(feature = "native", target_os = "linux"))] -pub(crate) const ABI_VERSION: u32 = 1; +pub(crate) const ABI_VERSION: u32 = 2; /// Result codes from `gpu_cuda.h`. pub(crate) mod code { @@ -32,6 +32,7 @@ pub(crate) mod code { pub const EIO: c_int = -8; pub const EBUSY: c_int = -9; pub const ELEAKED: c_int = -10; + pub const EUNSUPPORTED: c_int = -11; } /// Device capability bits (`RACER_GPU_FLAG_*`). @@ -46,6 +47,12 @@ pub mod flag { pub const GDR: u32 = 8; } +/// `racer_gpu_file_open` flags (`RACER_GPU_FILE_*`). +pub(crate) mod file_flag { + /// Set cuFile's `use_pci_p2pdma` before the driver opens. + pub const PCI_P2PDMA: u32 = 1; +} + /// One device. Matches `struct racer_gpu_device`. #[repr(C)] #[derive(Clone, Copy)] @@ -192,7 +199,7 @@ native_api! { verdict: unsafe extern "C" fn(*mut c_void, *mut RawVerdict) -> c_int, - file_open: unsafe extern "C" fn() -> c_int, + file_open: unsafe extern "C" fn(u32) -> c_int, file: unsafe extern "C" fn(c_int, Out) -> c_int, @@ -283,6 +290,7 @@ pub(crate) fn check(rc: c_int, last_driver_error: unsafe extern "C" fn() -> i32) code::EIO => Err(Error::Io), code::EBUSY => Err(Error::Busy), code::ELEAKED => Err(Error::Leaked), + code::EUNSUPPORTED => Err(Error::Unsupported(unsafe { last_driver_error() })), _ => Err(Error::Driver(rc)), } } @@ -292,6 +300,9 @@ pub(crate) struct Shared { pub(crate) api: Rc, pub(crate) leaks: Cell, + + /// `racer_gpu_file_open` flags. + pub(crate) file_flags: Cell, } impl Shared { diff --git a/cmd/racer-dataplane/gpu/src/lib.rs b/cmd/racer-dataplane/gpu/src/lib.rs index 39b3e9a548..58d2db26c6 100644 --- a/cmd/racer-dataplane/gpu/src/lib.rs +++ b/cmd/racer-dataplane/gpu/src/lib.rs @@ -36,6 +36,7 @@ //! - Neither: [`Gpu::load`] returns [`Error::Unavailable`]. //! //! Hardware-free tests: `cargo test -p racer-gpu --features simulation`. +//! Tests on a real GPU: `make racer-gpu-test` (`tests/hardware.rs`). mod ffi; pub mod gcm; @@ -89,6 +90,10 @@ pub enum Error { /// The call failed and a resource was leaked on purpose. Leaked, + + /// cuFile does not support this file, its filesystem, or its device, + /// with this cuFile status. Files on other devices may still work. + Unsupported(i32), } impl std::fmt::Display for Error { @@ -96,6 +101,9 @@ impl std::fmt::Display for Error { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Driver(code) => write!(f, "GPU driver call failed with status {code}"), + Self::Unsupported(code) => { + write!(f, "cuFile does not support the file (status {code})") + } other => write!(f, "{other:?}"), } } @@ -166,15 +174,7 @@ impl DeviceInfo { /// The UUID in `GPU-xxxxxxxx-xxxx-...` form, as `nvidia-smi` prints it. pub fn uuid_string(&self) -> String { - let hex: String = self.uuid.iter().map(|b| format!("{b:02x}")).collect(); - format!( - "GPU-{}-{}-{}-{}-{}", - &hex[..8], - &hex[8..12], - &hex[12..16], - &hex[16..20], - &hex[20..] - ) + uuid_string(&self.uuid) } /// From the C record. @@ -218,6 +218,7 @@ impl Gpu { let shared = Rc::new(Shared { api, leaks: Cell::new(0), + file_flags: Cell::new(0), }); shared.check(unsafe { (shared.api.init)() })?; Ok(Self { shared }) @@ -258,7 +259,7 @@ impl Gpu { /// Open the cuFile driver if needed and register an `O_DIRECT` file. pub fn open_file(&self, fd: OwnedFd) -> Result { let shared = &self.shared; - shared.check(unsafe { (shared.api.file_open)() })?; + shared.check(unsafe { (shared.api.file_open)(shared.file_flags.get()) })?; let fd = cufile_descriptor(fd); let raw = create(shared, |out| unsafe { (shared.api.file)(fd.as_raw_fd(), out) @@ -275,12 +276,27 @@ impl Gpu { self.shared.leaks.get() } + /// Ask cuFile to use PCI P2PDMA (`use_pci_p2pdma`) instead of + /// nvidia-fs. Takes effect only if set before the cuFile driver opens, + /// and applies to [`FileReader`]s made after this call. When on, a + /// cuFile that cannot set it fails to open. + pub fn set_pci_p2pdma(&self, enabled: bool) { + let flags = self.shared.file_flags.get() & !ffi::file_flag::PCI_P2PDMA; + let bit = if enabled { + ffi::file_flag::PCI_P2PDMA + } else { + 0 + }; + self.shared.file_flags.set(flags | bit); + } + /// A handle that runs [`FileRead`]s on any thread. It does not keep the /// backend loaded: drop every reader thread's work before this [`Gpu`] /// and the objects the reads name. pub fn file_reader(&self) -> FileReader { FileReader { file_open: self.shared.api.file_open, + file_flags: self.shared.file_flags.get(), file_read: self.shared.api.file_read, last_driver_error: self.shared.api.last_driver_error, #[cfg(any(test, feature = "simulation"))] @@ -292,7 +308,9 @@ impl Gpu { /// Runs cuFile reads off the helper thread. Cheap to clone. #[derive(Clone)] pub struct FileReader { - file_open: unsafe extern "C" fn() -> c_int, + file_open: unsafe extern "C" fn(u32) -> c_int, + + file_flags: u32, file_read: unsafe extern "C" fn(*mut c_void, *mut c_void, u64, u64, u64) -> i64, @@ -310,7 +328,10 @@ impl FileReader { pub fn open_driver(&self) -> Result<()> { #[cfg(any(test, feature = "simulation"))] let _environment = self.enter(); - ffi::check(unsafe { (self.file_open)() }, self.last_driver_error) + ffi::check( + unsafe { (self.file_open)(self.file_flags) }, + self.last_driver_error, + ) } /// Run `read` and wait for it. Errors are those of [`GpuFile::read`]. @@ -551,7 +572,7 @@ impl StagingBuffer { return Ok(()); } let shared = &self.0.context.0.shared; - shared.check(unsafe { (shared.api.file_open)() })?; + shared.check(unsafe { (shared.api.file_open)(shared.file_flags.get()) })?; shared.check(unsafe { (shared.api.staging_register)(self.0.raw.as_ptr()) })?; self.0.registered.set(true); Ok(()) @@ -643,6 +664,19 @@ pub enum Verdict { NotRun, } +/// A GPU UUID in `GPU-xxxxxxxx-xxxx-...` form, as `nvidia-smi` prints it. +pub fn uuid_string(uuid: &[u8; 16]) -> String { + let hex: String = uuid.iter().map(|b| format!("{b:02x}")).collect(); + format!( + "GPU-{}-{}-{}-{}-{}", + &hex[..8], + &hex[8..12], + &hex[12..16], + &hex[16..20], + &hex[20..] + ) +} + /// A CUDA stream. Work is visible only after [`Stream::synchronize`]. pub struct Stream { context: Context, diff --git a/cmd/racer-dataplane/gpu/src/simulation.rs b/cmd/racer-dataplane/gpu/src/simulation.rs index d35b2996ae..e212eb0b62 100644 --- a/cmd/racer-dataplane/gpu/src/simulation.rs +++ b/cmd/racer-dataplane/gpu/src/simulation.rs @@ -267,6 +267,9 @@ struct World { file_driver: bool, + /// Whether the open that set `file_driver` asked for PCI P2PDMA. + pci_p2pdma: bool, + no_kernels: bool, last_error: i32, @@ -292,6 +295,13 @@ impl Simulation { } } + /// Whether the cuFile driver is open with PCI P2PDMA on. `None` before + /// it opens. + pub fn pci_p2pdma(&self) -> Option { + let w = self.world(); + w.file_driver.then_some(w.pci_p2pdma) + } + /// Make this the simulation [`crate::Gpu::load`] uses on this thread. pub fn enter(&self) -> Environment { let previous = CURRENT.with(|current| current.borrow_mut().replace(self.clone())); @@ -372,6 +382,9 @@ fn with + Copy + From>( if rc == code::EDRIVER { world.last_error = 999; } + if rc == code::EUNSUPPORTED { + world.last_error = 5008; // CU_FILE_IO_NOT_SUPPORTED + } (R::from(rc), 0) } None => f(&mut world), @@ -794,8 +807,12 @@ unsafe extern "C" fn verdict(stream: *mut c_void, out: *mut RawVerdict) -> c_int }) } -unsafe extern "C" fn file_open() -> c_int { +unsafe extern "C" fn file_open(flags: u32) -> c_int { with(Operation::FileOpen, |w| { + // As in the shim, flags apply only to the open that loads cuFile. + if !w.file_driver { + w.pci_p2pdma = flags & ffi::file_flag::PCI_P2PDMA != 0; + } w.file_driver = true; (code::OK, 0) }) diff --git a/cmd/racer-dataplane/gpu/tests/hardware.rs b/cmd/racer-dataplane/gpu/tests/hardware.rs new file mode 100644 index 0000000000..414a29d329 --- /dev/null +++ b/cmd/racer-dataplane/gpu/tests/hardware.rs @@ -0,0 +1,581 @@ +//! Tests on a real GPU through `libracer_gpu_cuda.so.1`. +//! +//! Opt in with `RACER_GPU_TESTS=1`; `make racer-gpu-test` builds the shim +//! with its kernels and runs these. Without the variable every test returns +//! early, so the normal CI build compiles them and they pass as no-ops. With +//! `RACER_GPU_TESTS=1` a missing shim, driver or GPU is reported as a skip; +//! with `RACER_GPU_TESTS=require` it fails instead, for dedicated GPU runners. +//! cuFile tests skip when cuFile cannot open unless `RACER_GPU_TESTS_CUFILE` +//! is `require`. +#![cfg(all(feature = "native", target_os = "linux"))] + +use racer_gpu::{ + ComputeMode, Context, DecryptRequest, Destination, Error, Gpu, StagingBuffer, Verdict, flag, +}; +use std::{ + ffi::{CStr, c_int, c_uint, c_void}, + fs::OpenOptions, + io::Write, + os::{ + fd::{FromRawFd, OwnedFd}, + unix::fs::OpenOptionsExt, + }, +}; + +/// Whether hardware is mandatory for this run. +fn required() -> bool { + std::env::var("RACER_GPU_TESTS").as_deref() == Ok("require") +} + +/// Skip, or fail when hardware is required. +fn skip(test: &str, why: impl std::fmt::Display) { + assert!(!required(), "{test}: RACER_GPU_TESTS=require but {why}"); + eprintln!("SKIP {test}: {why}"); +} + +/// Load the backend and pick the first usable GPU, or skip. +fn setup(test: &str) -> Option<(Gpu, i32)> { + match std::env::var("RACER_GPU_TESTS").as_deref() { + Ok("1") | Ok("require") => {} + _ => { + eprintln!("SKIP {test}: set RACER_GPU_TESTS=1 to run GPU hardware tests"); + return None; + } + } + let gpu = match Gpu::load() { + Ok(gpu) => gpu, + Err(err) => { + skip(test, format_args!("no GPU backend ({err})")); + return None; + } + }; + match gpu.usable_devices() { + Ok(devices) if !devices.is_empty() => { + let ordinal = devices[0].ordinal; + Some((gpu, ordinal)) + } + Ok(_) => { + skip(test, "no usable GPU"); + None + } + Err(err) => { + skip(test, format_args!("device discovery failed ({err})")); + None + } + } +} + +/// Deterministic xorshift bytes. +fn fill(seed: u64, out: &mut [u8]) { + let mut state = seed | 1; + for chunk in out.chunks_mut(8) { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + chunk.copy_from_slice(&state.to_le_bytes()[..chunk.len()]); + } +} + +/// A page sealed by `racer-crypto`, as the dataplane stores it. +struct Page { + key: [u8; 32], + nonce: [u8; 24], + aad: Vec, + plaintext: Vec, + sealed: Vec, +} + +impl Page { + fn new(seed: u64, length: usize) -> Self { + let mut key = [0; 32]; + let mut nonce = [0; 24]; + let mut aad = vec![0; 97]; + let mut plaintext = vec![0; length]; + fill(seed, &mut key); + fill(seed + 1, &mut nonce); + fill(seed + 2, &mut aad); + fill(seed + 3, &mut plaintext); + let mut sealed = vec![0; length + racer_crypto::TAG_LEN]; + racer_crypto::seal(&key, &nonce, &aad, &plaintext, &mut sealed).unwrap(); + Self { + key, + nonce, + aad, + plaintext, + sealed, + } + } + + /// The CPU decrypt, the reference the GPU must match. + fn cpu_open(&self) -> Vec { + let mut out = vec![0; self.plaintext.len()]; + racer_crypto::open(&self.key, &self.nonce, &self.aad, &self.sealed, &mut out).unwrap(); + out + } + + /// Decrypt `sealed` (already at `src` offset 0) into `dst` and return the verdict. + fn gpu_decrypt( + &self, + context: &Context, + src: &StagingBuffer, + dst: Destination<'_>, + dst_offset: u64, + expected_crc: Option, + ) -> Verdict { + let derived = racer_crypto::xaes_derive(&self.key, &self.nonce); + let stream = context.stream().unwrap(); + let request = DecryptRequest { + key: derived.key(), + nonce: derived.nonce(), + aad: &self.aad, + length: self.plaintext.len() as u64, + expected_crc, + }; + match stream.decrypt(&request, src, 0, dst, dst_offset) { + Err(Error::KernelUnavailable) => { + panic!( + "the shim has no decrypt kernels: build it with nvcc (RACER_GPU_REQUIRE_PTX=1)" + ) + } + other => other.unwrap(), + } + stream.synchronize().unwrap(); + stream.verdict().unwrap() + } +} + +#[test] +/// The shim loads and CUDA initializes. +fn shim_loads_and_initializes() { + let Some((gpu, _)) = setup("shim_loads_and_initializes") else { + return; + }; + assert_eq!(gpu.leaked(), 0); +} + +#[test] +/// Discovery reports sane ordinals, UUIDs, PCI addresses and compute modes. +fn devices_are_enumerated() { + let Some((gpu, _)) = setup("devices_are_enumerated") else { + return; + }; + let devices = gpu.devices().unwrap(); + assert!(!devices.is_empty()); + let mut uuids = std::collections::HashSet::new(); + for (index, device) in devices.iter().enumerate() { + eprintln!( + "GPU {} {} pci={} numa={:?} mode={:?} flags={:#x} memory={}", + device.ordinal, + device.uuid_string(), + device.pci_bus_id, + device.numa_node, + device.compute_mode, + device.flags, + device.total_memory + ); + assert_eq!(device.ordinal, index as i32); + assert_ne!(device.uuid, [0; 16], "GPU {index} has no UUID"); + assert!(uuids.insert(device.uuid), "duplicate UUID"); + assert_eq!(device.pci_bus_id.len(), 12, "{}", device.pci_bus_id); + assert_eq!(device.pci_bus_id, device.pci_bus_id.to_lowercase()); + assert!(device.total_memory > 0); + assert!(!matches!(device.compute_mode, ComputeMode::Other(_))); + assert_eq!( + device.usable(), + device.compute_mode == ComputeMode::Default + && device.flags & (flag::VMM | flag::POSIX_FD) == flag::VMM | flag::POSIX_FD + ); + } + assert!(gpu.context(devices.len() as i32).is_err()); +} + +#[test] +/// A client allocation exported as a POSIX fd imports, accepts copies and +/// decrypts inside its window, and rejects writes outside it. +fn exported_allocation_imports_and_receives_decrypt() { + let Some((gpu, ordinal)) = setup("exported_allocation_imports_and_receives_decrypt") else { + return; + }; + let cuda = match Cuda::load() { + Ok(cuda) => cuda, + Err(why) => return skip("exported_allocation_imports_and_receives_decrypt", why), + }; + let page = Page::new(11, 70_001); + let allocation = cuda.allocate(ordinal, 4 << 20); + let size = allocation.size; + let window = 4096..4096 + page.plaintext.len() as u64; + + let context = gpu.context(ordinal).unwrap(); + let import = context + .import(allocation.export(), size, window.clone()) + .unwrap(); + assert_eq!(import.size(), size); + assert_eq!(import.window(), window); + + let staging = context.staging(page.sealed.len() as u64).unwrap(); + let stream = context.stream().unwrap(); + stream + .copy_from_host(Destination::Staging(&staging), 0, page.sealed.clone()) + .unwrap(); + assert_eq!( + stream.copy_from_host(Destination::Import(&import), 0, vec![1u8; 16]), + Err(Error::InvalidRange) + ); + stream.synchronize().unwrap(); + + let crc = racer_crypto::crc64(&page.sealed); + let verdict = page.gpu_decrypt( + &context, + &staging, + Destination::Import(&import), + window.start, + Some(crc), + ); + assert_eq!(verdict, Verdict::Passed { crc }); + let mut out = vec![0; page.plaintext.len()]; + allocation.read(window.start, &mut out); + assert!( + out == page.cpu_open(), + "client allocation differs from CPU decrypt" + ); + drop(import); + drop(context); + assert_eq!(gpu.leaked(), 0); +} + +#[test] +/// GPU decrypt matches the CPU decrypt for several lengths, and a tampered +/// page is rejected without writing. +fn gpu_decrypt_matches_cpu_decrypt() { + let Some((gpu, ordinal)) = setup("gpu_decrypt_matches_cpu_decrypt") else { + return; + }; + let context = gpu.context(ordinal).unwrap(); + for (seed, length) in [ + (1u64, 0usize), + (2, 1), + (3, 4097), + (4, (1 << 20) + 13), + (5, 16 << 20), + ] { + let page = Page::new(seed, length); + let sealed_len = page.sealed.len() as u64; + let src = context.staging(sealed_len).unwrap(); + let dst = context.staging(length.max(1) as u64).unwrap(); + let stream = context.stream().unwrap(); + stream + .copy_from_host(Destination::Staging(&src), 0, page.sealed.clone()) + .unwrap(); + stream + .copy_from_host(Destination::Staging(&dst), 0, vec![0xa5u8; length.max(1)]) + .unwrap(); + stream.synchronize().unwrap(); + + let crc = racer_crypto::crc64(&page.sealed); + let verdict = page.gpu_decrypt(&context, &src, Destination::Staging(&dst), 0, Some(crc)); + assert_eq!(verdict, Verdict::Passed { crc }, "length {length}"); + let verdict = page.gpu_decrypt(&context, &src, Destination::Staging(&dst), 0, None); + assert_eq!(verdict, Verdict::Passed { crc }, "length {length}"); + if length > 0 { + let mut out = vec![0; length]; + stream.read_staging(&dst, 0, &mut out).unwrap(); + assert!( + out == page.cpu_open(), + "length {length}: GPU differs from CPU" + ); + + // Flip one ciphertext bit: tag check fails, destination untouched. + let mut tampered = page.sealed.clone(); + tampered[length / 2] ^= 0x10; + stream + .copy_from_host(Destination::Staging(&src), 0, tampered) + .unwrap(); + stream + .copy_from_host(Destination::Staging(&dst), 0, vec![0xa5u8; length]) + .unwrap(); + stream.synchronize().unwrap(); + let verdict = page.gpu_decrypt(&context, &src, Destination::Staging(&dst), 0, None); + assert!( + matches!(verdict, Verdict::TagMismatch { .. }), + "length {length}: {verdict:?}" + ); + let verdict = + page.gpu_decrypt(&context, &src, Destination::Staging(&dst), 0, Some(crc)); + assert!( + matches!(verdict, Verdict::CrcMismatch { .. }), + "length {length}: {verdict:?}" + ); + stream.read_staging(&dst, 0, &mut out).unwrap(); + assert!( + out.iter().all(|&b| b == 0xa5), + "length {length}: tampered page was written" + ); + } + } + assert_eq!(gpu.leaked(), 0); +} + +#[test] +/// cuFile reads a sealed page from an `O_DIRECT` file into staging memory, +/// and the GPU decrypts it to the CPU result. +fn cufile_reads_sealed_page_into_gpu() { + const TEST: &str = "cufile_reads_sealed_page_into_gpu"; + let Some((gpu, ordinal)) = setup(TEST) else { + return; + }; + let cufile_required = std::env::var("RACER_GPU_TESTS_CUFILE").as_deref() == Ok("require"); + let skip_cufile = |why: String| { + assert!( + !cufile_required, + "{TEST}: RACER_GPU_TESTS_CUFILE=require but {why}" + ); + eprintln!("SKIP {TEST}: {why}"); + }; + + // A 64 KiB sealed page at a 4 KiB-aligned record offset, as in a slab. + let page = Page::new(21, (64 << 10) - racer_crypto::TAG_LEN); + let record_offset = 8192u64; + let dir = std::env::var_os("RACER_GPU_TEST_DIR") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| env!("CARGO_TARGET_TMPDIR").into()); + let path = dir.join(format!("racer-gpu-cufile-{}.slab", std::process::id())); + { + let mut file = std::fs::File::create(&path).unwrap(); + file.write_all(&vec![0u8; record_offset as usize]).unwrap(); + file.write_all(&page.sealed).unwrap(); + file.write_all(&vec![0u8; 4096]).unwrap(); + file.sync_all().unwrap(); + } + let opened = OpenOptions::new() + .read(true) + .custom_flags(libc::O_DIRECT) + .open(&path); + let _ = std::fs::remove_file(&path); + let file = match opened { + Ok(file) => OwnedFd::from(file), + Err(err) => { + return skip_cufile(format!( + "O_DIRECT open of {} failed ({err}); set RACER_GPU_TEST_DIR", + dir.display() + )); + } + }; + let gpu_file = match gpu.open_file(file) { + Ok(file) => file, + Err(err) => return skip_cufile(format!("cuFile is unavailable ({err})")), + }; + + let context = gpu.context(ordinal).unwrap(); + let length = page.sealed.len() as u64; + let staging = context.staging(length).unwrap(); + staging.register_cufile().unwrap(); + gpu_file.read(&staging, 0, record_offset, length).unwrap(); + let stream = context.stream().unwrap(); + let mut readback = vec![0; page.sealed.len()]; + stream.read_staging(&staging, 0, &mut readback).unwrap(); + assert!(readback == page.sealed, "cuFile read the wrong bytes"); + + // A detached read on another thread lands the same bytes. + let reader = gpu.file_reader(); + let other = context.staging(length).unwrap(); + let read = unsafe { gpu_file.detach(&other, 0, record_offset, length) }.unwrap(); + std::thread::scope(|scope| scope.spawn(|| reader.read(read)).join().unwrap()).unwrap(); + stream.read_staging(&other, 0, &mut readback).unwrap(); + assert!( + readback == page.sealed, + "detached cuFile read the wrong bytes" + ); + + let dst = context.staging(page.plaintext.len() as u64).unwrap(); + let crc = racer_crypto::crc64(&page.sealed); + let verdict = page.gpu_decrypt(&context, &staging, Destination::Staging(&dst), 0, Some(crc)); + assert_eq!(verdict, Verdict::Passed { crc }); + let mut out = vec![0; page.plaintext.len()]; + stream.read_staging(&dst, 0, &mut out).unwrap(); + assert!(out == page.cpu_open()); + assert_eq!( + gpu_file.read(&staging, 0, record_offset, length + 1), + Err(Error::InvalidRange) + ); +} + +/// The few CUDA driver calls a client makes to export an allocation, loaded +/// from `libcuda.so.1`, so the import path is tested against real exports. +struct Cuda { + init: unsafe extern "C" fn(c_uint) -> c_int, + device_get: unsafe extern "C" fn(*mut c_int, c_int) -> c_int, + retain: unsafe extern "C" fn(*mut *mut c_void, c_int) -> c_int, + set_current: unsafe extern "C" fn(*mut c_void) -> c_int, + granularity: unsafe extern "C" fn(*mut usize, *const AllocationProp, c_int) -> c_int, + create: unsafe extern "C" fn(*mut u64, usize, *const AllocationProp, u64) -> c_int, + export: unsafe extern "C" fn(*mut c_void, u64, c_int, u64) -> c_int, + reserve: unsafe extern "C" fn(*mut u64, usize, usize, u64, u64) -> c_int, + map: unsafe extern "C" fn(u64, usize, usize, u64, u64) -> c_int, + set_access: unsafe extern "C" fn(u64, usize, *const AccessDesc, usize) -> c_int, + copy_d2h: unsafe extern "C" fn(*mut c_void, u64, usize) -> c_int, + unmap: unsafe extern "C" fn(u64, usize) -> c_int, + address_free: unsafe extern "C" fn(u64, usize) -> c_int, + release: unsafe extern "C" fn(u64) -> c_int, +} + +/// `CUmemAllocationProp`. +#[repr(C)] +struct AllocationProp { + kind: c_int, + handle_types: c_int, + location_type: c_int, + location_id: c_int, + win32: *mut c_void, + flags: [u8; 8], +} + +/// `CUmemAccessDesc`. +#[repr(C)] +struct AccessDesc { + location_type: c_int, + location_id: c_int, + flags: c_int, +} + +impl Cuda { + /// Resolve the driver entry points. Field types give the signatures. + #[allow(clippy::missing_transmute_annotations)] + fn load() -> Result { + let library = unsafe { libc::dlopen(c"libcuda.so.1".as_ptr(), libc::RTLD_NOW) }; + if library.is_null() { + return Err("libcuda.so.1 not found".into()); + } + let symbol = |name: &CStr| -> Result<*mut c_void, String> { + let ptr = unsafe { libc::dlsym(library, name.as_ptr()) }; + if ptr.is_null() { + Err(format!("libcuda.so.1 lacks {name:?}")) + } else { + Ok(ptr) + } + }; + macro_rules! load { + ($name:literal) => { + unsafe { std::mem::transmute::<*mut c_void, _>(symbol($name)?) } + }; + } + let cuda = Self { + init: load!(c"cuInit"), + device_get: load!(c"cuDeviceGet"), + retain: load!(c"cuDevicePrimaryCtxRetain"), + set_current: load!(c"cuCtxSetCurrent"), + granularity: load!(c"cuMemGetAllocationGranularity"), + create: load!(c"cuMemCreate"), + export: load!(c"cuMemExportToShareableHandle"), + reserve: load!(c"cuMemAddressReserve"), + map: load!(c"cuMemMap"), + set_access: load!(c"cuMemSetAccess"), + copy_d2h: load!(c"cuMemcpyDtoH_v2"), + unmap: load!(c"cuMemUnmap"), + address_free: load!(c"cuMemAddressFree"), + release: load!(c"cuMemRelease"), + }; + check("cuInit", unsafe { (cuda.init)(0) }); + Ok(cuda) + } + + /// Make the GPU's primary context current on this thread. + fn bind(&self, ordinal: i32) { + let mut device = 0; + check("cuDeviceGet", unsafe { + (self.device_get)(&mut device, ordinal) + }); + let mut context = std::ptr::null_mut(); + check("cuDevicePrimaryCtxRetain", unsafe { + (self.retain)(&mut context, device) + }); + check("cuCtxSetCurrent", unsafe { (self.set_current)(context) }); + } + + /// Allocate exportable device memory, as the Python client does. + fn allocate(&self, ordinal: i32, size: usize) -> Allocation<'_> { + self.bind(ordinal); + let prop = AllocationProp { + kind: 1, // CU_MEM_ALLOCATION_TYPE_PINNED + handle_types: 1, // CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + location_type: 1, + location_id: ordinal, + win32: std::ptr::null_mut(), + flags: [0; 8], + }; + let mut granularity = 0; + check("cuMemGetAllocationGranularity", unsafe { + (self.granularity)(&mut granularity, &prop, 0) + }); + let size = size.div_ceil(granularity) * granularity; + let mut handle = 0; + check("cuMemCreate", unsafe { + (self.create)(&mut handle, size, &prop, 0) + }); + Allocation { + cuda: self, + ordinal, + handle, + size: size as u64, + } + } +} + +/// One exportable allocation. +struct Allocation<'a> { + cuda: &'a Cuda, + ordinal: i32, + handle: u64, + size: u64, +} + +impl Allocation<'_> { + /// Export as a POSIX fd. + fn export(&self) -> OwnedFd { + let mut fd: c_int = -1; + check("cuMemExportToShareableHandle", unsafe { + (self.cuda.export)((&mut fd as *mut c_int).cast(), self.handle, 1, 0) + }); + assert!(fd >= 0); + unsafe { OwnedFd::from_raw_fd(fd) } + } + + /// Map the allocation and copy `out.len()` bytes at `offset` to the host. + fn read(&self, offset: u64, out: &mut [u8]) { + let cuda = self.cuda; + cuda.bind(self.ordinal); + let size = self.size as usize; + let mut ptr = 0; + check("cuMemAddressReserve", unsafe { + (cuda.reserve)(&mut ptr, size, 0, 0, 0) + }); + check("cuMemMap", unsafe { + (cuda.map)(ptr, size, 0, self.handle, 0) + }); + let access = AccessDesc { + location_type: 1, + location_id: self.ordinal, + flags: 3, // CU_MEM_ACCESS_FLAGS_PROT_READWRITE + }; + check("cuMemSetAccess", unsafe { + (cuda.set_access)(ptr, size, &access, 1) + }); + check("cuMemcpyDtoH", unsafe { + (cuda.copy_d2h)(out.as_mut_ptr().cast(), ptr + offset, out.len()) + }); + check("cuMemUnmap", unsafe { (cuda.unmap)(ptr, size) }); + check("cuMemAddressFree", unsafe { + (cuda.address_free)(ptr, size) + }); + } +} + +impl Drop for Allocation<'_> { + fn drop(&mut self) { + unsafe { (self.cuda.release)(self.handle) }; + } +} + +/// Panic on a CUDA driver error. +fn check(call: &str, rc: c_int) { + assert_eq!(rc, 0, "{call} failed with CUresult {rc}"); +} diff --git a/cmd/racer-dataplane/gpu/tests/simulation.rs b/cmd/racer-dataplane/gpu/tests/simulation.rs index 56fd76ac81..54a514c2a3 100644 --- a/cmd/racer-dataplane/gpu/tests/simulation.rs +++ b/cmd/racer-dataplane/gpu/tests/simulation.rs @@ -362,6 +362,66 @@ fn failed_frees_leak_and_count() { assert_eq!(gpu.leaked(), 3); } +#[test] +/// A file cuFile does not support fails with its cuFile status, apart from +/// driver failures. +fn unsupported_file_is_its_own_error() { + let sim = simulation(); + let _env = sim.enter(); + let gpu = Gpu::load().unwrap(); + sim.fault(Operation::File, -11); + assert_eq!( + gpu.open_file(file_with(&[1; 16])).err(), + Some(Error::Unsupported(5008)) + ); + sim.fault(Operation::File, -5); + assert_eq!( + gpu.open_file(file_with(&[1; 16])).err(), + Some(Error::Driver(999)) + ); + assert!(gpu.open_file(file_with(&[1; 16])).is_ok()); +} + +#[test] +/// PCI P2PDMA is off by default and reaches the driver open when set, from +/// either the helper's own open or a reader's. +fn pci_p2pdma_reaches_driver_open() { + let sim = simulation(); + let _env = sim.enter(); + let gpu = Gpu::load().unwrap(); + assert_eq!(sim.pci_p2pdma(), None); + assert!(gpu.open_file(file_with(&[1; 16])).is_ok()); + assert_eq!(sim.pci_p2pdma(), Some(false)); + + let sim = simulation(); + let _env = sim.enter(); + let gpu = Gpu::load().unwrap(); + gpu.set_pci_p2pdma(true); + let reader = gpu.file_reader(); + // Turning it off after the reader is made does not change that reader. + gpu.set_pci_p2pdma(false); + reader.open_driver().unwrap(); + assert_eq!(sim.pci_p2pdma(), Some(true)); + // A later open does not change the driver's setting. + assert!(gpu.open_file(file_with(&[1; 16])).is_ok()); + assert_eq!(sim.pci_p2pdma(), Some(true)); +} + +#[test] +/// A failed open with PCI P2PDMA is reported, as any failed open is. +fn pci_p2pdma_open_failure_is_reported() { + let sim = simulation(); + let _env = sim.enter(); + let gpu = Gpu::load().unwrap(); + gpu.set_pci_p2pdma(true); + sim.fault(Operation::FileOpen, -3); + assert_eq!( + gpu.file_reader().open_driver().err(), + Some(Error::Unavailable) + ); + assert_eq!(sim.pci_p2pdma(), None); +} + #[test] /// cuFile reads land in staging, respect bounds, and wait for streams. fn cufile_reads_into_staging() { diff --git a/cmd/racer-dataplane/object-wire/src/client.rs b/cmd/racer-dataplane/object-wire/src/client.rs index fd060280af..889d0f5b30 100644 --- a/cmd/racer-dataplane/object-wire/src/client.rs +++ b/cmd/racer-dataplane/object-wire/src/client.rs @@ -46,6 +46,10 @@ pub enum ReadKind { /// The size the client passed to `cuMemCreate`. allocation_size: u64, + + /// The UUID of the GPU that holds the allocation, from + /// `Racer-Gpu-Uuid`. Absent from older clients. + gpu_uuid: Option<[u8; 16]>, }, } @@ -183,6 +187,7 @@ impl RequestParser { let mut destination_offset = None; let mut destination_length = None; let mut allocation_size = None; + let mut gpu_uuid = None; let mut seen = HashSet::new(); let mut host = None; let mut pin = None; @@ -275,6 +280,12 @@ impl RequestParser { "racer-allocation-size" if direct == Some(DirectDelivery::Gpu) => { once(&mut allocation_size, parse_decimal(value)?)?; } + "racer-gpu-uuid" if direct == Some(DirectDelivery::Gpu) => { + if gpu_uuid.is_some() { + return Err(Error::InvalidRequest); + } + gpu_uuid = Some(parse_gpu_uuid(trim_ows(value))?); + } "racer-metadata" => { validate_opaque(value)?; metadata = Some(value); @@ -322,6 +333,7 @@ impl RequestParser { range, destination, allocation_size, + gpu_uuid, } } } @@ -410,6 +422,32 @@ fn parse_max_age(value: &[u8], max_age: &mut Option) -> Result<()> { } /// Store a destination field, rejecting a repeated one. +/// Parse a GPU UUID as `nvidia-smi` prints it, `GPU-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx`. +/// The `GPU-` prefix is optional and hex digits may be either case. +pub fn parse_gpu_uuid(value: &[u8]) -> Result<[u8; 16]> { + let value = match value.get(..4) { + Some(prefix) if prefix.eq_ignore_ascii_case(b"GPU-") => &value[4..], + _ => value, + }; + if value.len() != 36 { + return Err(Error::InvalidRequest); + } + let mut uuid = [0; 16]; + let mut digits = 0; + for (i, &b) in value.iter().enumerate() { + if matches!(i, 8 | 13 | 18 | 23) { + if b != b'-' { + return Err(Error::InvalidRequest); + } + continue; + } + let nibble = (b as char).to_digit(16).ok_or(Error::InvalidRequest)? as u8; + uuid[digits / 2] |= if digits % 2 == 0 { nibble << 4 } else { nibble }; + digits += 1; + } + Ok(uuid) +} + fn once(slot: &mut Option, value: u64) -> Result<()> { if slot.replace(value).is_some() { return Err(Error::InvalidRequest); diff --git a/cmd/racer-dataplane/object-wire/tests/direct.rs b/cmd/racer-dataplane/object-wire/tests/direct.rs index 8e2e5fe917..1a3b86f3a2 100644 --- a/cmd/racer-dataplane/object-wire/tests/direct.rs +++ b/cmd/racer-dataplane/object-wire/tests/direct.rs @@ -91,6 +91,7 @@ fn gpu_read_requires_allocation_that_holds_the_window() { length: 100 }, allocation_size: 4196, + gpu_uuid: None, } ); assert_eq!(parsed.direct(), Some(DirectDelivery::Gpu)); @@ -113,6 +114,83 @@ fn gpu_read_requires_allocation_that_holds_the_window() { } } +#[test] +fn gpu_read_parses_an_optional_gpu_uuid() { + let mut expected = [0u8; 16]; + expected[0] = 0x01; + expected[15] = 0xef; + for value in [ + "GPU-01000000-0000-0000-0000-0000000000ef", + "gpu-01000000-0000-0000-0000-0000000000EF", + "01000000-0000-0000-0000-0000000000ef", + " GPU-01000000-0000-0000-0000-0000000000ef ", + ] { + let parsed = parse(&direct( + "POST", + "gpu-read", + &[ + DEST[0], + DEST[1], + ("Racer-Allocation-Size", "4196"), + ("Racer-Gpu-Uuid", value), + ], + )) + .unwrap(); + let ReadKind::GpuRead { gpu_uuid, .. } = parsed else { + panic!("{parsed:?}"); + }; + assert_eq!(gpu_uuid, Some(expected), "{value}"); + } + for value in [ + "", + "GPU-", + "GPU-01000000-0000-0000-0000-0000000000e", + "GPU-01000000-0000-0000-0000-0000000000eff", + "GPU-0100000000-00-0000-0000-0000000000ef", + "GPU-01000000-0000-0000-0000-0000000000eg", + "MIG-01000000-0000-0000-0000-0000000000ef", + ] { + assert_eq!( + parse(&direct( + "POST", + "gpu-read", + &[ + DEST[0], + DEST[1], + ("Racer-Allocation-Size", "4196"), + ("Racer-Gpu-Uuid", value), + ], + )), + Err(Error::InvalidRequest), + "{value:?}" + ); + } + let uuid = ("Racer-Gpu-Uuid", "GPU-01000000-0000-0000-0000-0000000000ef"); + assert_eq!( + parse(&direct( + "POST", + "gpu-read", + &[ + DEST[0], + DEST[1], + ("Racer-Allocation-Size", "4196"), + uuid, + uuid + ], + )), + Err(Error::InvalidRequest) + ); + // Ignored on host-read, like the allocation size. + assert!(matches!( + parse(&direct( + "POST", + "host-read", + &[DEST[0], DEST[1], ("Racer-Gpu-Uuid", "x")] + )), + Ok(ReadKind::HostRead { .. }) + )); +} + #[test] fn destination_fields_are_required_canonical_unique_and_bounded() { let max = i64::MAX.to_string(); diff --git a/cmd/racer-dataplane/src/app.rs b/cmd/racer-dataplane/src/app.rs index 7d109682ab..85fb4215f6 100644 --- a/cmd/racer-dataplane/src/app.rs +++ b/cmd/racer-dataplane/src/app.rs @@ -131,8 +131,10 @@ pub struct Application { /// graph crosses a thread. Worker zero alone drives enrollment/control reloads. pub struct NodeState { devices: Option, - /// The process-wide GPU helper, when `gpu-read` is enabled. - gpu: Option>, + /// The process-wide GPU helper slot. In `RACER_ENABLE_GPU=auto` it + /// starts on the first `gpu-read`; read it with `get()`, which never + /// starts it. + gpu: Arc, inventory: Arc, /// Startup-fixed selector and claimed disks for `/debug/hardware`. @@ -211,7 +213,7 @@ impl NodeState { Ok(Self { peer_admission, devices: None, - gpu: None, + gpu: Default::default(), peer_receive: Default::default(), inventory: Arc::new(crate::rdma::Inventory::default()), hardware: Default::default(), @@ -286,19 +288,9 @@ impl Application { self.config.limits.retained_snapshots.get(), ))); self.node.inventory.update(self.discovered_nics.clone())?; - Arc::get_mut(&mut self.node).ok_or(Error::Internal)?.gpu = - crate::gpu::GpuHelper::start(self.config.enable_gpu, self.config.gpu_staging_bytes)?; - // GPU decrypt and RDMA together: peers' NICs may write gpu-read pages - // straight into GPU landing buffers (GPUDirect RDMA). Activation - // registers them per rail and leaves a rail host-only on any failure. - if self.config.enable_rdma - && let Some(gpu) = &self.node.gpu - { - self.node - .native - .link_node() - .set_gpu_landing(gpu.clone(), self.config.gpu_landing_bytes); - } + let gpu = self.gpu_slot()?; + self.node.metrics[0].1.gpu().attach_slot(gpu.clone()); + Arc::get_mut(&mut self.node).ok_or(Error::Internal)?.gpu = gpu; if self.config.enable_rdma { self.node.native.place(&plan)?; self.node @@ -1382,8 +1374,10 @@ impl WorkerApplication { )); let mut responses = Responses::new(io.clone(), delivery) .with_observer(admission.policy().observer().clone()); - if let Some(gpu) = &node.gpu { + if let Some(gpu) = node.gpu.get() { responses = responses.with_gpu(Rc::new(crate::gpu::GpuClient::new(gpu.clone()))); + } else if node.gpu.standby() { + responses = responses.with_lazy_gpu(node.gpu.clone()); } let responses = Rc::new(responses); let clients = ClientListeners::new( @@ -2078,6 +2072,20 @@ impl WorkerApplication { } impl Application { + /// Resolve `RACER_ENABLE_GPU` into the process GPU slot; see + /// [`gpu_slot`]. + fn gpu_slot(&self) -> Result> { + gpu_slot( + &self.config, + cfg!(feature = "cuda"), + self.node.metrics[0].1.gpu_shared(), + self.config + .enable_rdma + .then(|| self.node.native.link_node()), + Box::new(racer_gpu::Gpu::load), + ) + } + /// Allocate paired worker resources once before building service lanes. fn prepare_workers(&self) -> Result<()> { let resources = crate::worker::Resources::new( @@ -3415,3 +3423,57 @@ impl WorkerPressure { } } } + +/// Resolve `RACER_ENABLE_GPU` into the process GPU slot. `built` says +/// whether the binary has the `cuda` feature; `loader` loads the driver. +/// +/// - `false`: off; nothing GPU-related runs. +/// - `true`: an explicit request. Start the helper now and fail startup when +/// no GPU is usable. +/// - `auto` with `cuda`: touch no GPU now. The helper (driver load, cuInit, +/// contexts, decrypt probe, cuFile open) starts on the first `gpu-read` +/// request head, once per process; if that start fails or finds no usable +/// GPU, GPU mode stays off (501) and is never retried. +/// - `auto` without `cuda`: off, logged as `reason=not-built`. +/// +/// Once a helper exists, metrics report it and, with RDMA on (`links`), +/// peers' NICs may write `gpu-read` pages into GPU landing buffers +/// (GPUDirect RDMA): activation registers them per rail, and rails already +/// active when a lazy start finishes get them on the owner worker's next +/// turn. +pub(crate) fn gpu_slot( + config: &Config, + built: bool, + telemetry: Arc, + links: Option>, + loader: crate::gpu::Loader, +) -> Result> { + use crate::config::GpuMode; + let landing_bytes = config.gpu_landing_bytes; + let on_start = move |gpu: &Arc| { + telemetry.attach_helper(gpu.clone()); + if let Some(links) = links { + links.set_gpu_landing(gpu.clone(), landing_bytes); + telemetry.attach_links(links); + } + }; + let (mode, staging, p2pdma) = ( + config.enable_gpu, + config.gpu_staging_bytes, + config.gpu_pci_p2pdma, + ); + if mode == GpuMode::Auto && built { + eprintln!("racer-dataplane: stage=gpu state=standby start=first-gpu-read"); + return Ok(Arc::new(crate::gpu_slot::GpuSlot::lazy( + Box::new(move || { + crate::gpu::GpuHelper::start_built(mode, built, staging, p2pdma, loader) + }), + Box::new(on_start), + ))); + } + let gpu = crate::gpu::GpuHelper::start_built(mode, built, staging, p2pdma, loader)?; + if let Some(gpu) = &gpu { + on_start(gpu); + } + Ok(Arc::new(crate::gpu_slot::GpuSlot::ready(gpu))) +} diff --git a/cmd/racer-dataplane/src/app/native.rs b/cmd/racer-dataplane/src/app/native.rs index 48dd3f2974..01f037658a 100644 --- a/cmd/racer-dataplane/src/app/native.rs +++ b/cmd/racer-dataplane/src/app/native.rs @@ -166,6 +166,22 @@ impl WorkerApplication { } } self.actual_rails = links.active(); + // The GPU helper started lazily after these rails activated: attach + // its landing buffers to them now. Only the owner registers memory. + if self.native_owner() + && self.native_task.is_none() + && !self.actual_rails.is_empty() + && links.node().take_landing_pending() + { + let rails = self.actual_rails.clone(); + let turn = scope(Duration::from_secs(5))?; + self.native_task = Some(Box::pin(async move { + links.attach_landings(&rails, &turn).await; + Ok(rails) + })); + cx.waker().wake_by_ref(); + return Ok(()); + } if !self.native_owner() || self.native_task.is_some() || !self.actual_rails.is_empty() diff --git a/cmd/racer-dataplane/src/client.rs b/cmd/racer-dataplane/src/client.rs index c9d0520679..a90a1dfdb7 100644 --- a/cmd/racer-dataplane/src/client.rs +++ b/cmd/racer-dataplane/src/client.rs @@ -37,6 +37,7 @@ use http1::StartLine; use racer_control_wire::CacheId; use std::collections::BTreeMap; use std::rc::Rc; +use std::sync::Arc; use std::task::Context; use std::task::Poll; use std::time::Duration; @@ -246,7 +247,12 @@ pub struct Responses { delivery: Rc, - gpu: Option>, + /// The `gpu-read` backend once known. Filled from `lazy_gpu` when the + /// first `gpu-read` on this worker starts (or finds) the GPU helper. + gpu: std::cell::RefCell>>, + + /// The process GPU slot while it waits for its first `gpu-read`. + lazy_gpu: Option>, } /// Duplex subscription framing. Delivered leases stay pinned until exact release. @@ -845,23 +851,70 @@ impl Responses { io, delivery, observer: Observer::default(), - gpu: None, + gpu: Default::default(), + lazy_gpu: None, } } /// Serve `gpu-read` with this backend. Without one, `gpu-read` is 501. pub fn with_gpu(mut self, gpu: Rc) -> Self { - self.gpu = Some(gpu); + self.gpu = std::cell::RefCell::new(Some(gpu)); + self + } + + /// Serve `gpu-read` with the helper in `slot`, starting it on the first + /// `gpu-read` (see [`Self::prepare_gpu`]). + pub fn with_lazy_gpu(mut self, slot: Arc) -> Self { + self.lazy_gpu = Some(slot); self } + /// Before opening a `gpu-read` sink: start the GPU helper if this is the + /// process's first `gpu-read`, or pick up one another worker started. + /// Awaits without blocking the worker. Other requests never call this, + /// so nothing else touches a GPU. + /// + /// The wait honors `scope`: an expired or cancelled request stops waiting + /// and releases its connection, while the shared start keeps running for + /// later requests. + pub(crate) async fn prepare_gpu(&self, scope: &RequestScope) -> Result<()> { + if self.gpu.borrow().is_some() { + return Ok(()); + } + let Some(slot) = &self.lazy_gpu else { + return Ok(()); + }; + let cancellation = scope.cancellation.subscribe()?; + let mut start = std::pin::pin!(slot.get_or_start()); + let started = std::future::poll_fn(|cx| { + cancellation.register(cx.waker()); + scope.check()?; + start.as_mut().poll(cx).map(Ok::<_, Error>) + }) + .await?; + let Some(helper) = started else { + return Ok(()); + }; + let mut gpu = self.gpu.borrow_mut(); + if gpu.is_none() { + *gpu = Some(Rc::new(crate::gpu::GpuClient::new(helper))); + } + Ok(()) + } + + /// Whether this worker serves `gpu-read`. + pub(crate) fn has_gpu(&self) -> bool { + self.gpu.borrow().is_some() + } + /// Check the descriptors that came with a request and open its sink. pub(crate) fn open_sink( &self, kind: &ReadKind, received: http1::connection::ReceivedDescriptors, ) -> Result>> { - direct::open_sink(kind, received, &self.delivery, self.gpu.as_ref()) + let gpu = self.gpu.borrow().clone(); + direct::open_sink(kind, received, &self.delivery, gpu.as_ref()) } /// Route detailed delivery failures to the worker's observer. diff --git a/cmd/racer-dataplane/src/client/direct.rs b/cmd/racer-dataplane/src/client/direct.rs index 61c648804d..60c94ffd5b 100644 --- a/cmd/racer-dataplane/src/client/direct.rs +++ b/cmd/racer-dataplane/src/client/direct.rs @@ -113,11 +113,13 @@ pub trait GpuService { /// request. Called after the request is validated and before any read. /// Return [`Error::NotImplemented`] when the allocation is on a GPU the /// dataplane cannot use; [`Error::InvalidRequest`] for a bad handle. + /// `gpu_uuid` names the client's GPU when the client sent it. fn open( &self, descriptor: OwnedFd, destination: Destination, allocation_size: u64, + gpu_uuid: Option<[u8; 16]>, ) -> Result>; } @@ -225,11 +227,13 @@ pub(crate) fn open_sink( ( DirectDelivery::Gpu, ReadKind::GpuRead { - allocation_size, .. + allocation_size, + gpu_uuid, + .. }, ) => gpu .ok_or(Error::NotImplemented)? - .open(descriptor, destination, *allocation_size) + .open(descriptor, destination, *allocation_size, *gpu_uuid) .map(Some), _ => Err(Error::Internal), } diff --git a/cmd/racer-dataplane/src/client/listener.rs b/cmd/racer-dataplane/src/client/listener.rs index 24ebef69ad..e57d10cb33 100644 --- a/cmd/racer-dataplane/src/client/listener.rs +++ b/cmd/racer-dataplane/src/client/listener.rs @@ -964,9 +964,20 @@ async fn serve_connection( deadline.set(scope.deadline.0); let scope = &scope; // Every descriptor must be claimed by a direct read; the rest close here. + let mut gpu_read = false; let parsed = received.value.and_then(|received| { let request = parser.parse(cache, received.head)?; - let sink = responses.open_sink(&request.kind, received.descriptors)?; + gpu_read = matches!(request.kind, crate::client::ReadKind::GpuRead { .. }); + Ok((request, received.descriptors)) + }); + // A `gpu-read` head is the only request that may start the GPU + // helper (`RACER_ENABLE_GPU=auto`); see `Responses::prepare_gpu`. + let parsed = match parsed { + Ok(parsed) if gpu_read => responses.prepare_gpu(scope).await.map(|()| parsed), + other => other, + }; + let parsed = parsed.and_then(|(request, descriptors)| { + let sink = responses.open_sink(&request.kind, descriptors)?; Ok((request, sink)) }); let parsed = match parsed { @@ -976,6 +987,14 @@ async fn serve_connection( }, other => other, }; + if gpu_read { + metrics + .gpu() + .record_request(crate::gpu_telemetry::GpuReadOutcome::classify( + parsed.as_ref().map(|_| ()).map_err(|error| *error), + responses.has_gpu(), + )); + } let (request, sink) = match parsed { Ok(parsed) => parsed, Err(error) => { diff --git a/cmd/racer-dataplane/src/client/tests/direct.rs b/cmd/racer-dataplane/src/client/tests/direct.rs index 8d9c0480ae..fdeb25eb33 100644 --- a/cmd/racer-dataplane/src/client/tests/direct.rs +++ b/cmd/racer-dataplane/src/client/tests/direct.rs @@ -95,15 +95,24 @@ fn direct_request(suffix: &str, fields: &str) -> Vec { /// A fixture whose fake origin serves `body` as version "v1". fn direct_fixture(body: Vec, gpu: Option>) -> Fixture { + direct_fixture_with(body, |responses| match gpu { + Some(gpu) => responses.with_gpu(gpu), + None => responses, + }) +} + +/// A fixture whose fake origin serves `body`, with `setup` applied to the +/// worker's responses. +fn direct_fixture_with(body: Vec, setup: impl FnOnce(Responses) -> Responses) -> Fixture { let mut fixture = Fixture::new(); let delivery = Rc::new(Delivery::new( fixture.reactor.clone(), Duration::from_secs(2), )); - let mut responses = Responses::new(fixture.listeners.io.clone(), delivery.clone()); - if let Some(gpu) = gpu { - responses = responses.with_gpu(gpu); - } + let responses = setup(Responses::new( + fixture.listeners.io.clone(), + delivery.clone(), + )); fixture.listeners.responses = Rc::new(responses); let worker = ReadWorker::new( definition(), @@ -406,12 +415,20 @@ fn gpu_read_is_validated_then_501_without_a_backend() { settle(&fixture); assert_eq!(memfd.open_copies(), 1); } + use crate::gpu_telemetry::GpuReadOutcome; + let gpu = fixture.listeners.metrics.gpu(); + assert_eq!(gpu.requests(GpuReadOutcome::NotEnabled), 1); + // A head the parser rejects is not known to be a gpu-read: only the + // missing descriptor counts. + assert_eq!(gpu.requests(GpuReadOutcome::Invalid), 1); + assert_eq!(gpu.requests(GpuReadOutcome::Gpu), 0); } /// Records what a GPU backend would be asked to do. #[derive(Default)] struct Recorded { opened: Vec<(Destination, u64)>, + uuids: Vec>, writes: Vec<(u64, Vec)>, finished: usize, } @@ -426,11 +443,12 @@ impl GpuService for FakeGpu { descriptor: OwnedFd, destination: Destination, allocation_size: u64, + gpu_uuid: Option<[u8; 16]>, ) -> Result> { - self.0 - .borrow_mut() - .opened - .push((destination, allocation_size)); + let mut recorded = self.0.borrow_mut(); + recorded.opened.push((destination, allocation_size)); + recorded.uuids.push(gpu_uuid); + drop(recorded); Ok(Box::new(FakeSink(self.0.clone(), descriptor))) } } @@ -505,11 +523,54 @@ fn gpu_read_delivers_pages_through_the_backend_seam() { ); assert_eq!(recorded.writes, [(100, b"llo".to_vec())]); assert_eq!(recorded.finished, 1); + // An older client sends no GPU UUID. + assert_eq!(recorded.uuids, [None]); } settle(&fixture); assert_eq!(memfd.open_copies(), 1); } +#[test] +fn gpu_read_passes_the_client_gpu_uuid_to_the_backend() { + let recorded = Rc::new(RefCell::new(Recorded::default())); + let fixture = direct_fixture(b"hello".to_vec(), Some(Rc::new(FakeGpu(recorded.clone())))); + let memfd = Memfd::new(); + let output = exchange( + &fixture, + &direct_request( + "gpu-read", + "Racer-Destination-Offset: 0\r\nRacer-Destination-Length: 5\r\n\ + Racer-Allocation-Size: 4096\r\n\ + Racer-Gpu-Uuid: GPU-03030303-0303-0303-0303-030303030303\r\n", + ), + &[memfd.file.as_raw_fd()], + ); + let (head, _) = split(&output); + assert!(head.starts_with("http/1.1 200"), "{head}"); + assert_eq!(recorded.borrow().uuids, [Some([3; 16])]); + settle(&fixture); +} + +#[test] +fn gpu_read_rejects_a_malformed_gpu_uuid() { + let recorded = Rc::new(RefCell::new(Recorded::default())); + let fixture = direct_fixture(b"hello".to_vec(), Some(Rc::new(FakeGpu(recorded.clone())))); + let memfd = Memfd::new(); + let output = exchange( + &fixture, + &direct_request( + "gpu-read", + "Racer-Destination-Offset: 0\r\nRacer-Destination-Length: 5\r\n\ + Racer-Allocation-Size: 4096\r\nRacer-Gpu-Uuid: 0\r\n", + ), + &[memfd.file.as_raw_fd()], + ); + let (head, _) = split(&output); + assert!(head.starts_with("http/1.1 400"), "{head}"); + assert!(recorded.borrow().opened.is_empty()); + settle(&fixture); +} + mod gpu; /// A plain HEAD pipelined ahead of a host-read on one connection: one receive diff --git a/cmd/racer-dataplane/src/client/tests/direct/gpu.rs b/cmd/racer-dataplane/src/client/tests/direct/gpu.rs index 3fd10e79e6..e89d6deceb 100644 --- a/cmd/racer-dataplane/src/client/tests/direct/gpu.rs +++ b/cmd/racer-dataplane/src/client/tests/direct/gpu.rs @@ -90,6 +90,70 @@ fn gpu_read_writes_the_range_at_the_destination_offset() { assert_no_body_leases(&fixture); } +/// A UUID header as `nvidia-smi` prints it for simulated GPU `index`. +fn uuid_field(index: u8) -> String { + format!( + "Racer-Gpu-Uuid: {}\r\n", + racer_gpu::uuid_string(&[index; 16]) + ) +} + +/// Serve a five-byte `gpu-read` into an allocation on GPU `ordinal` of four, +/// sending `extra` fields. Returns the status line and the import attempts. +fn read_on_four_gpus(ordinal: i32, extra: &str) -> (String, usize, Vec) { + let helper = start(GpuMode::Required, 4).unwrap(); + let fixture = gpu_fixture(b"hello".to_vec(), &helper); + let allocation = allocate(&helper, ordinal, GRANULE); + let from = trace_len(&helper); + let output = exchange( + &fixture, + &direct_request("gpu-read", &format!("{}{extra}", fields("", 0, 5, GRANULE))), + &[allocation.as_raw_fd()], + ); + let (head, _) = split(&output); + let status = head.lines().next().unwrap().to_string(); + let imports = calls(&helper, from, Operation::Import); + settle(&fixture); + assert_released(&helper); + assert_no_body_leases(&fixture); + (status, imports, read_allocation(&allocation, 0, 5)) +} + +#[test] +fn gpu_read_uuid_selects_a_nonzero_ordinal_first() { + // The client's container sees this GPU as device 0; the dataplane sees it + // as device 3. The UUID is the same in both. + let (status, imports, written) = read_on_four_gpus(3, &uuid_field(3)); + assert!(status.starts_with("http/1.1 200"), "{status}"); + assert_eq!(imports, 1, "only the named GPU is tried"); + assert_eq!(written, b"hello"); +} + +#[test] +fn gpu_read_without_uuid_tries_each_gpu() { + let (status, imports, written) = read_on_four_gpus(3, ""); + assert!(status.starts_with("http/1.1 200"), "{status}"); + assert_eq!(imports, 4); + assert_eq!(written, b"hello"); +} + +#[test] +fn gpu_read_unknown_uuid_is_501_so_the_client_uses_the_host() { + let (status, imports, written) = read_on_four_gpus(3, &uuid_field(9)); + assert!(status.starts_with("http/1.1 501"), "{status}"); + assert_eq!(imports, 0, "nothing is imported"); + assert_eq!(written, [0; 5]); +} + +#[test] +fn gpu_read_wrong_uuid_still_finds_the_owning_gpu() { + let (status, imports, written) = read_on_four_gpus(2, &uuid_field(1)); + assert!(status.starts_with("http/1.1 200"), "{status}"); + // GPU 1 first, then 0, then 2. + assert_eq!(imports, 3); + assert_eq!(written, b"hello"); +} + #[test] fn gpu_read_spans_pages() { let length = 2 * PAGE_BYTES as usize + 7; @@ -158,6 +222,12 @@ fn gpu_read_errors_before_the_head_release_everything() { "nothing written" ); assert_no_body_leases(&fixture); + // The 412 comes after the import, so it counts as a GPU open. + use crate::gpu_telemetry::GpuReadOutcome; + let gpu = fixture.listeners.metrics.gpu(); + assert_eq!(gpu.requests(GpuReadOutcome::Gpu), 1); + assert_eq!(gpu.requests(GpuReadOutcome::Invalid), 2); + assert_eq!(gpu.requests(GpuReadOutcome::WrongDevice), 1); } #[test] @@ -236,9 +306,13 @@ fn gpu_mode_resolves_against_hardware() { // Without the cuda feature, auto never loads a backend. if !cfg!(feature = "cuda") { assert!( - GpuHelper::start(GpuMode::Auto, crate::config::DEFAULT_GPU_STAGING_BYTES) - .unwrap() - .is_none() + GpuHelper::start( + GpuMode::Auto, + crate::config::DEFAULT_GPU_STAGING_BYTES, + false + ) + .unwrap() + .is_none() ); } } @@ -648,6 +722,7 @@ fn gpu_read_without_decrypt_kernels_decrypts_on_the_cpu() { let fixture = gpu_fixture(data.clone(), &helper); arrange(&fixture, &helper, 1, 5000, false, true); let mark = trace_len(&helper); + let plaintext = counter(&helper.counters().plaintext); let (head, frames, allocation, _) = gpu_read(&fixture, &helper, None, 5000); assert!(head.starts_with("http/1.1 200"), "{head}"); assert_eq!(frames, [(1, 0, 0, 5000), (2, 1, 5000, 0)]); @@ -655,5 +730,357 @@ fn gpu_read_without_decrypt_kernels_decrypts_on_the_cpu() { assert_eq!(calls(&helper, mark, Operation::Decrypt), 0); assert_eq!(calls(&helper, mark, Operation::FileRead), 0); assert_eq!(worker(&fixture).cpu_decrypts(), 1); + assert_eq!(counter(&helper.counters().plaintext), plaintext + 1); + assert_eq!(counter(&helper.counters().decrypted), 0); assert_idle(&fixture, &helper, default_slots()); } + +/// A loader on `count` simulated GPUs that counts its calls. +fn counted_loader(count: u8, calls: &Arc) -> crate::gpu::Loader { + let calls = calls.clone(); + Box::new(move || { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let devices: Vec = (0..count).map(SimDevice::new).collect(); + let simulation = Simulation::new(devices); + let environment = simulation.enter(); + SIM.with(|sim| *sim.borrow_mut() = Some((simulation, environment))); + racer_gpu::Gpu::load() + }) +} + +/// The process GPU slot for `mode` in a `cuda` build, as the app makes it. +fn slot( + mode: GpuMode, + loader: crate::gpu::Loader, + telemetry: Arc, + links: Option>, +) -> crate::error::Result> { + let mut config = crate::test_support::cluster::config(links.is_some()); + config.enable_gpu = mode; + crate::app::gpu_slot(&config, true, telemetry, links, loader) +} + +fn lazy_fixture(body: Vec, slot: &Arc) -> Fixture { + let slot = slot.clone(); + direct_fixture_with(body, move |responses| responses.with_lazy_gpu(slot)) +} + +fn render(telemetry: &crate::gpu_telemetry::GpuTelemetry) -> String { + let mut out = String::new(); + telemetry.write(&mut out).unwrap(); + out +} + +fn calls_now(calls: &std::sync::atomic::AtomicUsize) -> usize { + calls.load(std::sync::atomic::Ordering::SeqCst) +} + +#[test] +/// `auto` touches no GPU at startup, for host reads, or for metrics. The +/// first `gpu-read` starts the helper once, and later reads reuse it. +fn auto_starts_the_gpu_on_the_first_gpu_read_only() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let telemetry = Arc::new(crate::gpu_telemetry::GpuTelemetry::default()); + let links = Arc::new(crate::rdma::LinkNode::default()); + let slot = slot( + GpuMode::Auto, + counted_loader(1, &calls), + telemetry.clone(), + Some(links.clone()), + ) + .unwrap(); + telemetry.attach_slot(slot.clone()); + assert!(slot.standby()); + assert!(slot.get().is_none()); + let fixture = lazy_fixture(b"hello".to_vec(), &slot); + // A host read, and a metrics scrape, never start it. + let memfd = Memfd::new(); + let output = exchange( + &fixture, + &direct_request( + "host-read", + "Racer-Destination-Offset: 0\r\nRacer-Destination-Length: 5\r\nRacer-Page-Credits: 1\r\n", + ), + &[memfd.file.as_raw_fd()], + ); + assert!(split(&output).0.starts_with("http/1.1 200")); + settle(&fixture); + let text = render(&telemetry); + assert!(text.contains("racer_gpu_standby 1\n"), "{text}"); + assert!(text.contains("racer_gpu_enabled 0\n"), "{text}"); + assert_eq!(calls_now(&calls), 0); + assert!(!links.gpu_landing_configured()); + // The first gpu-read starts it. A host buffer is not a GPU allocation, + // so the import rejects this one as invalid. + let output = exchange( + &fixture, + &direct_request("gpu-read", &fields("", 0, 5, GRANULE)), + &[memfd.file.as_raw_fd()], + ); + let (head, _) = split(&output); + assert!(head.starts_with("http/1.1 400"), "{head}"); + settle(&fixture); + assert_eq!(calls_now(&calls), 1); + assert!(!slot.standby()); + let helper = slot.get().unwrap().clone(); + let text = render(&telemetry); + assert!(text.contains("racer_gpu_standby 0\n"), "{text}"); + assert!(text.contains("racer_gpu_enabled 1\n"), "{text}"); + // RDMA landing was configured for the rails, and marked pending. + assert!(links.gpu_landing_configured()); + assert!(links.take_landing_pending()); + use crate::gpu_telemetry::GpuReadOutcome; + let gpu = fixture.listeners.metrics.gpu(); + assert_eq!(gpu.requests(GpuReadOutcome::Invalid), 1); + // Later reads use the started helper. + let allocation = allocate(&helper, 0, GRANULE); + let output = exchange( + &fixture, + &direct_request("gpu-read", &fields("", 0, 5, GRANULE)), + &[allocation.as_raw_fd()], + ); + let (head, _) = split(&output); + assert!(head.starts_with("http/1.1 200"), "{head}"); + assert_eq!(read_allocation(&allocation, 0, 5), b"hello"); + settle(&fixture); + assert_released(&helper); + assert_eq!(calls_now(&calls), 1); + assert_eq!(gpu.requests(GpuReadOutcome::Gpu), 1); +} + +#[test] +/// Concurrent first `gpu-read` requests share one start. +fn concurrent_first_gpu_reads_start_the_gpu_once() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let slot = slot( + GpuMode::Auto, + counted_loader(1, &calls), + Default::default(), + None, + ) + .unwrap(); + let fixture = lazy_fixture(b"hello".to_vec(), &slot); + let memfd = Memfd::new(); + let request = direct_request("gpu-read", &fields("", 0, 5, GRANULE)); + let mut sockets: Vec<_> = (0..4) + .map(|_| { + let socket = fixture.connect(); + send_with_fds(&socket, &request, &[memfd.file.as_raw_fd()]); + socket + }) + .collect(); + for socket in &mut sockets { + let output = fixture.receive(socket, true); + let (head, _) = split(&output); + // A host buffer is not a GPU allocation; the helper did start. + assert!(head.starts_with("http/1.1 400"), "{head}"); + } + drop(sockets); + settle(&fixture); + assert_eq!(calls_now(&calls), 1); + assert!(slot.get().is_some()); +} + +#[test] +/// A failed start or a node with no usable GPU leaves `gpu-read` off (501) +/// for good, so clients use host reads, and it is never retried. +fn auto_start_failure_is_501_and_not_retried() { + let failing = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let counted = failing.clone(); + let failed: crate::gpu::Loader = Box::new(move || { + counted.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Err(racer_gpu::Error::Unavailable) + }); + let empty = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + for (loader, calls) in [(failed, &failing), (counted_loader(0, &empty), &empty)] { + let telemetry = Arc::new(crate::gpu_telemetry::GpuTelemetry::default()); + let links = Arc::new(crate::rdma::LinkNode::default()); + let slot = slot( + GpuMode::Auto, + loader, + telemetry.clone(), + Some(links.clone()), + ) + .unwrap(); + telemetry.attach_slot(slot.clone()); + let fixture = lazy_fixture(b"hello".to_vec(), &slot); + let memfd = Memfd::new(); + for _ in 0..3 { + let output = exchange( + &fixture, + &direct_request("gpu-read", &fields("", 0, 5, GRANULE)), + &[memfd.file.as_raw_fd()], + ); + let (head, _) = split(&output); + assert!(head.starts_with("http/1.1 501"), "{head}"); + settle(&fixture); + } + assert_eq!(calls_now(calls), 1); + assert!(!slot.standby()); + assert!(slot.get().is_none()); + assert!(!links.gpu_landing_configured()); + use crate::gpu_telemetry::GpuReadOutcome; + let gpu = fixture.listeners.metrics.gpu(); + assert_eq!(gpu.requests(GpuReadOutcome::NotEnabled), 3); + let text = render(&telemetry); + assert!(text.contains("racer_gpu_standby 0\n"), "{text}"); + assert!(text.contains("racer_gpu_enabled 0\n"), "{text}"); + } +} + +#[test] +/// `true` starts eagerly and fails startup without a usable GPU; `false` +/// and an `auto` build without `cuda` never load a driver. +fn required_starts_eagerly_and_off_modes_never_load() { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let links = Arc::new(crate::rdma::LinkNode::default()); + let telemetry = Arc::new(crate::gpu_telemetry::GpuTelemetry::default()); + let required = slot( + GpuMode::Required, + counted_loader(1, &calls), + telemetry.clone(), + Some(links.clone()), + ) + .unwrap(); + assert_eq!(calls_now(&calls), 1); + assert!(!required.standby()); + assert!(required.get().is_some()); + assert!(links.gpu_landing_configured()); + assert!(render(&telemetry).contains("racer_gpu_enabled 1\n")); + { + let count = 0; + let none = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + assert!( + slot( + GpuMode::Required, + counted_loader(count, &none), + Default::default(), + None + ) + .is_err() + ); + assert_eq!(calls_now(&none), 1); + } + assert!( + slot( + GpuMode::Required, + Box::new(|| Err(racer_gpu::Error::Unavailable)), + Default::default(), + None + ) + .is_err() + ); + let off = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let disabled = slot( + GpuMode::Disabled, + counted_loader(1, &off), + Default::default(), + None, + ) + .unwrap(); + assert!(!disabled.standby()); + assert!(disabled.get().is_none()); + let mut config = crate::test_support::cluster::config(false); + config.enable_gpu = GpuMode::Auto; + let unbuilt = crate::app::gpu_slot( + &config, + false, + Default::default(), + None, + counted_loader(1, &off), + ) + .unwrap(); + assert!(!unbuilt.standby()); + assert!(futures::executor::block_on(unbuilt.get_or_start()).is_none()); + assert_eq!(calls_now(&off), 0); +} + +/// A loader like [`counted_loader`] on one GPU that first waits for `gate`, +/// standing in for a CUDA initialization that stalls. +fn gated_loader( + calls: &Arc, + gate: std::sync::mpsc::Receiver<()>, +) -> crate::gpu::Loader { + let load = counted_loader(1, calls); + Box::new(move || { + let _ = gate.recv(); + load() + }) +} + +#[test] +/// A `gpu-read` that expires or is cancelled while the GPU start stalls +/// stops waiting and releases its connection; the shared start keeps going +/// and a later `gpu-read` uses the helper it started. +fn stalled_gpu_start_honors_request_expiry_and_cancellation() { + for mode in ["expire", "cancel"] { + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let (open, gate) = std::sync::mpsc::channel(); + let slot = slot( + GpuMode::Auto, + gated_loader(&calls, gate), + Default::default(), + None, + ) + .unwrap(); + let mut fixture = lazy_fixture(b"hello".to_vec(), &slot); + let timeout = Duration::from_millis(200); + if mode == "expire" { + fixture.listeners = fixture.listeners.with_request_timeout(timeout); + } + let memfd = Memfd::new(); + let mut socket = fixture.connect(); + send_with_fds( + &socket, + &direct_request("gpu-read", &fields("", 0, 5, GRANULE)), + &[memfd.file.as_raw_fd()], + ); + let started = Instant::now(); + for _ in 0..16 { + fixture.pump(16); + } + // The start is pending, so the request is still waiting on it. + assert!(slot.standby(), "{mode}"); + assert_eq!(fixture.listeners.active_connections(), 1, "{mode}"); + if mode == "expire" { + std::thread::sleep( + (started + timeout + Duration::from_millis(20)) + .saturating_duration_since(Instant::now()), + ); + } else { + fixture.listeners.cancel_cache(&definition().id).unwrap(); + } + let output = fixture.receive(&mut socket, true); + let (head, _) = split(&output); + assert!(head.starts_with("http/1.1 503"), "{mode}: {head}"); + settle(&fixture); + // Giving up did not cancel the shared start. + assert!(slot.standby(), "{mode}"); + open.send(()).unwrap(); + let deadline = Instant::now() + Duration::from_secs(5); + while slot.standby() { + assert!(Instant::now() < deadline, "{mode}: start did not finish"); + std::thread::yield_now(); + } + assert_eq!(calls_now(&calls), 1, "{mode}"); + let helper = slot.get().unwrap().clone(); + if mode == "expire" { + // A later request gets the helper the stalled start produced. + fixture.listeners = fixture + .listeners + .with_request_timeout(Duration::from_secs(30)); + let allocation = allocate(&helper, 0, GRANULE); + let output = exchange( + &fixture, + &direct_request("gpu-read", &fields("", 0, 5, GRANULE)), + &[allocation.as_raw_fd()], + ); + let (head, _) = split(&output); + assert!(head.starts_with("http/1.1 200"), "{head}"); + assert_eq!(read_allocation(&allocation, 0, 5), b"hello"); + settle(&fixture); + assert_released(&helper); + assert_eq!(calls_now(&calls), 1); + } + } +} diff --git a/cmd/racer-dataplane/src/config.rs b/cmd/racer-dataplane/src/config.rs index 812c219bc4..7ea5880f87 100644 --- a/cmd/racer-dataplane/src/config.rs +++ b/cmd/racer-dataplane/src/config.rs @@ -131,12 +131,18 @@ pub struct Config { /// `gpu-read` support. Resolved against hardware at startup, on the GPU /// helper thread, because GPU handles cannot move between threads. pub enable_gpu: GpuMode, - /// Per-GPU cap on dataplane-owned ciphertext staging memory. Parsed and - /// validated now; plaintext `gpu-read` does not allocate staging. + /// Per-GPU size of the dataplane-owned staging buffer that `gpu-read` + /// decrypts sealed (ciphertext) pages through: two scratch slots plus + /// slots for cuFile reads. Allocated on a GPU's first sealed page; pages + /// that arrive as plaintext are copied straight to the client and never + /// touch it. Below two slots, GPU decrypt is off and the CPU decrypts. pub gpu_staging_bytes: u64, /// Per-GPU GPU memory that peers' NICs write `gpu-read` ciphertext into /// over RDMA (GPUDirect RDMA landing). 0 turns landing off. pub gpu_landing_bytes: u64, + /// Ask cuFile to read local NVMe into GPU memory over PCI P2PDMA + /// (`use_pci_p2pdma`) instead of nvidia-fs. Defaults off. + pub gpu_pci_p2pdma: bool, /// Opt into experimental opaque HTTP transit; materialized relay is the default. pub opaque_relay: bool, /// Peer-only TCP_NODELAY. Peer exchanges are strictly request/response with @@ -286,6 +292,7 @@ impl Config { enable_gpu, gpu_staging_bytes, gpu_landing_bytes, + gpu_pci_p2pdma: raw.gpu_pci_p2pdma.unwrap_or(false), opaque_relay: raw.opaque_relay.unwrap_or(false), peer_tcp_nodelay: raw.peer_tcp_nodelay.unwrap_or(true), memory_pressure: raw.memory_pressure.unwrap_or(true), @@ -719,6 +726,7 @@ const ENV_NAMES: &[&str] = &[ "ENABLE_GPU", "GPU_STAGING_BYTES", "GPU_LANDING_BYTES", + "GPU_PCI_P2PDMA", "ALLOW_SMT", "OPAQUE_RELAY", "PEER_TCP_NODELAY", @@ -799,6 +807,7 @@ struct RawConfig { gpu_staging_bytes: Option, #[serde(default, deserialize_with = "decimal")] gpu_landing_bytes: Option, + gpu_pci_p2pdma: Option, allow_smt: Option, opaque_relay: Option, peer_tcp_nodelay: Option, @@ -1096,7 +1105,7 @@ mod tests { .map(|name| { let value = match *name { "ALLOW_SMT" | "OPAQUE_RELAY" | "PEER_TCP_NODELAY" | "MEMORY_PRESSURE" - | "PPROF_ENABLED" => "true", + | "PPROF_ENABLED" | "GPU_PCI_P2PDMA" => "true", _ => "1", }; ((*name).to_owned(), value.to_owned()) @@ -1520,6 +1529,25 @@ mod tests { } } + #[test] + fn gpu_pci_p2pdma_parses_and_validates() { + assert!(!parse(&[]).unwrap().gpu_pci_p2pdma); + for (value, enabled) in [("true", true), ("false", false)] { + assert_eq!( + parse(&[("RACER_GPU_PCI_P2PDMA", value)]) + .unwrap() + .gpu_pci_p2pdma, + enabled + ); + } + for value in ["TRUE", "1", "yes", "", " true"] { + assert!( + parse(&[("RACER_GPU_PCI_P2PDMA", value)]).is_err(), + "{value:?}" + ); + } + } + #[test] fn native_auto_requires_build_and_available_library() { for requested in [None, Some("auto"), Some("false")] { diff --git a/cmd/racer-dataplane/src/gpu.rs b/cmd/racer-dataplane/src/gpu.rs index 39570e52a3..549da927ea 100644 --- a/cmd/racer-dataplane/src/gpu.rs +++ b/cmd/racer-dataplane/src/gpu.rs @@ -38,15 +38,19 @@ //! //! # RDMA landing //! -//! When RDMA is active, each rail may get one GPU's landing buffer +//! When RDMA is active, each decrypting GPU may get one landing buffer //! ([`GpuHelper::landing`]): separate GPU memory, capped by -//! `RACER_GPU_LANDING_BYTES`, exported once as a dma-buf and registered with -//! the rail's NIC as an `rdma_verbs` device pool. A peer's NIC writes a -//! `gpu-read` page's ciphertext straight into it, and the GPU decrypts from -//! there ([`LandedCiphertext`]). The CPU never reads or copies those bytes. -//! Landing buffers are never freed while the process runs, and are leaked on -//! purpose when the helper exits, because a NIC registration may still cover -//! them. +//! `RACER_GPU_LANDING_BYTES`, exported once as a dma-buf. Disjoint slices of +//! it ([`GpuLanding::slice`]) are registered with the rails' NICs as +//! `rdma_verbs` device pools. A peer's NIC writes a `gpu-read` page's +//! ciphertext straight into it, and the GPU decrypts from there +//! ([`LandedCiphertext`]). The CPU never reads or copies those bytes. +//! The buffer is freed when the last [`GpuLanding`] referring to it drops: +//! the NIC registrations and received pages hold one, so this happens only +//! after every registration is gone. A registration `rdma_verbs` leaks on +//! purpose keeps its reference, so that memory is never freed. Buffers still +//! referenced when the helper exits are leaked on purpose for the same +//! reason. //! //! At startup the helper probes each GPU with a zero-length decrypt. A GPU //! whose shim has no decrypt kernels, or whose staging cap cannot hold the two @@ -73,6 +77,7 @@ use std::ops::Range; use std::os::fd::OwnedFd; use std::sync::Arc; use std::sync::Mutex; +use std::sync::Weak; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::AtomicUsize; @@ -204,30 +209,74 @@ impl Drop for StagedCiphertext { /// value works: the address is local to the memory registration. const LANDING_IOVA: u64 = 1 << 40; -/// One GPU's RDMA landing buffer, shared with `rdma_verbs` as a dma-buf. -/// See the module docs, "RDMA landing". -pub struct GpuLanding { - /// Export of the whole landing buffer. - fd: OwnedFd, +/// One GPU landing allocation. Dropping the last reference asks the helper +/// to free the GPU memory. +struct LandingAllocation { + /// Helper-side key of the memory. + id: u64, ordinal: i32, + /// Export of the whole allocation. `None` only during drop. + fd: Option, + jobs: mpsc::Sender, +} + +impl Drop for LandingAllocation { + fn drop(&mut self) { + if let Some(fd) = self.fd.take() { + // The helper closes the export before it frees the memory. If + // it is gone, the memory is leaked with it. + let _ = self.jobs.send(Job::FreeLanding { id: self.id, fd }); + } + } +} + +/// A slice of one GPU's RDMA landing buffer, shared with `rdma_verbs` as a +/// dma-buf. See the module docs, "RDMA landing". +pub struct GpuLanding { + allocation: Arc, + /// Start of the slice in the allocation. + offset: u64, length: u64, } impl GpuLanding { /// The GPU the memory is on. pub fn ordinal(&self) -> i32 { - self.ordinal + self.allocation.ordinal + } + + /// `length` bytes at `offset` of this slice, as their own slice that + /// keeps the GPU memory alive. `None` unless both are 4 KiB aligned, + /// `length` is nonzero, and the range fits. + pub fn slice(&self, offset: u64, length: u64) -> Option> { + let aligned = offset.is_multiple_of(4096) && length.is_multiple_of(4096); + (aligned && length != 0 && offset.checked_add(length)? <= self.length).then(|| { + Arc::new(GpuLanding { + allocation: self.allocation.clone(), + offset: self.offset + offset, + length, + }) + }) + } + + /// True if `other` is a slice of the same GPU memory. + pub fn same_allocation(&self, other: &GpuLanding) -> bool { + Arc::ptr_eq(&self.allocation, &other.allocation) } } impl rdma_verbs::DmabufMemory for GpuLanding { fn fd(&self) -> std::os::fd::BorrowedFd<'_> { use std::os::fd::AsFd; - self.fd.as_fd() + self.allocation + .fd + .as_ref() + .expect("landing export lives until drop") + .as_fd() } fn offset(&self) -> u64 { - 0 + self.offset } fn length(&self) -> usize { @@ -235,7 +284,7 @@ impl rdma_verbs::DmabufMemory for GpuLanding { } fn iova(&self) -> u64 { - LANDING_IOVA * (self.ordinal as u64 + 1) + LANDING_IOVA * (self.allocation.ordinal as u64 + 1) + self.offset } } @@ -245,6 +294,9 @@ pub struct LandedCiphertext { /// Holds the slot. `None` only for test pages written without a NIC. _page: Option, ordinal: i32, + /// Helper-side key of the landing allocation. + allocation: u64, + /// Byte offset in the allocation. offset: u64, length: u64, } @@ -257,8 +309,9 @@ impl LandedCiphertext { return Err(Error::Internal); } Ok(Self { - ordinal: landing.ordinal, - offset: page.offset(), + ordinal: landing.ordinal(), + allocation: landing.allocation.id, + offset: landing.offset + page.offset(), length: page.len() as u64, _page: Some(page), }) @@ -271,8 +324,9 @@ impl LandedCiphertext { assert!(offset + length <= landing.length); Self { _page: None, - ordinal: landing.ordinal, - offset, + ordinal: landing.ordinal(), + allocation: landing.allocation.id, + offset: landing.offset + offset, length, } } @@ -292,7 +346,7 @@ impl LandedCiphertext { self.length == 0 } - /// Byte offset in the landing buffer. + /// Byte offset in the landing allocation. pub(crate) fn offset(&self) -> u64 { self.offset } @@ -318,6 +372,48 @@ pub struct GpuCounters { /// Sealed pages decrypted from an RDMA landing buffer. pub landed: AtomicU64, + + /// Verified plaintext pages copied from host memory (decrypted by the + /// CPU or already plaintext in cache). + pub plaintext: AtomicU64, + + /// Devices (`st_dev`) cuFile rejected as unsupported. Each device is + /// counted once; cuFile stays on for other devices. + pub cufile_unsupported: AtomicU64, + + /// Times cuFile was turned off, indexed by [`CufileDisable`]. + pub cufile_disabled: [AtomicU64; CufileDisable::ALL.len()], +} + +/// Why cuFile was turned off. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum CufileDisable { + /// The cuFile driver failed to open at startup. + DriverOpen, + /// A staging buffer could not be registered with cuFile. + StagingRegister, + /// A file could not be registered with cuFile. + FileOpen, +} + +impl CufileDisable { + pub const ALL: [Self; 3] = [Self::DriverOpen, Self::StagingRegister, Self::FileOpen]; + + /// Metric label value. + pub fn label(self) -> &'static str { + match self { + Self::DriverOpen => "driver_open", + Self::StagingRegister => "staging_register", + Self::FileOpen => "file_open", + } + } +} + +impl GpuCounters { + /// Count one cuFile turn-off. + pub(crate) fn cufile_disabled(&self, reason: CufileDisable) { + self.cufile_disabled[reason as usize].fetch_add(1, Ordering::Relaxed); + } } /// Per-GPU state shared with I/O threads. @@ -337,6 +433,8 @@ enum Job { fd: OwnedFd, allocation_size: u64, window: Range, + /// The client's `Racer-Gpu-Uuid`, when it sent one. + gpu_uuid: Option<[u8; 16]>, reply: oneshot::Sender>, }, /// Read ciphertext from a file straight into a staging slot. @@ -359,10 +457,13 @@ enum Job { }, /// Allocate a GPU's RDMA landing buffer and export it as a dma-buf. Landing { + id: u64, ordinal: i32, length: u64, reply: oneshot::Sender>, }, + /// Close a landing buffer's export, then free it: nothing refers to it. + FreeLanding { id: u64, fd: OwnedFd }, #[cfg(test)] Run(Box), } @@ -379,20 +480,39 @@ pub struct GpuHelper { cufile: Arc, counters: Arc, kernels_off: Arc>>, - /// Landing buffers made so far, one per GPU at most. - landings: Mutex>>, + /// Live landing buffers, one per GPU at most. Weak: a buffer is freed + /// once nothing registers or reads it, and made again on demand. + landings: Mutex>>, } impl GpuHelper { /// Resolve `mode` against hardware. Returns `None` when GPU mode is off. - pub fn start(mode: GpuMode, staging_bytes: u64) -> Result>> { + pub fn start(mode: GpuMode, staging_bytes: u64, pci_p2pdma: bool) -> Result>> { + Self::start_built( + mode, + cfg!(feature = "cuda"), + staging_bytes, + pci_p2pdma, + Box::new(racer_gpu::Gpu::load), + ) + } + + /// [`Self::start`] with `built` standing in for the `cuda` feature and a + /// custom backend loader. `loader` is not called when GPU mode is off. + pub fn start_built( + mode: GpuMode, + built: bool, + staging_bytes: u64, + pci_p2pdma: bool, + loader: Loader, + ) -> Result>> { match mode { GpuMode::Disabled => Ok(None), - GpuMode::Auto if !cfg!(feature = "cuda") => { + GpuMode::Auto if !built => { eprintln!("racer-dataplane: stage=gpu state=disabled reason=not-built"); Ok(None) } - _ => Self::start_with_staging(mode, staging_bytes, Box::new(racer_gpu::Gpu::load)), + _ => Self::start_with_options(mode, staging_bytes, pci_p2pdma, loader), } } @@ -406,6 +526,17 @@ impl GpuHelper { mode: GpuMode, staging_bytes: u64, loader: Loader, + ) -> Result>> { + Self::start_with_options(mode, staging_bytes, false, loader) + } + + /// Start the helper with a custom backend loader, a per-GPU staging cap, + /// and whether cuFile should use PCI P2PDMA. + pub fn start_with_options( + mode: GpuMode, + staging_bytes: u64, + pci_p2pdma: bool, + loader: Loader, ) -> Result>> { if mode == GpuMode::Disabled { return Ok(None); @@ -423,6 +554,7 @@ impl GpuHelper { counters: counters.clone(), kernels_off: kernels_off.clone(), slots, + pci_p2pdma, }; std::thread::Builder::new() .name("racer-gpu".into()) @@ -458,10 +590,11 @@ impl GpuHelper { }) .collect::>(); eprintln!( - "racer-dataplane: stage=gpu state=enabled usable={} decrypt={} staging_slots={}", + "racer-dataplane: stage=gpu state=enabled usable={} decrypt={} staging_slots={} pci_p2pdma={}", devices.len(), shared.iter().filter(|device| device.decrypt).count(), slots, + pci_p2pdma, ); Ok(Some(Arc::new(Self { jobs, @@ -563,28 +696,36 @@ impl GpuHelper { .landings .lock() .ok() - .and_then(|l| l.get(&ordinal).cloned()) + .and_then(|l| l.get(&ordinal).and_then(Weak::upgrade)) { return Ok(landing); } + let id = self.next.fetch_add(1, Ordering::Relaxed); let (reply, mut receiver) = oneshot::channel(); self.send(Job::Landing { + id, ordinal, length, reply, })?; - let fd = answer(&mut receiver, scope).await??; + // Dropped on any failure below, which frees the memory again. + let made = Arc::new(GpuLanding { + allocation: Arc::new(LandingAllocation { + id, + ordinal, + fd: Some(answer(&mut receiver, scope).await??), + jobs: self.jobs.clone(), + }), + offset: 0, + length, + }); let mut landings = self.landings.lock().map_err(|_| Error::Internal)?; - Ok(landings - .entry(ordinal) - .or_insert_with(|| { - Arc::new(GpuLanding { - fd, - ordinal, - length, - }) - }) - .clone()) + // A concurrent caller may have made one first: keep that one. + if let Some(landing) = landings.get(&ordinal).and_then(Weak::upgrade) { + return Ok(landing); + } + landings.insert(ordinal, Arc::downgrade(&made)); + Ok(made) } /// Where a request on GPU `ordinal` should have sealed pages delivered, @@ -645,6 +786,7 @@ impl GpuService for GpuClient { descriptor: OwnedFd, destination: Destination, allocation_size: u64, + gpu_uuid: Option<[u8; 16]>, ) -> Result> { let window = destination.offset..destination.end(); if window.start >= window.end || window.end > allocation_size { @@ -669,6 +811,7 @@ impl GpuService for GpuClient { fd: descriptor, allocation_size, window, + gpu_uuid, reply, })?; sink.import = Some(answer); @@ -808,6 +951,8 @@ struct HelperState { kernels_off: Arc>>, /// Staging slots per GPU. slots: u64, + /// Ask cuFile to use PCI P2PDMA. + pci_p2pdma: bool, } /// A registered cuFile handle and the reads using it. @@ -872,6 +1017,8 @@ fn reader_main( struct Helper { gpu: racer_gpu::Gpu, contexts: Vec, + /// Each context's GPU UUID, in the same order as `contexts`. + uuids: Vec<[u8; 16]>, /// Each GPU's staging buffer, allocated on first sealed page. staging: HashMap, /// Registered files by (device, inode). The slab descriptors are few and @@ -880,6 +1027,9 @@ struct Helper { /// Registered files by id, including retired ones that reads still use. open: HashMap, next_file: u64, + /// Devices (`st_dev`) whose files cuFile rejected as unsupported. Reads + /// of files on them use the host path without asking cuFile again. + unsupported: std::collections::HashSet, /// Reads for the reader threads. `None` once they are told to stop. reads: Option>, /// Reads the reader threads finished, so their files can be released. @@ -888,9 +1038,10 @@ struct Helper { /// Opens the cuFile driver ahead of the first sealed page. opener: Option>, requests: HashMap, - /// RDMA landing buffers by GPU. Never freed: a NIC may hold a - /// registration over them. - landing: HashMap, + /// RDMA landing buffers by allocation id. Freed by `FreeLanding` once + /// nothing refers to them; leaked if still referenced at exit, because + /// a NIC may hold a registration over them. + landing: HashMap, state: HelperState, } @@ -904,6 +1055,9 @@ fn helper_main( ) { let mut helper = match discover(loader) { Ok((gpu, contexts, usable)) => { + // Before any reader or open: the flag applies only to the open + // that loads cuFile. + gpu.set_pci_p2pdma(state.pci_p2pdma); let decrypt = contexts.iter().map(probe_decrypt).collect(); let (reads, queue) = mpsc::channel(); let queue = Arc::new(Mutex::new(queue)); @@ -924,15 +1078,18 @@ fn helper_main( readers.push(thread); } } - let opener = open_cufile(&gpu, &state.cufile); + let opener = open_cufile(&gpu, &state.cufile, &state.counters); + let uuids = usable.iter().map(|device| device.uuid).collect(); let _ = ready.send(Ok((usable, decrypt))); Helper { gpu, contexts, + uuids, staging: HashMap::new(), files: HashMap::new(), open: HashMap::new(), next_file: 0, + unsupported: std::collections::HashSet::new(), reads: (!readers.is_empty()).then_some(reads), done, readers, @@ -978,17 +1135,20 @@ fn helper_main( fn open_cufile( gpu: &racer_gpu::Gpu, cufile: &Arc, + counters: &Arc, ) -> Option> { if !cufile.load(Ordering::Relaxed) { return None; } let (reader, cufile) = (gpu.file_reader(), cufile.clone()); + let counters = counters.clone(); std::thread::Builder::new() .name("racer-gpu-cufile".into()) .spawn(move || { if let Err(error) = reader.open_driver() && cufile.swap(false, Ordering::Relaxed) { + counters.cufile_disabled(CufileDisable::DriverOpen); eprintln!("racer-dataplane: stage=gpu state=cufile-disabled reason={error}"); } }) @@ -1101,9 +1261,10 @@ impl Helper { fd, allocation_size, window, + gpu_uuid, reply, } => { - let result = self.import(id, fd, allocation_size, window); + let result = self.import(id, fd, allocation_size, window, gpu_uuid); self.note_leaks(); let _ = reply.send(result); } @@ -1140,22 +1301,38 @@ impl Helper { } } Job::Landing { + id, ordinal, length, reply, } => { - let result = self.landing(ordinal, length); + let result = self.landing(id, ordinal, length); + if result.is_err() { + self.landing.remove(&id); + } + self.note_leaks(); + // The caller stopped waiting, so no FreeLanding will come: + // close the export and free the memory now. + if let Err(Ok(fd)) = reply.send(result) { + drop(fd); + self.landing.remove(&id); + } + } + Job::FreeLanding { id, fd } => { + drop(fd); + if self.landing.remove(&id).is_some() { + eprintln!("racer-dataplane: stage=gpu-landing allocation={id} state=freed"); + } self.note_leaks(); - let _ = reply.send(result); } #[cfg(test)] Job::Run(f) => f(), } } - /// Allocate GPU `ordinal`'s landing buffer, once, and export it. - fn landing(&mut self, ordinal: i32, length: u64) -> Result { - if !self.landing.contains_key(&ordinal) { + /// Allocate landing buffer `id` on GPU `ordinal`, and export it. + fn landing(&mut self, id: u64, ordinal: i32, length: u64) -> Result { + if !self.landing.contains_key(&id) { let context = self .contexts .iter() @@ -1167,9 +1344,9 @@ impl Helper { ); Error::Overloaded })?; - self.landing.insert(ordinal, buffer); + self.landing.insert(id, buffer); } - let buffer = self.landing.get(&ordinal).ok_or(Error::Internal)?; + let buffer = self.landing.get(&id).ok_or(Error::Internal)?; if buffer.size() != length { return Err(Error::InvalidConfiguration); } @@ -1181,13 +1358,45 @@ impl Helper { }) } - /// Try each usable GPU until one owns the allocation. - fn import(&mut self, id: u64, fd: OwnedFd, size: u64, window: Range) -> Result { + /// Import the allocation on the GPU that owns it. With the client's GPU + /// UUID, that GPU is tried first: device numbers differ between the + /// client's container and the dataplane, but UUIDs do not. An unknown + /// UUID is a GPU the dataplane cannot use, so the read falls back to the + /// host (501). Without a UUID (older clients), each usable GPU is tried + /// until one owns the allocation. + fn import( + &mut self, + id: u64, + fd: OwnedFd, + size: u64, + window: Range, + gpu_uuid: Option<[u8; 16]>, + ) -> Result { + let mut order: Vec = (0..self.contexts.len()).collect(); + if let Some(uuid) = gpu_uuid { + let Some(first) = self.uuids.iter().position(|known| *known == uuid) else { + eprintln!( + "racer-dataplane: stage=gpu state=import-skipped reason=unknown-gpu-uuid uuid={}", + racer_gpu::uuid_string(&uuid) + ); + return Err(Error::NotImplemented); + }; + order.retain(|&index| index != first); + order.insert(0, first); + } let mut last = racer_gpu::Error::WrongDevice; - for context in &self.contexts { - let attempt = fd.try_clone().map_err(|_| Error::Unavailable)?; - match context.import(attempt, size, window.clone()) { + for (attempt, index) in order.into_iter().enumerate() { + let context = &self.contexts[index]; + let handle = fd.try_clone().map_err(|_| Error::Unavailable)?; + match context.import(handle, size, window.clone()) { Ok(import) => { + if attempt > 0 && gpu_uuid.is_some() { + eprintln!( + "racer-dataplane: stage=gpu state=uuid-mismatch ordinal={} uuid={}", + context.ordinal(), + racer_gpu::uuid_string(&self.uuids[index]) + ); + } let stream = context.stream().map_err(import_error)?; self.requests.insert(id, Request { stream, import }); return Ok(context.ordinal()); @@ -1214,7 +1423,12 @@ impl Helper { .copy_from_host(destination, offset, plaintext.clone()); // Synchronize even after a failed queue call: part of it may run. let synchronized = request.stream.synchronize(); - queued.and(synchronized).map_err(gpu_error) + queued.and(synchronized).map_err(gpu_error)?; + self.state + .counters + .plaintext + .fetch_add(1, Ordering::Relaxed); + Ok(()) } /// This GPU's staging buffer, allocated on first use. @@ -1254,6 +1468,9 @@ impl Helper { Ok(identity) => identity, Err(error) => return Err((error, reply)), }; + if self.unsupported.contains(&identity.0) { + return Err((Error::Unavailable, reply)); + } if let Err(error) = self.staging(staged.ordinal) { return Err((error, reply)); } @@ -1265,7 +1482,7 @@ impl Helper { .counters .file_failures .fetch_add(1, Ordering::Relaxed); - self.disable_cufile(error); + self.disable_cufile(error, CufileDisable::StagingRegister); return Err((Error::Unavailable, reply)); } let id = match self.files.get(&identity) { @@ -1296,7 +1513,11 @@ impl Helper { .counters .file_failures .fetch_add(1, Ordering::Relaxed); - self.disable_cufile(error); + if let racer_gpu::Error::Unsupported(_) = error { + self.unsupported_device(identity.0, error); + return Err((Error::Unavailable, reply)); + } + self.disable_cufile(error, CufileDisable::FileOpen); return Err((Error::Unavailable, reply)); } } @@ -1366,9 +1587,25 @@ impl Helper { } } + /// Stop trying cuFile for files on `device` after cuFile rejected one of + /// them. Other devices keep using cuFile. + fn unsupported_device(&mut self, device: u64, error: racer_gpu::Error) { + if self.unsupported.insert(device) { + self.state + .counters + .cufile_unsupported + .fetch_add(1, Ordering::Relaxed); + let (major, minor) = (libc::major(device), libc::minor(device)); + eprintln!( + "racer-dataplane: stage=gpu state=cufile-unsupported dev={major}:{minor} reason={error}" + ); + } + } + /// Stop trying cuFile after a failure that is not about one read. - fn disable_cufile(&self, error: racer_gpu::Error) { + fn disable_cufile(&self, error: racer_gpu::Error, reason: CufileDisable) { if self.state.cufile.swap(false, Ordering::Relaxed) { + self.state.counters.cufile_disabled(reason); eprintln!("racer-dataplane: stage=gpu state=cufile-disabled reason={error}"); } } @@ -1417,7 +1654,10 @@ impl Helper { if landed.ordinal() != ordinal || landed.len() != length + 16 { return Err(Error::Internal); } - let buffer = self.landing.get(&ordinal).ok_or(Error::Internal)?; + let buffer = self + .landing + .get(&landed.allocation) + .ok_or(Error::Internal)?; (buffer, landed.offset()) } }; @@ -1517,6 +1757,11 @@ pub(crate) mod testing { SIM.with(|sim| sim.borrow().as_ref().unwrap().0.clone()) } + /// Live simulated staging and landing buffers. Run on the helper thread. + pub(crate) fn live_staging() -> usize { + simulation().live().staging + } + /// A test scope with a few seconds on the real clock. pub(crate) fn scope() -> RequestScope { RequestScope::new( @@ -1539,8 +1784,9 @@ pub(crate) mod testing { /// a NIC would. pub(crate) fn nic_write(landing: &GpuLanding, at: u64, bytes: &[u8]) { use std::os::unix::fs::FileExt; - let file = std::fs::File::from(landing.fd.try_clone().unwrap()); - file.write_all_at(bytes, at).unwrap(); + let fd = landing.allocation.fd.as_ref().unwrap(); + let file = std::fs::File::from(fd.try_clone().unwrap()); + file.write_all_at(bytes, landing.offset + at).unwrap(); } /// Import a fresh allocation on GPU `ordinal` as request `id`. @@ -1554,6 +1800,7 @@ pub(crate) mod testing { fd: fd.try_clone().unwrap(), allocation_size: size, window: 0..size, + gpu_uuid: None, reply, }) .unwrap(); @@ -1600,6 +1847,15 @@ pub(crate) mod testing { pub(crate) fn helper_with( numa: &[Option], setup: impl FnOnce(&racer_gpu::simulation::Simulation) + Send + 'static, + ) -> Arc { + helper_with_options(numa, false, setup) + } + + /// A helper as [`helper_with`] makes, with cuFile PCI P2PDMA set. + pub(crate) fn helper_with_options( + numa: &[Option], + pci_p2pdma: bool, + setup: impl FnOnce(&racer_gpu::simulation::Simulation) + Send + 'static, ) -> Arc { let devices: Vec = numa .iter() @@ -1610,8 +1866,10 @@ pub(crate) mod testing { device }) .collect(); - GpuHelper::start_with( + GpuHelper::start_with_options( GpuMode::Required, + crate::config::DEFAULT_GPU_STAGING_BYTES, + pci_p2pdma, Box::new(move || { let simulation = racer_gpu::simulation::Simulation::new(devices); setup(&simulation); @@ -1862,6 +2120,71 @@ mod tests { assert_eq!(helper.counters().staged.load(Ordering::Relaxed), 2); } + #[test] + /// cuFile rejecting one file as unsupported keeps cuFile on, remembers + /// the file's device so it is not registered again, and leaves files on + /// other devices on cuFile. + fn unsupported_file_skips_its_device_only() { + use racer_gpu::simulation::Operation; + let helper = helper(&[None]); + let fd = slab(&[7; 8192]); + helper.run(|| simulation().fault(Operation::File, -11)); + assert_eq!(stage(&helper, &fd, 0, 4096).err(), Some(Error::Unavailable)); + assert!(helper.cufile()); + assert_eq!(helper.counters().file_failures.load(Ordering::Relaxed), 1); + assert_eq!( + helper.counters().cufile_unsupported.load(Ordering::Relaxed), + 1 + ); + assert_eq!(operations(&helper, Operation::File), 1); + // Same device: no second registration, still the host path. + let other = slab(&[8; 4096]); + assert_eq!( + stage(&helper, &other, 0, 4096).err(), + Some(Error::Unavailable) + ); + assert_eq!(operations(&helper, Operation::File), 1); + assert_eq!(helper.counters().file_failures.load(Ordering::Relaxed), 1); + // A file on another device still goes through cuFile. + let elsewhere: OwnedFd = std::fs::File::open("/proc/self/exe").unwrap().into(); + stage(&helper, &elsewhere, 0, 4096).unwrap(); + assert_eq!(operations(&helper, Operation::File), 2); + assert_eq!(helper.counters().staged.load(Ordering::Relaxed), 1); + assert!(helper.cufile()); + // The device is counted once, and cuFile was not turned off. + assert_eq!( + helper.counters().cufile_unsupported.load(Ordering::Relaxed), + 1 + ); + assert_eq!( + helper.counters().cufile_disabled[CufileDisable::FileOpen as usize] + .load(Ordering::Relaxed), + 0 + ); + } + + #[test] + /// A driver failure while registering a file still turns cuFile off for + /// the whole node. + fn driver_failure_on_register_turns_cufile_off() { + use racer_gpu::simulation::Operation; + let helper = helper(&[None]); + let fd = slab(&[7; 4096]); + helper.run(|| simulation().fault(Operation::File, -5)); + assert_eq!(stage(&helper, &fd, 0, 4096).err(), Some(Error::Unavailable)); + assert!(!helper.cufile()); + assert_eq!(helper.counters().file_failures.load(Ordering::Relaxed), 1); + assert_eq!( + helper.counters().cufile_disabled[CufileDisable::FileOpen as usize] + .load(Ordering::Relaxed), + 1 + ); + assert_eq!( + helper.counters().cufile_unsupported.load(Ordering::Relaxed), + 0 + ); + } + /// Wait for the helper's background work to reach `done`. fn eventually(mut done: impl FnMut() -> bool) { let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); @@ -1883,6 +2206,32 @@ mod tests { assert_eq!(helper.counters().staged.load(Ordering::Relaxed), 1); } + #[test] + /// `RACER_GPU_PCI_P2PDMA` reaches the startup cuFile driver open, and is + /// off by default. + fn cufile_driver_opens_with_pci_p2pdma_setting() { + use racer_gpu::simulation::Operation; + for enabled in [false, true] { + let helper = helper_with_options(&[None], enabled, |_| {}); + eventually(|| operations(&helper, Operation::FileOpen) == 1); + assert_eq!(helper.run(|| simulation().pci_p2pdma()), Some(enabled)); + assert!(helper.cufile()); + } + } + + #[test] + /// A driver open with PCI P2PDMA that fails turns cuFile off, so pages + /// use the host path. + fn cufile_pci_p2pdma_open_failure_turns_cufile_off() { + use racer_gpu::simulation::Operation; + let helper = helper_with_options(&[None], true, |simulation| { + simulation.fault(Operation::FileOpen, -3) + }); + eventually(|| !helper.cufile()); + assert_eq!(helper.run(|| simulation().pci_p2pdma()), None); + assert_eq!(helper.counters().staged.load(Ordering::Relaxed), 0); + } + #[test] /// A driver that fails to open at startup turns cuFile off before any /// sealed page asks for it. diff --git a/cmd/racer-dataplane/src/gpu_slot.rs b/cmd/racer-dataplane/src/gpu_slot.rs new file mode 100644 index 0000000000..c153bd5887 --- /dev/null +++ b/cmd/racer-dataplane/src/gpu_slot.rs @@ -0,0 +1,316 @@ +//! The process-wide GPU helper slot. +//! +//! `RACER_ENABLE_GPU=auto` must not touch a GPU until a user asks for one, so +//! the helper starts on the first `gpu-read` request head instead of at +//! startup. [`GpuSlot::get`] never starts it; only [`GpuSlot::get_or_start`] +//! does, and only the `gpu-read` path calls that. +//! +//! The start runs once per process on its own thread: concurrent first +//! requests all wait on that one start, and no I/O worker blocks on it. The +//! outcome is final. A start that fails or finds no usable GPU leaves GPU +//! mode off for the life of the process; it is not retried, so a bad node +//! pays the probe cost once and every later `gpu-read` gets 501 at once. + +use crate::gpu::GpuHelper; +use futures::channel::oneshot; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; + +/// Starts the helper. `Ok(None)` and `Err` both mean GPU mode is off. +pub type Starter = Box crate::error::Result>> + Send>; + +/// Runs once with the started helper, before any waiter sees it. +pub type OnStart = Box) + Send>; + +enum Pending { + /// Not started yet; holds what starts it. + Idle(Starter, Option), + /// The start thread is running; these wait for it. + Starting(Vec>), + /// Resolved; see `helper`. + Done, +} + +/// Shared with the start thread. +struct Inner { + helper: OnceLock>>, + pending: Mutex, +} + +impl Inner { + /// Publish the outcome, then wake every waiter. + fn finish(&self, helper: Option>) { + let _ = self.helper.set(helper); + let waiters = match std::mem::replace(&mut *self.lock(), Pending::Done) { + Pending::Starting(waiters) => waiters, + _ => Vec::new(), + }; + for waiter in waiters { + let _ = waiter.send(()); + } + } + + fn lock(&self) -> std::sync::MutexGuard<'_, Pending> { + self.pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +/// The process-wide GPU helper, started now or on first use. +pub struct GpuSlot { + inner: Arc, +} + +impl Default for GpuSlot { + /// GPU mode off. + fn default() -> Self { + Self::ready(None) + } +} + +impl GpuSlot { + /// A slot whose outcome is already known (eager start, off, or tests). + pub fn ready(helper: Option>) -> Self { + Self { + inner: Arc::new(Inner { + helper: OnceLock::from(helper), + pending: Mutex::new(Pending::Done), + }), + } + } + + /// A slot that runs `start` on the first [`Self::get_or_start`], then + /// `on_start` with the helper when one started. + pub fn lazy(start: Starter, on_start: OnStart) -> Self { + Self { + inner: Arc::new(Inner { + helper: OnceLock::new(), + pending: Mutex::new(Pending::Idle(start, Some(on_start))), + }), + } + } + + /// The helper, if it has started. Never starts it. + pub fn get(&self) -> Option<&Arc> { + self.inner.helper.get().and_then(Option::as_ref) + } + + /// Whether the slot still waits for its first `gpu-read`. + pub fn standby(&self) -> bool { + self.inner.helper.get().is_none() + } + + /// The helper, starting it on first use. Waits without blocking the + /// calling thread; concurrent callers share one start. + pub async fn get_or_start(&self) -> Option> { + if let Some(helper) = self.inner.helper.get() { + return helper.clone(); + } + let waiter = { + let mut pending = self.inner.lock(); + // Checked under the lock: `finish` publishes before it takes the + // lock to drain waiters, so no waiter is missed. + if let Some(helper) = self.inner.helper.get() { + return helper.clone(); + } + let (sender, receiver) = oneshot::channel(); + match std::mem::replace(&mut *pending, Pending::Done) { + Pending::Idle(start, on_start) => { + *pending = Pending::Starting(vec![sender]); + drop(pending); + self.spawn(start, on_start); + } + Pending::Starting(mut waiters) => { + // Drop waiters whose request gave up, so a stalled + // start does not grow this without bound. + waiters.retain(|waiter| !waiter.is_canceled()); + waiters.push(sender); + *pending = Pending::Starting(waiters); + } + // Done is only set together with `helper`; unreachable. + Pending::Done => return None, + } + receiver + }; + let _ = waiter.await; + self.inner.helper.get().cloned().flatten() + } + + /// Run the start on its own thread: it blocks on driver calls. + fn spawn(&self, start: Starter, on_start: Option) { + let inner = self.inner.clone(); + let spawned = std::thread::Builder::new() + .name("racer-gpu-start".into()) + .spawn(move || { + let helper = start().unwrap_or_else(|error| { + eprintln!("racer-dataplane: stage=gpu state=disabled reason={error}"); + None + }); + if let (Some(helper), Some(on_start)) = (&helper, on_start) { + on_start(helper); + } + inner.finish(helper); + }); + if spawned.is_err() { + eprintln!("racer-dataplane: stage=gpu state=disabled reason=spawn-failed"); + self.inner.finish(None); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::future::Future; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + /// A lazy slot whose start counts its calls, waits for `gate`, and + /// returns `helper`. + fn counted( + calls: &Arc, + gate: std::sync::mpsc::Receiver<()>, + helper: Option>, + ) -> GpuSlot { + let calls = calls.clone(); + GpuSlot::lazy( + Box::new(move || { + calls.fetch_add(1, Ordering::SeqCst); + let _ = gate.recv(); + Ok(helper) + }), + Box::new(|_| {}), + ) + } + + #[test] + /// Reading a lazy slot never starts it. + fn get_never_starts() { + let calls = Arc::new(AtomicUsize::new(0)); + let (_open, gate) = std::sync::mpsc::channel(); + let slot = counted(&calls, gate, None); + for _ in 0..3 { + assert!(slot.get().is_none()); + assert!(slot.standby()); + } + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + + #[test] + /// Concurrent first callers wait on one start and see the same helper. + fn concurrent_callers_share_one_start() { + let calls = Arc::new(AtomicUsize::new(0)); + let (open, gate) = std::sync::mpsc::channel(); + let helper = crate::gpu::testing::helper(&[None]); + let slot = Arc::new(counted(&calls, gate, Some(helper.clone()))); + let waiters: Vec<_> = (0..8) + .map(|_| { + let slot = slot.clone(); + std::thread::spawn(move || futures::executor::block_on(slot.get_or_start())) + }) + .collect(); + // Let every waiter queue up behind the one start. + std::thread::sleep(std::time::Duration::from_millis(50)); + assert!(slot.standby()); + open.send(()).unwrap(); + for waiter in waiters { + let got = waiter.join().unwrap().unwrap(); + assert!(Arc::ptr_eq(&got, &helper)); + } + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(Arc::ptr_eq(slot.get().unwrap(), &helper)); + assert!(!slot.standby()); + } + + #[test] + /// A start that fails or finds nothing is final: no retry. + fn failure_is_final() { + for result in [Ok(None), Err(crate::error::Error::Unavailable)] { + let calls = Arc::new(AtomicUsize::new(0)); + let counted = calls.clone(); + let hooked = Arc::new(AtomicUsize::new(0)); + let hook = hooked.clone(); + let mut result = Some(result); + let slot = GpuSlot::lazy( + Box::new(move || { + counted.fetch_add(1, Ordering::SeqCst); + result.take().unwrap() + }), + Box::new(move |_| { + hook.fetch_add(1, Ordering::SeqCst); + }), + ); + for _ in 0..3 { + assert!(futures::executor::block_on(slot.get_or_start()).is_none()); + } + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(hooked.load(Ordering::SeqCst), 0); + assert!(!slot.standby()); + } + } + + #[test] + /// The start hook runs once, before any waiter sees the helper. + fn on_start_runs_before_waiters_resume() { + let helper = crate::gpu::testing::helper(&[None]); + let hooked = Arc::new(AtomicUsize::new(0)); + let hook = hooked.clone(); + let started = helper.clone(); + let slot = GpuSlot::lazy( + Box::new(move || Ok(Some(started))), + Box::new(move |_| { + std::thread::sleep(std::time::Duration::from_millis(20)); + hook.fetch_add(1, Ordering::SeqCst); + }), + ); + assert!(futures::executor::block_on(slot.get_or_start()).is_some()); + assert_eq!(hooked.load(Ordering::SeqCst), 1); + assert!(futures::executor::block_on(slot.get_or_start()).is_some()); + assert_eq!(hooked.load(Ordering::SeqCst), 1); + } + + #[test] + /// Waiters that give up during a stalled start are pruned, and a live + /// waiter still gets the result. + fn canceled_waiters_are_pruned() { + let calls = Arc::new(AtomicUsize::new(0)); + let (open, gate) = std::sync::mpsc::channel(); + let helper = crate::gpu::testing::helper(&[None]); + let slot = counted(&calls, gate, Some(helper.clone())); + let waiting = |slot: &GpuSlot| match &*slot.inner.lock() { + Pending::Starting(waiters) => waiters.len(), + _ => panic!("start is not pending"), + }; + for _ in 0..100 { + let mut request = Box::pin(slot.get_or_start()); + let waker = futures::task::noop_waker(); + let mut cx = std::task::Context::from_waker(&waker); + assert!(request.as_mut().poll(&mut cx).is_pending()); + // Dropping the request cancels its waiter. + drop(request); + assert!(waiting(&slot) <= 1); + } + let mut live = Box::pin(slot.get_or_start()); + let waker = futures::task::noop_waker(); + let mut cx = std::task::Context::from_waker(&waker); + assert!(live.as_mut().poll(&mut cx).is_pending()); + assert_eq!(waiting(&slot), 1); + open.send(()).unwrap(); + let got = futures::executor::block_on(live).unwrap(); + assert!(Arc::ptr_eq(&got, &helper)); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[test] + fn ready_slots_are_resolved() { + let off = GpuSlot::default(); + assert!(!off.standby()); + assert!(off.get().is_none()); + assert!(futures::executor::block_on(off.get_or_start()).is_none()); + let helper = crate::gpu::testing::helper(&[None]); + let on = GpuSlot::ready(Some(helper.clone())); + assert!(Arc::ptr_eq(on.get().unwrap(), &helper)); + } +} diff --git a/cmd/racer-dataplane/src/gpu_telemetry.rs b/cmd/racer-dataplane/src/gpu_telemetry.rs new file mode 100644 index 0000000000..195ccba867 --- /dev/null +++ b/cmd/racer-dataplane/src/gpu_telemetry.rs @@ -0,0 +1,412 @@ +//! Prometheus metrics for `gpu-read`, cuFile, and GPU RDMA landing. +//! +//! The GPU helper and the RDMA link node keep their own counters and state. +//! [`GpuTelemetry`] holds handles to both, attached once at startup, and reads +//! them when `/metrics` renders. Request outcomes are counted here because the +//! client listener sees them, including on nodes with no GPU helper. +//! +//! Labels are bounded: request outcomes and cuFile reasons are fixed sets, GPU +//! labels come from the GPUs found at startup, and rail labels from the active +//! RDMA rails. + +use crate::error::Error; +use crate::gpu::CufileDisable; +use crate::gpu::GpuHelper; +use crate::rdma::LinkNode; +use std::fmt::Write; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +/// How the dataplane answered a `gpu-read` request head. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum GpuReadOutcome { + /// The allocation was imported on a GPU and the read proceeds. + Gpu, + /// `gpu-read` is off on this node (501). + NotEnabled, + /// The allocation is on a GPU the dataplane cannot use, or the + /// `Racer-Gpu-Uuid` header names a GPU it did not find (501). + WrongDevice, + /// Too many open requests or no GPU memory (503). + Busy, + /// The request or its handle was invalid (400). + Invalid, + /// Any other failure. + Failed, +} + +impl GpuReadOutcome { + pub const ALL: [Self; 6] = [ + Self::Gpu, + Self::NotEnabled, + Self::WrongDevice, + Self::Busy, + Self::Invalid, + Self::Failed, + ]; + + /// Metric label value. + pub fn label(self) -> &'static str { + match self { + Self::Gpu => "gpu", + Self::NotEnabled => "not_enabled", + Self::WrongDevice => "wrong_device", + Self::Busy => "busy", + Self::Invalid => "invalid", + Self::Failed => "failed", + } + } + + /// Classify the result of opening a `gpu-read` sink. `gpu` says whether + /// this node serves `gpu-read` at all. + pub fn classify(result: Result<(), Error>, gpu: bool) -> Self { + match result { + Ok(()) => Self::Gpu, + Err(Error::NotImplemented) if !gpu => Self::NotEnabled, + Err(Error::NotImplemented) => Self::WrongDevice, + Err(Error::Overloaded) => Self::Busy, + Err(Error::InvalidRequest | Error::InvalidRange) => Self::Invalid, + Err(_) => Self::Failed, + } + } +} + +/// Samples written without a GPU helper: one per request outcome, plus +/// `racer_gpu_enabled` and `racer_gpu_standby`. +pub const BASE_SERIES: usize = GpuReadOutcome::ALL.len() + 2; + +/// GPU metric sources shared by every worker's metrics handle. +#[derive(Default)] +pub struct GpuTelemetry { + helper: OnceLock>, + slot: OnceLock>, + links: OnceLock>, + requests: [AtomicU64; GpuReadOutcome::ALL.len()], +} + +impl GpuTelemetry { + /// Report this GPU helper. Later calls are ignored. + pub fn attach_helper(&self, helper: Arc) { + let _ = self.helper.set(helper); + } + + /// Report the helper in `slot` once it starts, and standby until the + /// first `gpu-read` resolves it. Reading never starts it. Later calls are + /// ignored. + pub fn attach_slot(&self, slot: Arc) { + let _ = self.slot.set(slot); + } + + /// Report GPU landing state from this link node. Later calls are ignored. + pub fn attach_links(&self, links: Arc) { + let _ = self.links.set(links); + } + + /// Count one `gpu-read` request head. + pub fn record_request(&self, outcome: GpuReadOutcome) { + let counter = &self.requests[outcome as usize]; + let _ = counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| { + Some(v.saturating_add(1)) + }); + } + + /// Requests counted with `outcome`. + pub fn requests(&self, outcome: GpuReadOutcome) -> u64 { + self.requests[outcome as usize].load(Ordering::Relaxed) + } + + /// Write every GPU series in Prometheus text format. + pub fn write(&self, out: &mut impl Write) -> std::fmt::Result { + writeln!(out, "# TYPE racer_gpu_read_requests_total counter")?; + for outcome in GpuReadOutcome::ALL { + writeln!( + out, + "racer_gpu_read_requests_total{{outcome=\"{}\"}} {}", + outcome.label(), + self.requests(outcome) + )?; + } + let slot = self.slot.get(); + let helper = self + .helper + .get() + .or_else(|| slot.and_then(|slot| slot.get())); + writeln!( + out, + "# TYPE racer_gpu_enabled gauge\nracer_gpu_enabled {}", + u8::from(helper.is_some()) + )?; + // 1 while `RACER_ENABLE_GPU=auto` waits for the first `gpu-read` + // before touching any GPU; 0 once that start resolves either way. + writeln!( + out, + "# TYPE racer_gpu_standby gauge\nracer_gpu_standby {}", + u8::from(slot.is_some_and(|slot| slot.standby())) + )?; + let Some(helper) = helper else { + return Ok(()); + }; + writeln!(out, "# TYPE racer_gpu_decrypt_available gauge")?; + for device in helper.devices() { + writeln!( + out, + "racer_gpu_decrypt_available{{ordinal=\"{}\",uuid=\"{}\"}} {}", + device.ordinal, + device.uuid_string(), + u8::from(helper.decrypts(device.ordinal)) + )?; + } + writeln!(out, "# TYPE racer_gpu_staging_slots_free gauge")?; + for device in helper.devices() { + writeln!( + out, + "racer_gpu_staging_slots_free{{ordinal=\"{}\"}} {}", + device.ordinal, + helper.free_slots(device.ordinal) + )?; + } + let counters = helper.counters(); + let load = |counter: &AtomicU64| counter.load(Ordering::Relaxed); + writeln!( + out, + "# TYPE racer_gpu_cufile_enabled gauge\nracer_gpu_cufile_enabled {}", + u8::from(helper.cufile()) + )?; + writeln!(out, "# TYPE racer_gpu_cufile_disabled_total counter")?; + for reason in CufileDisable::ALL { + writeln!( + out, + "racer_gpu_cufile_disabled_total{{reason=\"{}\"}} {}", + reason.label(), + load(&counters.cufile_disabled[reason as usize]) + )?; + } + writeln!(out, "# TYPE racer_gpu_read_pages_total counter")?; + for (decrypt, value) in [ + ("gpu", load(&counters.decrypted)), + ("cpu", load(&counters.plaintext)), + ] { + writeln!( + out, + "racer_gpu_read_pages_total{{decrypt=\"{decrypt}\"}} {value}" + )?; + } + for (name, kind, value) in [ + ( + "racer_gpu_pages_rejected_total", + "counter", + load(&counters.rejected), + ), + ( + "racer_gpu_cufile_staged_pages_total", + "counter", + load(&counters.staged), + ), + ( + "racer_gpu_cufile_failures_total", + "counter", + load(&counters.file_failures), + ), + ( + "racer_gpu_staging_exhausted_total", + "counter", + load(&counters.exhausted), + ), + ( + "racer_gpu_landed_pages_decrypted_total", + "counter", + load(&counters.landed), + ), + ( + "racer_gpu_cufile_unsupported_total", + "counter", + load(&counters.cufile_unsupported), + ), + ( + "racer_gpu_open_requests", + "gauge", + helper.open_requests() as u64, + ), + ("racer_gpu_leaked_resources", "gauge", helper.leaked()), + ] { + writeln!(out, "# TYPE {name} {kind}\n{name} {value}")?; + } + if let Some(links) = self.links.get() + && links.gpu_landing_configured() + { + let rails = links.landing_state(); + writeln!(out, "# TYPE racer_gpu_landing_enabled gauge")?; + for (rail, state) in &rails { + writeln!( + out, + "racer_gpu_landing_enabled{{rail=\"{rail}\"}} {}", + u8::from(!state.is_empty()) + )?; + } + writeln!(out, "# TYPE racer_gpu_landing_slots gauge")?; + for (rail, state) in &rails { + for (ordinal, slots) in state { + writeln!( + out, + "racer_gpu_landing_slots{{rail=\"{rail}\",ordinal=\"{ordinal}\"}} {slots}" + )?; + } + } + writeln!(out, "# TYPE racer_gpu_landing_exhausted_total counter")?; + for ((rail, ordinal), count) in links.landing_exhausted() { + writeln!( + out, + "racer_gpu_landing_exhausted_total{{rail=\"{rail}\",ordinal=\"{ordinal}\"}} {count}" + )?; + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn render(telemetry: &GpuTelemetry) -> String { + let mut out = String::new(); + telemetry.write(&mut out).unwrap(); + out + } + + #[test] + fn classify_maps_errors_to_bounded_outcomes() { + use GpuReadOutcome as O; + assert_eq!(O::classify(Ok(()), true), O::Gpu); + assert_eq!( + O::classify(Err(Error::NotImplemented), false), + O::NotEnabled + ); + assert_eq!( + O::classify(Err(Error::NotImplemented), true), + O::WrongDevice + ); + assert_eq!(O::classify(Err(Error::Overloaded), true), O::Busy); + assert_eq!(O::classify(Err(Error::InvalidRequest), true), O::Invalid); + assert_eq!(O::classify(Err(Error::InvalidRange), true), O::Invalid); + assert_eq!(O::classify(Err(Error::Unavailable), true), O::Failed); + } + + #[test] + fn without_a_helper_only_requests_and_enabled_render() { + let telemetry = GpuTelemetry::default(); + telemetry.record_request(GpuReadOutcome::NotEnabled); + telemetry.record_request(GpuReadOutcome::NotEnabled); + let text = render(&telemetry); + assert!(text.contains("racer_gpu_read_requests_total{outcome=\"not_enabled\"} 2\n")); + assert!(text.contains("racer_gpu_read_requests_total{outcome=\"gpu\"} 0\n")); + assert!(text.contains("racer_gpu_enabled 0\n")); + assert!(text.contains("racer_gpu_standby 0\n")); + assert!(!text.contains("racer_gpu_cufile_enabled")); + assert!(!text.contains("racer_gpu_landing")); + assert_eq!(text.matches("# TYPE").count(), 3); + assert_eq!( + text.lines().filter(|l| !l.starts_with('#')).count(), + BASE_SERIES + ); + } + + #[test] + /// A lazy slot reports standby without starting, then the helper once + /// it starts. + fn lazy_slot_reports_standby_then_helper() { + use std::sync::atomic::AtomicUsize; + let calls = Arc::new(AtomicUsize::new(0)); + let counted = calls.clone(); + let slot = Arc::new(crate::gpu_slot::GpuSlot::lazy( + Box::new(move || { + counted.fetch_add(1, Ordering::Relaxed); + Ok(Some(crate::gpu::testing::helper(&[None]))) + }), + Box::new(|_| {}), + )); + let telemetry = GpuTelemetry::default(); + telemetry.attach_slot(slot.clone()); + let text = render(&telemetry); + assert!(text.contains("racer_gpu_enabled 0\n")); + assert!(text.contains("racer_gpu_standby 1\n")); + assert_eq!( + text.lines().filter(|l| !l.starts_with('#')).count(), + BASE_SERIES + ); + assert_eq!(calls.load(Ordering::Relaxed), 0); + assert!(futures::executor::block_on(slot.get_or_start()).is_some()); + let text = render(&telemetry); + assert!(text.contains("racer_gpu_enabled 1\n")); + assert!(text.contains("racer_gpu_standby 0\n")); + assert!(text.contains("racer_gpu_cufile_enabled")); + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + + #[test] + fn request_counter_saturates() { + let telemetry = GpuTelemetry::default(); + telemetry.requests[0].store(u64::MAX, Ordering::Relaxed); + telemetry.record_request(GpuReadOutcome::Gpu); + assert_eq!(telemetry.requests(GpuReadOutcome::Gpu), u64::MAX); + } + + #[test] + fn with_a_helper_every_gpu_series_renders() { + let helper = crate::gpu::testing::helper(&[None]); + helper.counters().decrypted.fetch_add(3, Ordering::Relaxed); + helper.counters().plaintext.fetch_add(2, Ordering::Relaxed); + helper.counters().cufile_disabled(CufileDisable::FileOpen); + let telemetry = GpuTelemetry::default(); + telemetry.attach_helper(helper.clone()); + let links = Arc::new(LinkNode::default()); + telemetry.attach_links(links.clone()); + let text = render(&telemetry); + for name in [ + "racer_gpu_read_requests_total", + "racer_gpu_enabled", + "racer_gpu_standby", + "racer_gpu_decrypt_available", + "racer_gpu_staging_slots_free", + "racer_gpu_cufile_enabled", + "racer_gpu_cufile_disabled_total", + "racer_gpu_read_pages_total", + "racer_gpu_pages_rejected_total", + "racer_gpu_cufile_staged_pages_total", + "racer_gpu_cufile_failures_total", + "racer_gpu_staging_exhausted_total", + "racer_gpu_landed_pages_decrypted_total", + "racer_gpu_cufile_unsupported_total", + "racer_gpu_open_requests", + "racer_gpu_leaked_resources", + ] { + assert!(text.contains(&format!("# TYPE {name} ")), "{name}\n{text}"); + } + assert!(text.contains("racer_gpu_enabled 1\n")); + assert!(text.contains("racer_gpu_read_pages_total{decrypt=\"gpu\"} 3\n")); + assert!(text.contains("racer_gpu_read_pages_total{decrypt=\"cpu\"} 2\n")); + assert!(text.contains("racer_gpu_cufile_disabled_total{reason=\"file_open\"} 1\n")); + assert!(text.contains("racer_gpu_cufile_disabled_total{reason=\"driver_open\"} 0\n")); + let uuid = helper.devices()[0].uuid_string(); + assert!(text.contains(&format!( + "racer_gpu_decrypt_available{{ordinal=\"0\",uuid=\"{uuid}\"}} 1\n" + ))); + // Landing is not configured on this link node. + assert!(!text.contains("racer_gpu_landing_enabled")); + links.set_gpu_landing(helper, 1 << 30); + let text = render(&telemetry); + assert!(text.contains("# TYPE racer_gpu_landing_enabled gauge\n")); + assert!(text.contains("# TYPE racer_gpu_landing_slots gauge\n")); + assert!(text.contains("# TYPE racer_gpu_landing_exhausted_total counter\n")); + links.count_landing_exhausted(racer_control_wire::RailId(2), 0); + links.count_landing_exhausted(racer_control_wire::RailId(2), 0); + links.count_landing_exhausted(racer_control_wire::RailId(1), 3); + let text = render(&telemetry); + assert!(text.contains( + "racer_gpu_landing_exhausted_total{rail=\"1\",ordinal=\"3\"} 1\n\ + racer_gpu_landing_exhausted_total{rail=\"2\",ordinal=\"0\"} 2\n" + )); + } +} diff --git a/cmd/racer-dataplane/src/lib.rs b/cmd/racer-dataplane/src/lib.rs index 4b4014fa9f..175a73bc69 100644 --- a/cmd/racer-dataplane/src/lib.rs +++ b/cmd/racer-dataplane/src/lib.rs @@ -13,6 +13,8 @@ pub mod config; pub mod control; pub mod error; pub mod gpu; +pub mod gpu_slot; +pub mod gpu_telemetry; pub mod hardware; pub mod http; pub mod memory; diff --git a/cmd/racer-dataplane/src/peer/transport.rs b/cmd/racer-dataplane/src/peer/transport.rs index a23af0ce40..0bfd3dc1e6 100644 --- a/cmd/racer-dataplane/src/peer/transport.rs +++ b/cmd/racer-dataplane/src/peer/transport.rs @@ -409,7 +409,7 @@ impl Transfers { links .active() .iter() - .any(|rail| links.landing_ordinal(rail.rail) == Some(ordinal)) + .any(|rail| links.has_landing(rail.rail, ordinal)) }) } /// Receiver: when the hop plan is RDMA and the link is ready, grant a page @@ -455,7 +455,7 @@ impl Transfers { if matches!( request.request.operation, crate::peer::protocol::Operation::Page { .. } - ) && links.landing_ordinal(rail) == Some(ordinal) => + ) && links.has_landing(rail, ordinal) => { match links .grant_landing(&peer, rail, binding.transfer, ordinal, &native_scope(scope)) @@ -2115,7 +2115,7 @@ mod tests { let bytes = (crate::rdma::BUFFER_LEN as u64) .next_multiple_of(crate::config::GPU_STAGING_UNIT); crate::rdma::links::tests::attach_landing(&mut w, &helper, bytes); - assert_eq!(w.a.links.landing_ordinal(RailId(7)), Some(0)); + assert!(w.a.links.has_landing(RailId(7), 0)); helper }); // The receiver grants in-band and names the grant in its signed Accept. @@ -2286,7 +2286,7 @@ mod tests { // Read what the NIC wrote, through the landing dma-buf. use std::os::unix::fs::FileExt; assert_eq!(landed.envelope.ciphertext_length, 144); - let buffer = w.a.links.node().landing(RailId(7)).unwrap(); + let buffer = w.a.links.node().landing(RailId(7), 0).unwrap(); let file = std::fs::File::from( rdma_verbs::DmabufMemory::fd(&*buffer) .try_clone_to_owned() diff --git a/cmd/racer-dataplane/src/rdma/links.rs b/cmd/racer-dataplane/src/rdma/links.rs index 4ffa35544f..5bd43cfb3f 100644 --- a/cmd/racer-dataplane/src/rdma/links.rs +++ b/cmd/racer-dataplane/src/rdma/links.rs @@ -98,36 +98,58 @@ pub struct LinkNode { /// in GPU memory (GPUDirect RDMA). gpu_landing: Mutex, u64)>>, - /// The GPU landing buffer registered on each rail this round. - landings: Mutex>>, + /// Landing was configured after the current rails attached theirs, so + /// the owner must attach landings on the rails already active. Set and + /// cleared under the `gpu_landing` lock, so an activation that reads the + /// new configuration consumes it and rails never attach twice. + landing_pending: std::sync::atomic::AtomicBool, + + /// GPU landing slices registered this round, by rail and GPU. + landings: Mutex>>, + + /// NIC slots in each registered landing slice this round, by rail and GPU. + landing_slots: Mutex>, + + /// Landing grants that found no free slot, by rail and GPU. Kept across + /// rounds so the metric only grows. + landing_exhausted: Mutex>, } -/// Pick at most one GPU per rail for landing, and each GPU for at most one -/// rail. A rail first takes a free GPU on its NUMA node, then any free GPU. -/// `rails` and `gpus` give each one's NUMA node, if known. PCI distance is -/// not used: only NUMA locality is known on both sides. -pub fn assign_landing( +/// Split each GPU's landing buffer of `slots` page slots across `rails`. +/// Every GPU gets a slice on every rail when there are enough slots, so a +/// read on any GPU can land over whichever rail its peer link uses. With +/// fewer slots than rails, a GPU's rails on its own NUMA node come first, +/// then the rest in rail order. Slots are shared out evenly; the first +/// rails get the remainder. `rails` and `gpus` give each one's NUMA node, +/// if known. PCI distance is not used: only NUMA locality is known on both +/// sides. Returns `(ordinal, rail, slots)` for each slice. +pub fn plan_landing( rails: &[(RailId, Option)], gpus: &[(i32, Option)], -) -> Vec<(RailId, i32)> { - let mut free: Vec<(i32, Option)> = gpus.to_vec(); - let mut assigned = Vec::new(); - let mut later = Vec::new(); - for &(rail, numa) in rails { - match free + slots: usize, +) -> Vec<(i32, RailId, usize)> { + let mut plan = Vec::new(); + for &(ordinal, gpu_numa) in gpus { + let local = |numa: Option| gpu_numa.is_some() && numa == gpu_numa; + let mut ordered: Vec = rails .iter() - .position(|&(_, node)| numa.is_some() && node == numa) - { - Some(index) => assigned.push((rail, free.remove(index).0)), - None => later.push(rail), - } - } - for rail in later { - if !free.is_empty() { - assigned.push((rail, free.remove(0).0)); + .filter(|&&(_, numa)| local(numa)) + .map(|&(rail, _)| rail) + .collect(); + ordered.extend( + rails + .iter() + .filter(|&&(_, numa)| !local(numa)) + .map(|&(rail, _)| rail), + ); + ordered.truncate(slots); + let count = ordered.len(); + for (index, rail) in ordered.into_iter().enumerate() { + let share = slots / count + usize::from(index < slots % count); + plan.push((ordinal, rail, share)); } } - assigned + plan } /// One setup this node initiated: reserved before the owner answers, then @@ -186,19 +208,98 @@ impl LinkNode { } /// Let `gpu-read` pages land in GPU memory: on each activation, give - /// each rail one GPU's landing buffer of `bytes` (see - /// [`assign_landing`]). 0 bytes turns landing off. + /// each decrypting GPU a landing buffer of `bytes`, split across the + /// rails (see [`plan_landing`]). 0 bytes turns landing off. + /// + /// Configuring it while rails are already active (the GPU helper started + /// lazily on the first `gpu-read`) marks those rails pending; the owner + /// worker attaches them on its next turn (see [`Self::take_landing_pending`]). pub fn set_gpu_landing(&self, helper: Arc, bytes: u64) { - *Self::lock(&self.gpu_landing) = (bytes != 0).then_some((helper, bytes)); + let mut landing = Self::lock(&self.gpu_landing); + *landing = (bytes != 0).then_some((helper, bytes)); + self.landing_pending + .store(landing.is_some(), std::sync::atomic::Ordering::Release); + } + + /// Whether landing still needs attaching on the rails already active. + /// Clears the mark; the caller must then attach them. + pub fn take_landing_pending(&self) -> bool { + let _landing = Self::lock(&self.gpu_landing); + self.landing_pending + .swap(false, std::sync::atomic::Ordering::AcqRel) + } + + /// The landing configuration, consuming any pending mark: the caller + /// attaches with it. + fn take_gpu_landing(&self) -> Option<(Arc, u64)> { + let landing = Self::lock(&self.gpu_landing); + self.landing_pending + .store(false, std::sync::atomic::Ordering::Release); + landing.clone() + } + + /// GPU `ordinal`'s landing slice on `rail` this round, if any. + pub fn landing(&self, rail: RailId, ordinal: i32) -> Option> { + Self::lock(&self.landings).get(&(rail.0, ordinal)).cloned() + } + + /// Landing state of each active rail, for metrics: the rail, and the + /// `(ordinal, slots)` of each GPU landing slice registered on it. + pub fn landing_state(&self) -> Vec<(u16, Vec<(i32, usize)>)> { + let slots = Self::lock(&self.landing_slots); + let mut rails: Vec = Self::lock(&self.active) + .iter() + .map(|mapping| mapping.rail.0) + .collect(); + rails.sort_unstable(); + rails.dedup(); + rails + .into_iter() + .map(|rail| { + let mut gpus: Vec<(i32, usize)> = slots + .iter() + .filter(|((r, _), _)| *r == rail) + .map(|(&(_, ordinal), &count)| (ordinal, count)) + .collect(); + gpus.sort_unstable(); + (rail, gpus) + }) + .collect() + } + + /// Count one landing grant on `rail` for GPU `ordinal` that found no + /// free slot. + pub fn count_landing_exhausted(&self, rail: RailId, ordinal: i32) { + let mut counts = Self::lock(&self.landing_exhausted); + let count = counts.entry((rail.0, ordinal)).or_default(); + *count = count.saturating_add(1); + } + + /// Landing exhaustion counts for metrics: every slice registered this + /// round (0 if none yet) and every slice that ever exhausted, sorted by + /// rail then GPU ordinal. + pub fn landing_exhausted(&self) -> Vec<((u16, i32), u64)> { + let counts = Self::lock(&self.landing_exhausted); + let mut keys: Vec<(u16, i32)> = Self::lock(&self.landing_slots) + .keys() + .chain(counts.keys()) + .copied() + .collect(); + keys.sort_unstable(); + keys.dedup(); + keys.into_iter() + .map(|key| (key, counts.get(&key).copied().unwrap_or(0))) + .collect() } - /// The GPU landing buffer on `rail` this round, if any. - pub fn landing(&self, rail: RailId) -> Option> { - Self::lock(&self.landings).get(&rail.0).cloned() + /// Whether GPU landing is configured (GPU decrypt and RDMA both on). + pub fn gpu_landing_configured(&self) -> bool { + Self::lock(&self.gpu_landing).is_some() } fn reset(&self) { Self::lock(&self.landings).clear(); + Self::lock(&self.landing_slots).clear(); Self::lock(&self.active).clear(); self.settled .store(false, std::sync::atomic::Ordering::Release); @@ -483,12 +584,14 @@ impl Links { } /// Register GPU landing buffers on the activated `rails`. Failures only - /// log: pages then land in host memory as usual. - async fn attach_landings(&self, rails: &[RailMapping], scope: &RequestScope) { - let Some((helper, bytes)) = LinkNode::lock(&self.node.gpu_landing).clone() else { + /// log: pages then land in host memory as usual. A GPU whose buffer no + /// rail registered drops it again, which frees the GPU memory. + pub(crate) async fn attach_landings(&self, rails: &[RailMapping], scope: &RequestScope) { + let Some((helper, bytes)) = self.node.take_gpu_landing() else { return; }; - if bytes < BUFFER_LEN as u64 { + let slots = (bytes / BUFFER_LEN as u64) as usize; + if slots == 0 { eprintln!( "racer-dataplane: stage=gpu-landing state=disabled reason=too-small bytes={bytes}" ); @@ -502,40 +605,66 @@ impl Links { .collect(); let rails: Vec<(RailId, Option)> = rails.iter().map(|m| (m.rail, m.numa_node)).collect(); - for (rail, ordinal) in assign_landing(&rails, &gpus) { + let plan = plan_landing(&rails, &gpus, slots); + for &(ordinal, _) in &gpus { + let mine: Vec<(RailId, usize)> = plan + .iter() + .filter(|&&(o, _, _)| o == ordinal) + .map(|&(_, rail, slots)| (rail, slots)) + .collect(); + if mine.is_empty() { + continue; + } let landing = match helper.landing(ordinal, bytes, scope).await { Ok(landing) => landing, Err(error) => { eprintln!( - "racer-dataplane: stage=gpu-landing rail={} ordinal={ordinal} state=failed reason={error}", - rail.0 + "racer-dataplane: stage=gpu-landing ordinal={ordinal} state=failed reason={error}" ); continue; } }; - let memory = landing.clone() as rdma_verbs::DeviceMemory; - match self - .wait(self.client.attach_device(rail.0, memory), scope) - .await - { - Ok(slots) => { + let mut offset = 0u64; + for (rail, slots) in mine { + let length = (slots * BUFFER_LEN) as u64; + let slice = landing.slice(offset, length); + offset += length; + let Some(slice) = slice else { eprintln!( - "racer-dataplane: stage=gpu-landing rail={} ordinal={ordinal} state=enabled slots={slots}", + "racer-dataplane: stage=gpu-landing rail={} ordinal={ordinal} state=failed reason=bad-slice", rail.0 ); - LinkNode::lock(&self.node.landings).insert(rail.0, landing); + continue; + }; + let memory = slice.clone() as rdma_verbs::DeviceMemory; + match self + .wait( + self.client + .attach_device(rail.0, landing_pool(ordinal), memory), + scope, + ) + .await + { + Ok(slots) => { + eprintln!( + "racer-dataplane: stage=gpu-landing rail={} ordinal={ordinal} state=enabled slots={slots}", + rail.0 + ); + LinkNode::lock(&self.node.landings).insert((rail.0, ordinal), slice); + LinkNode::lock(&self.node.landing_slots).insert((rail.0, ordinal), slots); + } + Err(error) => eprintln!( + "racer-dataplane: stage=gpu-landing rail={} ordinal={ordinal} state=failed reason={error}", + rail.0 + ), } - Err(error) => eprintln!( - "racer-dataplane: stage=gpu-landing rail={} ordinal={ordinal} state=failed reason={error}", - rail.0 - ), } } } - /// The GPU whose landing buffer is on `rail`, if any. - pub fn landing_ordinal(&self, rail: RailId) -> Option { - self.node.landing(rail).map(|landing| landing.ordinal()) + /// Whether GPU `ordinal` has a landing slice on `rail` this round. + pub fn has_landing(&self, rail: RailId, ordinal: i32) -> bool { + self.node.landing(rail, ordinal).is_some() } /// Close the node: every link breaks and the owner drains. A later @@ -793,10 +922,10 @@ impl Links { }) } - /// Receiver: like [`Self::grant`], but the window is over a slot of the - /// GPU landing buffer on `rail`, which must be on GPU `ordinal`. Fails - /// with `Unavailable` if it is not, and `Overloaded` if no slot is free; - /// the caller then grants host memory. + /// Receiver: like [`Self::grant`], but the window is over a slot of GPU + /// `ordinal`'s landing slice on `rail`. Fails with `Unavailable` if the + /// rail has none, and `Overloaded` if no slot is free; the caller then + /// grants host memory. pub async fn grant_landing( &self, peer: &NodeId, @@ -806,25 +935,30 @@ impl Links { scope: &RequestScope, ) -> Result { scope.check()?; - let landing = self - .node - .landing(rail) - .filter(|landing| landing.ordinal() == ordinal) - .ok_or(Error::Unavailable)?; + let landing = self.node.landing(rail, ordinal).ok_or(Error::Unavailable)?; let (epoch, link) = self.ready(peer, rail).ok_or(Error::Unavailable)?; let length = MAX_CIPHERTEXT as u32; - let granted = self + let granted = match self .wait( self.client.grant_device( link_key(peer, rail), epoch, + landing_pool(ordinal), transfer.0, length, scope.deadline.0, ), scope, ) - .await?; + .await + { + Ok(granted) => granted, + Err(Error::Overloaded) => { + self.node.count_landing_exhausted(rail, ordinal); + return Err(Error::Overloaded); + } + Err(error) => return Err(error), + }; Ok(PendingReceive { metrics: self.metrics.clone(), retry: self.retry.clone(), @@ -902,6 +1036,11 @@ impl Links { } } +/// The `rdma_verbs` device pool id of GPU `ordinal`'s landing slice. +fn landing_pool(ordinal: i32) -> u32 { + ordinal as u32 +} + /// Lends a ciphertext page to the owner as write source bytes. struct PageSource(CiphertextPage); diff --git a/cmd/racer-dataplane/src/rdma/links/tests.rs b/cmd/racer-dataplane/src/rdma/links/tests.rs index cc6a1ccf8c..9d119bda0e 100644 --- a/cmd/racer-dataplane/src/rdma/links/tests.rs +++ b/cmd/racer-dataplane/src/rdma/links/tests.rs @@ -849,25 +849,47 @@ mod send_in_place { } #[test] -/// Rails take a free GPU on their NUMA node first, then any free GPU; each -/// GPU serves one rail at most. -fn assign_landing_prefers_numa_then_any() { +/// Every GPU gets a slice on every rail when slots allow, so no GPU is left +/// without landing on a node with more GPUs than rails. With fewer slots +/// than rails, NUMA-local rails come first. Slots are shared out evenly. +fn plan_landing_covers_every_gpu_and_prefers_numa() { let rails = [ (RailId(1), Some(1)), (RailId(2), Some(0)), (RailId(3), None), - (RailId(4), Some(0)), ]; - let gpus = [(0, Some(0)), (1, Some(1)), (2, None)]; + let gpus = [(0, Some(0)), (1, Some(1)), (2, None), (3, Some(0))]; + let plan = plan_landing(&rails, &gpus, 7); + for &(ordinal, _) in &gpus { + let mine: Vec<_> = plan.iter().filter(|p| p.0 == ordinal).collect(); + assert_eq!(mine.len(), 3, "GPU {ordinal} is on every rail"); + assert_eq!(mine.iter().map(|p| p.2).sum::(), 7); + } + assert_eq!( + plan.iter() + .filter(|p| p.0 == 0) + .copied() + .collect::>(), + [(0, RailId(2), 3), (0, RailId(1), 2), (0, RailId(3), 2)] + ); + // One slot per GPU: only the best rail. assert_eq!( - assign_landing(&rails, &gpus), - [(RailId(1), 1), (RailId(2), 0), (RailId(3), 2)] + plan_landing(&rails, &gpus, 1), + [ + (0, RailId(2), 1), + (1, RailId(1), 1), + (2, RailId(1), 1), + (3, RailId(2), 1), + ] ); - assert_eq!(assign_landing(&rails, &[]), []); + // Two slots, two rails: one each, local first. assert_eq!( - assign_landing(&[(RailId(9), Some(3))], &[(5, Some(1))]), - [(RailId(9), 5)] + plan_landing(&rails[..2], &[(5, Some(1))], 2), + [(5, RailId(1), 1), (5, RailId(2), 1)] ); + assert_eq!(plan_landing(&rails, &[], 4), []); + assert_eq!(plan_landing(&[], &gpus, 4), []); + assert_eq!(plan_landing(&rails, &gpus, 0), []); } /// Register a simulated GPU's landing buffer on `a`'s rail. @@ -901,8 +923,9 @@ fn landing_grant_receives_into_gpu_memory() { // One slot only. let bytes = (BUFFER_LEN as u64).next_multiple_of(crate::config::GPU_STAGING_UNIT); attach_landing(&mut w, &helper, bytes); - assert_eq!(w.a.links.landing_ordinal(RAIL), Some(0)); - assert_eq!(w.b.links.landing_ordinal(RAIL), None); + assert!(w.a.links.has_landing(RAIL, 0)); + assert!(!w.a.links.has_landing(RAIL, 1)); + assert!(!w.b.links.has_landing(RAIL, 0)); let (a, b) = (node_id(1), node_id(2)); let scope = scope(); assert_eq!( @@ -913,7 +936,7 @@ fn landing_grant_receives_into_gpu_memory() { ) .err(), Some(Error::Unavailable), - "landing buffer is on another GPU" + "GPU 1 has no landing buffer" ); let pending = drive!( w, @@ -941,7 +964,7 @@ fn landing_grant_receives_into_gpu_memory() { _ => panic!("expected a device page"), }; // The bytes are in the GPU buffer, read here through its dma-buf. - let landing = w.a.links.node().landing(RAIL).unwrap(); + let landing = w.a.links.node().landing(RAIL, 0).unwrap(); let mut bytes = [0; 35]; let file = std::fs::File::from( rdma_verbs::DmabufMemory::fd(&*landing) @@ -969,9 +992,69 @@ fn landing_grant_receives_into_gpu_memory() { drive!(w, poll_fn(|cx| again.arrival.poll(cx))), Ok(Arrival::Released) )); - // Closing the node forgets the landing buffer for the next round. + // Closing the node forgets the landing buffer for the next round, and + // once the node drains the GPU memory is freed. + drop((landing, pending, again)); w.a.links.close(); - assert_eq!(w.a.links.landing_ordinal(RAIL), None); + assert!(!w.a.links.has_landing(RAIL, 0)); + for _ in 0..100 { + pump!(w); + if helper.run(crate::gpu::testing::live_staging) == 0 { + return; + } + std::thread::sleep(Duration::from_millis(1)); + } + panic!("landing buffer was not freed after the node drained"); +} + +#[test] +/// Two GPUs on one rail each get a landing slice and a separate device +/// pool: a read on either GPU lands in its own GPU's memory, and one GPU's +/// full slice does not refuse the other. +fn every_gpu_lands_on_a_shared_rail() { + let helper = crate::gpu::testing::helper(&[None, None]); + let mut w = world(); + link_up(&mut w); + let bytes = (BUFFER_LEN as u64).next_multiple_of(crate::config::GPU_STAGING_UNIT); + attach_landing(&mut w, &helper, bytes); + assert!(w.a.links.has_landing(RAIL, 0)); + assert!(w.a.links.has_landing(RAIL, 1)); + let (a, b) = (node_id(1), node_id(2)); + let scope = scope(); + let mut held = Vec::new(); + for (ordinal, fill) in [(0, 0x10u8), (1, 0x20)] { + let transfer = TransferId([fill; 16]); + let pending = drive!( + w, + w.a.links.grant_landing(&b, RAIL, transfer, ordinal, &scope) + ) + .unwrap(); + let header = STANDARD.decode(pending.header_value()).unwrap(); + let grant = AuthenticatedGrant::decode(&header, RAIL, transfer).unwrap(); + let page = page(&w.b.admission.clone(), 20, fill); + drive!(w, w.b.links.send(&a, &grant, &page, &scope)).unwrap(); + let device = match drive!(w, poll_fn(|cx| pending.arrival.poll(cx))).unwrap() { + Arrival::Device(page) => page, + _ => panic!("expected a device page"), + }; + let landing = w.a.links.node().landing(RAIL, ordinal).unwrap(); + let landed = crate::gpu::LandedCiphertext::new(&landing, device).unwrap(); + assert_eq!(landed.ordinal(), ordinal); + held.push(landed); + } + // Each GPU's only slot is taken. + for ordinal in [0, 1] { + assert_eq!( + drive!( + w, + w.a.links + .grant_landing(&b, RAIL, TransferId([9; 16]), ordinal, &scope) + ) + .err(), + Some(Error::Overloaded) + ); + } + drop(held); } #[test] @@ -981,17 +1064,45 @@ fn landing_failures_leave_rails_host_only() { let helper = crate::gpu::testing::helper(&[None]); let mut w = world(); attach_landing(&mut w, &helper, (BUFFER_LEN as u64) / 2); - assert_eq!(w.a.links.landing_ordinal(RAIL), None); + assert!(!w.a.links.has_landing(RAIL, 0)); attach_landing(&mut w, &helper, 0); - assert_eq!(w.a.links.landing_ordinal(RAIL), None); + assert!(!w.a.links.has_landing(RAIL, 0)); w._sim.fault( rdma_verbs::simulation::Operation::RegisterDmabuf, rdma_verbs::simulation::Fault::Reject, ); let bytes = (BUFFER_LEN as u64).next_multiple_of(crate::config::GPU_STAGING_UNIT); attach_landing(&mut w, &helper, bytes); - assert_eq!(w.a.links.landing_ordinal(RAIL), None); - // Registration works on the next try, reusing the same GPU buffer. + assert!(!w.a.links.has_landing(RAIL, 0)); + // No rail holds the buffer, so its GPU memory is freed again. + assert_eq!(helper.run(crate::gpu::testing::live_staging), 0); + // Registration works on the next try, with a new GPU buffer. + attach_landing(&mut w, &helper, bytes); + assert!(w.a.links.has_landing(RAIL, 0)); + assert_eq!(helper.run(crate::gpu::testing::live_staging), 1); +} + +#[test] +/// A GPU helper that starts after the rails activated (lazy `gpu-read` +/// start) marks them pending; attaching on the active rails registers the +/// landing buffer and consumes the mark, so rails never attach twice. +fn late_gpu_landing_attaches_on_active_rails_once() { + let helper = crate::gpu::testing::helper(&[None]); + let mut w = world(); + link_up(&mut w); + let node = w.a.links.node().clone(); + assert!(!node.take_landing_pending()); + node.set_gpu_landing(helper.clone(), 0); + assert!(!node.take_landing_pending(), "0 bytes turns landing off"); + let bytes = (BUFFER_LEN as u64).next_multiple_of(crate::config::GPU_STAGING_UNIT); + node.set_gpu_landing(helper.clone(), bytes); + assert!(!w.a.links.has_landing(RAIL, 0)); + // `attach_landing` sets the landing again, then attaches like the owner + // worker does on its next turn. attach_landing(&mut w, &helper, bytes); - assert_eq!(w.a.links.landing_ordinal(RAIL), Some(0)); + assert!(w.a.links.has_landing(RAIL, 0)); + assert!(!node.take_landing_pending(), "the attach consumed the mark"); + node.set_gpu_landing(helper, bytes); + assert!(node.take_landing_pending()); + assert!(!node.take_landing_pending()); } diff --git a/cmd/racer-dataplane/src/read/fill.rs b/cmd/racer-dataplane/src/read/fill.rs index 55ec6b1409..610544ac51 100644 --- a/cmd/racer-dataplane/src/read/fill.rs +++ b/cmd/racer-dataplane/src/read/fill.rs @@ -157,6 +157,12 @@ fn stage_events(stage: DriverStage) -> Option<(Event, Event)> { /// starter's inherited scope. const LOCAL_PLAINTEXT_JOIN_RETRIES: usize = 2; +/// Providers a `gpu-read` asks, in visit order, to land a page in GPU +/// memory before it takes the usual path. Each is one CopyOnly request +/// against the read's acquisition budget; more would add a round trip per +/// miss before the usual path asks the same providers again. +const LANDING_PROVIDERS: usize = 2; + impl Fill { pub(super) async fn warm_local( &self, @@ -990,12 +996,13 @@ impl Fill { /// or create a shared single flight, and a landed page is never published /// to memory or persisted, because there are no host bytes to share. /// - /// Only the first provider in visit order is asked, with one CopyOnly - /// request. None means use the usual path: no landing buffer on the - /// request's GPU, a local copy exists, this node owns the page, or the - /// provider missed or failed. A page the provider sent to host memory - /// instead (no free landing slot, or HTTP) is used privately in the same - /// way, so it is not fetched twice. + /// Up to [`LANDING_PROVIDERS`] providers are asked in visit order, one + /// CopyOnly request each; a miss or failure moves on to the next. None + /// means use the usual path: no landing buffer on the request's GPU, a + /// local copy exists, this node owns the page, or every provider asked + /// missed or failed. A page a provider sent to host memory instead (no + /// free landing slot, or HTTP) is used privately in the same way, so it + /// is not fetched twice. async fn land_from_peer( &self, page: &PageId, @@ -1033,40 +1040,54 @@ impl Fill { return Ok(None); } let order = resolver.visit_order(&candidates, &page.version.object, page.number); - let Some(first) = order.ranks().next() else { - return Ok(None); - }; - let destination = candidates.ordered[first].clone(); - let delivered = resolver - .request_landing( - &membership, - &destination, - context, - page, - scope, - budget, - target.ordinal, - ) - .await?; - let landed = match delivered { - crate::peer::Delivered::Landed(landed) => landed, - // The provider sent the page to host memory (no free landing - // slot, or HTTP): use it privately rather than fetch it again. - crate::peer::Delivered::Response(response) => { - let Ok(copy) = response_copy(response.response(), page) else { - return Ok(None); - }; - return Ok(self - .seal( - page, - UnverifiedPage { - copy, - disk_token: None, - }, - ) - .ok() - .map(super::sealed::Fetched::Sealed)); + let ranks: Vec = order.ranks().take(LANDING_PROVIDERS).collect(); + let mut landed = None; + for rank in ranks { + let destination = candidates.ordered[rank].clone(); + let delivered = match resolver + .request_landing( + &membership, + &destination, + context, + page, + scope, + budget, + target.ordinal, + ) + .await + { + Ok(delivered) => delivered, + Err(error @ (Error::Cancelled | Error::DeadlineExceeded)) => return Err(error), + // A failed provider: ask the next one. + Err(_) => continue, + }; + match delivered { + crate::peer::Delivered::Landed(answer) => { + landed = Some(answer); + break; + } + // The provider sent the page to host memory (no free landing + // slot, or HTTP): use it privately rather than fetch it again. + crate::peer::Delivered::Response(response) => { + // A miss or unusable answer: ask the next provider. + let Ok(copy) = response_copy(response.response(), page) else { + continue; + }; + return Ok(self + .seal( + page, + UnverifiedPage { + copy, + disk_token: None, + }, + ) + .ok() + .map(super::sealed::Fetched::Sealed)); + } } + } + let Some(landed) = landed else { + return Ok(None); }; let crate::peer::transport::LandedResponse { metadata, diff --git a/cmd/racer-dataplane/src/read/tests/peer_copies.rs b/cmd/racer-dataplane/src/read/tests/peer_copies.rs index 4128c680c2..d2984e137a 100644 --- a/cmd/racer-dataplane/src/read/tests/peer_copies.rs +++ b/cmd/racer-dataplane/src/read/tests/peer_copies.rs @@ -2994,6 +2994,12 @@ mod landing { /// the landing buffer, as a NIC would; otherwise the request goes to the /// usual peer script. fn landing(ordinal: i32, lands: bool) -> Landing { + landing_on(ordinal, lands.then_some(1)) + } + + /// Like [`landing`], but only the `call`th landing request (1-based) + /// lands; the others go to the usual peer script. + fn landing_on(ordinal: i32, call: Option) -> Landing { let mut f = fixture(); let copy = encrypted_copy(&mut f); let (peers, ordered) = install_peers(&mut f, None); @@ -3007,6 +3013,7 @@ mod landing { let seen = asked.clone(); let land: crate::peer::ScriptedLanding = Box::new(move |request| { seen.set(seen.get() + 1); + let lands = call == Some(seen.get()); assert!(matches!( request.operation, PeerOperation::Page { @@ -3201,14 +3208,63 @@ mod landing { { let mut replies = l.peers.replies.borrow_mut(); replies.push_back((l.ordered[0].clone(), Reply::Unauthorized)); + replies.push_back((l.ordered[1].clone(), Reply::Copy(l.copy.clone()))); + } + let fetched = acquire_sealed(&mut l); + assert!(fetched.is_ok(), "{:?}", fetched.err()); + // The failed first provider is followed by a landing request to the + // second, whose host answer is used once. + assert_eq!(l.asked.get(), 2); + assert_eq!(l.peers.calls.borrow().len(), 2); + assert_eq!(landed_count(&l.helper), 0); + } + + #[test] + /// The first provider misses: the next provider in visit order is asked + /// to land the page, and its page lands in GPU memory. + fn next_provider_lands_after_a_miss() { + let queue = Rc::new(uring_runtime::drivers::DriverQueue::new(1024)); + let _owner = queue.enter(); + let mut l = landing_on(0, Some(2)); + l.peers + .replies + .borrow_mut() + .push_back((l.ordered[0].clone(), Reply::Miss)); + let Fetched::Sealed(sealed) = acquire_sealed(&mut l).unwrap() else { + panic!("expected a sealed page"); + }; + assert!(matches!(sealed.source, SealedSource::Landed(_))); + assert_eq!(l.asked.get(), 2); + assert_eq!( + l.peers.calls.borrow().len(), + 1, + "only the miss went to the script" + ); + let fd = crate::gpu::testing::import(&l.helper, 1, 0); + crate::gpu::testing::write(&l.helper, 1, sealed, 0..3).unwrap(); + assert_eq!(racer_gpu::simulation::read_allocation(&fd, 0, 3), b"abc"); + assert_eq!(landed_count(&l.helper), 1); + } + + #[test] + /// Landing asks at most `LANDING_PROVIDERS` providers; after that the + /// usual path runs. + fn landing_stops_after_its_provider_limit() { + let queue = Rc::new(uring_runtime::drivers::DriverQueue::new(1024)); + let _owner = queue.enter(); + let mut l = landing_on(0, None); + { + let mut replies = l.peers.replies.borrow_mut(); + replies.push_back((l.ordered[0].clone(), Reply::Miss)); + replies.push_back((l.ordered[1].clone(), Reply::Miss)); for node in &l.ordered { replies.push_back((node.clone(), Reply::Copy(l.copy.clone()))); } } let fetched = acquire_sealed(&mut l); assert!(fetched.is_ok(), "{:?}", fetched.err()); - assert_eq!(l.asked.get(), 1); - assert!(l.peers.calls.borrow().len() >= 2); + assert_eq!(l.asked.get(), 2); + assert!(l.peers.calls.borrow().len() >= 3); assert_eq!(landed_count(&l.helper), 0); } } diff --git a/cmd/racer-dataplane/src/telemetry.rs b/cmd/racer-dataplane/src/telemetry.rs index e589c5bb5d..be817da441 100644 --- a/cmd/racer-dataplane/src/telemetry.rs +++ b/cmd/racer-dataplane/src/telemetry.rs @@ -3017,6 +3017,9 @@ pub struct Metrics { memory: Arc, cache_disks: ::telemetry::SnapshotShards>, + /// GPU, cuFile, and GPU landing sources, shared by every worker. + gpu: Arc, + shard: usize, } @@ -3698,6 +3701,7 @@ impl Metrics { process: ::telemetry::SnapshotShards::new(1), }); let cache_disks = ::telemetry::SnapshotShards::new(count); + let gpu = Arc::new(crate::gpu_telemetry::GpuTelemetry::default()); Ok(::telemetry::Metrics::shards(count) .into_iter() .enumerate() @@ -3709,6 +3713,7 @@ impl Metrics { disks: disks.clone(), memory: memory.clone(), cache_disks: cache_disks.clone(), + gpu: gpu.clone(), shard, }) .collect()) @@ -3734,6 +3739,14 @@ impl Metrics { self.record(sum, u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX)); self.record(count, 1); } + /// GPU metric sources shared by every worker. + pub fn gpu(&self) -> &crate::gpu_telemetry::GpuTelemetry { + &self.gpu + } + /// A shared handle to the GPU metric sources, for the lazy GPU start. + pub fn gpu_shared(&self) -> Arc { + self.gpu.clone() + } pub(crate) fn writer_shard(&self) -> usize { self.shard } @@ -3772,6 +3785,7 @@ impl Metrics { self.write_memory(out)?; self.write_worker_pools(out)?; self.write_cache_disks(out)?; + self.gpu.write(out)?; // Quota labels use runtime worker IDs. Read the authority's actual charge, // including pooled ciphertext capacity, without sampling on worker polls. // A stalled worker therefore remains observable from another worker. @@ -6336,6 +6350,7 @@ pub(crate) mod tests { + RETENTION_METRICS.len() + HELD_METRICS.len() + DISK_CLASS_METRICS.len() * 2 + + crate::gpu_telemetry::BASE_SERIES ); assert_only_disk_class_labels(&output); assert!(output.len() < 64 * 1024); @@ -6408,6 +6423,7 @@ pub(crate) mod tests { + RETENTION_METRICS.len() + HELD_METRICS.len() + DISK_CLASS_METRICS.len() * 2 + + crate::gpu_telemetry::BASE_SERIES ); assert_only_disk_class_labels(&output); } @@ -6431,7 +6447,12 @@ pub(crate) mod tests { } fn assert_only_disk_class_labels(output: &str) { - let labeled: Vec<_> = output.lines().filter(|line| line.contains('{')).collect(); + let labeled: Vec<_> = output + .lines() + .filter(|line| { + line.contains('{') && !line.starts_with("racer_gpu_read_requests_total{") + }) + .collect(); assert_eq!(labeled.len(), DISK_CLASS_METRICS.len() * 2); for line in labeled { let (name, rest) = line.split_once('{').unwrap(); diff --git a/cmd/racer-dataplane/src/test_support.rs b/cmd/racer-dataplane/src/test_support.rs index 300e2686e5..ebd80edf7b 100644 --- a/cmd/racer-dataplane/src/test_support.rs +++ b/cmd/racer-dataplane/src/test_support.rs @@ -313,6 +313,7 @@ pub mod cluster { enable_gpu: crate::config::GpuMode::Disabled, gpu_staging_bytes: crate::config::DEFAULT_GPU_STAGING_BYTES, gpu_landing_bytes: crate::config::DEFAULT_GPU_LANDING_BYTES, + gpu_pci_p2pdma: false, control_endpoint: Some("https://control.invalid".into()), managed_control: None, enrollment_directory: None, diff --git a/cmd/racer-dataplane/verbs/src/link.rs b/cmd/racer-dataplane/verbs/src/link.rs index ccc7d8be85..3257610d21 100644 --- a/cmd/racer-dataplane/verbs/src/link.rs +++ b/cmd/racer-dataplane/verbs/src/link.rs @@ -421,9 +421,10 @@ pub(crate) struct Link { free_windows: Vec, - /// Free windows tied to the device pool. A window stays tied to the - /// first region it binds, so device grants never reuse host windows. - free_device_windows: Vec, + /// Free windows tied to each device pool, by pool id. A window stays + /// tied to the first region it binds, so device grants never reuse host + /// windows or another device pool's windows. + free_device_windows: BTreeMap>, free_slots: Vec, @@ -503,7 +504,7 @@ impl Link { free_windows: (0..windows.len()).rev().collect(), windows, device: device.clone(), - free_device_windows: Vec::new(), + free_device_windows: BTreeMap::new(), control, epoch, role, @@ -617,14 +618,14 @@ impl Link { /// Bind a window over a free buffer and queue a `Grant`. `done` gets the /// tag and arrival op once the `Grant` is posted. With `device`, the - /// buffer comes from the device pool and the arrival is + /// buffer comes from that device pool and the arrival is /// [`Arrival::Device`]: fails with [`Error::Unavailable`] if the rail has - /// no device pool and [`Error::Overloaded`] if none is free. + /// no such device pool and [`Error::Overloaded`] if none is free. #[allow(clippy::too_many_arguments)] pub(crate) fn grant( &mut self, pool: &Pool, - device: bool, + device: Option, transfer: [u8; 16], length: u32, deadline: Instant, @@ -646,17 +647,18 @@ impl Link { return; } - if device && pool.device_shared().is_none() { + if device.is_some_and(|id| pool.device_shared(id).is_none()) { let _ = done.complete(Err(Error::Unavailable)); return; } - let window = if device { - match self.free_device_windows.pop() { + let window = if let Some(id) = device { + match self.free_device_windows.get_mut(&id).and_then(Vec::pop) { Some(window) => window, None => match MemoryWindow::new(&self.device) { - // At most one device window per concurrent grant, so - // device windows never exceed the grant cap. + // At most one device window per concurrent grant per + // device pool, so each pool's windows never exceed the + // grant cap. Ok(window) => { self.windows.push(window); self.windows.len() - 1 @@ -671,8 +673,8 @@ impl Link { self.free_windows.pop().expect("a window per grant") }; - let taken = if device { - pool.take_device(BufferState::Binding) + let taken = if let Some(id) = device { + pool.take_device(id, BufferState::Binding) } else { pool.take(BufferState::Binding) }; @@ -702,9 +704,9 @@ impl Link { } /// Return `window` to its free list. - fn free_window(&mut self, device: bool, window: usize) { - if device { - self.free_device_windows.push(window); + fn free_window(&mut self, device: Option, window: usize) { + if let Some(id) = device { + self.free_device_windows.entry(id).or_default().push(window); } else { self.free_windows.push(window); } @@ -938,7 +940,7 @@ impl Link { InFlight::Invalidate(tag) => { let mut entry = self.grants.remove(&tag).ok_or(Break::Error)?; let buffer = entry.buffer.take().ok_or(Break::Error)?; - self.free_window(buffer.is_device(), entry.window); + self.free_window(buffer.device(), entry.window); let GrantState::Invalidating(outcome) = entry.state else { return Err(Break::Error); }; diff --git a/cmd/racer-dataplane/verbs/src/link/tests.rs b/cmd/racer-dataplane/verbs/src/link/tests.rs index 8bed8cb418..d43366c6a6 100644 --- a/cmd/racer-dataplane/verbs/src/link/tests.rs +++ b/cmd/racer-dataplane/verbs/src/link/tests.rs @@ -142,7 +142,7 @@ fn poll(op: &Op) -> Option> { fn grant(side: &mut Side, length: u32, deadline: Instant, now: Instant) -> Op { let (granted, done) = op(); side.link - .grant(&side.pool, false, [7; 16], length, deadline, done, now); + .grant(&side.pool, None, [7; 16], length, deadline, done, now); granted } diff --git a/cmd/racer-dataplane/verbs/src/node.rs b/cmd/racer-dataplane/verbs/src/node.rs index 23dac0c2b8..6e93bb9f1d 100644 --- a/cmd/racer-dataplane/verbs/src/node.rs +++ b/cmd/racer-dataplane/verbs/src/node.rs @@ -159,6 +159,7 @@ enum Command { AttachDevice { rail: u16, + id: u32, memory: DeviceMemory, done: Completer, }, @@ -166,7 +167,7 @@ enum Command { Grant { key: LinkKey, epoch: Epoch, - device: bool, + device: Option, transfer: [u8; 16], length: u32, deadline: Instant, @@ -198,8 +199,8 @@ struct NodeShared { pools: RwLock>>, - /// Device pools by rail, once attached. - device_pools: RwLock>>, + /// Device pools by rail and pool id, once attached. + device_pools: RwLock>>, /// Breaks that did not fit in a command queue. breaks: Mutex>, @@ -355,7 +356,7 @@ impl Client { self.submit(|done| Command::Grant { key, epoch, - device: false, + device: None, transfer, length, deadline, @@ -364,16 +365,18 @@ impl Client { } /// Receiver: like [`Self::grant`], but over a free buffer of the rail's - /// device pool (see [`Self::attach_device`]). The arrival is + /// device pool `id` (see [`Self::attach_device`]). The arrival is /// [`crate::Arrival::Device`]. The fence is the same: the page is handed /// out only after the peer's write completion and the window's signaled /// invalidation, so the peer can no longer write it. Fails with - /// `Unavailable` if the rail has no device pool, and `Overloaded` if no - /// device buffer is free. Nothing falls back to host memory. + /// `Unavailable` if the rail has no device pool `id`, and `Overloaded` + /// if no buffer of it is free. Nothing falls back to host memory. + #[allow(clippy::too_many_arguments)] pub fn grant_device( &self, key: LinkKey, epoch: Epoch, + id: u32, transfer: [u8; 16], length: u32, deadline: Instant, @@ -381,7 +384,7 @@ impl Client { self.submit(|done| Command::Grant { key, epoch, - device: true, + device: Some(id), transfer, length, deadline, @@ -390,12 +393,19 @@ impl Client { } /// Register `memory` (a dma-buf, such as GPU memory) on `rail` as its - /// device pool, split into page buffers of the node's buffer length. - /// Once per rail per round, after [`Self::activate`]. Resolves to the - /// number of device buffers. The memory stays registered until the node - /// drains; it is never read or written by the CPU. - pub fn attach_device(&self, rail: u16, memory: DeviceMemory) -> Op { - self.submit(|done| Command::AttachDevice { rail, memory, done }) + /// device pool `id`, split into page buffers of the node's buffer + /// length. A rail may hold several device pools with different ids (one + /// per GPU, say); each id once per rail per round, after + /// [`Self::activate`]. Resolves to the number of device buffers. The + /// memory stays registered until the node drains, and the node holds + /// `memory` until then; it is never read or written by the CPU. + pub fn attach_device(&self, rail: u16, id: u32, memory: DeviceMemory) -> Op { + self.submit(|done| Command::AttachDevice { + rail, + id, + memory, + done, + }) } /// Sender: write `source` into the peer's grant `tag` with immediate, @@ -452,11 +462,11 @@ impl Client { Some(pools.get(&rail)?.count_in(state)) } - /// Device page buffers of `rail` in `state`, for diagnostics. `None` - /// without a device pool. - pub fn device_buffers(&self, rail: u16, state: BufferState) -> Option { + /// Page buffers of `rail`'s device pool `id` in `state`, for + /// diagnostics. `None` without that device pool. + pub fn device_buffers(&self, rail: u16, id: u32, state: BufferState) -> Option { let pools = self.shared.device_pools.try_read().ok()?; - Some(pools.get(&rail)?.count_in(state)) + Some(pools.get(&(rail, id))?.count_in(state)) } /// Close the node: the owner sends `Close` on ready links, breaks every @@ -715,8 +725,13 @@ impl Owner { } => { let _ = done.complete(self.connect(key, epoch, remote, link_id, now)); } - Command::AttachDevice { rail, memory, done } => { - let _ = done.complete(self.attach_device(rail, memory)); + Command::AttachDevice { + rail, + id, + memory, + done, + } => { + let _ = done.complete(self.attach_device(rail, id, memory)); } Command::Grant { key, @@ -789,8 +804,8 @@ impl Owner { } } - /// Register `memory` as `rail`'s device pool. - fn attach_device(&mut self, rail: u16, memory: DeviceMemory) -> Result { + /// Register `memory` as `rail`'s device pool `id`. + fn attach_device(&mut self, rail: u16, id: u32, memory: DeviceMemory) -> Result { #[cfg(any(test, feature = "simulation"))] let _simulation = self .simulation @@ -799,20 +814,20 @@ impl Owner { let quota = self.quota.clone().ok_or(Error::InvalidRequest)?; let state = self.devices.get_mut(&rail).ok_or(Error::InvalidRequest)?; - if state.pool.device_shared().is_some() { + if state.pool.device_shared(id).is_some() { return Err(Error::InvalidRequest); } let region = NativeRegion::dmabuf(state.device.clone(), memory, quota)?; - state.pool.attach_device(region)?; + state.pool.attach_device(id, region)?; // The driver waker reaches the device pool on its next registration. - let shared = state.pool.device_shared().expect("just attached").clone(); + let shared = state.pool.device_shared(id).expect("just attached").clone(); let count = shared.count(); self.shared .device_pools .write() .unwrap_or_else(PoisonError::into_inner) - .insert(rail, shared); + .insert((rail, id), shared); Ok(count) } diff --git a/cmd/racer-dataplane/verbs/src/pool.rs b/cmd/racer-dataplane/verbs/src/pool.rs index 19127544ef..bdd76208c4 100644 --- a/cmd/racer-dataplane/verbs/src/pool.rs +++ b/cmd/racer-dataplane/verbs/src/pool.rs @@ -31,6 +31,7 @@ //! may use it. use std::{ + collections::BTreeMap, rc::Rc, sync::{ Arc, Mutex, MutexGuard, PoisonError, @@ -184,13 +185,18 @@ impl PoolShared { pub(crate) struct Buffer { index: u32, - /// In the device pool. - device: bool, + /// The device pool it is in, if any. + device: Option, } impl Buffer { /// True for a device pool buffer. pub(crate) fn is_device(&self) -> bool { + self.device.is_some() + } + + /// The device pool it is in, if any. + pub(crate) fn device(&self) -> Option { self.device } } @@ -214,8 +220,8 @@ pub(crate) struct Pool { /// every write source. `None` only during drop. implicit: Option>, - /// Device page buffers, once attached. - device: Option, + /// Device page buffers by caller id, once attached. + device: BTreeMap, quota: Arc, } @@ -261,15 +267,17 @@ impl Pool { }), region: Some(region), implicit: Some(implicit), - device: None, + device: BTreeMap::new(), quota, }) } /// Split the dma-buf `region` (from [`NativeRegion::dmabuf`]) into - /// device page buffers of the host buffer length. At most one per pool. - pub(crate) fn attach_device(&mut self, region: Rc) -> Result<()> { - if self.device.is_some() { + /// device page buffers of the host buffer length, as device pool `id`. + /// A rail may hold several device pools (one per GPU, say), each with + /// its own id. + pub(crate) fn attach_device(&mut self, id: u32, region: Rc) -> Result<()> { + if self.device.contains_key(&id) { return Err(Error::InvalidRequest); } @@ -283,41 +291,44 @@ impl Pool { return Err(Error::InvalidConfiguration); } - self.device = Some(DevicePool { - shared: Arc::new(PoolShared { - base: std::ptr::null_mut(), - memory: Some(memory), - buffer_len, - states: (0..count) - .map(|_| AtomicU8::new(BufferState::Free.code())) - .collect(), - free: Mutex::new((0..count as u32).rev().collect()), - owner: AtomicWaker::new(), - }), - region: Some(region), - }); + self.device.insert( + id, + DevicePool { + shared: Arc::new(PoolShared { + base: std::ptr::null_mut(), + memory: Some(memory), + buffer_len, + states: (0..count) + .map(|_| AtomicU8::new(BufferState::Free.code())) + .collect(), + free: Mutex::new((0..count as u32).rev().collect()), + owner: AtomicWaker::new(), + }), + region: Some(region), + }, + ); Ok(()) } - /// State shared with device pages, if a device pool is attached. - pub(crate) fn device_shared(&self) -> Option<&Arc> { - self.device.as_ref().map(|d| &d.shared) + /// State shared with device pages of pool `id`, if attached. + pub(crate) fn device_shared(&self, id: u32) -> Option<&Arc> { + self.device.get(&id).map(|d| &d.shared) } /// Register the owner's waker for buffers returned to either pool. pub(crate) fn register_owner(&self, waker: &Waker) { self.shared.register_owner(waker); - if let Some(device) = &self.device { + for device in self.device.values() { device.shared.register_owner(waker); } } /// State shared by `buffer`'s pool. fn shared_of(&self, buffer: &Buffer) -> &Arc { - if buffer.device { + if let Some(id) = buffer.device { &self .device - .as_ref() + .get(&id) .expect("device buffer without pool") .shared } else { @@ -327,22 +338,22 @@ impl Pool { /// The registered region `buffer` lives in. pub(crate) fn region_of(&self, buffer: &Buffer) -> &Rc { - if buffer.device { - let device = self.device.as_ref().expect("device buffer without pool"); + if let Some(id) = buffer.device { + let device = self.device.get(&id).expect("device buffer without pool"); device.region.as_ref().unwrap() } else { self.region() } } - /// Take a free device buffer straight into an owner `state`. `None` if - /// no device pool is attached or none is free. - pub(crate) fn take_device(&self, state: BufferState) -> Option { + /// Take a free buffer of device pool `id` straight into an owner + /// `state`. `None` if that pool is not attached or none is free. + pub(crate) fn take_device(&self, id: u32, state: BufferState) -> Option { assert!(state.owner_held() && state != BufferState::Quarantined); - let device = self.device.as_ref()?; + let device = self.device.get(&id)?; device.shared.pop(state).map(|index| Buffer { index, - device: true, + device: Some(id), }) } @@ -353,7 +364,7 @@ impl Pool { buffer: Buffer, len: usize, ) -> std::result::Result { - if !buffer.device || len == 0 || len > self.shared.buffer_len { + if buffer.device.is_none() || len == 0 || len > self.shared.buffer_len { return Err(buffer); } @@ -392,7 +403,7 @@ impl Pool { assert!(state.owner_held() && state != BufferState::Quarantined); self.shared.pop(state).map(|index| Buffer { index, - device: false, + device: None, }) } @@ -413,7 +424,7 @@ impl Pool { buffer: Buffer, len: usize, ) -> std::result::Result { - if buffer.device || len == 0 || len > self.shared.buffer_len { + if buffer.device.is_some() || len == 0 || len > self.shared.buffer_len { return Err(buffer); } @@ -437,8 +448,9 @@ impl Pool { Arc::strong_count(&self.shared) - 1 + self .device - .as_ref() - .map_or(0, |d| Arc::strong_count(&d.shared) - 1) + .values() + .map(|d| Arc::strong_count(&d.shared) - 1) + .sum::() } /// True if dropping now would deregister the region: no buffer is @@ -447,7 +459,7 @@ impl Pool { let free = |shared: &PoolShared| shared.count_in(BufferState::Free) == shared.count(); self.outstanding() == 0 && free(&self.shared) - && self.device.as_ref().is_none_or(|d| free(&d.shared)) + && self.device.values().all(|d| free(&d.shared)) } } @@ -460,7 +472,7 @@ impl Drop for Pool { std::mem::forget(self.quota.clone()); std::mem::forget(self.region.take()); std::mem::forget(self.implicit.take()); - if let Some(device) = &mut self.device { + for device in self.device.values_mut() { std::mem::forget(device.region.take()); } } @@ -757,18 +769,22 @@ mod tests { let mut pool = Pool::new(region, implicit, LEN, quota).unwrap(); let memory = crate::simulation::MemfdMemory::new(2 * LEN + 1, 0x9000).unwrap(); let (quota, _observer) = lifetime_tests::quota(); - assert!(pool.take_device(BufferState::Binding).is_none()); + assert!(pool.take_device(0, BufferState::Binding).is_none()); // Host regions are refused as device pools. let host = NativeRegion::with_access(device.clone(), LEN, Access::PAGES, quota.clone()).unwrap(); - assert_eq!(pool.attach_device(host), Err(Error::InvalidConfiguration)); + assert_eq!( + pool.attach_device(0, host), + Err(Error::InvalidConfiguration) + ); let region = NativeRegion::dmabuf(device.clone(), memory.clone(), quota.clone()).unwrap(); - pool.attach_device(region.clone()).unwrap(); - assert_eq!(pool.attach_device(region), Err(Error::InvalidRequest)); - let shared = pool.device_shared().unwrap().clone(); + pool.attach_device(0, region.clone()).unwrap(); + assert_eq!(pool.attach_device(0, region), Err(Error::InvalidRequest)); + let shared = pool.device_shared(0).unwrap().clone(); assert_eq!(shared.count(), 2); - let buffer = pool.take_device(BufferState::Binding).unwrap(); + let buffer = pool.take_device(0, BufferState::Binding).unwrap(); assert!(buffer.is_device()); + assert_eq!(buffer.device(), Some(0)); assert!(pool.region_of(&buffer).is_device()); assert_eq!(pool.region_of(&buffer).address(), 0x9000); pool.set(&buffer, BufferState::Invalidating); @@ -779,10 +795,11 @@ mod tests { let page = pool.own_device(buffer, 10).unwrap(); assert_eq!((page.offset(), page.len()), (0, 10)); assert_eq!(shared.state(0), BufferState::Owned); - let second = pool.take_device(BufferState::Binding).unwrap(); + let second = pool.take_device(0, BufferState::Binding).unwrap(); let second = pool.own_device(second, LEN).unwrap(); assert_eq!(second.offset(), LEN as u64); - assert!(pool.take_device(BufferState::Binding).is_none()); + assert!(pool.take_device(0, BufferState::Binding).is_none()); + assert!(pool.take_device(1, BufferState::Binding).is_none()); let host = pool.take(BufferState::Binding).unwrap(); let host = pool.own_device(host, 1).unwrap_err(); pool.release(host); @@ -799,4 +816,43 @@ mod tests { "leaked region keeps the owner" ); } + + #[test] + /// Several device pools on one rail keep their buffers apart: each id + /// hands out its own buffers over its own memory, and the pool is idle + /// only when every device pool is. + fn device_pools_are_kept_apart_by_id() { + let sim = Simulation::new(); + let device = self::device(&sim); + let (quota, _charge) = lifetime_tests::quota(); + let region = + NativeRegion::with_access(device.clone(), LEN, Access::PAGES, quota.clone()).unwrap(); + let implicit = NativeRegion::implicit(device.clone(), quota.clone()).unwrap(); + let mut pool = Pool::new(region, implicit, LEN, quota).unwrap(); + let (quota, _observer) = lifetime_tests::quota(); + let first = crate::simulation::MemfdMemory::new(LEN, 0x9000).unwrap(); + let second = crate::simulation::MemfdMemory::new(2 * LEN, 0xa000).unwrap(); + let attach = |pool: &mut Pool, id, memory| { + let region = NativeRegion::dmabuf(device.clone(), memory, quota.clone()).unwrap(); + pool.attach_device(id, region) + }; + attach(&mut pool, 3, first.clone()).unwrap(); + attach(&mut pool, 7, second.clone()).unwrap(); + assert_eq!(attach(&mut pool, 7, second), Err(Error::InvalidRequest)); + assert_eq!(pool.device_shared(3).unwrap().count(), 1); + assert_eq!(pool.device_shared(7).unwrap().count(), 2); + let a = pool.take_device(3, BufferState::Binding).unwrap(); + assert!(pool.take_device(3, BufferState::Binding).is_none()); + let b = pool.take_device(7, BufferState::Binding).unwrap(); + assert_eq!((a.device(), b.device()), (Some(3), Some(7))); + assert_eq!(pool.region_of(&a).address(), 0x9000); + assert_eq!(pool.region_of(&b).address(), 0xa000); + let page = pool.own_device(b, 5).unwrap(); + assert_eq!(page.memory().length(), 2 * LEN); + pool.release(a); + assert!(!pool.idle(), "a page of pool 7 is outstanding"); + drop(page); + assert!(pool.idle()); + drop(first); + } } diff --git a/cmd/racer-dataplane/verbs/tests/links.rs b/cmd/racer-dataplane/verbs/tests/links.rs index 5da2706af7..6c6eb78b3f 100644 --- a/cmd/racer-dataplane/verbs/tests/links.rs +++ b/cmd/racer-dataplane/verbs/tests/links.rs @@ -838,14 +838,17 @@ fn device_grants_land_in_dmabuf_memory() { let mut w = world(); let (ea, eb) = link_up(&mut w); let memory = device_memory(); - let attached = w.a.clients[0].attach_device(7, memory.clone()); + let attached = w.a.clients[0].attach_device(7, 0, memory.clone()); assert_eq!(wait(&mut w, &attached), Ok(2)); - assert_eq!(w.a.clients[0].device_buffers(7, BufferState::Free), Some(2)); - assert_eq!(w.b.clients[0].device_buffers(7, BufferState::Free), None); + assert_eq!( + w.a.clients[0].device_buffers(7, 0, BufferState::Free), + Some(2) + ); + assert_eq!(w.b.clients[0].device_buffers(7, 0, BufferState::Free), None); let deadline = now() + 5 * SECOND; let mut held = Vec::new(); for round in 0..2u8 { - let granted = w.a.clients[0].grant_device(TO_B, ea, [round; 16], 100, deadline); + let granted = w.a.clients[0].grant_device(TO_B, ea, 0, [round; 16], 100, deadline); let granted = wait(&mut w, &granted).unwrap(); let data = vec![0xa0 | round; 100]; let written = w.b.clients[1].write(TO_A, eb, granted.tag, bytes(data.clone()), deadline); @@ -861,12 +864,12 @@ fn device_grants_land_in_dmabuf_memory() { held.push(page); } assert_eq!( - w.a.clients[0].device_buffers(7, BufferState::Owned), + w.a.clients[0].device_buffers(7, 0, BufferState::Owned), Some(2) ); // Both device buffers are held: the next device grant is refused, and // nothing falls back to host memory. - let refused = w.a.clients[0].grant_device(TO_B, ea, [7; 16], 10, deadline); + let refused = w.a.clients[0].grant_device(TO_B, ea, 0, [7; 16], 10, deadline); assert_eq!(wait(&mut w, &refused).err(), Some(Error::Overloaded)); assert_eq!(w.a.clients[0].buffers(7, BufferState::Free), Some(BUFFERS)); // Host grants still use the host pool. @@ -880,12 +883,15 @@ fn device_grants_land_in_dmabuf_memory() { let page = held.pop().unwrap(); std::thread::spawn(move || drop(page)).join().unwrap(); pump(&mut w); - let granted = w.a.clients[0].grant_device(TO_B, ea, [9; 16], 10, deadline); + let granted = w.a.clients[0].grant_device(TO_B, ea, 0, [9; 16], 10, deadline); let granted = wait(&mut w, &granted).unwrap(); let released = w.b.clients[0].release(TO_A, eb, granted.tag); assert_eq!(wait(&mut w, &released), Ok(())); assert!(matches!(ready(&granted.arrival), Ok(Arrival::Released))); - assert_eq!(w.a.clients[0].device_buffers(7, BufferState::Free), Some(1)); + assert_eq!( + w.a.clients[0].device_buffers(7, 0, BufferState::Free), + Some(1) + ); // Shutdown waits for the held device page. w.a.clients[0].close(); pump(&mut w); @@ -901,6 +907,59 @@ fn device_grants_land_in_dmabuf_memory() { shutdown(w); } +#[test] +/// One rail holds several device pools (one per GPU, say). Each grant names +/// its pool: the write lands in that pool's memory only, one full pool does +/// not refuse grants on another, and an unknown pool id is unavailable. +fn device_pools_on_one_rail_are_independent() { + let mut w = world(); + let (ea, eb) = link_up(&mut w); + let first = MemfdMemory::new(LEN, 0x7000_0000).unwrap(); + let second = MemfdMemory::new(LEN, 0x7100_0000).unwrap(); + let attached = w.a.clients[0].attach_device(7, 0, first.clone()); + assert_eq!(wait(&mut w, &attached), Ok(1)); + let attached = w.a.clients[0].attach_device(7, 1, second.clone()); + assert_eq!(wait(&mut w, &attached), Ok(1)); + let deadline = now() + 5 * SECOND; + let mut held = Vec::new(); + for (id, memory, fill) in [(0u32, &first, 0x11u8), (1, &second, 0x22)] { + let granted = w.a.clients[0].grant_device(TO_B, ea, id, [fill; 16], 50, deadline); + let granted = wait(&mut w, &granted).unwrap(); + let written = w.b.clients[0].write(TO_A, eb, granted.tag, bytes(vec![fill; 50]), deadline); + assert_eq!(wait(&mut w, &written), Ok(())); + let page = device_page(ready(&granted.arrival)); + assert!(Arc::ptr_eq( + page.memory(), + &(memory.clone() as rdma_verbs::DeviceMemory) + )); + assert_eq!(memory.read(0, 50), vec![fill; 50]); + held.push(page); + } + assert_eq!( + w.a.clients[0].device_buffers(7, 0, BufferState::Owned), + Some(1) + ); + assert_eq!( + w.a.clients[0].device_buffers(7, 1, BufferState::Owned), + Some(1) + ); + let full = w.a.clients[0].grant_device(TO_B, ea, 0, [3; 16], 10, deadline); + assert_eq!(wait(&mut w, &full).err(), Some(Error::Overloaded)); + let unknown = w.a.clients[0].grant_device(TO_B, ea, 9, [3; 16], 10, deadline); + assert_eq!(wait(&mut w, &unknown).err(), Some(Error::Unavailable)); + // Freeing pool 1's page lets pool 1 grant again; pool 0 stays full. + drop(held.pop()); + pump(&mut w); + let again = w.a.clients[0].grant_device(TO_B, ea, 1, [4; 16], 10, deadline); + let again = wait(&mut w, &again).unwrap(); + let released = w.b.clients[0].release(TO_A, eb, again.tag); + assert_eq!(wait(&mut w, &released), Ok(())); + drop(held); + shutdown(w); + assert_eq!(Arc::strong_count(&first), 1); + assert_eq!(Arc::strong_count(&second), 1); +} + #[test] /// Attach is once per activated rail; a device grant without a device pool /// is unavailable; a failed registration leaves nothing behind and can be @@ -909,26 +968,26 @@ fn device_attach_rules_and_registration_faults() { let mut w = world(); let (ea, _) = link_up(&mut w); let deadline = now() + 5 * SECOND; - let missing = w.a.clients[0].grant_device(TO_B, ea, [1; 16], 10, deadline); + let missing = w.a.clients[0].grant_device(TO_B, ea, 0, [1; 16], 10, deadline); assert_eq!(wait(&mut w, &missing).err(), Some(Error::Unavailable)); let memory = device_memory(); - let wrong_rail = w.a.clients[0].attach_device(3, memory.clone()); + let wrong_rail = w.a.clients[0].attach_device(3, 0, memory.clone()); assert_eq!(wait(&mut w, &wrong_rail), Err(Error::InvalidRequest)); let live = w.sim.live_resources(); w.sim.fault(Operation::RegisterDmabuf, Fault::Reject); - let failed = w.a.clients[0].attach_device(7, memory.clone()); + let failed = w.a.clients[0].attach_device(7, 0, memory.clone()); assert_eq!(wait(&mut w, &failed), Err(Error::Unavailable)); assert_eq!(w.sim.live_resources(), live, "nothing registered stays"); assert_eq!(Arc::strong_count(&memory), 1); // Memory smaller than one buffer is refused after registration, and // the region is deregistered again. - let small = w.a.clients[0].attach_device(7, MemfdMemory::new(LEN - 1, 0).unwrap()); + let small = w.a.clients[0].attach_device(7, 0, MemfdMemory::new(LEN - 1, 0).unwrap()); assert_eq!(wait(&mut w, &small), Err(Error::InvalidConfiguration)); assert_eq!(w.sim.live_resources(), live); - let attached = w.a.clients[0].attach_device(7, memory.clone()); + let attached = w.a.clients[0].attach_device(7, 0, memory.clone()); assert_eq!(wait(&mut w, &attached), Ok(2)); assert_eq!(w.sim.live_resources(), live + 1); - let again = w.a.clients[0].attach_device(7, device_memory()); + let again = w.a.clients[0].attach_device(7, 0, device_memory()); assert_eq!(wait(&mut w, &again), Err(Error::InvalidRequest)); let attempts: Vec = w .sim @@ -949,13 +1008,13 @@ fn device_attach_rules_and_registration_faults() { fn device_grant_break_fences_before_reuse() { let mut w = world(); let (ea, eb) = link_up(&mut w); - let attached = w.a.clients[0].attach_device(7, device_memory()); + let attached = w.a.clients[0].attach_device(7, 0, device_memory()); assert_eq!(wait(&mut w, &attached), Ok(2)); let deadline = now() + 5 * SECOND; - let granted = w.a.clients[0].grant_device(TO_B, ea, [1; 16], 10, deadline); + let granted = w.a.clients[0].grant_device(TO_B, ea, 0, [1; 16], 10, deadline); let granted = wait(&mut w, &granted).unwrap(); assert_eq!( - w.a.clients[0].device_buffers(7, BufferState::Exposed), + w.a.clients[0].device_buffers(7, 0, BufferState::Exposed), Some(1) ); w.sim @@ -965,7 +1024,10 @@ fn device_grant_break_fences_before_reuse() { let _ = ready(&written); assert_eq!(ready(&granted.arrival).err(), Some(Error::Unavailable)); assert!(matches!(w.a.clients[0].link(TO_B), LinkState::Backoff(_))); - assert_eq!(w.a.clients[0].device_buffers(7, BufferState::Free), Some(2)); + assert_eq!( + w.a.clients[0].device_buffers(7, 0, BufferState::Free), + Some(2) + ); shutdown(w); } @@ -976,7 +1038,7 @@ fn failed_deregister_leaks_device_memory_owner() { let mut w = world(); link_up(&mut w); let memory = device_memory(); - let attached = w.a.clients[0].attach_device(7, memory.clone()); + let attached = w.a.clients[0].attach_device(7, 0, memory.clone()); assert_eq!(wait(&mut w, &attached), Ok(2)); w.sim.reject(Operation::Deregister, None, true); drop(w.a.owner.take()); diff --git a/deploy/racer/dataplane-config.yaml.tmpl b/deploy/racer/dataplane-config.yaml.tmpl index 15bdb02a4d..a3a6ecaa72 100644 --- a/deploy/racer/dataplane-config.yaml.tmpl +++ b/deploy/racer/dataplane-config.yaml.tmpl @@ -22,8 +22,10 @@ data: # containerd then fails the Pod on any node without an NVIDIA CDI spec. # The operator reads this key and changes the DaemonSet accordingly. # RACER_DATAPLANE_NVIDIA: "env" - # GPU reads need an image built with RACER_NATIVE_CUDA=true. auto turns GPU - # mode off when no usable GPU is found; true fails startup instead. + # GPU reads. The racer-dataplane image is built with CUDA. auto touches no + # GPU until the first gpu-read, then turns GPU mode on if a usable GPU is + # found; true sets up the GPUs at startup and fails startup when none is + # found; false turns it off. # RACER_ENABLE_GPU: "auto" # Cap on ciphertext staging memory per GPU, in bytes: a nonzero multiple of # 2 MiB, default 256 MiB. GPU reads stage ciphertext here for decryption. diff --git a/designs/racer-gpudirect.md b/designs/racer-gpudirect.md index 44fdb546ca..632297d2ce 100644 --- a/designs/racer-gpudirect.md +++ b/designs/racer-gpudirect.md @@ -37,6 +37,7 @@ If-Match: "" optional, as for Get Racer-Destination-Offset: required Racer-Destination-Length: required, > 0 Racer-Allocation-Size: required for gpu-read only +Racer-Gpu-Uuid: GPU- optional, gpu-read only ``` - The client sends exactly one file descriptor with the request head, as @@ -46,6 +47,13 @@ Racer-Allocation-Size: required for gpu-read only - For `gpu-read` the descriptor comes from `cuMemExportToShareableHandle` with `CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR`. `Racer-Allocation-Size` is the size passed to `cuMemCreate`. +- `Racer-Gpu-Uuid` is the UUID of the GPU that holds the allocation, from + `cuDeviceGetUuid`, in the `GPU-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx` form + `nvidia-smi` prints. The `GPU-` prefix is optional and hex is case + insensitive. Device numbers differ between the client's container and the + dataplane (a pod with one GPU sees it as device 0), but UUIDs do not, so the + dataplane uses the UUID to pick its own context for that GPU. Without it + (older clients), the dataplane tries each usable GPU in turn. - `Range`, `If-Match`, `Authorization` and `Racer-Metadata` mean the same as for Get. `Racer-Page-Credits` and `Racer-Byte-Credits` are ignored. @@ -57,14 +65,19 @@ The dataplane rejects the request with 400, before writing anything, when: - a destination header is missing, duplicated, malformed, or `offset + length` overflows; - for `host-read`, the descriptor is not a regular file; -- for `gpu-read`, `offset + length > Racer-Allocation-Size`. +- for `gpu-read`, `offset + length > Racer-Allocation-Size`; +- for `gpu-read`, `Racer-Gpu-Uuid` is malformed or duplicated. After resolving the range, the dataplane rejects with 400 if the destination length is smaller than the range length. Range and version errors keep their Get status codes (412, 416, and so on). -`gpu-read` returns 501 when the dataplane has no usable GPU, or when the -allocation is on a GPU the dataplane did not discover. +`gpu-read` returns 501 when the dataplane has no usable GPU, when +`Racer-Gpu-Uuid` names a GPU the dataplane cannot use (it logs +`stage=gpu state=import-skipped reason=unknown-gpu-uuid uuid=`), or when +the allocation is on a GPU the dataplane did not discover. If the named GPU +does not own the allocation, the dataplane tries the others and logs +`stage=gpu state=uuid-mismatch` when another one does. ### Response @@ -93,7 +106,10 @@ the kind 2 frame, as for Get. The destination may then hold partial data. 1. The dataplane receives the descriptor with `recvmsg` and `MSG_CMSG_CLOEXEC`. 2. For `gpu-read` it imports the handle, maps the whole allocation into its own - CUDA context, and grants itself read and write access. + CUDA context, and grants itself read and write access. Before mapping, the + shim checks that the allocation lives on the context's GPU. It resolves the + device the driver reports for the handle in the dataplane's own numbering + and compares GPU UUIDs, not ordinals. 3. It writes pages. 4. When the read ends for any reason, including a client disconnect, it waits for all GPU work that touches the allocation, then unmaps it, releases the @@ -117,8 +133,10 @@ Clients try modes in order: GPU, then host, then stream. `recv`. - A new dataplane without a usable GPU answers `gpu-read` with 501. -A client remembers an unsupported mode for the socket path until the process -exits: +A client remembers an unsupported mode for the socket path, logs one warning +with the status, and tries the mode again after a retry interval (300 seconds +by default in the Python client) or when the socket file is replaced by a +restarted dataplane: - 501 marks the mode unsupported at once. - 400, 404 and 405 can also mean a bad request, so the client falls back for @@ -235,7 +253,16 @@ above: - The helper opens the cuFile driver on another thread at startup. Opening it probes every GPU and can take most of a second, which would otherwise delay the first disk read. A cuFile driver or registration failure turns cuFile off - until restart. A failed read only drops that file's registration. + until restart. When registration fails because cuFile does not support the + file, its filesystem, or its device (`CU_FILE_IO_NOT_SUPPORTED`, + `CU_FILE_DEVICE_NOT_SUPPORTED`, `CU_FILE_INVALID_FILE_TYPE`), only that + device (`st_dev`) skips cuFile until restart. A failed read only drops that + file's registration. +- `RACER_GPU_PCI_P2PDMA=true` makes the shim set cuFile's `use_pci_p2pdma` + (`cuFileSetParameterBool`) before it opens the driver, so local NVMe reads + use PCI P2PDMA without `nvidia-fs`. If cuFile cannot set it, the open fails + and cuFile is off, as for any failed open. The flag is an argument of + `racer_gpu_file_open`, so the shim ABI is version 2. - At startup the helper runs a zero-length decrypt on each GPU. A GPU whose shim has no decrypt kernels, or a staging cap smaller than the two scratch slots, gets plaintext pages only. @@ -249,24 +276,37 @@ above: ### GPUDirect RDMA landing -When GPU decryption and RDMA are both on, each GPU can get one landing -buffer: GPU memory separate from the staging buffer, sized by -`RACER_GPU_LANDING_BYTES` (default 256 MiB per GPU, 0 turns landing off). -The GPU helper allocates it once, exports it once as a dma-buf, and keeps it -until the process exits. At each RDMA activation, every rail is matched with -at most one GPU, and every GPU with at most one rail: a GPU on the rail's -NUMA node first, then any free GPU (`assign_landing` in -`src/rdma/links.rs`). PCI distance is not used; only NUMA nodes are known on -both sides. The rail registers the landing buffer as its `rdma_verbs` device -pool, so the buffer is split into page-sized slots. If allocation, export, or -registration fails, the rail stays host-only. +When GPU decryption and RDMA are both on, each decrypting GPU gets one +landing buffer: GPU memory separate from the staging buffer, sized by +`RACER_GPU_LANDING_BYTES` (default 256 MiB per GPU, not per rail; 0 turns +landing off). At each RDMA activation the GPU helper allocates it and exports +it once as a dma-buf. `plan_landing` in `src/rdma/links.rs` splits it into +disjoint slices, one per rail, so every GPU can land over every rail and a +node with more GPUs than rails leaves no GPU out. If the buffer has fewer +page slots than there are rails, rails on the GPU's NUMA node come first, +then the rest in rail order. PCI distance is not used; only NUMA nodes are +known on both sides. Each slice is registered as its own `rdma_verbs` device +pool on its rail, keyed by GPU ordinal, so a rail can hold one pool per GPU. +If allocation, export, or registration fails, that rail and GPU stay +host-only. + +Lifetime: the buffer is freed when nothing refers to it. NIC registrations +and landed pages hold references, so it is freed after the RDMA node closes +and drains (shutdown, or before a new activation round), or right away when +no rail registered it. A registration `rdma_verbs` leaks on purpose (failed +deregistration) keeps its reference, so that memory is never freed under the +NIC. The next activation allocates and registers again. Buffers are not +resized at runtime: `RACER_GPU_LANDING_BYTES` comes from the ConfigMap, and +the operator restarts the pods when it changes. A `gpu-read` page then takes this path (`Fill::land_from_peer`): 1. If this node is not a candidate for the page, has no local copy, and has - a landing buffer on the request's GPU, it sends one CopyOnly request to the - first provider in visit order, carrying the GPU ordinal. -2. When the hop is RDMA and the rail's landing buffer is on that GPU, the + a landing buffer on the request's GPU, it sends a CopyOnly request, + carrying the GPU ordinal, to the first provider in visit order. On a miss + or failure it asks the next one, up to two providers + (`LANDING_PROVIDERS`), then takes the usual path. +2. When the hop is RDMA and the rail has a landing slice for that GPU, the receiver grants a landing slot instead of a host buffer. If no slot is free, it grants host memory as usual. 3. The sender writes the ciphertext into the slot. After the usual fence, the @@ -281,9 +321,19 @@ A landed page is used once. It does not join or create a single flight, is not published to memory, and is not persisted, because there are no host bytes to share. The slot is held until the GPU job finishes. If the provider answers with a host page (no free slot, or HTTP), that page is used once in -the same way, so it is not fetched twice. A miss, any failure, or a landing -buffer on another GPU falls back to the usual path. Host-only nodes and -non-GPU reads do not change. Relays never ask for landing. +the same way, so it is not fetched twice. If every provider asked misses or +fails, or the rail has no slice for that GPU, the read takes the usual path. +Host-only nodes and non-GPU reads do not change. Relays never ask for +landing. + +Not caching landed pages is a design choice, not a gap. The host memory and +disk tiers are the cache. Putting a landed page there would need a +device-to-host copy of the ciphertext, which is the cost landing exists to +avoid. Keeping landed pages in GPU memory would hold scarce slots that peers +need for new pages, and a page landed on one GPU is not usable by a read on +another without a GPU-to-GPU copy. Concurrent reads of one page on the same +GPU each land their own copy; that is rare, since a model's shards are +usually read once per GPU. `racer_rdma_pages_landed_total` counts pages that landed, and the GPU helper's `landed` counter counts landed pages it decrypted. A landed page the diff --git a/docs/content/guides/racer-vllm.md b/docs/content/guides/racer-vllm.md index cbffbe86ee..141c9fe9b9 100644 --- a/docs/content/guides/racer-vllm.md +++ b/docs/content/guides/racer-vllm.md @@ -22,34 +22,32 @@ client falls back to host memory, then to a plain stream. prefix `Llama-3.1-8B-Instruct`. - NVIDIA nodes set up by the unbounded agent. -## 1. Turn On GPU Mode +## 1. Check GPU Mode -Build the dataplane image with CUDA support and push it to the registry the -operator uses: +GPU mode is on by default. The `racer-dataplane` image is built with CUDA, and +`RACER_ENABLE_GPU` is `auto`. The dataplane does not touch the GPUs until the +first GPU read, so vLLM's first load turns GPU mode on. After a load, check a +dataplane log on the GPU node: ```bash -podman build --build-arg RACER_NATIVE_CUDA=true \ - -t /racer-dataplane: -f images/racer-dataplane/Containerfile . +kubectl -n unbounded-system logs | grep 'stage=gpu state=' ``` -The CUDA image also installs NVIDIA's cuFile library (`libcufile-12-9`) for -GPUDirect Storage reads, because the NVIDIA container runtime injects only the -driver libraries. cuFile uses its packaged default configuration, except that -Racer turns off cuFile's compatibility mode: without GPUDirect Storage -(`nvidia-fs` or PCIe P2PDMA) on the host, disk reads use Racer's own host -memory path instead of cuFile's slower POSIX fallback. Set -`CUFILE_ALLOW_COMPAT_MODE=true` on the dataplane to keep cuFile's fallback. -cuFile reads the filesystem type from udev, so the dataplane mounts the host's -`/run/udev` read-only. +`state=enabled` means GPU mode is on. `state=disabled` gives the reason. +`state=standby` means no GPU read has come in yet. -GPU mode is `auto` by default. To make the dataplane fail when it finds no GPU, -set it to `true`: +To make the dataplane set up the GPUs at startup and fail when it finds none, +so a broken setup is not silent, set `RACER_ENABLE_GPU` to `true`: ```bash kubectl -n unbounded-system patch configmap racer-dataplane-config \ --type merge -p '{"data":{"RACER_ENABLE_GPU":"true"}}' ``` +The image also has NVIDIA's cuFile library and a `cufile.json` for GPUDirect +Storage. Without GPUDirect Storage on the host, disk reads go through host +memory. + The dataplane gets the NVIDIA driver libraries without requesting a GPU. This is on by default (`RACER_DATAPLANE_NVIDIA: env` in `racer-dataplane-config`). See [GPU reads]({{< relref "reference/racer" >}}#gpu-reads). @@ -70,7 +68,7 @@ metadata: spec: containers: - name: vllm - image: + image: /racer-vllm: command: - sh - -c @@ -90,17 +88,54 @@ spec: The webhook mounts `/run/racer/models/client` read-only into the container and sets `RACER_SOCKET` to the socket in it. To limit the mount to some containers, -add `racer.unbounded-cloud.io/client-containers: vllm`. The webhook does not -change init containers, so run `racer-vllm-fetch` in the vLLM container. +add `racer.unbounded-cloud.io/client-containers: vllm`. Init containers get +the mount too, so `racer-vllm-fetch` can run in an init container that writes +to the shared `emptyDir`. ## 3. Install racer-vllm -Install both Python packages into the vLLM image: +Release and nightly builds publish the image `racer-vllm` next to the other +images, at the same registry and tag. It is `vllm/vllm-openai:v0.25.1` with +both packages installed, so it needs no changes. + +To use another vLLM image, install both Python packages into it. They are not +on PyPI. GitHub releases include their wheels and sdists (`racer_client-*` and +`racer_vllm-*`) as assets. Releases made before the packages were added do not. +Set `VERSION` to a release tag and point pip at it: + +```bash +VERSION=vX.Y.Z +pip install --no-index \ + --find-links "https://github.com/Azure/unbounded/releases/expanded_assets/${VERSION}" \ + racer-vllm +``` + +`--find-links` reads the asset links from the release page. `--no-index` keeps +pip off PyPI, so a package there with the same name cannot be installed by +mistake. This works because `racer-vllm` needs only `racer-client`, and +`racer-client` needs nothing else. + +You can also install the wheels by URL. The file names carry the package +version, not the release tag: + +```bash +pip install --no-index \ + "https://github.com/Azure/unbounded/releases/download/${VERSION}/racer_client-0.2.0-py3-none-any.whl" \ + "https://github.com/Azure/unbounded/releases/download/${VERSION}/racer_vllm-0.2.0-py3-none-any.whl" +``` + +To build an image without network access, download the wheels first and use +`--find-links `. + +From a checkout: ```bash pip install ./clients/python/racer ./clients/python/racer-vllm ``` +`make racer-python-dist` builds wheels for both packages into +`tmp/racer-python-dist`. + `racer-vllm` registers the `racer` load format with vLLM through a plugin entry point. No other code change is needed. @@ -111,10 +146,31 @@ takes two steps: 1. `racer-vllm-fetch racer://// ` copies `config.json`, the tokenizer files, and the safetensors index into ``. - It also writes `racer-model.json`, which records the `racer://` location. + It also writes `racer-model.json`, which records the `racer://` location and + the ETag of each fetched file and each weight file. The prefix is the directory that holds `model.safetensors.index.json` and the safetensors files. -2. `vllm serve --load-format racer` reads the weights through Racer. +2. `vllm serve --load-format racer` reads the weights through Racer, at + the ETags in `racer-model.json`. If the model changed at the origin since + the fetch, or changes during the load, the load fails with + `ModelChangedError`. It also fails if the checkpoint layout changed, for + example an index was added after a single-file fetch, because a weight file + it would read has no ETag in `racer-model.json`. Run `racer-vllm-fetch` + again to pick up the new version. + +After each load, vLLM's log shows one `racer: loaded ...` line with the file +count, bytes, time, throughput, and the number of reads in each delivery mode +(`gpu`, `host`, `stream`). When the dataplane refuses a mode, the client logs a +warning with the reason once, then retries the mode after +`mode_retry_seconds` (default 300) or when the dataplane restarts. + +With tensor parallelism, each rank reads every weight file of a regular +checkpoint. Ranks on one node share the dataplane, so repeat reads come from +its cache, not the origin. To have each rank read only its own files, save a +vLLM `sharded_state` checkpoint and pass +`--model-loader-extra-config '{"sharded_state": true}'`. `racer-vllm-fetch` +records every rank's files; pass `--pattern` if the checkpoint uses a custom +file name pattern. See the `racer-vllm` README for loader options such as `staging_bytes`. diff --git a/docs/content/reference/gpu/nvidia.md b/docs/content/reference/gpu/nvidia.md index a6c9621245..e42658d119 100644 --- a/docs/content/reference/gpu/nvidia.md +++ b/docs/content/reference/gpu/nvidia.md @@ -122,6 +122,54 @@ The image does **not** include NVIDIA kernel drivers or userspace libraries. Those are forwarded from the host at boot time via the bind-mount mechanism described above. +## GPUDirect Storage and GPUDirect RDMA + +The agent bind-mounts `/dev/nvidia-fs*` into the machine when the host has +them, the same as other NVIDIA device nodes. It does not install or load the +modules that create them. + +Both features need the NVIDIA **open** kernel module. GPUDirect Storage and +dma-buf export of CUDA memory do not work with the proprietary flavor, and the +open flavor supports Turing and later GPUs only. + +- **GPUDirect RDMA (dma-buf).** Needs the open kernel module, Linux 5.12 or + later, and rdma-core 34 or later. No extra module is needed. The Azure Linux + NVIDIA host image (`host-azurelinux3-nvidia`) installs `cuda-open-hwe`, + which is the open flavor. +- **GPUDirect Storage.** On local NVMe with ext4 or xfs, prefer the kernel's + PCI P2PDMA path: Linux 6.2 or later, the open driver, and cuFile 1.13 + (CUDA 12.8) or later. It needs no `nvidia-fs`. The `nvidia-fs` path to local + NVMe also needs NVMe patches from MLNX_OFED or DOCA. The Azure Linux host + images do not install `nvidia-fs`; its repository has no package for it. + cuFile uses P2PDMA only when asked: for Racer, set + `RACER_GPU_PCI_P2PDMA=true` in `racer-dataplane-config`. + +On Ubuntu Azure VMs, `hack/agent/azure-gpu-vm/install-nvidia-drivers.sh` +installs `nvidia-fs` when you pass `--gds`: + +```bash +./hack/agent/azure-gpu-vm/gpu-vm.sh setup-nvidia --size Standard_ND96isr_H100_v5 --gds +# or, on the VM: +sudo ./install-nvidia-drivers.sh Standard_ND96isr_H100_v5 --gds +``` + +The option installs `nvidia-fs-dkms` from NVIDIA's CUDA apt repository and +writes `/etc/modules-load.d/nvidia-fs.conf`. An apt pin keeps the driver on +Ubuntu's packages. It fails for SKUs that use the proprietary driver and when +Secure Boot is enabled, because the DKMS module is unsigned. After the reboot, +check the module: + +```bash +lsmod | grep nvidia_fs +ls -la /dev/nvidia-fs* +cat /proc/driver/nvidia-fs/stats +``` + +The agent lists device nodes when it writes the machine configuration, so load +`nvidia-fs` before that. The modules-load.d entry loads it at boot. A module +loaded later on a running node is not visible in the machine or the Racer +dataplane Pod until the configuration is written again and they restart. + ## Troubleshooting ### CUDA error 802 (`SYSTEM_DRIVER_MISMATCH`) diff --git a/docs/content/reference/racer.md b/docs/content/reference/racer.md index 2b15e489c1..f4f26b47b0 100644 --- a/docs/content/reference/racer.md +++ b/docs/content/reference/racer.md @@ -197,18 +197,26 @@ Tune `racer-dataplane-config` in the installation namespace (normally `unbounded The operator preserves administrator data and hashes edits into managed rollouts. Explicit workload environment wiring takes precedence over ConfigMap imports. Numeric values are unsigned decimal strings: `268435456`, not `256Mi`. +The operator checks each value with the same per-key rules the dataplane uses at +startup: enum and boolean values, decimal syntax, nonzero and range limits, the +2 MiB multiples of the GPU byte settings, and setting `RACER_IO_WORKERS` and +`RACER_CRYPTO_HELPERS` together. When a value is invalid, reconciliation fails +with an error that names every bad key, and the DaemonSet keeps its current +configuration until the ConfigMap is fixed. Limits that depend on the node, such +as CPU count or memory, are still checked only by the dataplane at startup. | Setting | Default | Scope / constraint | |---------|---------|--------------------| | `RACER_MAX_THREADS` | Unset (automatic) | Total I/O, crypto, and managed control threads; managed deployments need at least 3 | -| `RACER_IO_WORKERS` | Unset (automatic) | Exact I/O count; requires `RACER_CRYPTO_HELPERS` | +| `RACER_IO_WORKERS` | Unset (automatic) | Exact I/O count, 1..65,536; requires `RACER_CRYPTO_HELPERS` | | `RACER_CRYPTO_HELPERS` | Unset (automatic) | Exact helper count, from 1 through the exact I/O count; requires `RACER_IO_WORKERS` | | `RACER_ALLOW_SMT` | `false` | Use eligible physical cores by default; `true` opts into logical CPUs | | `RACER_ENABLE_RDMA` | `auto` | `auto`, `true`, or `false`; see [modes](#rdma) | | `RACER_ENABLE_GPU` | `auto` | `auto`, `true`, or `false`; see [GPU reads](#gpu-reads) | | `RACER_GPU_STAGING_BYTES` | `268435456` (256 MiB) | Cap on ciphertext staging memory per GPU, in bytes. Must be a nonzero multiple of 2 MiB. Below about 34 MiB the GPU does not decrypt and the CPU does; see [GPU reads](#gpu-reads) | | `RACER_GPU_LANDING_BYTES` | `268435456` (256 MiB) | GPU memory per GPU that peers' NICs write GPU-read pages into over RDMA, in bytes. A multiple of 2 MiB; `0` turns it off. Below one page slot (16 MiB plus 4 KiB) it is off; see [GPU reads](#gpu-reads) | -| `RACER_SHARES` | `4` | Reported shares; explicit Node `shares` wins | +| `RACER_GPU_PCI_P2PDMA` | `false` | `true` asks cuFile to read local NVMe into GPU memory over PCI P2PDMA instead of `nvidia-fs` (cuFile `use_pci_p2pdma`). Needs cuFile 1.13 (CUDA 12.8) or later; see [GPU reads](#gpu-reads) | +| `RACER_SHARES` | `4` | Reported shares, 1..4,294,967,295; explicit Node `shares` wins | | `RACER_PLAINTEXT_BYTES` / `RACER_CIPHERTEXT_BYTES` | Unset: `67108864` (64 MiB) / `134217728` (128 MiB) per final I/O worker | Independent plaintext/ciphertext admission budgets; an explicit value is node-wide and divided among I/O workers | | `RACER_DIRTY_BYTES` | Unset: `167772160` (160 MiB) per final I/O worker | Dirty-data budget; an explicit value is node-wide | | `RACER_REGISTERED_BYTES` | Unset: `100663296` (96 MiB) per final I/O worker | Native RDMA memory budget when enabled; an explicit value is node-wide | @@ -227,14 +235,14 @@ Numeric values are unsigned decimal strings: `268435456`, not `256Mi`. | `RACER_CHECKPOINT_BYTES` | `67108864` (64 MiB) | Node-wide checkpoint working-set budget; at most 512 MiB; device mode can raise it as described below | | `RACER_ACTIVE_PATH_SEARCHES` | `8` | **Per I/O worker** distinct unfinished path-search limit; 1..64, independent of cache size | | `RACER_SLAB_BYTES` | `1073741824` (1 GiB) | **Per I/O worker** file-backed slab geometry; positive multiple of segment size; does not limit device-mode capacity | -| `RACER_SEGMENT_BYTES` | `67108864` (64 MiB) | Per-worker segment size; must fit a page plus storage overhead | -| `RACER_FREE_SEGMENT_RESERVE` | `2` | Positive reserved segment count per I/O worker, less than that worker's total segments | +| `RACER_SEGMENT_BYTES` | `67108864` (64 MiB) | Per-worker segment size; must fit a page plus storage overhead (at least 16,797,712 bytes) | +| `RACER_FREE_SEGMENT_RESERVE` | `2` | Positive reserved segment count per I/O worker, less than that worker's total segments (at most 1,048,575) | | `RACER_SLAB_DIRECTORY` | `/var/lib/racer/slabs` | Slab files and checkpoints; still required and writable in device mode | | `RACER_DEVICE_DIRECTORY` | `/host/dev` | Device root; scan its `disk/by-id` directory; explicitly wired by managed workloads | | `RACER_DEVICE_BUSY_WAIT_SECS` | `120` | Startup bound for retrying selected devices whose exclusive open is busy (`EBUSY`); all busy devices share one deadline; 0..3,600 seconds, `0` skips busy devices immediately | -| `RACER_CONNECTIONS_PER_NEIGHBOR` | `2` | Per-worker neighbor connection cap; at most 1024 | +| `RACER_CONNECTIONS_PER_NEIGHBOR` | `2` | Per-worker neighbor connection cap; 1..1024 | | `RACER_ORIGIN_CONNECTIONS_PER_CACHE` | `32` | Per-cache, per-worker origin cap, independent of peer cap; at most 1024 and configured client connections (the default is lowered to fit), further clamped to the worker's outbound connection admission (a quarter of client connections). When that admission is at least 2, one slot is reserved for origin metadata requests across all caches, so page and peer connections use at most one fewer | -| `RACER_PEER_INFLIGHT_MAX` / `RACER_PEER_PER_NEIGHBOR_MAX` | `256` / `32` | Node-wide peer-exchange admission caps | +| `RACER_PEER_INFLIGHT_MAX` / `RACER_PEER_PER_NEIGHBOR_MAX` | `256` / `32` | Node-wide peer-exchange admission caps; each 1..65,536, and the per-neighbor cap at most the in-flight cap | | `RACER_REPLICA_SELECTION` | `spread` | Which page owner a non-owner asks first. `spread`: direct neighbors first, then a per-node hash spreads pages across all owners. `primary`: always the first-ranked owner | | `RACER_REQUEST_TIMEOUT_MS` / `RACER_PEER_ATTEMPT_TIMEOUT_MS` | Each `30000` | Independent request / local peer-exchange deadlines; 1..86,400,000 ms | | `RACER_READER_STALL_TIMEOUT_MS` | `10000` | At least 1 ms and no greater than request timeout | @@ -400,8 +408,13 @@ Racer-Destination-Length: - The descriptor comes from `cuMemExportToShareableHandle` with the POSIX file descriptor handle type. `Racer-Allocation-Size` is required and must be the size passed to `cuMemCreate`. +- `Racer-Gpu-Uuid` is optional: the UUID of the GPU that holds the memory, as + `nvidia-smi` prints it (`GPU-xxxxxxxx-...`). Device numbers in the client's + container can differ from the dataplane's, so the dataplane uses the UUID to + find the GPU. Without it, the dataplane tries each GPU. The Python client + sends it. - The answer is 501 when GPU mode is off or the memory is on a GPU the - dataplane did not find. It is 400 when the size does not match the memory or + dataplane did not find, including an unknown `Racer-Gpu-Uuid`. It is 400 when the size does not match the memory or the window does not fit in it. - A success has `Racer-Delivery: gpu`. A page frame means the GPU copy of that page is complete. @@ -788,6 +801,7 @@ only when a cache exists again. | Setting in `racer-dataplane-config` | Default | Purpose | |------------------------------------|---------|---------| | `RACER_CONTROL_SEATS` | `auto` | Choose the number of control seats automatically, or request 1 to 16 | +| `RACER_DATAPLANE_NVIDIA` | `env` | How the dataplane asks for NVIDIA driver libraries: `off`, `env`, or `cdi`; see [GPU reads](#gpu-reads) | Set `RACER_CONTROL_SEATS` in `racer-dataplane-config` to `auto` (the default) or an integer from 1 to 16. The active count cannot exceed the eligible node count. @@ -915,6 +929,45 @@ network access, especially with host networking. The diagnostics listener uses port 9090 by default. Persist scrape configuration on workload templates or supported overrides, not only live Pods. +### GPU metrics + +`/metrics` also reports `gpu-read` state. Labels have fixed values or come +from the GPUs the dataplane found and the active RDMA rails. + +| Metric | Type | Meaning | +|--------|------|---------| +| `racer_gpu_read_requests_total{outcome}` | counter | `gpu-read` requests by how the dataplane opened them: `gpu` (imported on a GPU; the read can still fail later), `not_enabled` (501, `gpu-read` is off), `wrong_device` (501, the memory is on a GPU the dataplane cannot use, or the `Racer-Gpu-Uuid` header names a GPU it did not find), `busy` (503), `invalid` (400), or `failed`. Requests the parser rejects before it knows the mode are not counted. | +| `racer_gpu_enabled` | gauge | 1 when `gpu-read` is on and found a usable GPU, else 0. | +| `racer_gpu_standby` | gauge | 1 while `auto` mode waits for the first `gpu-read` before it touches a GPU, else 0. | + +The rest are reported only when `racer_gpu_enabled` is 1: + +| Metric | Type | Meaning | +|--------|------|---------| +| `racer_gpu_decrypt_available{ordinal,uuid}` | gauge | 1 when the GPU can decrypt sealed pages. 0 means pages for it are decrypted by the CPU. | +| `racer_gpu_staging_slots_free{ordinal}` | gauge | Free cuFile staging slots on the GPU. | +| `racer_gpu_cufile_enabled` | gauge | 1 while cuFile reads are on. A driver-level failure turns it off until restart. | +| `racer_gpu_cufile_disabled_total{reason}` | counter | Times cuFile was turned off for the node: `driver_open`, `staging_register`, or `file_open` (a file registration error other than unsupported). | +| `racer_gpu_cufile_unsupported_total` | counter | Devices (`st_dev`) whose files cuFile reported as unsupported. Files on such a device are read through host memory; cuFile stays on for other devices. | +| `racer_gpu_read_pages_total{decrypt}` | counter | Pages written to GPU memory: `gpu` (decrypted on the GPU) or `cpu` (plaintext from host memory, decrypted by the CPU or already cached). | +| `racer_gpu_pages_rejected_total` | counter | Sealed pages the GPU rejected for a CRC or tag mismatch. | +| `racer_gpu_cufile_staged_pages_total` | counter | Pages cuFile read from disk straight into GPU staging. | +| `racer_gpu_cufile_failures_total` | counter | cuFile reads or registrations that failed. The page is read through host memory instead. | +| `racer_gpu_staging_exhausted_total` | counter | Disk reads that found no free staging slot and went through host memory. | +| `racer_gpu_landed_pages_decrypted_total` | counter | Pages decrypted on the GPU from an RDMA landing buffer. | +| `racer_gpu_open_requests` | gauge | `gpu-read` requests that hold GPU resources now. | +| `racer_gpu_leaked_resources` | gauge | GPU resources leaked on purpose after a failed release. | + +When GPU landing is configured (GPU decrypt and RDMA both on), each active +rail also reports the following. Each decrypting GPU's landing buffer is split +into one slice per rail, so the `{rail,ordinal}` series are per slice. + +| Metric | Type | Meaning | +|--------|------|---------| +| `racer_gpu_landing_enabled{rail}` | gauge | 1 when at least one GPU landing slice is registered on the rail. | +| `racer_gpu_landing_slots{rail,ordinal}` | gauge | NIC slots in GPU `ordinal`'s landing slice on the rail. | +| `racer_gpu_landing_exhausted_total{rail,ordinal}` | counter | Landing grants for that slice that found no free slot. The page lands in host memory instead. | + ## Upgrade and recovery Keep the operator, dataplane, SDK, and origin adapters on matching @@ -1003,7 +1056,7 @@ directory read-only and tells the app where the socket is. |-----|---------| | Label `racer.unbounded-cloud.io/client-socket: enabled` | Turns on the mount | | Annotation `racer.unbounded-cloud.io/client-cache` | ClusterCache name. Required | -| Annotation `racer.unbounded-cloud.io/client-containers` | Optional comma-separated app container names. Default: every app container. Init containers are never changed | +| Annotation `racer.unbounded-cloud.io/client-containers` | Optional comma-separated init or app container names. Default: every init and app container, except the injected `racer-object` sidecar | | Annotation `racer.unbounded-cloud.io/client-socket-injected` | Set by the webhook to the cache name; do not edit | Each selected container gets: @@ -1012,7 +1065,9 @@ Each selected container gets: mount still allows `connect` and `sendmsg` on the socket. - `RACER_SOCKET=/run/racer//client/socket`. -The pod is placed on Linux nodes without the `exclude` label, like sidecar pods. +Init containers, including native sidecars (`restartPolicy: Always`), get the +same mount, so an init container can fetch through the socket before the app +starts. The pod is placed on Linux nodes without the `exclude` label, like sidecar pods. The hostPath type is `Directory`, so the pod waits until the dataplane has created the cache directory. Host networking is allowed. @@ -1026,14 +1081,19 @@ container already sets `RACER_SOCKET` or mounts that path. GPU reads let a client on the node read an object straight into its own NVIDIA GPU memory. They need: -- A dataplane image built with `--build-arg RACER_NATIVE_CUDA=true`. +- A dataplane image built with CUDA. The released `racer-dataplane` image is, + so no setup is needed. It runs on nodes without GPUs too. An image built with + `--build-arg RACER_NATIVE_CUDA=false` has no GPU reads. - An NVIDIA node set up by the unbounded agent, which installs the NVIDIA container runtime and a CDI spec. -- `RACER_ENABLE_GPU` in `racer-dataplane-config`. `auto` turns GPU mode off - when no usable GPU is found, or when the image was built without CUDA. `true` - makes startup fail instead. `false` turns it off. - -At startup the dataplane logs each GPU it finds: ordinal, UUID, PCI address, +- `RACER_ENABLE_GPU` in `racer-dataplane-config`. `auto`, the default, does + not touch any GPU until a client sends the first `gpu-read`. Only then does + the dataplane load CUDA, look for GPUs, and turn GPU mode on, or off if it + finds no usable GPU. It does not try again until the pod restarts. `true` + sets up the GPUs at startup and makes startup fail when there is no usable + GPU. `false` turns it off. + +When it sets up the GPUs, the dataplane logs each GPU it finds: ordinal, UUID, PCI address, NUMA node, compute mode, and whether it can use it. It skips GPUs in `EXCLUSIVE_PROCESS` or `PROHIBITED` mode. One helper thread does all GPU work for the process, except cuFile disk reads, which run on four reader threads. At most 256 GPU reads run at once; more get 503. @@ -1050,11 +1110,19 @@ Where the bytes come from decides who decrypts them: - A page on local disk is read straight from the disk into GPU memory with cuFile when it can be. Otherwise the CPU reads it and copies it to the GPU. This needs GPUDirect Storage on the host (`nvidia-fs`, or PCIe P2PDMA). + For PCI P2PDMA, set `RACER_GPU_PCI_P2PDMA=true`. cuFile then does not fall + back to `nvidia-fs`, and if it cannot turn P2PDMA on, the dataplane logs + `stage=gpu state=cufile-disabled` and reads through host memory. The + `stage=gpu state=enabled` line reports the setting as `pci_p2pdma=`. `racer_gpu_direct_disk_reads_total` and `racer_gpu_direct_disk_read_bytes_total` count these reads. After a cuFile driver or registration failure the dataplane logs `stage=gpu state=cufile-disabled` and reads through host memory until it - restarts. + restarts. When cuFile only rejects a file because its filesystem or device + is not supported, the dataplane logs + `stage=gpu state=cufile-unsupported dev=: reason=` once + for that device, reads files on that device through host memory until it + restarts, and keeps cuFile on for other devices. - With RDMA on, a page from a peer can be written by the peer's NIC straight into GPU memory on this node (GPUDirect RDMA). See below. @@ -1062,15 +1130,20 @@ Where the bytes come from decides who decrypts them: If the GPU finds a damaged page, the read fails and the dataplane drops that copy from memory or disk, as it does when the CPU finds one. -When both GPU reads and RDMA are on, each GPU can get a landing buffer of -`RACER_GPU_LANDING_BYTES`, split into page slots. Each RDMA rail gets at most -one GPU's landing buffer, preferring a GPU on the rail's NUMA node. A GPU -read that needs a page from a peer over that rail, on that GPU, asks the peer -to write the encrypted page into a free slot. The GPU then checks and -decrypts it. Such a page is used only by that read: it is not kept in memory, -written to disk, or shared with other reads. When no slot is free, the GPUs -do not match, or anything fails, the page comes through host memory as -before. The dataplane logs `stage=gpu-landing` lines at RDMA activation, and +When both GPU reads and RDMA are on, each GPU that can decrypt gets a landing +buffer of `RACER_GPU_LANDING_BYTES` (per GPU, not per rail), split into page +slots. The buffer is divided across the RDMA rails, so every GPU can land +pages over every rail. If the buffer has fewer slots than there are rails, +the rails on the GPU's NUMA node come first. A GPU read that needs a page +from a peer asks up to two providers, in turn, to write the encrypted page +into a free slot of that GPU's share on the link's rail. The GPU then checks +and decrypts it. Such a page is used only by that read: it is not kept in +memory, written to disk, or shared with other reads, because the host never +has its bytes. When no slot is free or anything fails, the page comes through +host memory as before. Landing buffers are registered when GPU mode turns +on and at each RDMA activation, and freed when the RDMA node closes or when no rail could register +them. They are not resized at runtime; changing the setting restarts the +pods. The dataplane logs `stage=gpu-landing` lines when it registers them, and `racer_rdma_pages_landed_total` counts pages that landed in GPU memory. `RACER_GPU_STAGING_BYTES` is the GPU memory the dataplane keeps on each GPU for @@ -1096,14 +1169,51 @@ The operator reads it and changes the DaemonSet: | `RACER_DATAPLANE_NVIDIA` | Behavior | |--------------------------|----------| -| `env` (default) | Sets `NVIDIA_VISIBLE_DEVICES=all` and `NVIDIA_DRIVER_CAPABILITIES=compute,utility` on the dataplane container. The NVIDIA runtime injects the libraries. Other runtimes ignore these variables, so nodes without NVIDIA are not affected | -| `cdi` | Also adds the Pod annotation `cdi.k8s.io/racer-dataplane: nvidia.com/gpu=all`. containerd fails to create the container on any node that cannot resolve this CDI device, so use it only when every dataplane node has an NVIDIA CDI spec | -| `off` | Adds nothing | +| `env` (default) | Sets `NVIDIA_VISIBLE_DEVICES=all` and `NVIDIA_DRIVER_CAPABILITIES=compute,utility` on the dataplane container, mounts the host `/run/udev` read-only at `/run/udev` (cuFile reads each slab device's filesystem type from the udev database), and mounts a 64 MiB `emptyDir` at `/var/log/cufile` for the cuFile log. The NVIDIA runtime injects the libraries. Other runtimes ignore these variables, so nodes without NVIDIA are not affected | +| `cdi` | Everything `env` adds, plus the Pod annotation `cdi.k8s.io/racer-dataplane: nvidia.com/gpu=all`. containerd fails to create the container on any node that cannot resolve this CDI device, so use it only when every dataplane node has an NVIDIA CDI spec | +| `off` | Adds none of these: no NVIDIA variables, no CDI annotation, no `/run/udev` mount, and no cuFile log volume. Without driver libraries there is no GPU, so cuFile is not used | + +The dataplane image ships `/etc/racer/cufile.json` and sets +`CUFILE_ENV_PATH_JSON` to it. It turns off cuFile compatibility mode, so +cuFile fails instead of copying through host memory and Racer uses its own +host read path, and it logs errors to `/var/log/cufile/cufile.log`. Other +cuFile settings keep their built-in defaults. To use your own file, build an +image from the dataplane image that adds it, and set `CUFILE_ENV_PATH_JSON` +in `racer-dataplane-config`. Setting `CUFILE_ALLOW_COMPAT_MODE=true` +turns compatibility mode back on. On nodes set up by the unbounded agent, the NVIDIA runtime is the default containerd runtime and runs in CDI mode, so `env` is enough there. These env entries are operator wiring: workload overrides cannot change them. +### Testing on GPUs + +Regular CI has no GPU. It compiles the decrypt kernels with `nvcc` and runs +them on the CPU through an emulation, but never loads them on a GPU. On a +machine with an NVIDIA driver, `nvcc`, and Rust, run: + +```sh +make racer-gpu-test +``` + +This builds `libracer_gpu_cuda.so.1` with its kernels and runs the hardware +tests in `cmd/racer-dataplane/gpu/tests/hardware.rs`: shim load, device +discovery, import of an exported client allocation, GPU decrypt checked +against the CPU, and cuFile reads. It fails when no GPU is found; pass +`RACER_GPU_TESTS=1` to skip instead. cuFile tests skip when cuFile cannot open +or the directory does not support `O_DIRECT` (set `RACER_GPU_TEST_DIR`); set +`RACER_GPU_TESTS_CUFILE=require` to make that a failure. To also test +`gpu-read` end to end from the Python client, set `RACER_GPU_E2E_SOCKET` to a +running dataplane's client socket and `RACER_GPU_E2E_KEY` to the hex key of an +object it can serve (`RACER_GPU_E2E_METADATA` and +`RACER_GPU_E2E_AUTHORIZATION` are optional). Without the socket, that test is +skipped; set `RACER_GPU_E2E=require` to make it a failure. + +The `racer gpu` workflow (`.github/workflows/racer-gpu.yaml`) runs this target +nightly and on demand. It runs only when the repository variable +`RACER_GPU_RUNNER` holds the GPU runner labels as a JSON array, for example +`["self-hosted", "linux", "gpu"]`. + See [Load vLLM models from Racer]({{< relref "guides/racer-vllm" >}}). ## RDMA diff --git a/hack/agent/azure-gpu-vm/gpu-vm.sh b/hack/agent/azure-gpu-vm/gpu-vm.sh index e7488e6c36..509077a7be 100755 --- a/hack/agent/azure-gpu-vm/gpu-vm.sh +++ b/hack/agent/azure-gpu-vm/gpu-vm.sh @@ -25,6 +25,10 @@ # -i, --image OS image (default: Ubuntu2204) # -u, --admin-user Admin username (default: azureuser) # +# Setup Options: +# --gds Also install the GPUDirect Storage kernel module +# (nvidia-fs). Open kernel module SKUs only. +# # Delete Options: # -y, --yes Skip confirmation prompt # @@ -54,6 +58,7 @@ VM_SIZE="Standard_NC6s_v3" IMAGE="Ubuntu2204" ADMIN_USER="azureuser" SKIP_CONFIRM=false +GDS=false # --------------------------------------------------------------------------- # # Helpers @@ -80,12 +85,17 @@ Create Options: -i, --image OS image (default: Ubuntu2204) -u, --admin-user Admin username (default: azureuser) +Setup Options: + --gds Also install the GPUDirect Storage kernel module + (nvidia-fs). Open kernel module SKUs only. + Delete Options: -y, --yes Skip confirmation prompt Examples: gpu-vm.sh create gpu-vm.sh setup-nvidia + gpu-vm.sh setup-nvidia --size Standard_ND96isr_H100_v5 --gds gpu-vm.sh create --size Standard_NC24ads_A100_v4 --location eastus2 gpu-vm.sh delete -y EOF @@ -166,8 +176,16 @@ cmd_setup-nvidia() { echo "==> Installing NVIDIA drivers on VM '${VM_NAME}'" echo " VM SKU: ${VM_SIZE}" echo " Resource Group: ${RESOURCE_GROUP}" + echo " GDS: ${GDS}" echo "" + # Use name=value so az does not parse "--gds" as its own option. For + # Linux scripts the values are passed to the script positionally. + local params=("sku=${VM_SIZE}") + if [[ "${GDS}" == true ]]; then + params+=("gds=--gds") + fi + echo "==> Uploading and executing driver install script..." local result result=$(az vm run-command invoke \ @@ -176,7 +194,7 @@ cmd_setup-nvidia() { --name "${VM_NAME}" \ --command-id RunShellScript \ --scripts @"${driver_script}" \ - --parameters "${VM_SIZE}" \ + --parameters "${params[@]}" \ --output json) local stdout stderr @@ -191,6 +209,10 @@ cmd_setup-nvidia() { echo "==> Driver installation initiated. The VM will reboot automatically." echo " After reboot, SSH into the VM and verify with:" echo " ls -la /dev/nvidia*" + if [[ "${GDS}" == true ]]; then + echo " lsmod | grep nvidia_fs" + echo " ls -la /dev/nvidia-fs*" + fi } # --------------------------------------------------------------------------- # @@ -244,6 +266,7 @@ while [[ $# -gt 0 ]]; do -i|--image) IMAGE="$2"; shift 2 ;; -u|--admin-user) ADMIN_USER="$2"; shift 2 ;; -y|--yes) SKIP_CONFIRM=true; shift ;; + --gds) GDS=true; shift ;; -h|--help) usage ;; *) die "Unknown option: $1" ;; esac diff --git a/hack/agent/azure-gpu-vm/install-nvidia-drivers.sh b/hack/agent/azure-gpu-vm/install-nvidia-drivers.sh index 9541e7b74d..21ee822d55 100755 --- a/hack/agent/azure-gpu-vm/install-nvidia-drivers.sh +++ b/hack/agent/azure-gpu-vm/install-nvidia-drivers.sh @@ -30,16 +30,27 @@ # A systemd oneshot service is installed to run `nvidia-smi -pm 1` on boot, # keeping the driver loaded even when no GPU processes are active. # +# GPUDirect Storage (optional, --gds): +# Installs nvidia-fs-dkms from NVIDIA's CUDA apt repository and loads the +# nvidia-fs module at boot through /etc/modules-load.d/nvidia-fs.conf. The +# CUDA repository is pinned so that only nvidia-fs and the dkms version it +# requires come from it; the driver keeps coming from Ubuntu. GPUDirect +# Storage and dma-buf export need the open kernel module, and the DKMS build +# needs the driver sources in /usr/src, so --gds requires an open SKU and +# Secure Boot disabled. nvidia-fs alone does not enable GPUDirect reads from +# local NVMe on a stock kernel; see docs/content/reference/gpu/nvidia.md. +# # After a successful install the VM must be rebooted for the kernel module to # load and /dev/nvidia* devices to appear. # # Usage: -# sudo ./install-nvidia-drivers.sh +# sudo ./install-nvidia-drivers.sh [--gds] # # Examples: # sudo ./install-nvidia-drivers.sh Standard_NC6s_v3 # V100 -> proprietary 580 # sudo ./install-nvidia-drivers.sh Standard_NC24ads_A100_v4 # A100 -> open 580 # sudo ./install-nvidia-drivers.sh Standard_ND96isr_H100_v5 # H100 -> open 580 +# sudo ./install-nvidia-drivers.sh Standard_ND96isr_H100_v5 --gds # plus nvidia-fs set -euo pipefail @@ -246,19 +257,114 @@ UNIT log "nvidia-persistence-mode.service installed and enabled." } +# Print the NVIDIA CUDA apt repository directory for an os-release file and +# machine architecture, for example "ubuntu2404/x86_64". +cuda_repo_path() { + local os_release="$1" machine="$2" + local id version_id arch + # shellcheck disable=SC1090 + id=$(. "${os_release}" && echo "${ID:-}") + # shellcheck disable=SC1090 + version_id=$(. "${os_release}" && echo "${VERSION_ID:-}") + [[ "${id}" == "ubuntu" ]] || die "GPUDirect Storage install supports Ubuntu only, found '${id}'." + [[ "${version_id}" =~ ^[0-9]+\.[0-9]+$ ]] || die "Unexpected Ubuntu VERSION_ID '${version_id}'." + case "${machine}" in + x86_64) arch="x86_64" ;; + aarch64) arch="sbsa" ;; + *) die "GPUDirect Storage install does not support architecture '${machine}'." ;; + esac + echo "ubuntu${version_id//./}/${arch}" +} + +# Check that the GPUDirect Storage module can be installed for this driver. +check_gds_supported() { + local driver_type="$1" secure_boot="$2" + [[ "${driver_type}" == "open" ]] || die "--gds needs the NVIDIA open kernel module; this SKU uses the ${driver_type} driver." + [[ "${secure_boot}" == "no" ]] || die "--gds builds nvidia-fs with DKMS, which Secure Boot rejects unless you enroll a MOK key. Disable Secure Boot or install nvidia-fs yourself." +} + +# Write the apt pin that limits the CUDA repository to nvidia-fs and dkms. +write_cuda_pin() { + local path="$1" + cat > "${path}" <<'PIN' +# Written by install-nvidia-drivers.sh --gds. Take only nvidia-fs and the dkms +# version it requires from the CUDA repository; keep the Ubuntu driver. +Package: * +Pin: origin developer.download.nvidia.com +Pin-Priority: 100 + +Package: nvidia-fs nvidia-fs-dkms dkms +Pin: origin developer.download.nvidia.com +Pin-Priority: 600 +PIN +} + +# Write the modules-load.d entry that loads nvidia-fs at boot. modprobe also +# loads nvidia, which nvidia-fs depends on. +write_gds_modules_load() { + local path="$1" + echo "nvidia-fs" > "${path}" +} + +# Install the GPUDirect Storage kernel module (nvidia-fs) through DKMS. +install_gds() { + local repo + repo=$(cuda_repo_path /etc/os-release "$(uname -m)") + log "Installing GPUDirect Storage kernel module (nvidia-fs-dkms) from CUDA repo ${repo}..." + + local tmpdir + tmpdir=$(mktemp -d) + curl -fsSL -o "${tmpdir}/cuda-keyring.deb" \ + "https://developer.download.nvidia.com/compute/cuda/repos/${repo}/cuda-keyring_1.1-1_all.deb" + dpkg -i "${tmpdir}/cuda-keyring.deb" + rm -rf "${tmpdir}" + + # cuda-keyring ships /etc/apt/preferences.d/cuda-repository-pin-600, which + # pins every "NVIDIA CUDA" package to 600. APT takes the maximum of + # matching generic-form pins (apt_preferences(5)), so a separate lower pin + # cannot override it. Replace that file with our policy instead. It is a + # dpkg conffile, so later cuda-keyring upgrades keep our edit. + write_cuda_pin /etc/apt/preferences.d/cuda-repository-pin-600 + apt-get update -qq + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends nvidia-fs-dkms + + write_gds_modules_load /etc/modules-load.d/nvidia-fs.conf + dkms status nvidia-fs || true + + log "nvidia-fs installed; it loads at boot." +} + # --------------------------------------------------------------------------- # # Main # --------------------------------------------------------------------------- # -[[ $# -ge 1 ]] || die "Usage: $0 " +main() { +[[ $# -ge 1 ]] || die "Usage: $0 [--gds]" +VM_SKU="$1" +INSTALL_GDS=no +shift +while [[ $# -gt 0 ]]; do + case "$1" in + --gds) INSTALL_GDS=yes ;; + *) die "Unknown option: $1" ;; + esac + shift +done [[ $(id -u) -eq 0 ]] || die "This script must be run as root (use sudo)." -VM_SKU="$1" log "VM SKU: ${VM_SKU}" read -r GPU_NAME DRIVER_TYPE DRIVER_VERSION NEEDS_FABRIC <<< "$(resolve_sku "${VM_SKU}")" log "Detected GPU: ${GPU_NAME}, driver type: ${DRIVER_TYPE}, driver version: ${DRIVER_VERSION}, fabric manager: ${NEEDS_FABRIC}" +SECURE_BOOT=no if is_secure_boot; then + SECURE_BOOT=yes +fi +if [[ "${INSTALL_GDS}" == "yes" ]]; then + check_gds_supported "${DRIVER_TYPE}" "${SECURE_BOOT}" +fi + +if [[ "${SECURE_BOOT}" == "yes" ]]; then log "Secure Boot is enabled - using pre-signed kernel module packages." install_prebuilt_azure "${GPU_NAME}" "${DRIVER_VERSION}" "${DRIVER_TYPE}" else @@ -270,6 +376,14 @@ else esac fi +# Install nvidia-fs after the driver so DKMS builds it against the driver's +# sources and the new nv-p2p interface. Run it before Fabric Manager: +# install_fabric_manager uses dpkg --force-depends, which leaves unmet +# dependencies that make any later apt-get install abort. +if [[ "${INSTALL_GDS}" == "yes" ]]; then + install_gds +fi + # Install Fabric Manager for NVSwitch-based multi-GPU systems (H100, H200, # ND A100). Without it, CUDA initialization fails with error 802. if [[ "${NEEDS_FABRIC}" == "yes" ]]; then @@ -287,6 +401,18 @@ log "A reboot is required for the kernel module to load." log "After reboot, verify with:" log " ls -la /dev/nvidia*" log " nvidia-smi" +if [[ "${INSTALL_GDS}" == "yes" ]]; then + log " lsmod | grep nvidia_fs" + log " ls -la /dev/nvidia-fs*" + log " cat /proc/driver/nvidia-fs/stats" +fi log "" log "Rebooting now..." shutdown -r +0 +} + +# Tests set INSTALL_NVIDIA_DRIVERS_SOURCED=1 to source the helpers without +# running main. Do not test BASH_SOURCE: run-command may pipe the script. +if [[ "${INSTALL_NVIDIA_DRIVERS_SOURCED:-0}" != "1" ]]; then + main "$@" +fi diff --git a/hack/agent/azure-gpu-vm/install-nvidia-drivers_test.sh b/hack/agent/azure-gpu-vm/install-nvidia-drivers_test.sh new file mode 100755 index 0000000000..e39b2b9136 --- /dev/null +++ b/hack/agent/azure-gpu-vm/install-nvidia-drivers_test.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# Tests for the GPUDirect Storage helpers in install-nvidia-drivers.sh. +# They run without root, apt, or a GPU. +# +# Usage: ./hack/agent/azure-gpu-vm/install-nvidia-drivers_test.sh + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +INSTALL_NVIDIA_DRIVERS_SOURCED=1 +# shellcheck source=hack/agent/azure-gpu-vm/install-nvidia-drivers.sh +. "${SCRIPT_DIR}/install-nvidia-drivers.sh" +set +e + +TMP="$(mktemp -d)" +trap 'rm -rf "${TMP}"' EXIT +failures=0 + +expect_eq() { + local name="$1" want="$2" got="$3" + if [[ "${want}" != "${got}" ]]; then + echo "FAIL ${name}: want '${want}', got '${got}'" + failures=$((failures + 1)) + else + echo "ok ${name}" + fi +} + +# Run a command in a subshell and report whether it exited non-zero. +expect_fail() { + local name="$1" + shift + if ("$@") >/dev/null 2>&1; then + echo "FAIL ${name}: expected failure" + failures=$((failures + 1)) + else + echo "ok ${name}" + fi +} + +os_release() { + printf 'ID=%s\nVERSION_ID="%s"\n' "$1" "$2" > "${TMP}/os-release" + echo "${TMP}/os-release" +} + +expect_eq "repo ubuntu 22.04 x86_64" "ubuntu2204/x86_64" "$(cuda_repo_path "$(os_release ubuntu 22.04)" x86_64)" +expect_eq "repo ubuntu 24.04 x86_64" "ubuntu2404/x86_64" "$(cuda_repo_path "$(os_release ubuntu 24.04)" x86_64)" +expect_eq "repo ubuntu 24.04 aarch64" "ubuntu2404/sbsa" "$(cuda_repo_path "$(os_release ubuntu 24.04)" aarch64)" +expect_fail "repo rejects non-Ubuntu" cuda_repo_path "$(os_release azurelinux 3.0)" x86_64 +expect_fail "repo rejects bad version" cuda_repo_path "$(os_release ubuntu noble)" x86_64 +expect_fail "repo rejects unknown arch" cuda_repo_path "$(os_release ubuntu 24.04)" riscv64 + +expect_eq "gds open without Secure Boot" "0" "$( (check_gds_supported open no) >/dev/null 2>&1; echo $?)" +expect_fail "gds rejects proprietary driver" check_gds_supported proprietary no +expect_fail "gds rejects Secure Boot" check_gds_supported open yes + +write_gds_modules_load "${TMP}/nvidia-fs.conf" +expect_eq "modules-load entry" "nvidia-fs" "$(cat "${TMP}/nvidia-fs.conf")" + +write_cuda_pin "${TMP}/pin" +expect_eq "pin lowers CUDA repo" "1" "$(grep -c '^Pin-Priority: 100$' "${TMP}/pin")" +expect_eq "pin raises nvidia-fs and dkms" "Package: nvidia-fs nvidia-fs-dkms dkms" "$(grep '^Package: nvidia-fs' "${TMP}/pin")" + +expect_fail "main rejects unknown option" main Standard_ND96isr_H100_v5 --bogus +expect_fail "main requires a SKU" main + +if [[ "${failures}" -ne 0 ]]; then + echo "${failures} test(s) failed" + exit 1 +fi +echo "all tests passed" diff --git a/hack/cmd/release-bom/main.go b/hack/cmd/release-bom/main.go index 83f251f93c..121160a1b3 100644 --- a/hack/cmd/release-bom/main.go +++ b/hack/cmd/release-bom/main.go @@ -43,6 +43,7 @@ var releaseImageNames = []string{ "orca", "racer-dataplane", "racer-object", + "racer-vllm", "token-refresher", "unbounded-net-controller", "unbounded-net-node", diff --git a/hack/cmd/release-bom/main_test.go b/hack/cmd/release-bom/main_test.go index 45b939e820..7a990828dd 100644 --- a/hack/cmd/release-bom/main_test.go +++ b/hack/cmd/release-bom/main_test.go @@ -10,8 +10,11 @@ import ( "fmt" "os" "path/filepath" + "slices" "testing" + "sigs.k8s.io/yaml" + "github.com/Azure/unbounded/pkg/agent/goalstates" ) @@ -157,3 +160,47 @@ func TestWriteBOMCreatesOutputDirectory(t *testing.T) { t.Fatalf("BOM = %#v, want %#v", got, want) } } + +// TestReleaseBOMCoversReleaseWorkflowImages fails when release.yaml builds and +// signs an image matrix entry that the BOM does not list. An image missing +// from the BOM is never digest-pinned or signature-checked by +// hack/release/verify-release-artifacts.sh. +func TestReleaseBOMCoversReleaseWorkflowImages(t *testing.T) { + raw, err := os.ReadFile(filepath.Join("..", "..", "..", ".github", "workflows", "release.yaml")) + if err != nil { + t.Fatalf("read release workflow: %v", err) + } + + var workflow struct { + Jobs map[string]struct { + Strategy struct { + Matrix struct { + Component []struct { + Name string `json:"name"` + } `json:"component"` + } `json:"matrix"` + } `json:"strategy"` + } `json:"jobs"` + } + if err := yaml.Unmarshal(raw, &workflow); err != nil { + t.Fatalf("parse release workflow: %v", err) + } + + jobs := []string{"component-images", "netboot-image"} + for _, job := range jobs { + components := workflow.Jobs[job].Strategy.Matrix.Component + if len(components) == 0 { + t.Fatalf("release workflow job %q has no component matrix", job) + } + + for _, component := range components { + if !slices.Contains(releaseImageNames, component.Name) { + t.Errorf("release workflow job %q builds %q, which releaseImageNames does not list", job, component.Name) + } + } + } + + if !slices.Contains(releaseImageNames, "racer-vllm") { + t.Error("releaseImageNames must list the amd64-only racer-vllm image") + } +} diff --git a/images/racer-dataplane/Containerfile b/images/racer-dataplane/Containerfile index 15f655e19a..2b100517ce 100644 --- a/images/racer-dataplane/Containerfile +++ b/images/racer-dataplane/Containerfile @@ -7,10 +7,14 @@ ARG RUNTIME_IMAGE=ubuntu:noble # Build on the target platform: both Rust and the verbs adapter use its ABI. FROM ${RUST_IMAGE} AS builder ARG RACER_NATIVE_RDMA=true -# Build the CUDA shim (libracer_gpu_cuda.so.1) and the `cuda` feature. The -# shim dlopens the driver at runtime, so only nvcc (for the decrypt kernel PTX) -# is installed here; the runtime image gets the driver from the NVIDIA runtime. -ARG RACER_NATIVE_CUDA=false +# Build the CUDA shim (libracer_gpu_cuda.so.1) and the `cuda` feature. On by +# default so one image serves GPU and CPU nodes. The shim dlopens the driver at +# runtime, so only nvcc (for the decrypt kernel PTX) is installed here; the +# runtime image gets the driver from the NVIDIA runtime. With +# RACER_ENABLE_GPU=auto (the default), the dataplane touches no GPU until the +# first gpu-read, and turns GPU reads off then if no GPU is usable. Set false +# for a smaller image without GPU reads. +ARG RACER_NATIVE_CUDA=true ARG CUDA_NVCC_PACKAGE=cuda-nvcc-12-9 ENV RUSTFLAGS="-C force-frame-pointers=yes" WORKDIR /src @@ -107,7 +111,7 @@ RUN set -eu; \ # symbol versions, so the binary and verbs adapter run on noble's newer ones. FROM ${RUNTIME_IMAGE} ARG RACER_NATIVE_RDMA=true -ARG RACER_NATIVE_CUDA=false +ARG RACER_NATIVE_CUDA=true # cuFile user-space library for GPUDirect Storage. The NVIDIA container runtime # injects driver libraries only, not the toolkit's libcufile.so.0, so the CUDA # image ships it. The repo matches RUNTIME_IMAGE (ubuntu:noble). cuFile loads @@ -162,8 +166,12 @@ RUN set -eu; \ # adapter and CUDA shim under /usr/local/lib where the dynamic loader finds it by soname. COPY --from=builder /out/ / COPY LICENSE NOTICE /unbounded/licenses/ +# cuFile reads its configuration from CUFILE_ENV_PATH_JSON. Images built +# without CUDA never load cuFile, so the file and variable are inert there. +COPY images/racer-dataplane/cufile.json /etc/racer/cufile.json RUN ldconfig +ENV CUFILE_ENV_PATH_JSON=/etc/racer/cufile.json ENV PATH="/unbounded/bin:${PATH}" WORKDIR /unbounded diff --git a/images/racer-dataplane/cufile.json b/images/racer-dataplane/cufile.json new file mode 100644 index 0000000000..ef368da118 --- /dev/null +++ b/images/racer-dataplane/cufile.json @@ -0,0 +1,22 @@ +{ + // cuFile configuration for the Racer dataplane. The image sets + // CUFILE_ENV_PATH_JSON to this file. Keys left out keep cuFile's built-in + // defaults, which suit Racer's 16 MiB page reads (cuFile's default maximum + // direct IO size is also 16 MiB). To use another file, build an image on this + // one and set CUFILE_ENV_PATH_JSON to it. + "logging": { + // The dataplane root filesystem is read-only. The operator mounts an + // emptyDir here whenever RACER_DATAPLANE_NVIDIA is not off. + "dir": "/var/log/cufile", + // Errors only. Racer logs its own stage=gpu lines to stderr. + "level": "ERROR" + }, + "properties": { + // Compatibility mode does a POSIX read into a bounce buffer and a copy on + // the GPU helper thread. Racer's host read path does that work on I/O + // threads instead, so cuFile must fail rather than fall back. The CUDA + // shim also sets it at runtime: off, or on when CUFILE_ALLOW_COMPAT_MODE + // is true (cmd/racer-dataplane/gpu/native/gpu_cuda.c). + "allow_compat_mode": false + } +} diff --git a/images/racer-vllm/Containerfile b/images/racer-vllm/Containerfile new file mode 100644 index 0000000000..e04b1233d6 --- /dev/null +++ b/images/racer-vllm/Containerfile @@ -0,0 +1,33 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# vLLM's OpenAI server with racer-client and racer-vllm installed, so +# `vllm serve --load-format racer` and `racer-vllm-fetch` work as is. +# racer-vllm targets vLLM v0.25.x; keep this tag in that series. +ARG VLLM_IMAGE=docker.io/vllm/vllm-openai:v0.25.1 + +# Build stage +# Both packages are pure Python, so the wheels do not depend on the base image. +FROM docker.io/library/python:3.12-slim AS builder +WORKDIR /src +RUN pip install --no-cache-dir build==1.3.0 +COPY clients/python/racer ./racer +COPY clients/python/racer-vllm ./racer-vllm +RUN python -m build --wheel --outdir /wheels ./racer \ + && python -m build --wheel --outdir /wheels ./racer-vllm + +# Runtime stage +FROM ${VLLM_IMAGE} +# Install both wheels together so racer-vllm's racer-client dependency comes +# from them, not from an index. vLLM images install packages with uv; fall +# back to pip for bases without it. +RUN --mount=type=bind,from=builder,source=/wheels,target=/wheels \ + set -eu; \ + if command -v uv >/dev/null; then \ + uv pip install --system --no-cache --no-index --find-links /wheels racer-client racer-vllm; \ + else \ + python3 -m pip install --no-cache-dir --no-index --find-links /wheels racer-client racer-vllm; \ + fi; \ + python3 -c "import racer, racer_vllm"; \ + command -v racer-vllm-fetch +COPY LICENSE NOTICE /unbounded/licenses/ diff --git a/internal/operator/components/racer/config.go b/internal/operator/components/racer/config.go new file mode 100644 index 0000000000..2ac57a0c48 --- /dev/null +++ b/internal/operator/components/racer/config.go @@ -0,0 +1,204 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +package racer + +import ( + "fmt" + "math" + "slices" + "sort" + "strconv" + "strings" + "unicode" + + "github.com/Azure/unbounded/internal/racer/workload" +) + +// gpuStagingUnit mirrors GPU_STAGING_UNIT in cmd/racer-dataplane/src/config.rs. +const gpuStagingUnit = 2 << 20 + +// Constants mirrored from cmd/racer-dataplane/src: MIB, MAX_BYTES, and +// MAX_ENTRIES in config.rs, PAGE_BYTES in model, DUPLICATE_BYTES in +// read/candidates.rs, and the storage record size checked in config.rs +// (a page, a 16-byte tag, and MAX_HEADER_BYTES from store.rs). +const ( + mib = 1 << 20 + maxBytes = 64 << 30 + maxEntries = 1 << 20 + pageBytes = 16 * mib + duplicateBytes = 2*pageBytes + 16 + recordBytes = pageBytes + 16 + 16384 + 4096 +) + +// The checks below mirror the startup parsing in +// cmd/racer-dataplane/src/config.rs so that a value the dataplane would refuse +// fails the operator plan instead of rolling Pods into a crash loop. They +// cover per-key syntax and ranges, not cross-key or hardware policy. +var ( + // Enumerated keys and their accepted values. + dataplaneEnumKeys = map[string][]string{ + "RACER_ENABLE_RDMA": {"auto", "true", "false"}, + "RACER_ENABLE_GPU": {"auto", "true", "false"}, + "RACER_REPLICA_SELECTION": {"primary", "spread"}, + "RACER_ADMISSION_MODE": {"disabled", "second-sight"}, + } + + // Booleans parse with Rust's bool FromStr, which accepts only true and false. + dataplaneBoolKeys = []string{ + "RACER_ALLOW_SMT", "RACER_OPAQUE_RELAY", "RACER_PEER_TCP_NODELAY", + "RACER_MEMORY_PRESSURE", "RACER_PPROF_ENABLED", "RACER_GPU_PCI_P2PDMA", + } + + // Inclusive single-key ranges from Config::from_lookup and the validate + // methods it calls. Cross-key rules, such as RACER_CRYPTO_HELPERS at most + // RACER_IO_WORKERS or plaintext budgets that must fit the range window, + // are left to the dataplane. + dataplaneRangeKeys = map[string][2]uint64{ + "RACER_GPU_LANDING_BYTES": {0, math.MaxUint64}, + "RACER_IO_WORKERS": {1, 65536}, + "RACER_CRYPTO_HELPERS": {1, 65536}, + "RACER_SHARES": {1, math.MaxUint32}, + "RACER_DISK_PAGE_ENTRIES": {1, maxEntries}, + "RACER_ADMISSION_HISTORY_BYTES": {32, 512 * mib}, + "RACER_ADMISSION_PERIOD_SECS": {1, 86400}, + "RACER_CHECKPOINT_BYTES": {1, 512 * mib}, + "RACER_DEVICE_BUSY_WAIT_SECS": {0, 3600}, + "RACER_WARM_MAX_TTL_SECS": {1, 86400}, + "RACER_WARM_FILL_TIMEOUT_SECS": {1, 86400}, + "RACER_WARM_MAX_JOBS": {1, 1024}, + "RACER_WARM_WINDOW_PAGES": {1, 64}, + "RACER_RETAINED_SEGMENT_PERCENT": {1, 100}, + "RACER_RANGE_SPILL_PAGES": {0, 1024}, + "RACER_PAGE_HEDGE_DELAY_MS": {1, 30000}, + "RACER_PAGE_HEDGE_SLOTS": {0, 32}, + "RACER_PAGE_HEDGE_BYTES": {0, 32 * duplicateBytes}, + "RACER_PEER_INFLIGHT_MAX": {1, 65536}, + "RACER_PEER_PER_NEIGHBOR_MAX": {1, 65536}, + "RACER_PEER_RECEIVE_MAX": {0, 65536}, + "RACER_PEER_RECEIVE_QUEUE_MAX": {1, 65536}, + "RACER_PEER_RECEIVE_WAIT_MS": {1, 30000}, + "RACER_SLAB_BYTES": {1, math.MaxInt64}, + "RACER_SEGMENT_BYTES": {recordBytes, math.MaxInt64}, + "RACER_FREE_SEGMENT_RESERVE": {1, maxEntries - 1}, + "RACER_PLAINTEXT_BYTES": {1, maxBytes}, + "RACER_CIPHERTEXT_BYTES": {1, maxBytes}, + "RACER_DIRTY_BYTES": {pageBytes + 16, maxBytes}, + "RACER_REGISTERED_BYTES": {1, maxBytes}, + "RACER_REQUEST_CONTEXT_BYTES": {128 * 1024, maxBytes}, + "RACER_FLIGHTS": {1, 65536}, + "RACER_WAITERS_PER_FLIGHT": {1, 4096}, + "RACER_QUEUE_ENTRIES": {2, 65536}, + "RACER_CONNECTIONS_PER_NEIGHBOR": {1, 1024}, + "RACER_CLIENT_CONNECTIONS": {1, 65536}, + "RACER_PIPES": {1, 65536}, + "RACER_RANGE_WINDOW_PAGES": {1, 64}, + "RACER_HEADER_BYTES": {1024, 32 * 1024}, + "RACER_PLACEMENT_CACHE_BYTES": {1024, 512 * mib}, + "RACER_CACHED_PATHS": {1, maxEntries}, + "RACER_PATH_CACHE_BYTES": {1, 512 * mib}, + "RACER_ACTIVE_PATH_SEARCHES": {1, 64}, + "RACER_RETAINED_SNAPSHOTS": {2, 64}, + "RACER_METADATA_ENTRIES": {1, maxEntries}, + "RACER_RELAY_TRANSFERS": {1, 65536}, + "RACER_ORIGIN_CONNECTIONS_PER_CACHE": {1, 1024}, + "RACER_REQUEST_TIMEOUT_MS": {1, 86400000}, + "RACER_PEER_ATTEMPT_TIMEOUT_MS": {1, 86400000}, + "RACER_READER_STALL_TIMEOUT_MS": {1, 86400000}, + "RACER_SHUTDOWN_TIMEOUT_MS": {1, 3600000}, + } +) + +// validateDataplaneConfig reports every racer-dataplane-config value the +// dataplane or operator would reject, sorted by key, in one error. +func validateDataplaneConfig(data map[string]string) error { + var problems []string + + add := func(key, rule string) { + problems = append(problems, key+" "+rule) + } + + for key, allowed := range dataplaneEnumKeys { + if value, ok := data[key]; ok && !slices.Contains(allowed, value) { + add(key, "must be one of "+strings.Join(allowed, ", ")) + } + } + + for _, key := range dataplaneBoolKeys { + if value, ok := data[key]; ok && value != "true" && value != "false" { + add(key, "must be true or false") + } + } + + for key, bounds := range dataplaneRangeKeys { + if value, ok := data[key]; ok { + if n, valid := dataplaneDecimal(value); !valid || n < bounds[0] || n > bounds[1] { + add(key, fmt.Sprintf("must be an integer from %d to %d", bounds[0], bounds[1])) + } + } + } + + if value, ok := data["RACER_GPU_STAGING_BYTES"]; ok { + if n, valid := dataplaneDecimal(value); !valid || n < gpuStagingUnit || n%gpuStagingUnit != 0 { + add("RACER_GPU_STAGING_BYTES", "must be a nonzero multiple of 2097152 (2 MiB)") + } + } + + if value, ok := data["RACER_GPU_LANDING_BYTES"]; ok { + if n, valid := dataplaneDecimal(value); valid && n%gpuStagingUnit != 0 { + add("RACER_GPU_LANDING_BYTES", "must be 0 or a multiple of 2097152 (2 MiB)") + } + } + + _, ioSet := data["RACER_IO_WORKERS"] + if _, helpersSet := data["RACER_CRYPTO_HELPERS"]; ioSet != helpersSet { + add("RACER_IO_WORKERS", "and RACER_CRYPTO_HELPERS must be set together") + } + + if value, ok := data[workload.ControlSeatsKey]; ok && value != "auto" { + // The runtime accepts only the canonical decimal form, so reject + // leading zeros, signs, and whitespace before Pods roll. + seats, err := strconv.ParseUint(value, 10, 8) + if err != nil || strconv.FormatUint(seats, 10) != value || seats < 1 || seats > workload.MaxControlSeats { + add(workload.ControlSeatsKey, fmt.Sprintf("must be auto or an integer from 1 to %d", workload.MaxControlSeats)) + } + } + + if value, ok := data[workload.NVIDIAKey]; ok { + if mode := workload.NVIDIAInjection(value); mode != workload.NVIDIAOff && mode != workload.NVIDIAEnv && mode != workload.NVIDIACDI { + add(workload.NVIDIAKey, "must be off, env, or cdi") + } + } + + if value, ok := data["RACER_MAX_THREADS"]; ok { + if threads, valid := dataplaneDecimal(value); !valid || threads < minThreads { + add("RACER_MAX_THREADS", fmt.Sprintf("must be an integer of at least %d, including the control plane thread", minThreads)) + } + } + + if len(problems) == 0 { + return nil + } + + sort.Strings(problems) + + return fmt.Errorf("ConfigMap %s is invalid, so the dataplane DaemonSet keeps its current configuration: %s", dataplaneConfigName, strings.Join(problems, "; ")) +} + +// dataplaneDecimal mirrors the dataplane's text and decimal rules: nonempty, +// at most 4096 bytes, no control characters, ASCII digits only, and within +// the u64 range. +func dataplaneDecimal(value string) (uint64, bool) { + if value == "" || len(value) > 4096 || strings.IndexFunc(value, func(r rune) bool { return r < '0' || r > '9' || unicode.IsControl(r) }) >= 0 { + return 0, false + } + + trimmed := strings.TrimLeft(value, "0") + if trimmed == "" { + return 0, true + } + + n, err := strconv.ParseUint(trimmed, 10, 64) + + return n, err == nil +} diff --git a/internal/operator/components/racer/config_test.go b/internal/operator/components/racer/config_test.go new file mode 100644 index 0000000000..62ac4bb00b --- /dev/null +++ b/internal/operator/components/racer/config_test.go @@ -0,0 +1,208 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +package racer + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" +) + +func TestValidateDataplaneConfigAccepts(t *testing.T) { + for name, data := range map[string]map[string]string{ + "empty": {}, + "gpu": { + "RACER_ENABLE_GPU": "true", "RACER_GPU_STAGING_BYTES": "2097152", + "RACER_GPU_LANDING_BYTES": "0", "RACER_DATAPLANE_NVIDIA": "cdi", + "RACER_GPU_PCI_P2PDMA": "true", + }, + "gpu pci p2pdma off": {"RACER_GPU_PCI_P2PDMA": "false"}, + "gpu auto with large landing": {"RACER_ENABLE_GPU": "auto", "RACER_GPU_LANDING_BYTES": "268435456"}, + "booleans and enums": { + "RACER_ALLOW_SMT": "true", "RACER_PPROF_ENABLED": "false", "RACER_ENABLE_RDMA": "false", + "RACER_REPLICA_SELECTION": "primary", "RACER_ADMISSION_MODE": "disabled", + }, + "leading zeros and bounds": { + "RACER_MAX_THREADS": "0004", "RACER_WARM_MAX_JOBS": "1024", "RACER_RETAINED_SEGMENT_PERCENT": "1", + "RACER_RANGE_SPILL_PAGES": "0", "RACER_SHARES": "4294967295", + }, + "range minimums": { + "RACER_IO_WORKERS": "1", "RACER_CRYPTO_HELPERS": "1", "RACER_SHARES": "1", + "RACER_ADMISSION_HISTORY_BYTES": "32", "RACER_DEVICE_BUSY_WAIT_SECS": "0", + "RACER_PAGE_HEDGE_SLOTS": "0", "RACER_PEER_RECEIVE_MAX": "0", + "RACER_SLAB_BYTES": "1", "RACER_SEGMENT_BYTES": "16797712", "RACER_FREE_SEGMENT_RESERVE": "1", + "RACER_DIRTY_BYTES": "16777232", "RACER_REQUEST_CONTEXT_BYTES": "131072", + "RACER_QUEUE_ENTRIES": "2", "RACER_HEADER_BYTES": "1024", "RACER_PLACEMENT_CACHE_BYTES": "1024", + "RACER_RETAINED_SNAPSHOTS": "2", "RACER_RANGE_WINDOW_PAGES": "1", "RACER_REQUEST_TIMEOUT_MS": "1", + }, + "range maximums": { + "RACER_IO_WORKERS": "65536", "RACER_CRYPTO_HELPERS": "65536", "RACER_DISK_PAGE_ENTRIES": "1048576", + "RACER_ADMISSION_HISTORY_BYTES": "536870912", "RACER_ADMISSION_PERIOD_SECS": "86400", + "RACER_CHECKPOINT_BYTES": "536870912", "RACER_DEVICE_BUSY_WAIT_SECS": "3600", + "RACER_RANGE_SPILL_PAGES": "1024", "RACER_PAGE_HEDGE_DELAY_MS": "30000", "RACER_PAGE_HEDGE_SLOTS": "32", + "RACER_PAGE_HEDGE_BYTES": "1073742336", "RACER_PEER_INFLIGHT_MAX": "65536", + "RACER_PEER_RECEIVE_MAX": "65536", "RACER_PEER_RECEIVE_QUEUE_MAX": "65536", + "RACER_PEER_RECEIVE_WAIT_MS": "30000", "RACER_SLAB_BYTES": "9223372036854775807", + "RACER_FREE_SEGMENT_RESERVE": "1048575", "RACER_PLAINTEXT_BYTES": "68719476736", + "RACER_FLIGHTS": "65536", "RACER_WAITERS_PER_FLIGHT": "4096", "RACER_QUEUE_ENTRIES": "65536", + "RACER_CONNECTIONS_PER_NEIGHBOR": "1024", "RACER_PIPES": "65536", "RACER_RANGE_WINDOW_PAGES": "64", + "RACER_HEADER_BYTES": "32768", "RACER_PLACEMENT_CACHE_BYTES": "536870912", + "RACER_CACHED_PATHS": "1048576", "RACER_ACTIVE_PATH_SEARCHES": "64", "RACER_RETAINED_SNAPSHOTS": "64", + "RACER_ORIGIN_CONNECTIONS_PER_CACHE": "1024", "RACER_REQUEST_TIMEOUT_MS": "86400000", + "RACER_SHUTDOWN_TIMEOUT_MS": "3600000", + }, + "thread split": {"RACER_IO_WORKERS": "2", "RACER_CRYPTO_HELPERS": "1", "RACER_CONTROL_SEATS": "auto"}, + "unknown keys are ignored": {"RACER_SOMETHING_ELSE": "x"}, + } { + t.Run(name, func(t *testing.T) { + require.NoError(t, validateDataplaneConfig(data)) + }) + } +} + +func TestValidateDataplaneConfigRejects(t *testing.T) { + for _, tc := range []struct{ key, value string }{ + {"RACER_ENABLE_GPU", "TRUE"}, + {"RACER_ENABLE_GPU", " auto"}, + {"RACER_ENABLE_GPU", ""}, + {"RACER_ENABLE_RDMA", "yes"}, + {"RACER_GPU_STAGING_BYTES", "0"}, + {"RACER_GPU_STAGING_BYTES", "1048576"}, + {"RACER_GPU_STAGING_BYTES", "3145728"}, + {"RACER_GPU_STAGING_BYTES", "2MiB"}, + {"RACER_GPU_STAGING_BYTES", "18446744073709551616"}, + {"RACER_GPU_LANDING_BYTES", "1048576"}, + {"RACER_GPU_LANDING_BYTES", "-2097152"}, + {"RACER_GPU_LANDING_BYTES", "+2097152"}, + {"RACER_DATAPLANE_NVIDIA", "on"}, + {"RACER_ALLOW_SMT", "1"}, + {"RACER_GPU_PCI_P2PDMA", "TRUE"}, + {"RACER_GPU_PCI_P2PDMA", "1"}, + {"RACER_GPU_PCI_P2PDMA", ""}, + {"RACER_REPLICA_SELECTION", "random"}, + {"RACER_ADMISSION_MODE", "on"}, + {"RACER_SHARES", "0"}, + {"RACER_SHARES", "4294967296"}, + {"RACER_SHARES", "18446744073709551615"}, + {"RACER_IO_WORKERS", "0"}, + {"RACER_IO_WORKERS", "65537"}, + {"RACER_CRYPTO_HELPERS", "0"}, + {"RACER_DISK_PAGE_ENTRIES", "1048577"}, + {"RACER_ADMISSION_HISTORY_BYTES", "31"}, + {"RACER_ADMISSION_HISTORY_BYTES", "536870913"}, + {"RACER_ADMISSION_PERIOD_SECS", "0"}, + {"RACER_ADMISSION_PERIOD_SECS", "86401"}, + {"RACER_CHECKPOINT_BYTES", "536870913"}, + {"RACER_DEVICE_BUSY_WAIT_SECS", "3601"}, + {"RACER_RANGE_SPILL_PAGES", "1025"}, + {"RACER_PAGE_HEDGE_DELAY_MS", "0"}, + {"RACER_PAGE_HEDGE_DELAY_MS", "30001"}, + {"RACER_PAGE_HEDGE_SLOTS", "33"}, + {"RACER_PAGE_HEDGE_BYTES", "1073742337"}, + {"RACER_PEER_INFLIGHT_MAX", "0"}, + {"RACER_PEER_INFLIGHT_MAX", "65537"}, + {"RACER_PEER_PER_NEIGHBOR_MAX", "0"}, + {"RACER_PEER_RECEIVE_MAX", "65537"}, + {"RACER_PEER_RECEIVE_QUEUE_MAX", "0"}, + {"RACER_PEER_RECEIVE_QUEUE_MAX", "65537"}, + {"RACER_PEER_RECEIVE_WAIT_MS", "0"}, + {"RACER_PEER_RECEIVE_WAIT_MS", "30001"}, + {"RACER_SLAB_BYTES", "0"}, + {"RACER_SLAB_BYTES", "9223372036854775808"}, + {"RACER_SEGMENT_BYTES", "0"}, + {"RACER_SEGMENT_BYTES", "16797711"}, + {"RACER_FREE_SEGMENT_RESERVE", "0"}, + {"RACER_FREE_SEGMENT_RESERVE", "1048576"}, + {"RACER_PLAINTEXT_BYTES", "68719476737"}, + {"RACER_DIRTY_BYTES", "16777231"}, + {"RACER_REQUEST_CONTEXT_BYTES", "131071"}, + {"RACER_FLIGHTS", "65537"}, + {"RACER_WAITERS_PER_FLIGHT", "4097"}, + {"RACER_QUEUE_ENTRIES", "1"}, + {"RACER_CONNECTIONS_PER_NEIGHBOR", "1025"}, + {"RACER_CLIENT_CONNECTIONS", "65537"}, + {"RACER_PIPES", "65537"}, + {"RACER_RANGE_WINDOW_PAGES", "0"}, + {"RACER_RANGE_WINDOW_PAGES", "65"}, + {"RACER_HEADER_BYTES", "1023"}, + {"RACER_HEADER_BYTES", "32769"}, + {"RACER_PLACEMENT_CACHE_BYTES", "1023"}, + {"RACER_PLACEMENT_CACHE_BYTES", "536870913"}, + {"RACER_CACHED_PATHS", "1048577"}, + {"RACER_PATH_CACHE_BYTES", "536870913"}, + {"RACER_ACTIVE_PATH_SEARCHES", "65"}, + {"RACER_RETAINED_SNAPSHOTS", "1"}, + {"RACER_RETAINED_SNAPSHOTS", "65"}, + {"RACER_METADATA_ENTRIES", "1048577"}, + {"RACER_RELAY_TRANSFERS", "65537"}, + {"RACER_ORIGIN_CONNECTIONS_PER_CACHE", "1025"}, + {"RACER_REQUEST_TIMEOUT_MS", "0"}, + {"RACER_REQUEST_TIMEOUT_MS", "86400001"}, + {"RACER_PEER_ATTEMPT_TIMEOUT_MS", "86400001"}, + {"RACER_READER_STALL_TIMEOUT_MS", "0"}, + {"RACER_SHUTDOWN_TIMEOUT_MS", "3600001"}, + {"RACER_SLAB_BYTES", "1 GiB"}, + {"RACER_WARM_MAX_TTL_SECS", "0"}, + {"RACER_WARM_WINDOW_PAGES", "65"}, + {"RACER_RETAINED_SEGMENT_PERCENT", "101"}, + {"RACER_MAX_THREADS", "2"}, + {"RACER_CONTROL_SEATS", "17"}, + {"RACER_CONTROL_SEATS", "01"}, + {"RACER_SLAB_BYTES", strings.Repeat("1", 4097)}, + } { + t.Run(tc.key+"="+tc.value, func(t *testing.T) { + err := validateDataplaneConfig(map[string]string{tc.key: tc.value}) + require.ErrorContains(t, err, tc.key) + require.ErrorContains(t, err, dataplaneConfigName) + }) + } +} + +func TestValidateDataplaneConfigReportsEveryProblem(t *testing.T) { + err := validateDataplaneConfig(map[string]string{ + "RACER_ENABLE_GPU": "maybe", "RACER_GPU_STAGING_BYTES": "1", "RACER_IO_WORKERS": "4", + }) + require.ErrorContains(t, err, "RACER_ENABLE_GPU") + require.ErrorContains(t, err, "RACER_GPU_STAGING_BYTES") + require.ErrorContains(t, err, "RACER_CRYPTO_HELPERS must be set together") +} + +// An invalid edit fails the plan, so the DaemonSet keeps its last good hash. +func TestInvalidDataplaneConfigKeepsDaemonSet(t *testing.T) { + env := testEnv(t, cacheObject("cache")) + initialize(t, env) + + before := liveDataplane(t, env).Spec.Template.Annotations[dataplaneConfigHash] + + config := &corev1.ConfigMap{} + require.NoError(t, env.Client.Get(t.Context(), objectKey(env, dataplaneConfigName), config)) + config.Data["RACER_GPU_STAGING_BYTES"] = "1048576" + require.NoError(t, env.Client.Update(t.Context(), config)) + + _, _, err := (Component{}).Plan(t.Context(), env, nil) + require.ErrorContains(t, err, "RACER_GPU_STAGING_BYTES") + require.Equal(t, before, liveDataplane(t, env).Spec.Template.Annotations[dataplaneConfigHash]) +} + +// One image serves GPU and CPU nodes, so no ConfigMap key selects another +// image. An unknown key such as RACER_DATAPLANE_CUDA is ignored. +func TestDataplaneImage(t *testing.T) { + env := testEnv(t, cacheObject("cache")) + initialize(t, env) + + image := func() string { return liveDataplane(t, env).Spec.Template.Spec.Containers[0].Image } + require.Equal(t, "mirror.test/unbounded/racer-dataplane:v1", image()) + + for _, value := range []string{"true", "false", "yes"} { + config := &corev1.ConfigMap{} + require.NoError(t, env.Client.Get(t.Context(), objectKey(env, dataplaneConfigName), config)) + config.Data["RACER_DATAPLANE_CUDA"] = value + require.NoError(t, env.Client.Update(t.Context(), config)) + + persist(t, env, planPass(t, env)) + require.Equal(t, "mirror.test/unbounded/racer-dataplane:v1", image(), "value %q", value) + } +} diff --git a/internal/operator/components/racer/dataplane.go b/internal/operator/components/racer/dataplane.go index ec65777919..f6538c5e89 100644 --- a/internal/operator/components/racer/dataplane.go +++ b/internal/operator/components/racer/dataplane.go @@ -8,7 +8,6 @@ import ( "errors" "fmt" "reflect" - "strconv" corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" @@ -74,26 +73,8 @@ func planDataplaneConfig(ctx context.Context, env *component.Env, plan *componen return nil, err } - if value, ok := config.Data[workload.ControlSeatsKey]; ok && value != "auto" { - // The runtime accepts only the canonical decimal form, so reject - // leading zeros, signs, and whitespace before Pods roll. - seats, err := strconv.ParseUint(value, 10, 8) - if err != nil || strconv.FormatUint(seats, 10) != value || seats < 1 || seats > workload.MaxControlSeats { - return nil, fmt.Errorf("ConfigMap %s %s must be auto or an integer from 1 to %d", dataplaneConfigName, workload.ControlSeatsKey, workload.MaxControlSeats) - } - } - - if value, ok := config.Data[workload.NVIDIAKey]; ok { - if mode := workload.NVIDIAInjection(value); mode != workload.NVIDIAOff && mode != workload.NVIDIAEnv && mode != workload.NVIDIACDI { - return nil, fmt.Errorf("ConfigMap %s %s must be off, env, or cdi", dataplaneConfigName, workload.NVIDIAKey) - } - } - - if value, ok := config.Data["RACER_MAX_THREADS"]; ok { - threads, err := strconv.ParseUint(value, 10, 64) - if err != nil || threads < minThreads { - return nil, fmt.Errorf("ConfigMap %s RACER_MAX_THREADS must be an integer of at least %d, including the control plane thread", dataplaneConfigName, minThreads) - } + if err := validateDataplaneConfig(config.Data); err != nil { + return nil, err } return config, nil diff --git a/internal/racer/workload/cufile_test.go b/internal/racer/workload/cufile_test.go new file mode 100644 index 0000000000..46ae950663 --- /dev/null +++ b/internal/racer/workload/cufile_test.go @@ -0,0 +1,47 @@ +// Copyright (c) Microsoft Corporation. +// SPDX-License-Identifier: Apache-2.0 + +package workload + +import ( + "encoding/json" + "os" + "regexp" + "testing" +) + +// The image's cufile.json must parse and agree with the workload wiring. +func TestImageCUFileConfig(t *testing.T) { + raw, err := os.ReadFile("../../../images/racer-dataplane/cufile.json") + if err != nil { + t.Fatal(err) + } + + // cuFile accepts // line comments; no value in the file contains "//". + stripped := regexp.MustCompile(`(?m)^\s*//.*$`).ReplaceAll(raw, nil) + + var config struct { + Logging struct { + Dir string `json:"dir"` + Level string `json:"level"` + } `json:"logging"` + Properties struct { + AllowCompatMode *bool `json:"allow_compat_mode"` + } `json:"properties"` + } + if err := json.Unmarshal(stripped, &config); err != nil { + t.Fatalf("cufile.json: %v", err) + } + + if config.Logging.Dir != CUFileLogDirectory { + t.Fatalf("logging.dir = %q, want %q", config.Logging.Dir, CUFileLogDirectory) + } + + if config.Logging.Level != "ERROR" { + t.Fatalf("logging.level = %q, want ERROR", config.Logging.Level) + } + + if config.Properties.AllowCompatMode == nil || *config.Properties.AllowCompatMode { + t.Fatal("properties.allow_compat_mode must be false") + } +} diff --git a/internal/racer/workload/workload.go b/internal/racer/workload/workload.go index ebd00878c4..867b4327c4 100644 --- a/internal/racer/workload/workload.go +++ b/internal/racer/workload/workload.go @@ -58,10 +58,11 @@ const ( type NVIDIAInjection string const ( - // NVIDIAOff adds nothing. + // NVIDIAOff adds no NVIDIA env, CDI annotation, /run/udev mount, or cuFile log volume. NVIDIAOff NVIDIAInjection = "off" - // NVIDIAEnv sets NVIDIA_VISIBLE_DEVICES and NVIDIA_DRIVER_CAPABILITIES. - // The NVIDIA container runtime acts on them; runc ignores them, so they + // NVIDIAEnv sets NVIDIA_VISIBLE_DEVICES and NVIDIA_DRIVER_CAPABILITIES and + // mounts /run/udev and a cuFile log emptyDir for cuFile. + // The NVIDIA container runtime acts on the variables; runc ignores them, so they // are harmless on nodes without NVIDIA. NVIDIAEnv NVIDIAInjection = "env" // NVIDIACDI also adds the CDI pod annotation. containerd fails container @@ -81,6 +82,15 @@ const ( CDIAllGPUs = "nvidia.com/gpu=all" ) +const ( + // CUFileLogDirectory is the logging.dir of the cufile.json in the + // dataplane image (images/racer-dataplane/cufile.json). + CUFileLogDirectory = "/var/log/cufile" + // CUFileLogSizeLimit bounds the cuFile log emptyDir. cuFile logs at + // ERROR level only, so this is far above normal use. + CUFileLogSizeLimit = "64Mi" +) + // ErrInvalid reports an invalid workload configuration. var ErrInvalid = errors.New("invalid racer dataplane configuration") @@ -287,6 +297,8 @@ func dataplaneVolumes(c Config) []corev1.Volume { // database and refuses files it cannot classify. if c.nvidiaMode() != NVIDIAOff { volumes = append(volumes, corev1.Volume{Name: "udev", VolumeSource: corev1.VolumeSource{HostPath: &corev1.HostPathVolumeSource{Path: "/run/udev", Type: ptr.To(corev1.HostPathDirectoryOrCreate)}}}) + // The root filesystem is read-only, so cuFile writes cufile.log here. + volumes = append(volumes, corev1.Volume{Name: "cufile-log", VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{SizeLimit: ptr.To(resource.MustParse(CUFileLogSizeLimit))}}}) } return volumes @@ -305,7 +317,7 @@ func dataplaneContainer(c Config) corev1.Container { {Name: "devices", MountPath: "/host/dev", ReadOnly: true}, } if c.nvidiaMode() != NVIDIAOff { - mounts = append(mounts, corev1.VolumeMount{Name: "udev", MountPath: "/run/udev", ReadOnly: true}) + mounts = append(mounts, corev1.VolumeMount{Name: "udev", MountPath: "/run/udev", ReadOnly: true}, corev1.VolumeMount{Name: "cufile-log", MountPath: CUFileLogDirectory}) } return corev1.Container{ @@ -371,7 +383,9 @@ func dataplaneEnvironment(c Config) []corev1.EnvVar { } // Ask the NVIDIA runtime for driver libraries only. The privileged - // container already sees /dev/nvidia*; no GPU resource is requested. + // container already sees /dev/nvidia*, including the GPUDirect Storage + // /dev/nvidia-fs* nodes that exist when the Pod starts, so NVIDIA_GDS is + // not needed. No GPU resource is requested. if c.nvidiaMode() != NVIDIAOff { env = append(env, corev1.EnvVar{Name: "NVIDIA_VISIBLE_DEVICES", Value: "all"}, diff --git a/internal/racer/workload/workload_test.go b/internal/racer/workload/workload_test.go index 8665aeb720..cee02433c1 100644 --- a/internal/racer/workload/workload_test.go +++ b/internal/racer/workload/workload_test.go @@ -231,6 +231,16 @@ func TestNVIDIAInjectionModes(t *testing.T) { t.Fatalf("CDI annotation = %q (present %v), want present %v", value, hasCDI, tc.wantCDI) } + // /dev/nvidia-fs* reaches the container through privileged + // mode, not the runtime's NVIDIA_GDS hook. + if _, ok := env["NVIDIA_GDS"]; ok { + t.Fatal("NVIDIA_GDS must not be set") + } + + if container.SecurityContext == nil || container.SecurityContext.Privileged == nil || !*container.SecurityContext.Privileged { + t.Fatal("dataplane must be privileged to see /dev/nvidia-fs*") + } + // cuFile classifies slab devices through the host udev database. hasUdevVolume := false @@ -252,6 +262,25 @@ func TestNVIDIAInjectionModes(t *testing.T) { t.Fatalf("udev volume/mount = %v/%v, want %v", hasUdevVolume, hasUdevMount, tc.wantEnv) } + // cuFile writes its log to an emptyDir on the read-only root. + hasLogVolume, hasLogMount := false, false + + for _, volume := range template.Spec.Volumes { + if volume.Name == "cufile-log" { + hasLogVolume = volume.EmptyDir != nil && volume.EmptyDir.SizeLimit != nil && volume.EmptyDir.SizeLimit.String() == CUFileLogSizeLimit + } + } + + for _, mount := range container.VolumeMounts { + if mount.Name == "cufile-log" { + hasLogMount = mount.MountPath == CUFileLogDirectory && !mount.ReadOnly + } + } + + if hasLogVolume != tc.wantEnv || hasLogMount != tc.wantEnv { + t.Fatalf("cufile-log volume/mount = %v/%v, want %v", hasLogVolume, hasLogMount, tc.wantEnv) + } + // GPU libraries must never come from a GPU resource request. for _, list := range []corev1.ResourceList{container.Resources.Requests, container.Resources.Limits} { if _, ok := list["nvidia.com/gpu"]; ok {