feat: optional CUDA host-memory pinning of the ring buffers

.gitignore

@@ -23,6 +23,9 @@ benches/*.png
 # Tooling
 .claude
 .worktrees/
+CLAUDE.md
+/docs/agents/
+.scratch/
 
 # Local-only files (personal scripts, scratch, etc.)
 .local/

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
 

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"] }

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
 

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<CudaApi>`, 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.

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

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<u8>`, 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<CudaApi>` 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<u8>` stays.
+
 ## What this type does not do
 
 - Track which slots are in use (the `Store` does that via the write/read

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