feat(racer): GPUDirect reads into client host and GPU memory

.github/workflows/ci.yaml

@@ -355,6 +355,69 @@ jobs:
           if [ -n "$VERBS_FEATURES" ]; then features=(--features "$VERBS_FEATURES"); fi
           timeout --signal=TERM --kill-after=10s 300s cargo +1.96.0 test --locked --manifest-path cmd/racer-dataplane/Cargo.toml -p rdma-verbs --no-default-features "${features[@]}"
 
+  racer-gpu:
+    name: Racer GPU (${{ matrix.configuration }})
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - configuration: no-default-features
+            features: ""
+            native: false
+          - configuration: simulation
+            features: simulation
+            native: false
+          - configuration: native
+            features: native
+            native: true
+          - configuration: native-and-simulation
+            features: native,simulation
+            native: true
+    env:
+      CARGO_TARGET_DIR: ${{ github.workspace }}/bin/racer-gpu-cargo
+      CARGO_BUILD_JOBS: "2"
+      GPU_FEATURES: ${{ matrix.features }}
+      LD_LIBRARY_PATH: ${{ github.workspace }}/bin/racer-gpu-native
+    steps:
+      - name: Checkout
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+
+      - 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.
+      - name: Compile CUDA shim without kernels (not hardware validation)
+        if: matrix.native
+        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)
+        if: matrix.native
+        run: timeout --signal=TERM --kill-after=10s 300s cargo +1.96.0 test --locked --manifest-path cmd/racer-dataplane/Cargo.toml -p racer-gpu --no-default-features --features "$GPU_FEATURES" --lib ffi::tests::required_native_shim_loads_abi_and_symbols -- --exact --ignored
+
+      - name: Run decrypt kernels on the CPU (kernel emulation)
+        if: matrix.native
+        run: timeout --signal=TERM --kill-after=10s 300s sh cmd/racer-dataplane/gpu/native/run-kernel-emulation.sh bin/racer-gpu-emulation
+
+      - name: Check GPU formatting
+        run: timeout --signal=TERM --kill-after=10s 300s cargo +1.96.0 fmt --manifest-path cmd/racer-dataplane/Cargo.toml -p racer-gpu --check
+
+      - name: Strict GPU Clippy
+        run: |
+          features=()
+          if [ -n "$GPU_FEATURES" ]; then features=(--features "$GPU_FEATURES"); fi
+          timeout --signal=TERM --kill-after=10s 300s cargo +1.96.0 clippy --locked --manifest-path cmd/racer-dataplane/Cargo.toml -p racer-gpu --all-targets --no-default-features "${features[@]}" -- -D warnings
+
+      # Includes the C AES-GCM/CRC core check against racer-crypto (16 MiB pages).
+      - name: Test GPU crate without requiring GPU hardware
+        run: |
+          features=()
+          if [ -n "$GPU_FEATURES" ]; then features=(--features "$GPU_FEATURES"); fi
+          timeout --signal=TERM --kill-after=10s 300s cargo +1.96.0 test --locked --manifest-path cmd/racer-dataplane/Cargo.toml -p racer-gpu --no-default-features "${features[@]}"
+
   runtime-miri:
     name: Runtime Miri (${{ matrix.group }})
     runs-on: ubuntu-24.04
@@ -392,6 +455,17 @@ jobs:
       - name: Provision envtest and run controller API-server race tests
         run: timeout --signal=TERM --kill-after=10s 300s make racer-envtest-ci
 
+  racer-python:
+    name: Racer Python Client
+    runs-on: ubuntu-24.04
+    timeout-minutes: 10
+    steps:
+      - name: Checkout
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+
+      - name: Run Python client and vLLM plugin tests
+        run: timeout --signal=TERM --kill-after=10s 540s make racer-python-test
+
   # ---------- Orca Integration Tests ----------
   # Spins up Garage and Azurite via testcontainers-go and runs the
   # orca in-process integration suite (internal/orca/inttest). Docker

Makefile

@@ -581,6 +581,15 @@ racer-server-test: ## Lint and race-test the Racer server
 
 racer-test: racer-server-test ## Check the Racer controller
 
+RACER_PY_VENV ?= $(CURDIR)/tmp/racer-python-venv
+
+.PHONY: racer-python-test
+racer-python-test: ## Run the Racer Python client and vLLM plugin tests (no GPU needed)
+	@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
+	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
+
 $(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)

NOTICE

@@ -48,6 +48,15 @@ notices:
         link: https://docs.rs/crate/base64/0.22.1/source/LICENSE-APACHE
       - name: MIT License
         link: https://docs.rs/crate/base64/0.22.1/source/LICENSE-MIT
+  - dependency: cc
+    ecosystem: cargo
+    copyright:
+      - Copyright (c) 2014 Alex Crichton
+    license:
+      - name: Apache License, Version 2.0
+        link: https://docs.rs/crate/cc/1.5.1/source/LICENSE-APACHE
+      - name: MIT License
+        link: https://docs.rs/crate/cc/1.5.1/source/LICENSE-MIT
   - dependency: chacha20poly1305
     ecosystem: cargo
     copyright:

clients/python/.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.egg-info/
+.pytest_cache/
+build/
+dist/

clients/python/racer-vllm/README.md

@@ -0,0 +1,83 @@
+# racer-vllm
+
+A vLLM plugin that loads model weights from a Racer cache. It adds the load
+format `racer`. It uses the `racer-client` package.
+
+## Install
+
+```
+pip install ./clients/python/racer ./clients/python/racer-vllm
+```
+
+vLLM finds the plugin through the `vllm.general_plugins` entry point. It was
+written for vLLM v0.25.x.
+
+## Model location
+
+A model lives at `racer://<cache>/<bucket>/<prefix>`. The plugin reads
+objects through `/run/racer/<cache>/client/socket`. Set `RACER_SOCKET` to use
+another socket. Object keys follow racer-object, so `<bucket>/<prefix>/<file>`
+names the S3 object.
+
+## Use
+
+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:
+
+```
+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.
+
+Then start vLLM on that directory:
+
+```
+vllm serve /models/llama-3-8b --load-format racer \
+  --model-loader-extra-config '{"staging_bytes": 268435456, "concurrency": 4}'
+```
+
+Extra config keys, all optional:
+
+| Key | Default | Meaning |
+|---|---|---|
+| `uri` | from `racer-model.json` | The `racer://` model location. |
+| `staging_bytes` | 256 MiB | Size of each GPU staging buffer. |
+| `concurrency` | 4 | Staging buffers, and reads in flight. |
+| `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`. |
+
+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
+`RACER_RANGE_WINDOW_PAGES` and `RACER_GPU_STAGING_BYTES` on the dataplane. See
+the GPU reads section of the Racer reference.
+
+## How it loads
+
+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
+   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
+   vLLM uses one buffer, the next range is read into another.
+
+A tensor bigger than a staging buffer gets its own GPU allocation.
+
+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.
+
+## Tests
+
+```
+make racer-python-test
+```

clients/python/racer-vllm/pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "racer-vllm"
+version = "0.1.0"
+description = "vLLM model loader that reads weights from a Racer cache"
+readme = "README.md"
+license = { text = "Apache-2.0" }
+requires-python = ">=3.10"
+dependencies = ["racer-client"]
+
+[project.optional-dependencies]
+test = ["pytest"]
+
+[project.entry-points."vllm.general_plugins"]
+racer = "racer_vllm:register"
+
+[project.scripts]
+racer-vllm-fetch = "racer_vllm.model:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

clients/python/racer-vllm/src/racer_vllm/__init__.py

@@ -0,0 +1,8 @@
+# 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"]

clients/python/racer-vllm/src/racer_vllm/loader.py

@@ -0,0 +1,97 @@
+# 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://<cache>/<bucket>/<prefix>; 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