diff --git a/.gitignore b/.gitignore index 18e7158..27b1d61 100644 --- a/.gitignore +++ b/.gitignore @@ -23,6 +23,9 @@ benches/*.png # Tooling .claude .worktrees/ +CLAUDE.md +/docs/agents/ +.scratch/ # Local-only files (personal scripts, scratch, etc.) .local/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 44b2b64..1e36d39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,24 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] +## [0.2.0] - 2026-08-03 + +### Added + +- `Server(..., pin_host_memory=True)` CUDA-page-locks the ring buffers, so a + downstream host-to-device copy of a sampled batch is a real DMA transfer on the + copy engine instead of a chunked, driver-mediated staging copy. Off by default. + Measured locally, `cudaMemcpyAsync` out of pageable memory holds the calling + thread for 97% of the copy; out of page-locked memory it returns in 0.002 ms. +- If the constructor returns, every ring buffer is page-locked. Any inability to + deliver that raises `RuntimeError` naming every path probed for the CUDA + runtime and the CUDA error symbolically. There is no silent no-op. +- The CUDA runtime is resolved through a three-rung ladder — already-mapped + libraries, then the installed CUDA wheels, then sonames (versioned before + unversioned) — so no soname symlink or `LD_LIBRARY_PATH` entry is needed. +- New guide: [Host-memory pinning](https://instadeepai.github.io/echo/guides/host-memory-pinning/), + covering the mechanism, the constructor guarantee, the unswappable footprint + arithmetic, and how the CUDA runtime is located. ## [0.1.1] - 2026-05-26 diff --git a/Cargo.lock b/Cargo.lock index 49e9145..08dd1de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -215,7 +215,7 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "echo" -version = "0.1.1" +version = "0.2.0" dependencies = [ "arc-swap", "criterion", @@ -226,6 +226,7 @@ dependencies = [ "numpy", "parking_lot", "pyo3", + "tempfile", "tokio", "tokio-test", ] @@ -236,12 +237,39 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "futures-core" version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + [[package]] name = "half" version = "2.7.1" @@ -317,6 +345,12 @@ version = "0.2.182" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "lock_api" version = "0.4.14" @@ -571,6 +605,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rawpointer" version = "0.2.1" @@ -641,6 +681,19 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustversion" version = "1.0.22" @@ -738,6 +791,19 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "tinytemplate" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index 028bfc6..197014a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "echo" -version = "0.1.1" +version = "0.2.0" edition = "2021" description = "A fast distributed replay buffer for reinforcement learning." license = "Apache-2.0" @@ -38,5 +38,6 @@ detailed-metrics = ["dep:hdrhistogram"] [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } +tempfile = "3" tokio-test = "0.4" tokio = { version = "1", features = ["macros", "rt-multi-thread", "time"] } diff --git a/README.md b/README.md index 61663c6..dbc084b 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,7 @@ pytree-agnostic. - GIL released while waiting for batches - FIFO sampling (with more strategies planned) - Detailed metrics exposed per batch +- Optional CUDA [host-memory pinning](https://instadeepai.github.io/echo/guides/host-memory-pinning/) so H2D copies are real DMA transfers, not driver-staged ## Example diff --git a/docs/src/design/host-pinning.md b/docs/src/design/host-pinning.md new file mode 100644 index 0000000..03f7da2 --- /dev/null +++ b/docs/src/design/host-pinning.md @@ -0,0 +1,99 @@ +# Host-memory pinning + +`src/host_pinning/` optionally CUDA-page-locks the ring buffers, so a +host-to-device copy of a sampled batch is a DMA transfer rather than a chunked +staging copy. For the mechanism and the user-facing contract, see the +[host-memory pinning guide](../guides/host-memory-pinning.md). + +| File | Role | +|---|---| +| `mod.rs` | What both halves need: `CudaApi`, `Region`, `PinError`. | +| `resolve.rs` | Find the CUDA runtime. | +| `register.rs` | Page-lock memory with it, and roll back cleanly. | + +`resolve` and `register` meet at one point: `resolve` returns a `CudaApi` and +`register` takes one. Beyond `dlsym`, `resolve` makes no CUDA calls, so most of +the module is testable without a GPU. + +## Why resolution is a ladder + +The feature was a silent no-op because it did one thing: +`dlopen("libcudart.so")`. The pip CUDA runtime wheels ship only the versioned +soname, with no unversioned symlink and no ldconfig entry, so that call returned +NULL and every pin became a no-op that reported nothing. + +Resolution now tries three rungs in order, accumulating every attempted path: + +1. **Already-loaded scan.** Parse `/proc/self/maps` for a mapped runtime and + `dlopen` its absolute path. Version- and path-agnostic, and it hits whenever + the framework has already initialised CUDA — the common case. `dlopen` on the + path of a mapped library reuses that mapping rather than loading a second copy. +2. **Installed-wheel search.** Search beneath the CUDA vendor package + directories, which `py_bindings` locates through Python's import machinery and + passes down. It has to be the vendor package rather than a named component + subpackage: CUDA 13 ships one consolidated wheel laid out as + `nvidia/cu13/lib/`, CUDA 12 one wheel per component as + `nvidia/cuda_runtime/lib/`, so naming the component finds nothing on a CUDA 13 + install. +3. **Soname load**, versioned before unversioned, each trying + `RTLD_NOLOAD` before a full load. + +`candidates()` builds the whole ladder as a pure function of the maps text and +the vendor roots, so its ordering is unit-tested. Reordering the rungs is how +this broke before, and it is the kind of thing an unrelated edit can change +without anything else failing. + +A rung that yields no candidate still reports itself, as a `(none)` line, so a +failure message distinguishes "searched and found nothing" from "never ran". + +`resolve` is the only place that knows about CUDA layouts; `py_bindings` supplies +the vendor roots because that is where a `Python` token exists, and keeping pyo3 +out of `host_pinning` lets `cargo test` exercise it with no interpreter. + +## Why registration is not in the constructor + +`PytreeRingBuf::new` and `Store::new` stay infallible. Page-locking is a separate +`pin_host_memory` call on a fully-constructed buffer. + +The reason is rollback. A constructor that returns `Err` never runs `Drop`, so +registering inside one forces a hand-written unregister loop on the error path — +code that is easy to get wrong and almost never exercised. Registering afterwards +lets `Drop` own rollback for both the failure path and normal teardown. +`PytreeRingBuf` holds `pinned_with: Option`, which is both the cue for +`Drop` and what keeps `Drop` off any process-global. + +Within one attempt it is all or nothing: if the *n*th buffer is rejected, +`pin_all` unregisters the preceding *n−1* before returning the error, so a failed +construction leaves nothing for a retry to accumulate. Registrations are still +leaked if a reference-counted `Store` outlives process shutdown, which is +accepted. + +## The injection seam + +`CudaApi` is a struct of function pointers passed explicitly to `pin_all` / +`unpin_all`, rather than a global the two reach into. + +Rollback is unsafe code whose only externally visible consequence is the absence +of leaked registrations, which cannot be observed from Python at all. With +injection, a test hands in a stub that fails on the third buffer and asserts +exactly two unregister calls, on a machine with no GPU. `tests/host_pinning.rs` +does that, and also drives a real `PytreeRingBuf` through registration and `Drop` +with the same stubs. + +The module is `pub`, like every other module in the crate, because `tests/` are +separate crates and can only see `pub` items — the same reason +`PytreeRingBuf::slot_mut` is public. `CudaApi`'s fields stay private behind +`unsafe fn CudaApi::new`, so being public does not let a caller assemble one from +arbitrary function pointers, and `pin_all`/`unpin_all` remain `unsafe fn` with +stated contracts. + +## No CUDA dependency at build time + +Nothing links CUDA and nothing is generated from CUDA headers: the four entry +points (`cudaHostRegister`, `cudaHostUnregister`, `cudaGetErrorName`, `cudaFree`) +are `dlsym`'d at pin time. One wheel installs on GPU and CPU-only hosts alike, and +with pinning off no library is loaded and no code here runs. + +The library is opened `RTLD_LOCAL` so echo cannot change how anything else in the +process resolves its symbols, and the handle is never `dlclose`d — the +registrations it backs must outlive it. diff --git a/docs/src/design/overview.md b/docs/src/design/overview.md index 1d239a8..9f4d750 100644 --- a/docs/src/design/overview.md +++ b/docs/src/design/overview.md @@ -56,6 +56,7 @@ client.send(pytree) | [`store.rs`](store.md) | Composes the ring buffer, sampler and remover. Does the CAS-reserve-then-memcpy dance. | | [`ring_buf.rs`](ring-buffer.md) | Raw pre-allocated buffer with `UnsafeCell` interior mutability; no synchronisation of its own. | | [`selector.rs`](selector.md) | `Sampler` and `Remover` traits, plus the FIFO implementation that does the commit-counter wake. | +| [`host_pinning/`](host-pinning.md) | Optional CUDA page-locking of the ring buffers. Off by default; resolves the runtime by `dlopen`, so no build-time CUDA dependency. | | `metrics.rs` | Per-drainer counters + optional hdrhistograms. See [Reading metrics](../guides/metrics.md). | | `py_bindings.rs` | PyO3 layer. Constructs the `Store`, picks the transport, releases the GIL during `sample()`. | | `array_spec.rs` | Tiny value type: shape + dtype size per leaf. | diff --git a/docs/src/design/ring-buffer.md b/docs/src/design/ring-buffer.md index e6ecd7a..bc24a7e 100644 --- a/docs/src/design/ring-buffer.md +++ b/docs/src/design/ring-buffer.md @@ -44,6 +44,28 @@ ever spans the wrap-around point. That property is what lets the `Contiguous { start, count }` sample result type exist at all; without it we'd need a scatter-gather variant. +## Optional host-memory pinning + +`pin_host_memory` CUDA-page-locks every backing `Vec`, so a host-to-device +copy of a sampled view is a DMA transfer rather than a chunked staging copy. The +module that does it is covered in [Host-memory pinning](host-pinning.md). + +What matters here is that the registration stays valid for the buffer's whole +life: the buffers are allocated once in `new` and never reallocated or resized. +Page-locking pins the physical pages behind specific addresses, so a growable +buffer would invalidate its own registration on the first reallocation. + +`Drop` owns the reverse. `pinned_with: Option` records the runtime the +buffers were registered with, and `Drop` unregisters before the `Vec`s are freed — +`Drop::drop` runs before a struct's fields are dropped, so the memory is still +valid there. It is also why pinning is a separate step after construction rather +than part of `new`; see +[Why registration is not in the constructor](host-pinning.md#why-registration-is-not-in-the-constructor). + +The buffers are page-*un*aligned in practice (glibc returns large allocations at a +small offset into a page). Measurement showed aligning them buys ~1.6% of copy +time and nothing on host-thread occupancy, so `Vec` stays. + ## What this type does not do - Track which slots are in use (the `Store` does that via the write/read diff --git a/docs/src/development.md b/docs/src/development.md index f2bae0f..5a80a25 100644 --- a/docs/src/development.md +++ b/docs/src/development.md @@ -22,6 +22,31 @@ cargo test # Rust unit + integration tests uv run pytest python/tests/ -v # Python tests ``` +Rust tests live in `tests/`, one file per module — including +`tests/host_pinning.rs`. Python tests live in `python/tests/`. + +### Working on host-memory pinning + +See [Host-memory pinning](design/host-pinning.md) for how the module is laid out, +and the [guide](guides/host-memory-pinning.md) for what it does for a user. + +Most of `tests/host_pinning.rs` needs no GPU: the resolution ladder is pure, and +registration and rollback run against injected stubs. + +To exercise the CUDA-runtime resolution ladder locally, install a runtime wheel +into the checkout's venv: + +```bash +uv pip install nvidia-cuda-runtime +``` + +Without it only rung 1 (already-mapped libraries) can hit, and on a machine with no +system CUDA install nothing resolves at all. The wheel is *not* a `dev` extra: +CI is CPU-only and should not download a CUDA runtime. + +`cargo test` has no Python interpreter to ask for the wheel's location, so the +Rust tests look under `.venv/lib/python*/site-packages/nvidia` in the checkout. + ## Benchmarks ```bash @@ -33,10 +58,31 @@ Benchmark results land in `benches/`. ## Docs ```bash -uv run --extra docs mkdocs serve # live-reload on http://127.0.0.1:8000 -uv run --extra docs mkdocs build # static site to site/ +just docs-serve # live-reload on http://127.0.0.1:8000/echo/ +just docs # static site to docs/site/ ``` +Note the **`/echo/` path prefix** — `mkdocs serve` honours `site_url`, so +`http://127.0.0.1:8000/` alone 404s. The address it prints on startup is correct. + +Working on a remote box, `docs-serve` binds to loopback only and is invisible from +your laptop. Either forward the port from the client side, which needs no change +here: + +```bash +ssh -L 8000:localhost:8000 you@remote-box # then browse http://localhost:8000/echo/ +``` + +or bind to an interface the client can reach: + +```bash +just docs-serve-on # 0.0.0.0:8000 — all interfaces +just docs-serve-on 10.0.0.5:8000 # just this one +``` + +Prefer the tunnel, or a specific private interface, over `0.0.0.0` on an untrusted +network: `mkdocs serve` is a development server with no authentication. + The docs are built and deployed by `.github/workflows/docs.yml` on every push to `main`. diff --git a/docs/src/guides/host-memory-pinning.md b/docs/src/guides/host-memory-pinning.md new file mode 100644 index 0000000..34ff5d6 --- /dev/null +++ b/docs/src/guides/host-memory-pinning.md @@ -0,0 +1,134 @@ +# Host-memory pinning + +`Server(..., pin_host_memory=True)` CUDA-page-locks echo's ring buffers, so that +copying a sampled batch to a GPU is a real DMA transfer instead of a chunked copy +the driver walks through by hand. + +It is off by default, and it is not free: page-locked memory is not swappable. + +```python +server = Server(example, batch_size=512, transport=TcpTransport(port=50051), + pin_host_memory=True) +``` + +## The guarantee + +> **If `Server(...)` returns, every ring buffer is page-locked.** + +There is no state in which you asked for pinning and silently did not get it. +Anything that would prevent it — no CUDA runtime found, no usable device, a +registration rejected — raises `RuntimeError` from the constructor, naming every +path echo probed and the CUDA error symbolically: + +```text +RuntimeError: pin_host_memory=True but no CUDA runtime could be loaded. Probed, in order: + - (none) [already-loaded scan]: this process has no CUDA runtime mapped + - (none) [installed-wheel search]: no CUDA vendor package is importable + - libcudart.so.13 [soname load]: libcudart.so.13: cannot open shared object file: No such file or directory + - libcudart.so.12 [soname load]: libcudart.so.12: cannot open shared object file: No such file or directory + - libcudart.so [soname load]: libcudart.so: cannot open shared object file: No such file or directory +Install a CUDA runtime (for example the nvidia-cuda-runtime wheel), or construct +the Server after your framework has initialised CUDA. +``` + +A `(none)` line means that rung ran and found nothing, so the message +distinguishes it from a rung that never ran at all. + +That is all the API there is: no status object, no report method, no warning. If +you would rather degrade than fail, catch `RuntimeError`. + +With the argument off, echo loads no CUDA library at all and behaves exactly like +a build without the feature. + +## Why a pageable "async" copy is not async + +Ordinary heap memory is *pageable* — the OS may move or swap it out, so the GPU +cannot safely DMA out of it. When you copy from pageable memory, the driver +cannot hand the transfer to the copy engine and walk away. Instead it: + +1. copies a chunk of your data into a small internal page-locked staging buffer, +2. DMAs that chunk to the device, +3. repeats, thousands of times for a large batch. + +Steps 1–3 run on the calling CPU thread and hold the driver lock, with two +consequences: + +- the bytes move slower, because every chunk pays a round-trip; +- **`cudaMemcpyAsync` stops being asynchronous.** The call does not return until + the staging loop is essentially done, and while it runs, the driver lock it + holds blocks the *same thread's* kernel launches. On a learner issuing tens of + thousands of small launches per step, that contention — not the bandwidth — is + what shows up as host-side dispatch dominating the step while the GPU sits idle. + +Page-locking the ring buffers removes the staging loop. The pages cannot move, so +the driver writes one DMA descriptor to the copy engine and returns immediately. + +Echo registers with `cudaHostRegisterPortable`, which makes the locked pages +valid in **every** CUDA context in the process, including contexts created later. +That is why there is no device argument: N servers across N GPUs in one process +each pin without any per-server configuration, and echo never selects a device or +allocates device memory. + +## Sizing the footprint + +The whole ring is locked, not one batch: + +```text +page-locked bytes = batch_size x num_buffers x sum(leaf.nbytes for leaf in example) +``` + +For a batch of 512 with `num_buffers=3` and 67 KB per sample, that is +512 x 3 x 67 KB ≈ 102 MB. + +- **It is not swappable.** Those pages are removed from the pool the kernel can + reclaim under pressure, for the whole life of the server. +- **It multiplies by the number of servers.** One `Server` per GPU in a single + process means N times the figure above, so size the host accordingly. + +Note that `VmLck` in `/proc/self/status` stays at **zero** even when pinning is +working — the driver's page-locking does not go through mlock accounting. So a +memory-lock limit (`ulimit -l`) does not bind here, and `VmLck` is not a way to +check that pinning engaged. + +## Construct the server after CUDA is initialised + +Registration needs an initialised CUDA runtime, so the supported order is: + +1. let your framework initialise CUDA, +2. then construct the `Server`. + +Echo does not depend on you getting this right — before registering it forces +initialisation best-effort, by freeing a null pointer. + +Echo never allocates device buffers of its own and never calls a +set-device function. + +If initialisation genuinely cannot happen, or no device is usable, the +constructor raises with the CUDA error named. It never degrades silently. + +## Finding the CUDA runtime + +Echo has no build-time or link-time CUDA dependency; one wheel installs on GPU +and CPU-only hosts alike. The runtime is located at pin time by trying, in order: + +1. **What the process already has mapped.** Version- and path-agnostic; hits + whenever your framework has already initialised CUDA. +2. **What the installed CUDA wheels ship.** Echo locates the `nvidia` vendor + package through Python's import machinery and searches beneath it. It searches + the vendor package rather than a named component because CUDA 13 ships one + consolidated wheel (`nvidia/cu13/lib/`) whose layout differs from CUDA 12's + per-component layout (`nvidia/cuda_runtime/lib/`). +3. **Sonames**, versioned first (`libcudart.so.13`, `libcudart.so.12`) and the + unversioned `libcudart.so` last. + +You should not need a soname symlink in your image or an `LD_LIBRARY_PATH` entry. +If you are carrying either as a workaround, remove it — and note that with them +in place, rung 3 hides whether rungs 1 and 2 work. + +The ladder and why each rung exists are covered in +[Host-memory pinning](../design/host-pinning.md) under Rust internals. + +## What is not covered + +- **Non-CUDA accelerators.** macOS wheels build and the default-off path is a + clean no-op there, but no equivalent functionality is provided. diff --git a/justfile b/justfile index f236ab7..e5d948f 100644 --- a/justfile +++ b/justfile @@ -31,10 +31,14 @@ install-telemetry-from-git BRANCH="main": develop: uv run maturin develop --features detailed-metrics -# Live-reload docs on http://127.0.0.1:8000 +# Live-reload docs on http://127.0.0.1:8000/echo/ docs-serve: uv run --extra docs mkdocs serve +# Live-reload docs bound so another machine can reach them (pass an interface to narrow it) +docs-serve-on ADDR="0.0.0.0:8000": + uv run --extra docs mkdocs serve -a {{ADDR}} + # Build static docs into site/ docs: uv run --extra docs mkdocs build diff --git a/mkdocs.yml b/mkdocs.yml index e32abae..d301a7e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -80,6 +80,7 @@ nav: - TrajectoryAccumulator: guides/trajectory-accumulator.md - In-process use: guides/in-process.md - Reading metrics: guides/metrics.md + - Host-memory pinning: guides/host-memory-pinning.md - Python API: - Server: api/server.md - Clients: api/clients.md @@ -93,4 +94,5 @@ nav: - Ingress & drainers: design/ingress.md - Selector: design/selector.md - Transport: design/transport.md + - Host-memory pinning: design/host-pinning.md - Development: development.md diff --git a/pyproject.toml b/pyproject.toml index 43b39cb..81bb457 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "id-echo" -version = "0.1.1" +version = "0.2.0" description = "A fast distributed replay buffer for reinforcement learning." readme = "README.md" license = { file = "LICENSE" } @@ -64,6 +64,9 @@ include = [{ path = "python/echo/py.typed", format = "wheel" }, { path = "python pythonpath = ["."] python_files = ["test_*.py", "bench_*.py"] testpaths = ["python/tests"] +markers = [ + "gpu: needs a usable CUDA device; skipped automatically when none is present", +] [tool.uv] # Rebuild package when any rust files change diff --git a/python/echo/echo.pyi b/python/echo/echo.pyi index cf045cc..0d6c69a 100644 --- a/python/echo/echo.pyi +++ b/python/echo/echo.pyi @@ -46,6 +46,7 @@ class _Server: num_buffers: int = 3, num_drainers: int = 8, producer_queue_size: int = 8, + pin_host_memory: bool = False, ) -> None: ... def start(self) -> None: ... def sample(self) -> tuple[list[np.ndarray[Any, np.dtype[np.uint8]]], SampleInfo] | None: ... diff --git a/python/echo/server.py b/python/echo/server.py index 1b39a7f..f9cea91 100644 --- a/python/echo/server.py +++ b/python/echo/server.py @@ -28,6 +28,12 @@ class Server: num_buffers: Number of ring buffer batches (min 2, default 3) num_drainers: Number of threads draining from producer queues producer_queue_size: Per-connection queue size + pin_host_memory: CUDA-page-lock the ring buffers, so a downstream H2D + copy is a real DMA transfer instead of a staging copy. + + Raises: + RuntimeError: If ``pin_host_memory`` is set and the buffers could not + be page-locked. Lifetime: the numpy arrays returned by ``sample()`` are views into Rust-owned ring-buffer memory. They are invalidated as soon as the next @@ -45,6 +51,7 @@ def __init__( num_buffers: int = 3, num_drainers: int = 8, producer_queue_size: int = 8, + pin_host_memory: bool = False, ): leaves, self._treedef = optree.tree_flatten(example) @@ -64,6 +71,7 @@ def __init__( num_buffers=num_buffers, num_drainers=num_drainers, producer_queue_size=producer_queue_size, + pin_host_memory=pin_host_memory, ) def start(self) -> None: diff --git a/python/tests/test_host_pinning.py b/python/tests/test_host_pinning.py new file mode 100644 index 0000000..075ff23 --- /dev/null +++ b/python/tests/test_host_pinning.py @@ -0,0 +1,337 @@ +"""Tests for ``Server(pin_host_memory=...)``. + +The contract under test is that *if the constructor returns, every ring buffer +is page-locked*. That is deliberately what echo offers instead of a status +object: construction succeeding is the assertion, and unlike an explicit check +it cannot be forgotten. + +So the confirmation here does not ask echo what it did. It asks the CUDA runtime +directly, through ``ctypes``, for the registration flags on the address behind a +sampled view. A test that trusted echo's own account of its work would pass just +as happily against the silent no-op this feature replaced. + +Note that ``VmLck`` in ``/proc/self/status`` is *not* a valid check: the NVIDIA +driver's page-locking does not go through mlock accounting, so it stays at zero +even when pinning demonstrably works. +""" +import ctypes +import glob +import os +import re +import subprocess +import sys +import textwrap + +import numpy as np +import pytest + +from echo import Server + +CUDA_SUCCESS = 0 +CUDA_HOST_REGISTER_PORTABLE = 0x01 + +# Rung 3 of echo's ladder, in the order echo tries them. +SONAMES = ("libcudart.so.13", "libcudart.so.12", "libcudart.so") + +EXAMPLE = { + "obs": np.zeros((16,), dtype=np.float32), + "reward": np.zeros((1,), dtype=np.float32), +} + + +def _cudart_paths() -> list[str]: + """Candidate CUDA runtime libraries, mirroring echo's resolution ladder. + + Independent of echo's implementation on purpose: the point of these tests is + to reach the runtime without going through the code under test. + """ + candidates = [] + with open("/proc/self/maps") as maps: + candidates += re.findall(r"\s(/\S*libcudart\.so[.\d]*)", maps.read()) + for site in sys.path: + candidates += glob.glob(os.path.join(site, "nvidia", "*", "lib", "libcudart.so*")) + candidates += list(SONAMES) + return [c for c in candidates if not c.endswith(".a")] + + +def _load_cudart() -> ctypes.CDLL | None: + """The CUDA runtime, or None on a machine without one.""" + for path in _cudart_paths(): + try: + return ctypes.CDLL(path) + except OSError: + continue + return None + + +CUDART = _load_cudart() + + +def _device_count(cudart: ctypes.CDLL) -> int: + count = ctypes.c_int(0) + if cudart.cudaGetDeviceCount(ctypes.byref(count)) != CUDA_SUCCESS: + return 0 + return count.value + + +HAS_GPU = CUDART is not None and _device_count(CUDART) > 0 +requires_gpu = pytest.mark.skipif(not HAS_GPU, reason="no CUDA device available") + + +def registration_flags(address: int) -> tuple[int, int]: + """The CUDA runtime's own account of how ``address`` is registered. + + Returns ``(cuda_error_code, flags)``; a non-zero code means the runtime does + not know the address as page-locked host memory. + """ + assert CUDART is not None + flags = ctypes.c_uint(0) + code = CUDART.cudaHostGetFlags(ctypes.byref(flags), ctypes.c_void_p(address)) + return code, flags.value + + +def batch_address(sample) -> int: + """Address of the first leaf of a sampled batch. + + The first batch starts at ring slot 0, so this is the base of the buffer. + """ + leaves = [sample.batch[key] for key in sorted(sample.batch)] + return leaves[0].ctypes.data + + +class TestDefaultOff: + def test_constructs_and_serves_batches(self): + """Off by default, and off is the unchanged behaviour — this runs on + CPU-only machines and CI.""" + server = Server(EXAMPLE, batch_size=2) + try: + for _ in range(2): + server.submit({k: v.copy() for k, v in EXAMPLE.items()}) + sample = server.sample() + assert sample is not None + assert sample.batch["obs"].shape == (2, 16) + finally: + server.close() + + @pytest.mark.skipif(not sys.platform.startswith("linux"), reason="needs procfs") + def test_loads_no_cuda_library(self): + """The default path must behave like a build without the feature. + + Checked in a subprocess because this module loads the CUDA runtime + through ``ctypes`` at import, which would mask a load by echo. Runs on + CPU-only machines too — there it asserts the absence stays an absence. + """ + script = textwrap.dedent( + """ + import re + import numpy as np + from echo import Server + + def mapped(): + with open("/proc/self/maps") as f: + return set(re.findall(r"\\s(/\\S*libcudart\\.so[.\\d]*)", f.read())) + + before = mapped() + server = Server({"obs": np.zeros((32,), dtype=np.float32)}, 4) + assert mapped() == before, "default-off loaded a CUDA library" + server.close() + print("ok") + """ + ) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="needs procfs") +def test_total_resolution_failure_names_every_rung_in_order(): + """The failure a misconfigured deployment gets: loud, at startup, specific. + + The message is the entire interface to a resolution failure, so it is checked + against a real one rather than a hand-built error. Forced in a subprocess with + every rung disabled: nothing mapped, ``nvidia`` unimportable, no loader path. + + Runs on GPU hosts too, since it needs resolution to fail rather than a device + to be absent. + + The soname probe has to happen in the subprocess: this module loads the + runtime by absolute path at import, and once it is loaded, ``dlopen`` by + soname succeeds against the already-loaded object. Probing in the parent + would report every rung as resolvable and quietly disable this test. + """ + script = textwrap.dedent( + """ + import re + import sys + + with open("/proc/self/maps") as f: + assert not re.search(r"libcudart\\.so", f.read()), "runtime already mapped" + + # If the loader can find a runtime by bare soname then rung 3 succeeds and + # a total failure is not reproducible here. A failed CDLL maps nothing, so + # rung 1 stays empty either way. + import ctypes + for soname in sys.argv[1:]: + try: + ctypes.CDLL(soname) + except OSError: + continue + print("SKIP", soname) + sys.exit(0) + + sys.modules["nvidia"] = None # rung 2: no importable vendor package + + import numpy as np + from echo import Server + + try: + Server({"obs": np.zeros((8,), dtype=np.float32)}, 4, pin_host_memory=True) + except RuntimeError as e: + print(e) + else: + raise AssertionError("pinning succeeded with every rung disabled") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, *SONAMES], + capture_output=True, + text=True, + env={k: v for k, v in os.environ.items() if k != "LD_LIBRARY_PATH"}, + ) + assert result.returncode == 0, result.stderr + if result.stdout.startswith("SKIP"): + pytest.skip(f"{result.stdout.split()[1]} resolves by soname, so rung 3 succeeds") + message = result.stdout + + # A rung that ran and found nothing reports itself as `(none)`, so the message + # cannot be mistaken for one where that rung never ran at all. Then every + # soname actually attempted, versioned before unversioned. Ordering is the + # part worth pinning down: reordering the rungs is how this feature broke. + expected = [ + "(none) [already-loaded scan]", + "(none) [installed-wheel search]", + *(f"{soname} [soname load]" for soname in SONAMES), + ] + positions = [message.find(fragment) for fragment in expected] + missing = [f for f, p in zip(expected, positions) if p < 0] + assert not missing, f"absent from the message: {missing}\n{message}" + assert positions == sorted(positions), f"rungs reported out of order:\n{message}" + assert "Install a CUDA runtime" in message, message + + +@requires_gpu +@pytest.mark.gpu +def test_pins_when_no_runtime_is_loaded_yet(): + """A server constructed *before* the framework touches CUDA must still pin. + + This test module loads the runtime through ``ctypes`` at import, so the + tests below resolve it from the process's own mappings. Here the subprocess + has nothing mapped, which leaves finding the pip-installed wheel — the case + the original defect broke, and the reason it needed no loader-path + environment variable to be set. + """ + script = textwrap.dedent( + """ + import re + import numpy as np + with open("/proc/self/maps") as f: + assert not re.search(r"libcudart\\.so", f.read()), "runtime already mapped" + + from echo import Server + server = Server({"obs": np.zeros((32,), dtype=np.float32)}, 4, pin_host_memory=True) + with open("/proc/self/maps") as f: + assert re.search(r"libcudart\\.so", f.read()), "nothing was loaded" + server.close() + print("ok") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env={k: v for k, v in os.environ.items() if k != "LD_LIBRARY_PATH"}, + ) + assert result.returncode == 0, result.stderr + assert "ok" in result.stdout + + +@requires_gpu +@pytest.mark.gpu +class TestPinningOnGpu: + def test_construction_page_locks_every_ring_buffer(self): + server = Server(EXAMPLE, batch_size=4, pin_host_memory=True) + try: + for _ in range(4): + server.submit({k: v.copy() for k, v in EXAMPLE.items()}) + sample = server.sample() + assert sample is not None + + for name in sorted(sample.batch): + address = sample.batch[name].ctypes.data + code, flags = registration_flags(address) + assert code == CUDA_SUCCESS, f"{name} is not registered host memory" + assert flags & CUDA_HOST_REGISTER_PORTABLE, ( + f"{name} is registered but not portable (flags={flags:#x})" + ) + finally: + server.close() + + def test_batch_is_bit_identical_with_and_without_pinning(self): + rng = np.random.default_rng(0) + samples = [ + { + "obs": rng.standard_normal(16).astype(np.float32), + "reward": rng.standard_normal(1).astype(np.float32), + } + for _ in range(4) + ] + + batches = [] + for pin in (False, True): + server = Server(EXAMPLE, batch_size=4, pin_host_memory=pin) + try: + for sample in samples: + server.submit(sample) + got = server.sample() + assert got is not None + batches.append({k: np.copy(v) for k, v in got.batch.items()}) + finally: + server.close() + + unpinned, pinned = batches + for name in unpinned: + assert unpinned[name].tobytes() == pinned[name].tobytes(), name + + def test_two_servers_in_one_process_both_pin(self): + """One server per GPU in one process needs no device configuration: + registration is portable across every context in the process.""" + first = Server(EXAMPLE, batch_size=4, pin_host_memory=True) + second = Server(EXAMPLE, batch_size=8, pin_host_memory=True) + try: + for server, batch_size in ((first, 4), (second, 8)): + for _ in range(batch_size): + server.submit({k: v.copy() for k, v in EXAMPLE.items()}) + sample = server.sample() + assert sample is not None + code, flags = registration_flags(batch_address(sample)) + assert code == CUDA_SUCCESS + assert flags & CUDA_HOST_REGISTER_PORTABLE + finally: + first.close() + second.close() + + def test_pinning_survives_server_teardown_and_reconstruction(self): + """A closed-and-dropped server unregisters, and the next one still pins: + a long-lived process must not accumulate or exhaust registrations.""" + for _ in range(3): + server = Server(EXAMPLE, batch_size=4, pin_host_memory=True) + try: + for _ in range(4): + server.submit({k: v.copy() for k, v in EXAMPLE.items()}) + sample = server.sample() + assert sample is not None + code, _ = registration_flags(batch_address(sample)) + assert code == CUDA_SUCCESS + finally: + server.close() + del server diff --git a/src/host_pinning/mod.rs b/src/host_pinning/mod.rs new file mode 100644 index 0000000..524804f --- /dev/null +++ b/src/host_pinning/mod.rs @@ -0,0 +1,129 @@ +//! Optional CUDA host-memory pinning (page-locking) of the ring buffers. +//! +//! Copying out of pageable memory is not a DMA transfer: the driver stages it +//! through a small internal pinned buffer in CPU-executed chunks, holding the +//! calling thread and the driver lock throughout. Page-locking makes the same +//! copy one descriptor on the copy engine. +//! +//! Either every buffer ends up locked or the caller gets an error saying why; +//! there is no best-effort path, because a silent no-op is indistinguishable +//! from pinning not helping. See `docs/src/guides/host-memory-pinning.md`. +//! +//! The runtime is `dlopen`ed at pin time, so there is no build-time or +//! link-time CUDA dependency and nothing loads unless pinning is requested. +//! +//! Two jobs, one submodule each: [`resolve`] finds the runtime, [`register`] +//! page-locks memory with it. This root holds only what both need — the entry +//! points, the regions they act on, and the error either can return. + +pub mod register; +pub mod resolve; + +pub use register::{pin_all, unpin_all}; +pub use resolve::api; + +use std::fmt; +use std::os::raw::{c_char, c_int, c_uint, c_void}; + +/// `cudaError_t cudaHostRegister(void*, size_t, unsigned int)` +pub type RegisterFn = unsafe extern "C" fn(*mut c_void, usize, c_uint) -> c_int; +/// `cudaError_t cudaHostUnregister(void*)` +pub type UnregisterFn = unsafe extern "C" fn(*mut c_void) -> c_int; +/// `const char* cudaGetErrorName(cudaError_t)` +pub type ErrorNameFn = unsafe extern "C" fn(c_int) -> *const c_char; +/// `cudaError_t cudaFree(void*)` +pub type FreeFn = unsafe extern "C" fn(*mut c_void) -> c_int; + +pub const CUDA_SUCCESS: c_int = 0; + +/// The CUDA entry points pinning needs. +/// +/// Passed explicitly rather than reached through a global so tests can inject +/// stubs and assert rollback without a GPU. +#[derive(Clone, Copy)] +pub struct CudaApi { + register: RegisterFn, + unregister: UnregisterFn, + error_name: ErrorNameFn, + free: FreeFn, +} + +impl CudaApi { + /// Build an api from raw entry points. `resolve` uses this after `dlsym`; + /// tests use it to inject stubs. + /// + /// # Safety + /// Each pointer must be a live function with the signature its type names. + pub unsafe fn new( + register: RegisterFn, + unregister: UnregisterFn, + error_name: ErrorNameFn, + free: FreeFn, + ) -> Self { + Self { + register, + unregister, + error_name, + free, + } + } +} + +/// A contiguous host allocation to page-lock. Owns nothing; the validity +/// invariants live on [`pin_all`] / [`unpin_all`]. +#[derive(Clone, Copy)] +pub struct Region { + pub ptr: *mut u8, + pub len: usize, +} + +/// Why pinning could not be delivered. On failure nothing is left registered. +#[derive(Debug)] +pub enum PinError { + /// No CUDA runtime could be loaded. One line per probed path, so a caller + /// can fix their environment without reading this source. + RuntimeUnavailable { probed: Vec }, + /// The runtime rejected a registration. `name` is the CUDA error symbol. + Registration { + name: String, + code: c_int, + /// Which of the caller's regions was rejected. + region_index: usize, + len: usize, + }, +} + +impl fmt::Display for PinError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + PinError::RuntimeUnavailable { probed } => { + write!( + f, + "pin_host_memory=True but no CUDA runtime could be loaded. Probed, in order:" + )?; + for line in probed { + write!(f, "\n - {line}")?; + } + write!( + f, + "\nInstall a CUDA runtime (for example the nvidia-cuda-runtime wheel), or \ + construct the Server after your framework has initialised CUDA." + ) + } + PinError::Registration { + name, + code, + region_index, + len, + } => write!( + f, + "pin_host_memory=True but cudaHostRegister failed with {name} ({code}) on ring \ + buffer {region_index} ({len} bytes). Registrations made before the failure have been \ + rolled back. A usable CUDA device must be visible to this process; construct the \ + Server after CUDA is initialised." + ), + } + } +} + +impl std::error::Error for PinError {} diff --git a/src/host_pinning/register.rs b/src/host_pinning/register.rs new file mode 100644 index 0000000..0b58a72 --- /dev/null +++ b/src/host_pinning/register.rs @@ -0,0 +1,90 @@ +//! Page-locking host memory, and rolling back cleanly when part of it fails. + +use std::ffi::CStr; +use std::os::raw::{c_int, c_uint, c_void}; + +use super::{CudaApi, PinError, Region, CUDA_SUCCESS}; + +/// `cudaHostRegisterPortable`: valid in every CUDA context in the process, +/// including ones created later. This is why pinning takes no device argument +/// and why N servers across N GPUs in one process need no configuration. +pub const CUDA_HOST_REGISTER_PORTABLE: c_uint = 0x01; + +/// Page-lock every region, or leave none of them locked. +/// +/// On failure the regions registered so far are unregistered first, so nothing +/// accumulates across retries. Normal teardown is `Drop`'s job. +/// +/// # Safety +/// - Every region's `ptr` must be valid for `len` bytes. +/// - That allocation must not be reallocated, resized, or moved while it stays +/// registered — the registration pins the physical pages behind *these* +/// addresses. +/// - Each region registered here must be passed to [`unpin_all`] before its +/// memory is freed. +pub unsafe fn pin_all(api: &CudaApi, regions: &[Region]) -> Result<(), PinError> { + // Registration against an uninitialised runtime fails, so force init. + // Freeing null frees nothing and never changes the current device, but it is + // not free: it creates the primary context on the current device, costing + // that context's memory (~128 MB measured). Constructing after the framework + // has initialised CUDA — the documented order — means it already exists. + // + // The result is ignored: portable registration makes the device irrelevant, + // so the only failure worth reporting is registration's own, below. + unsafe { (api.free)(std::ptr::null_mut()) }; + + for (index, region) in regions.iter().enumerate() { + if region.len == 0 { + continue; + } + // SAFETY: the caller guarantees `ptr` is valid for `len` bytes. + let code = unsafe { + (api.register)( + region.ptr as *mut c_void, + region.len, + CUDA_HOST_REGISTER_PORTABLE, + ) + }; + if code != CUDA_SUCCESS { + // SAFETY: these regions were just registered by this loop, so they + // satisfy `unpin_all`'s contract. + unsafe { unpin_all(api, ®ions[..index]) }; + return Err(PinError::Registration { + name: error_name(api, code), + code, + region_index: index, + len: region.len, + }); + } + } + Ok(()) +} + +/// Reverse of [`pin_all`]; must run before the memory is freed. +/// +/// # Safety +/// Every region must currently be registered, by a [`pin_all`] call through +/// this same `api`, and must still be valid for `len` bytes. +pub unsafe fn unpin_all(api: &CudaApi, regions: &[Region]) { + for region in regions { + if region.len == 0 { + continue; + } + // Errors are dropped: this runs from `Drop` and during rollback, and + // there is nothing useful either could do about a failure. + // SAFETY: the caller guarantees the region is registered and valid. + unsafe { (api.unregister)(region.ptr as *mut c_void) }; + } +} + +/// A CUDA error as its symbol (`cudaErrorInvalidValue`), so it can be looked up +/// rather than decoded from an integer. +fn error_name(api: &CudaApi, code: c_int) -> String { + let name = unsafe { (api.error_name)(code) }; + if name.is_null() { + return format!("unrecognised CUDA error {code}"); + } + unsafe { CStr::from_ptr(name) } + .to_string_lossy() + .into_owned() +} diff --git a/src/host_pinning/resolve.rs b/src/host_pinning/resolve.rs new file mode 100644 index 0000000..2ae6ae3 --- /dev/null +++ b/src/host_pinning/resolve.rs @@ -0,0 +1,301 @@ +//! Finding the CUDA runtime: a three-rung ladder, tried in order, accumulating +//! every attempted path so a total failure can say what it tried. + +use std::ffi::{CStr, CString}; +use std::os::raw::c_void; +use std::path::PathBuf; +use std::sync::OnceLock; + +use super::{CudaApi, ErrorNameFn, FreeFn, PinError, RegisterFn, UnregisterFn}; + +/// Current and previous CUDA major versions, then the unversioned soname. +/// +/// The unversioned name is last and rarely resolves: the pip CUDA runtime +/// wheels ship only the versioned soname, with no symlink and no ldconfig entry. +const SONAMES: [&str; 3] = ["libcudart.so.13", "libcudart.so.12", "libcudart.so"]; + +/// Which rung of the resolution ladder produced a candidate. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Rung { + /// Rung 1: a runtime the process already has mapped. + AlreadyLoaded, + /// Rung 2: a runtime shipped by an installed CUDA wheel. + InstalledWheel, + /// Rung 3: a soname, left to the system loader. + Soname, +} + +impl Rung { + pub fn label(self) -> &'static str { + match self { + Rung::AlreadyLoaded => "already-loaded scan", + Rung::InstalledWheel => "installed-wheel search", + Rung::Soname => "soname load", + } + } +} + +/// One thing to hand to `dlopen`, in ladder order. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Candidate { + /// An absolute path (rungs 1 and 2) or a bare soname (rung 3). + pub name: String, + pub rung: Rung, +} + +/// True for a CUDA runtime shared object, by file name. +fn is_runtime_lib(path: &str) -> bool { + let name = path.rsplit('/').next().unwrap_or(path); + // `libcudart.so`, `libcudart.so.13`, `libcudart.so.13.0.48` — but not + // `libcudart_static.a`. + name == "libcudart.so" || name.starts_with("libcudart.so.") +} + +/// Rung 1: runtimes the process already has mapped, in first-seen order. +/// +/// Hits whenever the framework has already initialised CUDA. `dlopen` on the +/// absolute path of a mapped library reuses that mapping rather than loading a +/// second copy. +pub fn scan_mapped_runtimes(maps: &str) -> Vec { + let mut found: Vec = Vec::new(); + for line in maps.lines() { + // The pathname is the last field and is absolute, so it starts at the + // first " /". Located this way, not by splitting on whitespace, because + // paths may contain spaces. Unlinked files get a " (deleted)" suffix. + let line = line.strip_suffix(" (deleted)").unwrap_or(line); + let Some(offset) = line.find(" /") else { + continue; + }; + let path = &line[offset + 1..]; + if is_runtime_lib(path) && !found.iter().any(|seen| seen == path) { + found.push(path.to_string()); + } + } + found +} + +/// `nvidia/cu13/lib/libcudart.so.13` is three deep; one spare for a relayout. +const WHEEL_SEARCH_DEPTH: usize = 4; + +/// Rung 2: runtimes shipped by installed wheels, newest major first. Lets the +/// server be constructed before the framework has loaded CUDA. +/// +/// `roots` are the CUDA *vendor* package directories, found by the caller +/// through Python's import machinery. Searching the whole vendor package rather +/// than a named component is load-bearing: CUDA 13 ships one consolidated wheel +/// (`nvidia/cu13/lib/`), CUDA 12 one per component (`nvidia/cuda_runtime/lib/`), +/// so naming the component finds nothing on a CUDA 13 install. +pub fn search_wheel_roots(roots: &[PathBuf]) -> Vec { + let mut found: Vec = Vec::new(); + let mut frontier: Vec<(PathBuf, usize)> = roots.iter().map(|r| (r.clone(), 0)).collect(); + + while let Some((dir, depth)) = frontier.pop() { + let Ok(entries) = std::fs::read_dir(&dir) else { + continue; // Missing or unreadable roots simply contribute nothing. + }; + for entry in entries.flatten() { + let path = entry.path(); + // `DirEntry::file_type` does not follow symlinks, so a symlinked + // directory is not descended into. Wheels do not put libraries behind + // directory symlinks, and it keeps the walk safe from link loops. + let Ok(file_type) = entry.file_type() else { + continue; + }; + if file_type.is_dir() { + if depth + 1 < WHEEL_SEARCH_DEPTH { + frontier.push((path, depth + 1)); + } + } else if path.to_str().is_some_and(is_runtime_lib) { + found.push(path); + } + } + } + + // Newest major version first, then by path so the result is deterministic + // regardless of directory iteration order. + found.sort_by(|a, b| { + let key = |p: &PathBuf| { + let major = p + .file_name() + .and_then(|n| n.to_str()) + .and_then(soname_major) + .unwrap_or(0); + std::cmp::Reverse(major) + }; + key(a).cmp(&key(b)).then_with(|| a.cmp(b)) + }); + found + .iter() + .filter_map(|p| p.to_str().map(str::to_owned)) + .collect() +} + +/// CUDA major version from a soname, e.g. `libcudart.so.13.0.48` -> 13. +fn soname_major(name: &str) -> Option { + name.strip_prefix("libcudart.so.")? + .split('.') + .next()? + .parse() + .ok() +} + +/// The full ladder, in the order it will be tried. Pure, so the ordering is +/// unit-testable without a GPU: reordering it is how this feature broke before. +pub fn candidates(maps: &str, vendor_roots: &[PathBuf]) -> Vec { + let rungs = [ + (Rung::AlreadyLoaded, scan_mapped_runtimes(maps)), + (Rung::InstalledWheel, search_wheel_roots(vendor_roots)), + ( + Rung::Soname, + SONAMES.iter().map(|s| (*s).to_owned()).collect(), + ), + ]; + + let mut ladder: Vec = Vec::new(); + for (rung, names) in rungs { + for name in names { + // The same file can turn up on two rungs (a mapped runtime that is + // also the one the wheel ships); probe it once, on the earlier rung. + if !ladder.iter().any(|c| c.name == name) { + ladder.push(Candidate { name, rung }); + } + } + } + ladder +} + +/// `dlopen` a candidate and resolve the four symbols pinning needs. +/// +/// The handle is never `dlclose`d: the registrations it backs must outlive it. +fn open(candidate: &Candidate) -> Result { + let name = CString::new(candidate.name.as_str()) + .map_err(|_| "path contains an interior NUL byte".to_string())?; + + // RTLD_LOCAL, so this does not change how any other library in the process + // resolves its symbols. + let flags = libc::RTLD_NOW | libc::RTLD_LOCAL; + let handle = unsafe { + // Rung 3 prefers a copy the process already has, under whatever path, + // over pulling a second one in from the loader path. + let noload = if candidate.rung == Rung::Soname { + libc::dlopen(name.as_ptr(), flags | libc::RTLD_NOLOAD) + } else { + std::ptr::null_mut() + }; + if noload.is_null() { + libc::dlopen(name.as_ptr(), flags) + } else { + noload + } + }; + if handle.is_null() { + return Err(dlerror().unwrap_or_else(|| "dlopen failed".to_string())); + } + + // SAFETY: each symbol is transmuted to the signature libcudart declares for + // it; a name that resolves in libcudart has that signature by definition. + unsafe { + Ok(CudaApi::new( + std::mem::transmute::<*mut c_void, RegisterFn>(symbol(handle, c"cudaHostRegister")?), + std::mem::transmute::<*mut c_void, UnregisterFn>(symbol( + handle, + c"cudaHostUnregister", + )?), + std::mem::transmute::<*mut c_void, ErrorNameFn>(symbol(handle, c"cudaGetErrorName")?), + std::mem::transmute::<*mut c_void, FreeFn>(symbol(handle, c"cudaFree")?), + )) + } +} + +/// # Safety +/// `handle` must be a live handle returned by `dlopen`. +unsafe fn symbol(handle: *mut c_void, name: &CStr) -> Result<*mut c_void, String> { + let sym = libc::dlsym(handle, name.as_ptr()); + if sym.is_null() { + return Err(format!("opened, but {} is missing", name.to_string_lossy())); + } + Ok(sym) +} + +fn dlerror() -> Option { + let err = unsafe { libc::dlerror() }; + if err.is_null() { + return None; + } + Some( + unsafe { CStr::from_ptr(err) } + .to_string_lossy() + .into_owned(), + ) +} + +/// Walk the ladder, returning the first runtime that opens or an error naming +/// every path probed. +fn resolve(vendor_roots: &[PathBuf]) -> Result { + // Absent (macOS, a container without procfs) just means rung 1 finds + // nothing. + let maps = std::fs::read_to_string("/proc/self/maps").unwrap_or_default(); + + let ladder = candidates(&maps, vendor_roots); + let mut probed = Vec::new(); + for rung in [Rung::AlreadyLoaded, Rung::InstalledWheel, Rung::Soname] { + let of_rung: Vec<&Candidate> = ladder.iter().filter(|c| c.rung == rung).collect(); + // A rung that produced no candidate at all would otherwise be invisible, + // leaving a reader unable to tell "searched and found nothing" from + // "never ran". + if of_rung.is_empty() { + probed.push(empty_rung_note(rung, vendor_roots)); + continue; + } + for candidate in of_rung { + match open(candidate) { + Ok(api) => return Ok(api), + Err(reason) => { + probed.push(format!("{} [{}]: {reason}", candidate.name, rung.label())) + } + } + } + } + Err(PinError::RuntimeUnavailable { probed }) +} + +/// What to report for a rung that contributed no candidate, so the error says +/// where it looked. +fn empty_rung_note(rung: Rung, vendor_roots: &[PathBuf]) -> String { + let detail = match rung { + Rung::AlreadyLoaded => "this process has no CUDA runtime mapped".to_string(), + Rung::InstalledWheel if vendor_roots.is_empty() => { + "no CUDA vendor package is importable".to_string() + } + Rung::InstalledWheel => format!( + "no runtime found beneath {}", + vendor_roots + .iter() + .map(|root| root.display().to_string()) + .collect::>() + .join(", ") + ), + // Unreachable: SONAMES is a non-empty const, and the dedup in + // `candidates` only ever compares absolute paths against bare sonames. + // Kept for exhaustiveness rather than panicking inside error reporting. + Rung::Soname => "no soname candidates".to_string(), + }; + format!("(none) [{}]: {detail}", rung.label()) +} + +static API: OnceLock = OnceLock::new(); + +/// The CUDA runtime, resolving it on first use. Only called when pinning was +/// asked for, so the default-off path loads nothing. +/// +/// Failures are not cached: a process that retries after its framework has +/// initialised CUDA should get the later, better answer. +pub fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError> { + if let Some(api) = API.get() { + return Ok(api); + } + let resolved = resolve(vendor_roots)?; + // Two threads racing here both resolve; `dlopen` is idempotent and + // reference-counted, so the loser just drops an identical set of pointers. + Ok(API.get_or_init(|| resolved)) +} diff --git a/src/lib.rs b/src/lib.rs index 37188bc..40f6a19 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ use pyo3::prelude::*; pub mod array_spec; +pub mod host_pinning; pub mod ingress; pub mod metrics; mod py_bindings; diff --git a/src/py_bindings.rs b/src/py_bindings.rs index eb706fb..1470638 100644 --- a/src/py_bindings.rs +++ b/src/py_bindings.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use std::sync::Arc; use numpy::npyffi::{ @@ -98,8 +99,10 @@ pub struct PyServer { #[pymethods] impl PyServer { #[new] - #[pyo3(signature = (shapes, dtype_sizes, batch_size, transport=None, num_buffers=3, num_drainers=8, producer_queue_size=8))] + #[pyo3(signature = (shapes, dtype_sizes, batch_size, transport=None, num_buffers=3, num_drainers=8, producer_queue_size=8, pin_host_memory=false))] + #[allow(clippy::too_many_arguments)] fn new( + py: Python<'_>, shapes: Vec>, dtype_sizes: Vec, batch_size: usize, @@ -107,6 +110,7 @@ impl PyServer { num_buffers: usize, num_drainers: usize, producer_queue_size: usize, + pin_host_memory: bool, ) -> PyResult { if shapes.len() != dtype_sizes.len() { return Err(PyValueError::new_err( @@ -123,7 +127,7 @@ impl PyServer { let capacity = batch_size * num_buffers; let metrics = Metrics::new(num_drainers); - let store = Arc::new(Store::new( + let mut store = Store::new( specs, batch_size, num_buffers, @@ -133,7 +137,19 @@ impl PyServer { Some(metrics.clone()), )), Box::new(FifoRemover::new()), - )); + ); + + // The guarantee: if this constructor returns, every ring buffer is + // registered. Anything short of that raises, so a workload can never be + // measured against a pinning path that silently did nothing. + if pin_host_memory { + // Holding the GIL across a CUDA call would also let echo deadlock + // against a thread that holds a CUDA lock and wants the GIL. + let roots = cuda_vendor_roots(py); + py.detach(|| store.pin_host_memory(&roots)) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + } + let store = Arc::new(store); // Drainer pool + transport are only needed when there's a network // transport; in-process submit() skips both and writes directly. @@ -250,6 +266,27 @@ impl PyServer { } } +/// CUDA vendor package directories, for rung 2 of the resolution ladder. This +/// is how pinning finds a pip-installed runtime with no loader-path variable set. +/// +/// `nvidia` is a namespace package, so `__path__` is an iterable of every +/// site-packages directory contributing to it. An absent or unusable package +/// just means that rung finds nothing. +fn cuda_vendor_roots(py: Python<'_>) -> Vec { + let Ok(vendor) = py.import("nvidia") else { + return Vec::new(); + }; + let Ok(path) = vendor.getattr("__path__") else { + return Vec::new(); + }; + let Ok(entries) = path.try_iter() else { + return Vec::new(); + }; + entries + .filter_map(|entry| entry.ok()?.extract::().ok()) + .collect() +} + /// Create a 1-D uint8 numpy array that is a view into existing memory. /// The array does NOT own the data (NPY_ARRAY_OWNDATA is not set). /// The caller must ensure the memory remains valid for the array's lifetime. diff --git a/src/ring_buf.rs b/src/ring_buf.rs index b820950..5715377 100644 --- a/src/ring_buf.rs +++ b/src/ring_buf.rs @@ -11,6 +11,9 @@ //! coordination. use std::cell::UnsafeCell; +use std::path::PathBuf; + +use crate::host_pinning::{self, CudaApi, PinError, Region}; pub struct PytreeRingBuf { /// One contiguous buffer per array in the flattened pytree. @@ -19,6 +22,9 @@ pub struct PytreeRingBuf { slot_bytes: Vec, /// Total number of slots. capacity: usize, + /// The runtime the buffers are registered with, once they are. `Some` is + /// `Drop`'s cue to unregister, and holding it here keeps `Drop` off a global. + pinned_with: Option, } impl PytreeRingBuf { @@ -34,7 +40,7 @@ impl PytreeRingBuf { ); assert!(!slot_bytes.is_empty(), "slot_bytes must not be empty"); - let buffers = slot_bytes + let buffers: Vec>> = slot_bytes .iter() .map(|&bytes| UnsafeCell::new(vec![0u8; bytes * capacity])) .collect(); @@ -43,9 +49,49 @@ impl PytreeRingBuf { buffers, slot_bytes, capacity, + pinned_with: None, } } + /// Page-lock every buffer, so a host-to-device copy of a sampled view is a + /// DMA transfer rather than a chunked staging copy. All or nothing: a partial + /// failure is rolled back before the error returns. + /// + /// Separate from `new` because a constructor returning `Err` never runs + /// `Drop`, so registering there would need a hand-written unregister loop on + /// the error path. Here `Drop` owns rollback and teardown alike. + /// + /// The buffers are contiguous and never reallocated, so a registration stays + /// valid for the buffer's whole life. `cuda_vendor_roots` comes from Python's + /// import machinery; see [`crate::host_pinning`]. + pub fn pin_host_memory(&mut self, cuda_vendor_roots: &[PathBuf]) -> Result<(), PinError> { + self.pin_with(*host_pinning::api(cuda_vendor_roots)?) + } + + /// [`Self::pin_host_memory`] against an already-resolved runtime. Split out + /// so tests can drive registration and teardown with stubs, without a GPU. + pub fn pin_with(&mut self, api: CudaApi) -> Result<(), PinError> { + // SAFETY: the regions are `self`'s own allocations, never reallocated, + // and `Drop` unregisters them before they are freed. + unsafe { host_pinning::pin_all(&api, &self.regions())? }; + self.pinned_with = Some(api); + Ok(()) + } + + /// Each backing buffer as a (pointer, length) pair for the CUDA runtime. + fn regions(&self) -> Vec { + self.buffers + .iter() + .map(|cell| { + let buf = unsafe { &*cell.get() }; + Region { + ptr: buf.as_ptr() as *mut u8, + len: buf.len(), + } + }) + .collect() + } + pub fn capacity(&self) -> usize { self.capacity } @@ -117,3 +163,16 @@ unsafe impl Sync for PytreeRingBuf {} // Safety: All data is heap-allocated and owned; transfer between threads is safe. unsafe impl Send for PytreeRingBuf {} + +impl Drop for PytreeRingBuf { + fn drop(&mut self) { + let Some(api) = self.pinned_with else { + return; + }; + // Drop::drop runs before the fields are dropped, so the memory is still + // valid here. + // SAFETY: `pinned_with` is Some only after `pin_with` registered exactly + // these regions through this same api. + unsafe { host_pinning::unpin_all(&api, &self.regions()) }; + } +} diff --git a/src/store.rs b/src/store.rs index 09743fa..4f54391 100644 --- a/src/store.rs +++ b/src/store.rs @@ -1,4 +1,5 @@ use std::cell::Cell; +use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; #[cfg(feature = "detailed-metrics")] use std::time::Instant; @@ -7,6 +8,7 @@ use crossbeam_utils::CachePadded; use tokio::sync::Notify; use crate::array_spec::ArraySpec; +use crate::host_pinning::PinError; use crate::metrics::DrainerMetrics; use crate::ring_buf::PytreeRingBuf; use crate::selector::{Remover, SampleResult, Sampler}; @@ -64,6 +66,14 @@ impl Store { } } + /// Page-lock the ring buffers; see [`PytreeRingBuf::pin_host_memory`]. + /// + /// `&mut self` puts this before the store is shared, so a caller that fails + /// to finish building unregisters by dropping it. + pub fn pin_host_memory(&mut self, cuda_vendor_roots: &[PathBuf]) -> Result<(), PinError> { + self.ring.pin_host_memory(cuda_vendor_roots) + } + pub fn batch_size(&self) -> usize { self.batch_size } diff --git a/tests/host_pinning.rs b/tests/host_pinning.rs new file mode 100644 index 0000000..472591a --- /dev/null +++ b/tests/host_pinning.rs @@ -0,0 +1,500 @@ +//! Tests for `host_pinning`: the CUDA-runtime resolution ladder, and +//! registration with its rollback. +//! +//! The resolution and stub-injected registration tests run anywhere. Rollback's +//! only visible consequence is the *absence* of leaked registrations, which +//! cannot be observed from Python, so it is checked through an injected API +//! rather than externally. The remaining tests need a real CUDA device and +//! report a skip when there isn't one — Rust has no native test skip, so they +//! pass rather than fail on CPU-only CI. + +use std::cell::RefCell; +use std::ffi::CString; +use std::os::raw::{c_char, c_int, c_uint, c_void}; +use std::path::PathBuf; + +use echo::host_pinning::register::CUDA_HOST_REGISTER_PORTABLE; +use echo::host_pinning::resolve::{ + api, candidates, scan_mapped_runtimes, search_wheel_roots, Rung, +}; +use echo::host_pinning::{pin_all, unpin_all, CudaApi, PinError, Region, RegisterFn, CUDA_SUCCESS}; +use echo::ring_buf::PytreeRingBuf; + +// =========================================================================== +// Resolution: the three-rung ladder +// =========================================================================== + +// --- rung 1: parsing the process's own mapped files --- + +/// A runtime mapped as several segments, unrelated libraries, anonymous and +/// special mappings, a path with a space in it, and a deleted mapping. +const MAPS_FIXTURE: &str = "\ +55a3c0000000-55a3c0021000 r--p 00000000 fd:01 1179651 /usr/bin/python3.11 +7f1a00000000-7f1a00021000 r--p 00000000 fd:01 2359310 /usr/lib/x86_64-linux-gnu/libc.so.6 +7f1a10000000-7f1a10800000 rw-p 00000000 00:00 0 +7f1a20000000-7f1a20a00000 r--p 00000000 fd:01 4194313 /venv/lib/python3.11/site-packages/nvidia/cu13/lib/libcudart.so.13 +7f1a20a00000-7f1a21400000 r-xp 00a00000 fd:01 4194313 /venv/lib/python3.11/site-packages/nvidia/cu13/lib/libcudart.so.13 +7f1a30000000-7f1a30100000 r-xp 00000000 fd:01 4194320 /opt/my libs/libcudart.so.12 +7f1a40000000-7f1a40100000 r-xp 00000000 fd:01 4194321 /tmp/stale/libcudart.so.11 (deleted) +7ffd00000000-7ffd00021000 rw-p 00000000 00:00 0 [stack] +ffffffffff600000-ffffffffff601000 --xp 00000000 00:00 0 [vsyscall] +"; + +#[test] +fn scan_finds_each_mapped_runtime_once_in_order() { + assert_eq!( + scan_mapped_runtimes(MAPS_FIXTURE), + vec![ + "/venv/lib/python3.11/site-packages/nvidia/cu13/lib/libcudart.so.13", + "/opt/my libs/libcudart.so.12", + "/tmp/stale/libcudart.so.11", + ] + ); +} + +#[test] +fn scan_ignores_unrelated_libraries() { + let maps = "\ +7f1a00000000-7f1a00021000 r-xp 00000000 fd:01 1 /usr/lib/libcudnn.so.9 +7f1a10000000-7f1a10021000 r-xp 00000000 fd:01 2 /usr/lib/libcublas.so.13 +7f1a20000000-7f1a20021000 r-xp 00000000 fd:01 3 /usr/lib/libcudart_static.a +"; + assert!(scan_mapped_runtimes(maps).is_empty()); +} + +#[test] +fn scan_of_a_process_with_no_runtime_yields_nothing() { + assert!(scan_mapped_runtimes("").is_empty()); +} + +// --- rung 2: searching beneath the CUDA vendor package --- + +fn touch(path: &std::path::Path) { + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(path, b"").unwrap(); +} + +#[test] +fn wheel_search_finds_the_consolidated_layout() { + // CUDA 13: one wheel, `nvidia/cu13/lib/`. + let root = tempfile::tempdir().unwrap(); + let nvidia = root.path().join("nvidia"); + touch(&nvidia.join("cu13/lib/libcudart.so.13")); + touch(&nvidia.join("cu13/lib/libcudart_static.a")); + touch(&nvidia.join("cudnn/lib/libcudnn.so.9")); + + assert_eq!( + search_wheel_roots(std::slice::from_ref(&nvidia)), + vec![nvidia + .join("cu13/lib/libcudart.so.13") + .to_string_lossy() + .into_owned()] + ); +} + +#[test] +fn wheel_search_finds_the_per_component_layout() { + // CUDA 12: one wheel per component, `nvidia/cuda_runtime/lib/`. + let root = tempfile::tempdir().unwrap(); + let nvidia = root.path().join("nvidia"); + touch(&nvidia.join("cuda_runtime/lib/libcudart.so.12")); + touch(&nvidia.join("cublas/lib/libcublas.so.12")); + + assert_eq!( + search_wheel_roots(std::slice::from_ref(&nvidia)), + vec![nvidia + .join("cuda_runtime/lib/libcudart.so.12") + .to_string_lossy() + .into_owned()] + ); +} + +#[test] +fn wheel_search_prefers_the_newest_major_version() { + let root = tempfile::tempdir().unwrap(); + let nvidia = root.path().join("nvidia"); + touch(&nvidia.join("cuda_runtime/lib/libcudart.so.9")); + touch(&nvidia.join("cu13/lib/libcudart.so.13")); + touch(&nvidia.join("cuda_runtime/lib/libcudart.so.12")); + + let found = search_wheel_roots(&[nvidia]); + let names: Vec<&str> = found + .iter() + .map(|p| p.rsplit('/').next().unwrap()) + .collect(); + assert_eq!( + names, + vec!["libcudart.so.13", "libcudart.so.12", "libcudart.so.9"] + ); +} + +#[test] +fn wheel_search_tolerates_missing_and_empty_roots() { + let root = tempfile::tempdir().unwrap(); + assert!(search_wheel_roots(&[]).is_empty()); + assert!(search_wheel_roots(&[root.path().join("does-not-exist")]).is_empty()); + assert!(search_wheel_roots(&[root.path().to_path_buf()]).is_empty()); +} + +// --- ladder ordering --- + +#[test] +fn ladder_tries_mapped_then_wheel_then_sonames() { + let root = tempfile::tempdir().unwrap(); + let nvidia = root.path().join("nvidia"); + touch(&nvidia.join("cu13/lib/libcudart.so.13")); + let wheel_lib = nvidia + .join("cu13/lib/libcudart.so.13") + .to_string_lossy() + .into_owned(); + + let ladder = candidates(MAPS_FIXTURE, &[nvidia]); + let rungs: Vec = ladder.iter().map(|c| c.rung).collect(); + let names: Vec<&str> = ladder.iter().map(|c| c.name.as_str()).collect(); + + assert_eq!( + rungs, + vec![ + Rung::AlreadyLoaded, + Rung::AlreadyLoaded, + Rung::AlreadyLoaded, + Rung::InstalledWheel, + Rung::Soname, + Rung::Soname, + Rung::Soname, + ] + ); + assert_eq!( + names, + vec![ + "/venv/lib/python3.11/site-packages/nvidia/cu13/lib/libcudart.so.13", + "/opt/my libs/libcudart.so.12", + "/tmp/stale/libcudart.so.11", + wheel_lib.as_str(), + // Versioned before unversioned: the wheels ship no unversioned + // symlink, which is what broke this before. + "libcudart.so.13", + "libcudart.so.12", + "libcudart.so", + ] + ); +} + +#[test] +fn ladder_probes_each_path_once() { + // A mapped runtime that is also the one the wheel ships. + let root = tempfile::tempdir().unwrap(); + let nvidia = root.path().join("nvidia"); + let lib = nvidia.join("cu13/lib/libcudart.so.13"); + touch(&lib); + let maps = format!( + "7f1a20000000-7f1a20a00000 r-xp 0 fd:01 1 {}\n", + lib.display() + ); + + let ladder = candidates(&maps, &[nvidia]); + let mut names: Vec<&str> = ladder.iter().map(|c| c.name.as_str()).collect(); + let before = names.len(); + names.sort_unstable(); + names.dedup(); + assert_eq!(names.len(), before, "ladder contains a duplicate path"); +} + +#[test] +fn unresolvable_runtime_reports_every_path_probed() { + let err = PinError::RuntimeUnavailable { + probed: vec![ + "/a/libcudart.so.13: boom".into(), + "libcudart.so: nope".into(), + ], + }; + let message = err.to_string(); + assert!(message.contains("/a/libcudart.so.13: boom"), "{message}"); + assert!(message.contains("libcudart.so: nope"), "{message}"); +} + +// =========================================================================== +// Registration, rollback and teardown +// =========================================================================== + +// --- registration and rollback, with the CUDA calls stubbed out --- + +thread_local! { + static REGISTERED: RefCell> = const { RefCell::new(Vec::new()) }; + static UNREGISTERED: RefCell> = const { RefCell::new(Vec::new()) }; + static INITIALISED: RefCell = const { RefCell::new(0) }; +} + +/// Fails on the third region, succeeds on every other. +unsafe extern "C" fn stub_register(ptr: *mut c_void, _len: usize, _flags: c_uint) -> c_int { + let nth = REGISTERED.with(|c| { + c.borrow_mut().push(ptr as usize); + c.borrow().len() + }); + if nth == 3 { + 1 // cudaErrorInvalidValue + } else { + CUDA_SUCCESS + } +} + +unsafe extern "C" fn stub_register_ok(ptr: *mut c_void, _len: usize, _flags: c_uint) -> c_int { + REGISTERED.with(|c| c.borrow_mut().push(ptr as usize)); + CUDA_SUCCESS +} + +unsafe extern "C" fn stub_unregister(ptr: *mut c_void) -> c_int { + UNREGISTERED.with(|c| c.borrow_mut().push(ptr as usize)); + CUDA_SUCCESS +} + +unsafe extern "C" fn stub_error_name(_code: c_int) -> *const c_char { + c"cudaErrorInvalidValue".as_ptr() +} + +unsafe extern "C" fn stub_free(_ptr: *mut c_void) -> c_int { + INITIALISED.with(|c| *c.borrow_mut() += 1); + CUDA_SUCCESS +} + +fn stub_api(register: RegisterFn) -> CudaApi { + REGISTERED.with(|c| c.borrow_mut().clear()); + UNREGISTERED.with(|c| c.borrow_mut().clear()); + INITIALISED.with(|c| *c.borrow_mut() = 0); + // SAFETY: the stubs below are live functions with these signatures. + unsafe { CudaApi::new(register, stub_unregister, stub_error_name, stub_free) } +} + +/// Dangling, but never dereferenced: only passed to the stub API. +fn fake_regions(n: usize) -> Vec { + (1..=n) + .map(|i| Region { + ptr: (i * 0x1000) as *mut u8, + len: 4096, + }) + .collect() +} + +#[test] +fn a_failure_on_the_third_region_rolls_back_exactly_the_first_two() { + let api = stub_api(stub_register); + let regions = fake_regions(5); + + // SAFETY: the stub API never dereferences these pointers. + let err = unsafe { pin_all(&api, ®ions) }.expect_err("registration should have failed"); + + let unregistered = UNREGISTERED.with(|c| c.borrow().clone()); + assert_eq!( + unregistered, + vec![0x1000, 0x2000], + "rollback must unregister the regions that succeeded, and only those" + ); + // Nothing past the failure is attempted. + assert_eq!(REGISTERED.with(|c| c.borrow().len()), 3); + assert!( + err.to_string().contains("cudaErrorInvalidValue"), + "the CUDA error must be named, not numeric: {err}" + ); +} + +#[test] +fn a_successful_pin_registers_every_region_and_rolls_back_nothing() { + let api = stub_api(stub_register_ok); + let regions = fake_regions(4); + + // SAFETY: the stub API never dereferences these pointers. + unsafe { pin_all(&api, ®ions) }.expect("registration should have succeeded"); + + assert_eq!( + REGISTERED.with(|c| c.borrow().clone()), + vec![0x1000, 0x2000, 0x3000, 0x4000] + ); + assert!(UNREGISTERED.with(|c| c.borrow().is_empty())); +} + +#[test] +fn pinning_forces_runtime_initialisation_first() { + let api = stub_api(stub_register_ok); + // SAFETY: the stub API never dereferences these pointers. + unsafe { pin_all(&api, &fake_regions(1)) }.unwrap(); + assert_eq!(INITIALISED.with(|c| *c.borrow()), 1); +} + +#[test] +fn dropping_a_pinned_ring_buffer_unregisters_every_buffer() { + // Teardown without a GPU: real ring buffer, stubbed CUDA. A leak here + // would be invisible from Python, hence the injection point. + let api = stub_api(stub_register_ok); + let mut ring = PytreeRingBuf::new(vec![64, 128], 8, 4); + ring.pin_with(api) + .expect("stubbed registration should succeed"); + + let registered = REGISTERED.with(|c| c.borrow().clone()); + assert_eq!(registered.len(), 2, "one registration per array"); + assert!(UNREGISTERED.with(|c| c.borrow().is_empty())); + + drop(ring); + assert_eq!( + UNREGISTERED.with(|c| c.borrow().clone()), + registered, + "drop must unregister exactly what was registered" + ); +} + +#[test] +fn dropping_an_unpinned_ring_buffer_touches_no_cuda_entry_point() { + let _api = stub_api(stub_register_ok); + drop(PytreeRingBuf::new(vec![64], 8, 4)); + assert!(REGISTERED.with(|c| c.borrow().is_empty())); + assert!(UNREGISTERED.with(|c| c.borrow().is_empty())); +} + +#[test] +fn a_ring_buffer_whose_registration_fails_leaves_nothing_registered() { + // The third array is rejected, so the two that took must be unregistered + // and `Drop` must then do nothing more. + let api = stub_api(stub_register); + let mut ring = PytreeRingBuf::new(vec![64, 64, 64, 64], 8, 4); + ring.pin_with(api).expect_err("the third array should fail"); + + let after_rollback = UNREGISTERED.with(|c| c.borrow().clone()); + assert_eq!(after_rollback.len(), 2); + drop(ring); + assert_eq!( + UNREGISTERED.with(|c| c.borrow().clone()), + after_rollback, + "a failed pin must leave Drop with nothing to reverse" + ); +} + +#[test] +fn unpin_all_unregisters_every_region() { + let api = stub_api(stub_register_ok); + // SAFETY: the stub API never dereferences these pointers. + unsafe { unpin_all(&api, &fake_regions(3)) }; + assert_eq!( + UNREGISTERED.with(|c| c.borrow().clone()), + vec![0x1000, 0x2000, 0x3000] + ); +} + +// --- against a real CUDA runtime, when the machine has one --- + +/// The dev checkout's CUDA vendor package: `cargo test` has no Python +/// interpreter to ask. See `docs/src/development.md` for installing the wheel. +fn dev_vendor_roots() -> Vec { + let venv = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(".venv/lib"); + let Ok(entries) = std::fs::read_dir(venv) else { + return Vec::new(); + }; + entries + .flatten() + .map(|e| e.path().join("site-packages/nvidia")) + .filter(|p| p.is_dir()) + .collect() +} + +/// `cudaError_t cudaHostGetFlags(unsigned int*, void*)` +type HostGetFlagsFn = unsafe extern "C" fn(*mut c_uint, *mut c_void) -> c_int; + +/// The runtime, or `None` on a machine without one. Rust has no native test +/// skip, so these report and pass rather than fail on CPU-only CI. +fn real_api() -> Option<&'static CudaApi> { + match api(&dev_vendor_roots()) { + Ok(api) => Some(api), + Err(e) => { + eprintln!("skipping: no CUDA runtime on this machine\n{e}"); + None + } + } +} + +/// `cudaHostGetFlags`, looked up separately from the [`CudaApi`] under test: +/// going through those pointers would only prove they believe they succeeded. +/// +/// Call after resolution has run, so this re-opens the mapped runtime +/// (`RTLD_NOLOAD`) rather than loading a second copy. +fn host_get_flags() -> HostGetFlagsFn { + let maps = std::fs::read_to_string("/proc/self/maps").unwrap_or_default(); + let mapped = scan_mapped_runtimes(&maps) + .into_iter() + .next() + .expect("resolution has run, so a runtime is mapped"); + let path = CString::new(mapped).unwrap(); + unsafe { + let handle = libc::dlopen( + path.as_ptr(), + libc::RTLD_NOW | libc::RTLD_LOCAL | libc::RTLD_NOLOAD, + ); + assert!(!handle.is_null(), "the mapped runtime should re-open"); + let sym = libc::dlsym(handle, c"cudaHostGetFlags".as_ptr()); + assert!(!sym.is_null(), "the runtime should export cudaHostGetFlags"); + std::mem::transmute::<*mut c_void, HostGetFlagsFn>(sym) + } +} + +#[test] +fn a_real_runtime_confirms_the_buffers_are_registered_portable() { + let Some(api) = real_api() else { return }; + let mut buffer = vec![0u8; 4 << 20]; + let region = Region { + ptr: buffer.as_mut_ptr(), + len: buffer.len(), + }; + + // SAFETY: `buffer` outlives the registration and is unregistered below. + unsafe { pin_all(api, std::slice::from_ref(®ion)) } + .expect("registration failed on a GPU host"); + + let mut flags: c_uint = 0; + let code = unsafe { (host_get_flags())(&mut flags, region.ptr as *mut c_void) }; + assert_eq!(code, CUDA_SUCCESS, "cudaHostGetFlags rejected the address"); + assert_eq!( + flags & CUDA_HOST_REGISTER_PORTABLE, + CUDA_HOST_REGISTER_PORTABLE + ); + + // SAFETY: just registered above, through this same api. + unsafe { unpin_all(api, std::slice::from_ref(®ion)) }; +} + +#[test] +fn a_real_ring_buffer_pins_and_unregisters_on_drop() { + if real_api().is_none() { + return; + } + let roots = dev_vendor_roots(); + let mut ring = PytreeRingBuf::new(vec![1024, 2048], 64, 8); + ring.pin_host_memory(&roots) + .expect("registration failed on a GPU host"); + + let (address, _) = ring.range_ptr(0, 0, 8); + let mut flags: c_uint = 0; + let code = unsafe { (host_get_flags())(&mut flags, address as *mut c_void) }; + assert_eq!(code, CUDA_SUCCESS); + assert_eq!( + flags & CUDA_HOST_REGISTER_PORTABLE, + CUDA_HOST_REGISTER_PORTABLE + ); + + // Drop unregisters, so the runtime must no longer know the address. + drop(ring); + let code = unsafe { (host_get_flags())(&mut flags, address as *mut c_void) }; + assert_ne!( + code, CUDA_SUCCESS, + "drop should have unregistered the ring buffer" + ); +} + +#[test] +fn pinning_twice_in_one_process_reuses_the_resolved_runtime() { + // Two servers in one process, no extra configuration. + if real_api().is_none() { + return; + } + let roots = dev_vendor_roots(); + let mut first = PytreeRingBuf::new(vec![4096], 32, 4); + let mut second = PytreeRingBuf::new(vec![4096], 32, 4); + first.pin_host_memory(&roots).expect("first ring failed"); + second.pin_host_memory(&roots).expect("second ring failed"); +} diff --git a/uv.lock b/uv.lock index 0b76d57..c510412 100644 --- a/uv.lock +++ b/uv.lock @@ -423,7 +423,7 @@ wheels = [ [[package]] name = "id-echo" -version = "0.1.0" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },