From b0a2c415265ce064c3194ccad19f89eb3475f7bd Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Thu, 25 Jun 2026 19:33:50 +0200 Subject: [PATCH 01/15] feat: optional CUDA host-memory pinning of ring buffers Page-lock each ring buffer via cudaHostRegister so a downstream jax.device_put of the numpy views is a fast, truly-async H2D DMA instead of a synchronous host->device staging copy. libcudart is resolved at runtime via dlopen (libc only) -- no CUDA toolchain / link-time dependency; the code is always compiled and is a graceful no-op unless ECHO_PIN_HOST_MEMORY=1 and libcudart is present. Pinning happens only at ring-buffer construction (PytreeRingBuf::new) and teardown (Drop), off the tokio runtime. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/host_pinning.rs | 83 +++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/ring_buf.rs | 23 ++++++++++++- 3 files changed, 106 insertions(+), 1 deletion(-) create mode 100644 src/host_pinning.rs diff --git a/src/host_pinning.rs b/src/host_pinning.rs new file mode 100644 index 0000000..72d697e --- /dev/null +++ b/src/host_pinning.rs @@ -0,0 +1,83 @@ +//! Optional CUDA host-memory pinning (page-locking) of the ring buffers. +//! +//! Enabled at runtime by `ECHO_PIN_HOST_MEMORY=1` (off by default). When enabled, +//! [`pin`] registers a buffer with `cudaHostRegister` so a downstream +//! `jax.device_put` of its numpy view becomes a fast, truly-async H2D DMA instead +//! of a synchronous host->device staging copy. libcudart is resolved at runtime +//! via `dlopen`, so this carries NO CUDA toolchain / link-time dependency: the +//! code is always compiled but is a graceful no-op when the env var is unset, on +//! CUDA-less hosts, or if registration fails. + +use std::os::raw::{c_char, c_int, c_uint, c_void}; +use std::sync::OnceLock; + +// cudaError_t cudaHostRegister(void* ptr, size_t size, unsigned int flags) +type RegisterFn = unsafe extern "C" fn(*mut c_void, usize, c_uint) -> c_int; +// cudaError_t cudaHostUnregister(void* ptr) +type UnregisterFn = unsafe extern "C" fn(*mut c_void) -> c_int; + +struct Api { + register: RegisterFn, + unregister: UnregisterFn, +} +// The dlopen handle outlives the process; the resolved fn pointers are immutable. +unsafe impl Sync for Api {} +unsafe impl Send for Api {} + +fn load() -> Option { + // Opt-in: only attempt to pin when explicitly enabled. + if std::env::var("ECHO_PIN_HOST_MEMORY").as_deref() != Ok("1") { + return None; + } + unsafe { + let handle = libc::dlopen( + b"libcudart.so\0".as_ptr() as *const c_char, + libc::RTLD_NOW | libc::RTLD_GLOBAL, + ); + if handle.is_null() { + return None; + } + let register = libc::dlsym(handle, b"cudaHostRegister\0".as_ptr() as *const c_char); + let unregister = libc::dlsym(handle, b"cudaHostUnregister\0".as_ptr() as *const c_char); + if register.is_null() || unregister.is_null() { + return None; + } + Some(Api { + register: std::mem::transmute::<*mut c_void, RegisterFn>(register), + unregister: std::mem::transmute::<*mut c_void, UnregisterFn>(unregister), + }) + } +} + +fn api() -> Option<&'static Api> { + static API: OnceLock> = OnceLock::new(); + API.get_or_init(load).as_ref() +} + +/// Page-lock `[ptr, ptr + len)` (portable across the process's CUDA contexts). +/// +/// No-op unless `ECHO_PIN_HOST_MEMORY=1` and libcudart is present. Best-effort: +/// registration errors leave the memory pageable. +pub(crate) fn pin(ptr: *mut u8, len: usize) { + if len == 0 { + return; + } + if let Some(api) = api() { + // cudaHostRegisterPortable = 0x01 (usable from every device/context). + unsafe { + (api.register)(ptr as *mut c_void, len, 0x01); + } + } +} + +/// Reverse of [`pin`]; must run before the buffer is freed. +pub(crate) fn unpin(ptr: *mut u8, len: usize) { + if len == 0 { + return; + } + if let Some(api) = api() { + unsafe { + (api.unregister)(ptr as *mut c_void); + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 37188bc..c447663 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ use pyo3::prelude::*; pub mod array_spec; +mod host_pinning; pub mod ingress; pub mod metrics; mod py_bindings; diff --git a/src/ring_buf.rs b/src/ring_buf.rs index b820950..dcfca56 100644 --- a/src/ring_buf.rs +++ b/src/ring_buf.rs @@ -34,11 +34,21 @@ 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(); + // Optionally page-lock each buffer (ECHO_PIN_HOST_MEMORY=1) so that a + // downstream `jax.device_put` of the numpy views is a fast, truly-async + // H2D DMA rather than a synchronous host->device staging copy. No-op + // without CUDA. Buffers are contiguous and never reallocated, so the + // registration stays valid for the buffer's life. + for cell in &buffers { + let buf = unsafe { &mut *cell.get() }; + crate::host_pinning::pin(buf.as_mut_ptr(), buf.len()); + } + Self { buffers, slot_bytes, @@ -117,3 +127,14 @@ 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) { + // Unregister (un-pin) before the backing Vecs are freed. Drop::drop runs + // before the struct's fields are dropped, so the memory is still valid. + for cell in &self.buffers { + let buf = unsafe { &mut *cell.get() }; + crate::host_pinning::unpin(buf.as_mut_ptr(), buf.len()); + } + } +} From 524bbc2dd09ce6ee20f06360ce3e191a0c159fb4 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Thu, 25 Jun 2026 19:33:50 +0200 Subject: [PATCH 02/15] chore: bump version to 0.1.2 Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.lock | 2 +- Cargo.toml | 2 +- pyproject.toml | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 49e9145..5d6da5b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -215,7 +215,7 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "echo" -version = "0.1.1" +version = "0.1.2" dependencies = [ "arc-swap", "criterion", diff --git a/Cargo.toml b/Cargo.toml index 028bfc6..0c8bb29 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "echo" -version = "0.1.1" +version = "0.1.2" edition = "2021" description = "A fast distributed replay buffer for reinforcement learning." license = "Apache-2.0" diff --git a/pyproject.toml b/pyproject.toml index 43b39cb..a28344f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "id-echo" -version = "0.1.1" +version = "0.1.2" description = "A fast distributed replay buffer for reinforcement learning." readme = "README.md" license = { file = "LICENSE" } From 2a13d09a0a5200f3a1560d6d26fd59278bd16645 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Fri, 31 Jul 2026 14:28:49 +0200 Subject: [PATCH 03/15] chore: ignore agent tooling and scratch directories Keeps CLAUDE.md, docs/agents/ and .scratch/ out of the repo. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) 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/ From 708bc38b0b1c4f40a2be21f455b3fb54b3e27105 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Fri, 31 Jul 2026 14:29:13 +0200 Subject: [PATCH 04/15] feat: pin_host_memory keyword argument with a page-locked guarantee MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewrites CUDA host-memory pinning so that it works, and so that it cannot silently not work. The previous implementation opened the unversioned `libcudart.so` soname, which the pip CUDA runtime wheels do not ship. The dlopen returned NULL, every pin call became a no-op, and nothing reported it — so a workload could be measured against a code path that never executed and be concluded not to benefit from pinning. Pinning is now one keyword argument, off by default: Server(example, batch_size, ..., pin_host_memory=False) The contract is that if `Server(...)` returns, every ring buffer is page-locked. Any inability to deliver that raises `RuntimeError` naming every path probed and the CUDA error symbolically. That guarantee is what replaces an observability API: construction succeeding *is* the assertion, and unlike an explicit check it cannot be forgotten. - Resolution is a three-rung ladder — already-mapped libraries, then the installed CUDA wheels (searching the `nvidia` vendor package rather than a named component, since CUDA 13's consolidated layout differs from CUDA 12's per-component one), then sonames, versioned first. No soname symlink or LD_LIBRARY_PATH entry is needed. - Page-locking moved out of `PytreeRingBuf::new` into an explicit fallible step on a constructed buffer, so the existing `Drop` owns rollback for both the failure path and normal teardown rather than a hand-written unregister loop. `new` stays infallible and the eight existing ring-buffer tests are untouched. - CUDA entry points go through an injectable struct of function pointers, so registration, rollback and drop are all exercised with stubs on a machine with no GPU. - Registration is portable-flag only. Page-aligned allocation and read-only registration were both measured and rejected: alignment gained 1.6% of copy time (bar was 5%), and read-only is unsupported on the test GPU and is documented as permission rather than optimisation. - `ECHO_PIN_HOST_MEMORY` is deleted; the argument is the only control. GPU tests carry a `gpu` marker and skip when no device is present. The new guide documents the mechanism, the measured numbers, the unswappable footprint arithmetic, and how to verify pinning engaged — including that `VmLck` stays at zero even when it is working, which is what made the original false negative believable. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 31 +- Cargo.lock | 68 +- Cargo.toml | 3 +- README.md | 3 + docs/src/design/ring-buffer.md | 30 + docs/src/development.md | 21 + docs/src/guides/host-memory-pinning.md | 262 ++++++ mkdocs.yml | 1 + pyproject.toml | 5 +- python/echo/echo.pyi | 1 + python/echo/server.py | 23 + python/tests/test_host_pinning.py | 271 +++++++ src/host_pinning.rs | 1031 ++++++++++++++++++++++-- src/py_bindings.rs | 45 +- src/ring_buf.rs | 81 +- src/store.rs | 13 + uv.lock | 2 +- 17 files changed, 1816 insertions(+), 75 deletions(-) create mode 100644 docs/src/guides/host-memory-pinning.md create mode 100644 python/tests/test_host_pinning.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 44b2b64..cb282ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,36 @@ 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] - unreleased + +### Added + +- `Server(..., pin_host_memory=False)` 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. 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, so a + throughput measurement can't be invalidated by the feature having done nothing. +- New guide: [Host-memory pinning](https://instadeepai.github.io/echo/guides/host-memory-pinning/), + covering the mechanism, the measured numbers, the unswappable footprint + arithmetic, and how to verify pinning engaged (and why `VmLck` cannot). + +### Changed + +- The CUDA runtime is now resolved through a three-rung ladder — already-mapped + libraries, then the installed CUDA wheels, then sonames (versioned before + unversioned). A soname symlink or `LD_LIBRARY_PATH` entry is no longer needed. + +### Removed + +- `ECHO_PIN_HOST_MEMORY`. Pinning is controlled only by the `pin_host_memory` + keyword argument, so the two can never disagree. Nothing released ever + responded to this variable — and it never worked: it opened the unversioned + `libcudart.so` soname, which the pip CUDA runtime wheels do not ship, so every + pin call was silently a no-op. ## [0.1.1] - 2026-05-26 diff --git a/Cargo.lock b/Cargo.lock index 5d6da5b..08dd1de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -215,7 +215,7 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "echo" -version = "0.1.2" +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 0c8bb29..197014a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "echo" -version = "0.1.2" +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..a1f7d88 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,9 @@ 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/) + of the ring buffers, so a host-to-device copy of a batch is a real DMA + transfer rather than a chunked staging copy that blocks the calling thread ## Example diff --git a/docs/src/design/ring-buffer.md b/docs/src/design/ring-buffer.md index e6ecd7a..096fc21 100644 --- a/docs/src/design/ring-buffer.md +++ b/docs/src/design/ring-buffer.md @@ -44,6 +44,36 @@ 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` (see the [guide](../guides/host-memory-pinning.md) for the +mechanism and the numbers) CUDA-page-locks every backing `Vec` so a +downstream host-to-device copy of a sampled view is a DMA transfer rather than a +chunked staging copy. Two design points are worth recording here. + +**It is a separate fallible step, not part of `new`.** `PytreeRingBuf::new` and +`Store::new` stay infallible; `pin_host_memory(&mut self, ..)` runs on a +fully-constructed buffer and returns `Result`. The reason is rollback +correctness, not taste: a constructor that returns `Err` never runs `Drop`, so +registering inside the constructor would force a hand-written unregister loop on +the error path. Registering afterwards lets the existing `Drop` own rollback for +both the failure path and normal teardown. A `pinned` flag says whether `Drop` +has anything to reverse. + +Within one attempt, either every buffer ends up registered or none does — +`pin_all` unregisters what succeeded before returning the error. Registrations +are still leaked if a reference-counted `Store` outlives process shutdown, which +is accepted. + +**The registration is valid for the buffer's life** because the buffers are +allocated once in `new` and never reallocated or resized. That property is what +makes registering the whole `Vec` up front sound; a growable buffer would +invalidate the registration on its first reallocation. + +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 at all 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..05ba0df 100644 --- a/docs/src/development.md +++ b/docs/src/development.md @@ -22,6 +22,27 @@ cargo test # Rust unit + integration tests uv run pytest python/tests/ -v # Python tests ``` +### Working on host-memory pinning + +The [pinning](guides/host-memory-pinning.md) tests that need a GPU carry the +`gpu` marker and skip automatically when no device is present, so CI needs no +per-runner configuration. + +To exercise the CUDA-runtime resolution ladder locally, install a runtime wheel +into the checkout's venv: + +```bash +uv pip install nvidia-cuda-runtime # CUDA 13; use nvidia-cuda-runtime-cu12 for CUDA 12 +``` + +Without it only rung 1 (already-mapped libraries) can hit, and on a machine with +no system CUDA install nothing resolves at all — the pinning tests then report +that they skipped rather than failing. The wheel is deliberately *not* a `dev` +extra: CI is CPU-only and should not download a CUDA runtime. + +The Rust tests find the wheel by looking under `.venv/lib/python*/site-packages/` +in the checkout, since `cargo test` has no Python interpreter to ask. + ## Benchmarks ```bash diff --git a/docs/src/guides/host-memory-pinning.md b/docs/src/guides/host-memory-pinning.md new file mode 100644 index 0000000..eb57ba4 --- /dev/null +++ b/docs/src/guides/host-memory-pinning.md @@ -0,0 +1,262 @@ +# 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 — the page-locked memory is not +swappable. This page covers what the mechanism actually is, what it measured, +how to size the footprint, and — the part that is easy to get wrong — **how to +confirm it engaged**. + +```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: + - 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. +``` + +That is deliberately all the API there is: no status object, no report method, no +warning. Construction succeeding *is* the assertion, and unlike an explicit check +you might add, it cannot be forgotten. 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. + +## The mechanism: 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. Two things +follow, and the second is usually the expensive one: + +- 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 waiting for work. + +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. + +## Measured + +Local microbenchmark on a **NVIDIA GeForce RTX 5060 Ti (16 GB), driver +595.71.05, CUDA runtime 13.3**, sized like a large-batch learner: +`batch_size=512` over four arrays (34.1 MB per batch, 102.2 MB ring), with the +copy measurement isolated to a single 52.4 MB buffer. 11 repeats; median and +interquartile range: + +Through echo's own ring buffers and sampled views (34.1 MB batch, 102.2 MB ring): + +| arm | copy time | IQR | throughput | host-write | registration | +|---|---|---|---|---|---| +| pageable | 2.441 ms | 2.440–2.442 | 14.0 GB/s | 12.32 GB/s | – | +| page-locked (portable) | 2.401 ms | 2.399–2.407 | 14.2 GB/s | 12.39 GB/s | 8.9 ms | + +And isolating one 52.4 MB buffer, to compare against driver-allocated pinned +memory and to time the copy call on its own: + +| source memory | copy time | IQR | throughput | `cudaMemcpyAsync` returns after | +|---|---|---|---|---| +| pageable | 3.751 ms | 3.748–3.752 | 14.0 GB/s | 3.645 ms — **97% of the copy** | +| page-locked | 3.715 ms | 3.714–3.719 | 14.1 GB/s | 0.002 ms — **0%** | +| `cudaHostAlloc` (reference ceiling) | 3.658 ms | 3.657–3.658 | 14.3 GB/s | 0.002 ms | + +Read the throughput and the return-latency columns separately, because they say +different things. + +**Bandwidth barely moves on this machine, and that is expected.** Both paths +saturate the host's link at ~14 GB/s; driver-allocated pinned memory only reaches +14.3 GB/s, so there is nothing more to win here. A host with more PCIe headroom +will show a larger gap. + +**Host-thread occupancy collapses by ~1800x.** That is the mechanism above, +measured: the same call goes from holding the calling thread for 3.645 ms to +0.002 ms. This is the component that scales into the driver-lock contention a +large-batch learner suffers, and it is the reason to turn pinning on. + +**Nothing on the write side pays for it.** Host-write throughput into the ring — +what a drainer costs — is unchanged at 12.3 GB/s, and registering the 102 MB ring +took 8.9 ms once, at construction. Pinning changes only how the memory is mapped, +not any code on the ingest or sample path, so there is no per-sample or per-drain +cost. A batch sampled with pinning on is bit-identical to the same batch with it +off. + +**Only the portable flag is used.** Two variants were measured and rejected. +Page-aligning the ring buffers gained 1.6% of copy time and nothing at all on +return latency, well under the bar set in advance, so the buffers stay plain +`Vec`. Read-only registration (`cudaHostRegisterReadOnly`) is not usable +here at all — `cudaHostRegister` returns `cudaErrorNotSupported` on this GPU, +whose `cudaDevAttrHostRegisterReadOnlySupported` is 0 — and the CUDA +documentation describes that flag as permission to register memory *mapped* +read-only rather than as a transfer optimisation, while saying nothing about host +writes to such a range. Echo's drainers write these pages continuously, so there +would be no documented basis for using it even where it is supported. + +**What this microbenchmark cannot tell you.** A single consumer GPU cannot +reproduce the driver-lock contention of a real learner issuing tens of thousands +of small kernel launches concurrently with the staging copy. The table above +measures the bandwidth component and the mechanism; it does *not* size the win on +a large-batch learner, where the contention component dominates and the effect is +correspondingly larger than the 1–3% bandwidth figure here. + +So do not read 1–3% as "pinning is worth 1–3%", and do not read the occupancy +column as a step-time prediction either. Measure your own workload: turn pinning +on, confirm it engaged (below), and compare step times. + +To reproduce the table: allocate a buffer with `malloc`, `cudaMalloc` a +destination, and time `cudaMemcpyAsync` + `cudaStreamSynchronize` before and +after `cudaHostRegister(ptr, size, cudaHostRegisterPortable)`. Time the +`cudaMemcpyAsync` call *on its own*, without the synchronize, to see the +occupancy column. + +## 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. + +Two things to hold onto: + +- **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. Size the host before the job does it + for you. + +Note that `VmLck` in `/proc/self/status` stays at **zero** even when pinning is +working, so a memory-lock resource limit (`ulimit -l`) does not bind here — see +[below](#verifying-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. That frees nothing and +never *changes* which device is current, so it cannot perturb your framework's +device selection. + +It is not entirely free, though, and this is the reason the ordering above is a +recommendation rather than a footnote: initialising the runtime creates the +primary CUDA context on whatever device is already current, and that context +costs device memory (~128 MB on the machine in the table above). Construct after +your framework has initialised CUDA and the context already exists, so echo adds +nothing. Construct before it, and echo creates the context first — which both +spends that memory early and may interact badly with a framework that +pre-allocates a fraction of *free* device memory. + +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. + +## Verifying that pinning engaged + +This section is the reason this page exists. Pinning was once believed not to +help a workload that had in fact never executed the pinning code path, and the +knowledge of how to check lived in one person's head. + +**In-process: the constructor.** `Server(..., pin_host_memory=True)` returning +*is* the assertion. There is nothing else to query. + +**From a profile: the copy's memory-source classification.** Profile the learner +and look at the host-to-device memcpy rows. A profiler reports the source kind +for each transfer; it must say the source is pinned/page-locked rather than +pageable. This is the authoritative external signal. + +**From a test, without a framework.** Ask the CUDA runtime directly for the flags +on the address behind a sampled batch — this is what echo's own test suite does, +precisely so that a test cannot pass merely because echo believes it worked: + +```python +import ctypes + +cudart = ctypes.CDLL("libcudart.so.13") +flags = ctypes.c_uint(0) +code = cudart.cudaHostGetFlags( + ctypes.byref(flags), ctypes.c_void_p(batch["obs"].ctypes.data) +) +assert code == 0 # 0 = cudaSuccess; non-zero means not registered +assert flags.value & 0x01 # cudaHostRegisterPortable +``` + +!!! warning "`VmLck` is not a valid check" + + `VmLck` in `/proc/self/status` stays at **zero** even when pinning + demonstrably works: the NVIDIA driver's page-locking does not go through + mlock accounting. Anyone who reads it as a check will conclude pinning is off + when it is on. The same goes for any tool built on mlock accounting. Use the + profile trace or `cudaHostGetFlags`. + +## 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. +- **Anything other than the ring buffers.** Producer queues, transport staging + buffers and accumulator storage are not registered: the ring buffers are the + only source of a host-to-device copy. +- **Allocating pinned memory directly** instead of registering the existing + buffers. Steady-state transfer performance would be identical, so it would + change only construction cost while making the CUDA runtime mandatory at + allocation time. diff --git a/mkdocs.yml b/mkdocs.yml index e32abae..201e682 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 diff --git a/pyproject.toml b/pyproject.toml index a28344f..81bb457 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "id-echo" -version = "0.1.2" +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..3e4faf3 100644 --- a/python/echo/server.py +++ b/python/echo/server.py @@ -28,6 +28,27 @@ 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 + 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. + + **If this constructor returns, every ring buffer is page-locked.** + Anything that would prevent that — no CUDA runtime found, no usable + device, a registration rejected — raises ``RuntimeError`` here, + naming every path probed and the CUDA error. There is no silent + fallback, so a throughput measurement can never be invalidated by + pinning having quietly done nothing. + + Construct the server *after* your framework has initialised CUDA. + The page-locked footprint is the full ring + (``batch_size * num_buffers * bytes_per_sample``), is not + swappable, and multiplies by the number of servers in the process. + See the [host-memory pinning guide](../guides/host-memory-pinning.md). + + 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 +66,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 +86,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..00e4462 --- /dev/null +++ b/python/tests/test_host_pinning.py @@ -0,0 +1,271 @@ +"""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 + +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 += ["libcudart.so.13", "libcudart.so.12", "libcudart.so"] + 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(HAS_GPU, reason="a usable CUDA device is present") +def test_requesting_pinning_without_cuda_raises_and_lists_probed_paths(): + """The failure a misconfigured deployment gets: loud, at startup, specific.""" + with pytest.raises(RuntimeError) as excinfo: + Server(EXAMPLE, batch_size=4, pin_host_memory=True) + + message = str(excinfo.value) + assert "libcudart.so" in message, message + # Every rung reports what it tried, so the user can fix their environment + # without reading echo's source. + assert "soname load" 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.rs b/src/host_pinning.rs index 72d697e..99dfdb3 100644 --- a/src/host_pinning.rs +++ b/src/host_pinning.rs @@ -1,83 +1,1012 @@ //! Optional CUDA host-memory pinning (page-locking) of the ring buffers. //! -//! Enabled at runtime by `ECHO_PIN_HOST_MEMORY=1` (off by default). When enabled, -//! [`pin`] registers a buffer with `cudaHostRegister` so a downstream -//! `jax.device_put` of its numpy view becomes a fast, truly-async H2D DMA instead -//! of a synchronous host->device staging copy. libcudart is resolved at runtime -//! via `dlopen`, so this carries NO CUDA toolchain / link-time dependency: the -//! code is always compiled but is a graceful no-op when the env var is unset, on -//! CUDA-less hosts, or if registration fails. +//! A host-to-device copy out of *pageable* memory is not the DMA transfer it +//! looks like: the driver stages it through a small internal pinned buffer in +//! CPU-executed chunks — thousands of driver round-trips for a large batch, +//! which both moves bytes slower and holds the driver lock long enough to +//! starve the consumer thread's own kernel launches. Page-locking the ring +//! buffers turns the same copy into one DMA descriptor on the copy engine. +//! +//! Off unless the caller asks for it, and when asked this module either +//! delivers fully page-locked buffers or says why it could not. There is +//! deliberately no partial-success or best-effort path: a silent no-op is +//! indistinguishable from "pinning doesn't help this workload", and that +//! confusion is the defect this module was rewritten to remove. +//! +//! The runtime is resolved by `dlopen` at pin time, so echo carries no +//! build-time or link-time CUDA dependency and one wheel installs on GPU and +//! CPU-only hosts alike. Nothing here runs unless pinning is requested. +use std::ffi::{CStr, CString}; +use std::fmt; use std::os::raw::{c_char, c_int, c_uint, c_void}; +use std::path::PathBuf; use std::sync::OnceLock; -// cudaError_t cudaHostRegister(void* ptr, size_t size, unsigned int flags) +/// `cudaError_t cudaHostRegister(void*, size_t, unsigned int)` type RegisterFn = unsafe extern "C" fn(*mut c_void, usize, c_uint) -> c_int; -// cudaError_t cudaHostUnregister(void* ptr) +/// `cudaError_t cudaHostUnregister(void*)` type UnregisterFn = unsafe extern "C" fn(*mut c_void) -> c_int; +/// `const char* cudaGetErrorName(cudaError_t)` +type ErrorNameFn = unsafe extern "C" fn(c_int) -> *const c_char; +/// `cudaError_t cudaFree(void*)` +type FreeFn = unsafe extern "C" fn(*mut c_void) -> c_int; -struct Api { +const CUDA_SUCCESS: c_int = 0; + +/// `cudaHostRegisterPortable`: the locked pages are valid in every CUDA context +/// in the process, including contexts created after registration. +/// +/// This is why echo needs no device parameter and never selects a device. N +/// servers across N GPUs in one process each register portably, and every +/// registration is DMA-fast for every device regardless of which context +/// happened to be current at the time. +const CUDA_HOST_REGISTER_PORTABLE: c_uint = 0x01; + +/// The CUDA entry points pinning needs, as plain function pointers. +/// +/// Passed explicitly to [`pin_all`] / [`unpin_all`] rather than reached through +/// a process-global so that registration and rollback can be exercised with +/// stubs on a machine with no GPU — the rollback path is unsafe code whose only +/// externally visible consequence is the *absence* of leaked registrations. +#[derive(Clone, Copy)] +pub(crate) struct CudaApi { register: RegisterFn, unregister: UnregisterFn, + error_name: ErrorNameFn, + free: FreeFn, } -// The dlopen handle outlives the process; the resolved fn pointers are immutable. -unsafe impl Sync for Api {} -unsafe impl Send for Api {} -fn load() -> Option { - // Opt-in: only attempt to pin when explicitly enabled. - if std::env::var("ECHO_PIN_HOST_MEMORY").as_deref() != Ok("1") { - return None; +/// A contiguous host allocation to page-lock. +/// +/// Just a (pointer, length) pair — it owns nothing and asserts nothing. The +/// invariants that make registering it sound live on [`pin_all`] and +/// [`unpin_all`], which are `unsafe` for that reason. +#[derive(Clone, Copy)] +pub(crate) struct Region { + pub ptr: *mut u8, + pub len: usize, +} + +/// Why pinning could not be delivered. Never returned for a partial success: +/// on failure nothing is left registered. +#[derive(Debug)] +pub(crate) enum PinError { + /// No CUDA runtime could be loaded. Carries one line per probed path so a + /// caller can fix their environment without reading echo's 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." + ), + } } - unsafe { - let handle = libc::dlopen( - b"libcudart.so\0".as_ptr() as *const c_char, - libc::RTLD_NOW | libc::RTLD_GLOBAL, - ); - if handle.is_null() { - return None; +} + +impl std::error::Error for PinError {} + +// --------------------------------------------------------------------------- +// Resolving the CUDA runtime +// --------------------------------------------------------------------------- + +/// Versioned sonames for the current and previous CUDA major versions, then the +/// unversioned one. +/// +/// The unversioned soname is last and on its own is near-useless: the pip CUDA +/// runtime wheels ship only the versioned soname, with no unversioned symlink +/// and no ldconfig entry. Opening `libcudart.so` alone is what made this +/// feature a silent no-op for months. +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)] +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 { + 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)] +struct Candidate { + /// An absolute path (rungs 1 and 2) or a bare soname (rung 3). + name: String, + 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: CUDA runtimes the process already has mapped, in first-seen order. +/// +/// Version- and path-agnostic, and it hits whenever the framework has already +/// initialised CUDA, which is the common case. `dlopen`-ing the absolute path +/// of an already-mapped library returns a handle to that same mapping rather +/// than loading a second copy. +fn scan_mapped_runtimes(maps: &str) -> Vec { + let mut found: Vec = Vec::new(); + for line in maps.lines() { + // The pathname is the last field of a `/proc//maps` line and is + // absolute, so it starts at the first " /" — found this way rather than + // by splitting on whitespace because paths may contain spaces. An + // unlinked file gets 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()); } - let register = libc::dlsym(handle, b"cudaHostRegister\0".as_ptr() as *const c_char); - let unregister = libc::dlsym(handle, b"cudaHostUnregister\0".as_ptr() as *const c_char); - if register.is_null() || unregister.is_null() { - return None; + } + found +} + +/// The deepest a runtime library sits below the vendor package +/// (`nvidia/cu13/lib/libcudart.so.13` is three), plus room for a +/// reorganisation. +const WHEEL_SEARCH_DEPTH: usize = 4; + +/// Rung 2: CUDA runtimes shipped by installed wheels, newest major first. +/// +/// `roots` are the CUDA *vendor* package directories, located through Python's +/// import machinery by the caller. Searching the whole vendor package instead +/// of a named component subpackage is load-bearing, not lazy: CUDA 13 ships one +/// consolidated wheel laid out as `nvidia/cu13/lib/`, whereas CUDA 12 ships one +/// wheel per component laid out as `nvidia/cuda_runtime/lib/`. A lookup that +/// names the component finds nothing on a CUDA 13 install. Searching beneath +/// the vendor package covers both layouts and survives the next one. +/// +/// This rung is what lets echo be constructed before the framework has loaded +/// CUDA, and it removes any need for loader-path environment variables. +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(); + 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 resolution ladder, in the order it will be tried. +/// +/// Kept as one pure function of its inputs so that the ordering — the thing a +/// future edit could silently break, reintroducing the original bug — is +/// unit-testable without a GPU or a CUDA install. +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 }); + } } - Some(Api { - register: std::mem::transmute::<*mut c_void, RegisterFn>(register), - unregister: std::mem::transmute::<*mut c_void, UnregisterFn>(unregister), + } + ladder +} + +/// `dlopen` a candidate and resolve the four symbols pinning needs. +/// +/// The handle is deliberately never `dlclose`d: the registrations it backs must +/// outlive it, and the runtime is process-wide state anyway. +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 rather than RTLD_GLOBAL: echo should 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 asks the loader for an already-present library first, so a + // bare soname prefers a copy the process has (under whatever path) + // over pulling in a second one 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 { + register: std::mem::transmute::<*mut c_void, RegisterFn>(symbol( + handle, + c"cudaHostRegister", + )?), + unregister: std::mem::transmute::<*mut c_void, UnregisterFn>(symbol( + handle, + c"cudaHostUnregister", + )?), + error_name: std::mem::transmute::<*mut c_void, ErrorNameFn>(symbol( + handle, + c"cudaGetErrorName", + )?), + free: std::mem::transmute::<*mut c_void, FreeFn>(symbol(handle, c"cudaFree")?), }) } } -fn api() -> Option<&'static Api> { - static API: OnceLock> = OnceLock::new(); - API.get_or_init(load).as_ref() +/// # 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(), + ) } -/// Page-lock `[ptr, ptr + len)` (portable across the process's CUDA contexts). +/// 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 rather than staying silent. +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(", ") + ), + 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. /// -/// No-op unless `ECHO_PIN_HOST_MEMORY=1` and libcudart is present. Best-effort: -/// registration errors leave the memory pageable. -pub(crate) fn pin(ptr: *mut u8, len: usize) { - if len == 0 { - return; - } - if let Some(api) = api() { - // cudaHostRegisterPortable = 0x01 (usable from every device/context). - unsafe { - (api.register)(ptr as *mut c_void, len, 0x01); +/// Only ever called when a caller asked for pinning, so the default-off path +/// loads nothing. A failure is not cached: a process that constructs a server +/// before its framework has initialised CUDA and retries later should get the +/// later, better answer. +pub(crate) 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)) +} + +// --------------------------------------------------------------------------- +// Registration +// --------------------------------------------------------------------------- + +/// Page-lock every region, or leave none of them locked. +/// +/// On failure the regions registered so far are unregistered before returning, +/// so a caller whose construction fails leaves nothing behind for a retry or a +/// long-lived process to accumulate. 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(crate) unsafe fn pin_all(api: &CudaApi, regions: &[Region]) -> Result<(), PinError> { + // Force runtime initialisation, best-effort: registration against an + // uninitialised runtime fails, and this is what gives the runtime a reason + // to set itself up. Freeing a null pointer frees nothing and does not + // *change* the current device — but note it is not free of consequence: it + // creates the primary CUDA context on whatever device is already current, + // which costs that context's device memory (order of 100 MB). When the + // caller follows the documented order and constructs the server after the + // framework has initialised CUDA, the context already exists and this costs + // nothing. + // + // The result is ignored on purpose: portable registration makes the choice + // of device irrelevant, so the only failure that matters is registration's + // own, which reports the real CUDA error a few lines 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(crate) 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 a reader can look +/// it up instead of decoding 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() } -/// Reverse of [`pin`]; must run before the buffer is freed. -pub(crate) fn unpin(ptr: *mut u8, len: usize) { - if len == 0 { - return; +#[cfg(test)] +mod tests { + use super::*; + use std::cell::RefCell; + + // --- rung 1: parsing the process's own mapped files --- + + /// Representative `/proc/self/maps` content: a versioned runtime mapped + /// several times (one line per segment), unrelated libraries, anonymous and + /// special mappings, a path containing a space, 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()); } - if let Some(api) = api() { + + // --- 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 sonames before the unversioned one: the pip CUDA + // runtime wheels ship no unversioned symlink, which is the + // original defect this ordering exists to prevent. + "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"); + } + + // --- 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); + CudaApi { + register, + unregister: stub_unregister, + error_name: stub_error_name, + free: stub_free, + } + } + + /// Dangling but never dereferenced: `pin_all` only passes them to the 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() { + // The success path's teardown, without a GPU: a real ring buffer, stubbed + // CUDA. Registrations leaked here would be invisible from Python, so this + // is checked through the injection point rather than externally. + let api = stub_api(stub_register_ok); + let mut ring = crate::ring_buf::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(crate::ring_buf::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() { + // Rollback across the real seam: 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 = crate::ring_buf::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, so rung 2 is exercisable from + /// `cargo test` (which 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*)` — the runtime's own + /// account of how an address is registered. + type HostGetFlagsFn = unsafe extern "C" fn(*mut c_uint, *mut c_void) -> c_int; + + /// Resolve the runtime, or `None` on a machine without one. Rust has no + /// native test skip, so these tests 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. + /// Confirming registration through echo's own function pointers would only + /// prove echo believes it succeeded. + /// + /// Call only after resolution has run, so the runtime is mapped and this + /// re-opens it (`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 { - (api.unregister)(ptr as *mut c_void); + 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 = symbol(handle, c"cudaHostGetFlags") + .expect("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 = crate::ring_buf::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; the same address must no longer be known to the + // runtime as registered host memory. + 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 must each pin with no extra configuration. + if real_api().is_none() { + return; } + let roots = dev_vendor_roots(); + let mut first = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); + let mut second = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); + first.pin_host_memory(&roots).expect("first ring failed"); + second.pin_host_memory(&roots).expect("second ring failed"); + } + + #[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}"); } } diff --git a/src/py_bindings.rs b/src/py_bindings.rs index eb706fb..48656bf 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,18 @@ impl PyServer { Some(metrics.clone()), )), Box::new(FifoRemover::new()), - )); + ); + + // The page-locked guarantee: if this constructor returns, every ring + // buffer is registered. Anything short of that raises, so a caller can + // never measure a workload against a pinning path that silently did + // nothing. Nothing is loaded at all when the caller didn't ask. + if pin_host_memory { + store + .pin_host_memory(&cuda_vendor_roots(py)) + .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 +265,30 @@ impl PyServer { } } +/// CUDA vendor package directories, located through Python's import machinery. +/// +/// Rung 2 of the CUDA resolution ladder searches beneath these, which is how +/// pinning works with a pip-installed CUDA runtime and without any loader-path +/// environment variable. `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; the other +/// two rungs still run, and total failure reports every path probed. +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 dcfca56..1d1f541 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,10 @@ pub struct PytreeRingBuf { slot_bytes: Vec, /// Total number of slots. capacity: usize, + /// The CUDA runtime the buffers are registered with, once they are. + /// `Some` is exactly the condition for `Drop` having registrations to + /// reverse, and holding it here keeps `Drop` off any process-global. + pinned_with: Option, } impl PytreeRingBuf { @@ -39,23 +46,63 @@ impl PytreeRingBuf { .map(|&bytes| UnsafeCell::new(vec![0u8; bytes * capacity])) .collect(); - // Optionally page-lock each buffer (ECHO_PIN_HOST_MEMORY=1) so that a - // downstream `jax.device_put` of the numpy views is a fast, truly-async - // H2D DMA rather than a synchronous host->device staging copy. No-op - // without CUDA. Buffers are contiguous and never reallocated, so the - // registration stays valid for the buffer's life. - for cell in &buffers { - let buf = unsafe { &mut *cell.get() }; - crate::host_pinning::pin(buf.as_mut_ptr(), buf.len()); - } - Self { buffers, slot_bytes, capacity, + pinned_with: None, } } + /// Page-lock every buffer, so a downstream host-to-device copy of a sampled + /// view is a DMA transfer rather than a chunked staging copy. + /// + /// Either every buffer ends up locked or none does — a partial failure is + /// rolled back before the error returns. That is why this is a separate + /// fallible step on a constructed buffer rather than part of `new`: 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 for both the failure path and normal teardown. + /// + /// `cuda_vendor_roots` are the CUDA vendor package directories located + /// through Python's import machinery; see [`crate::host_pinning`]. + /// + /// The buffers are contiguous and never reallocated, so a registration stays + /// valid for the buffer's whole life. + pub(crate) 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 stubbed CUDA + /// entry points, on a machine with no GPU. + pub(crate) fn pin_with(&mut self, api: CudaApi) -> Result<(), PinError> { + // SAFETY: the regions are this buffer's own allocations, which live as + // long as `self` and are 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 } @@ -130,11 +177,13 @@ unsafe impl Send for PytreeRingBuf {} impl Drop for PytreeRingBuf { fn drop(&mut self) { - // Unregister (un-pin) before the backing Vecs are freed. Drop::drop runs - // before the struct's fields are dropped, so the memory is still valid. - for cell in &self.buffers { - let buf = unsafe { &mut *cell.get() }; - crate::host_pinning::unpin(buf.as_mut_ptr(), buf.len()); - } + let Some(api) = self.pinned_with else { + return; + }; + // Unregister before the backing Vecs are freed. Drop::drop runs before + // the struct's fields are dropped, so the memory is still valid. + // SAFETY: `pinned_with` is Some only after `pin_with` registered exactly + // these regions through this same api, and they are still valid here. + unsafe { host_pinning::unpin_all(&api, &self.regions()) }; } } diff --git a/src/store.rs b/src/store.rs index 09743fa..b11e6f0 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,17 @@ impl Store { } } + /// Page-lock the ring buffers; see [`PytreeRingBuf::pin_host_memory`]. + /// + /// `&mut self` puts this between construction and sharing the store, so a + /// caller that then fails to finish building unregisters by dropping it. + pub(crate) 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/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'" }, From 34713a4b47aeb8d43b0ce2691a5b5f933b055f03 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Fri, 31 Jul 2026 14:43:38 +0200 Subject: [PATCH 05/15] refactor: move host_pinning tests to their own file, trim comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every other module keeps its tests under tests/; host_pinning was the only one with an in-file #[cfg(test)] module, and at 1012 lines was three times the next-largest source file. The tests move to src/host_pinning/tests.rs — in-crate rather than under tests/ because the module is private and the stub-injection tests need the internal CudaApi — leaving 488 lines of source. Also cuts the comments back towards the density of the surrounding code. The load-bearing "why" stays (the pageable staging mechanism, why the unversioned soname is not enough, why the vendor package is searched rather than a named component, why pinning is a post-construction step, what forcing runtime init actually costs); the restatements of the guide and the narrative asides go. Co-Authored-By: Claude Opus 5 (1M context) --- src/host_pinning.rs | 648 ++++---------------------------------- src/host_pinning/tests.rs | 483 ++++++++++++++++++++++++++++ src/py_bindings.rs | 20 +- src/ring_buf.rs | 42 +-- src/store.rs | 4 +- 5 files changed, 572 insertions(+), 625 deletions(-) create mode 100644 src/host_pinning/tests.rs diff --git a/src/host_pinning.rs b/src/host_pinning.rs index 99dfdb3..f1b69ff 100644 --- a/src/host_pinning.rs +++ b/src/host_pinning.rs @@ -1,21 +1,16 @@ //! Optional CUDA host-memory pinning (page-locking) of the ring buffers. //! -//! A host-to-device copy out of *pageable* memory is not the DMA transfer it -//! looks like: the driver stages it through a small internal pinned buffer in -//! CPU-executed chunks — thousands of driver round-trips for a large batch, -//! which both moves bytes slower and holds the driver lock long enough to -//! starve the consumer thread's own kernel launches. Page-locking the ring -//! buffers turns the same copy into one DMA descriptor on the copy engine. +//! 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. //! -//! Off unless the caller asks for it, and when asked this module either -//! delivers fully page-locked buffers or says why it could not. There is -//! deliberately no partial-success or best-effort path: a silent no-op is -//! indistinguishable from "pinning doesn't help this workload", and that -//! confusion is the defect this module was rewritten to remove. +//! 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 resolved by `dlopen` at pin time, so echo carries no -//! build-time or link-time CUDA dependency and one wheel installs on GPU and -//! CPU-only hosts alike. Nothing here runs unless pinning is requested. +//! 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. use std::ffi::{CStr, CString}; use std::fmt; @@ -34,21 +29,15 @@ type FreeFn = unsafe extern "C" fn(*mut c_void) -> c_int; const CUDA_SUCCESS: c_int = 0; -/// `cudaHostRegisterPortable`: the locked pages are valid in every CUDA context -/// in the process, including contexts created after registration. -/// -/// This is why echo needs no device parameter and never selects a device. N -/// servers across N GPUs in one process each register portably, and every -/// registration is DMA-fast for every device regardless of which context -/// happened to be current at the time. +/// `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. const CUDA_HOST_REGISTER_PORTABLE: c_uint = 0x01; -/// The CUDA entry points pinning needs, as plain function pointers. +/// The CUDA entry points pinning needs. /// -/// Passed explicitly to [`pin_all`] / [`unpin_all`] rather than reached through -/// a process-global so that registration and rollback can be exercised with -/// stubs on a machine with no GPU — the rollback path is unsafe code whose only -/// externally visible consequence is the *absence* of leaked registrations. +/// Passed explicitly rather than reached through a global so tests can inject +/// stubs and assert rollback without a GPU. #[derive(Clone, Copy)] pub(crate) struct CudaApi { register: RegisterFn, @@ -57,23 +46,19 @@ pub(crate) struct CudaApi { free: FreeFn, } -/// A contiguous host allocation to page-lock. -/// -/// Just a (pointer, length) pair — it owns nothing and asserts nothing. The -/// invariants that make registering it sound live on [`pin_all`] and -/// [`unpin_all`], which are `unsafe` for that reason. +/// A contiguous host allocation to page-lock. Owns nothing; the validity +/// invariants live on [`pin_all`] / [`unpin_all`]. #[derive(Clone, Copy)] pub(crate) struct Region { pub ptr: *mut u8, pub len: usize, } -/// Why pinning could not be delivered. Never returned for a partial success: -/// on failure nothing is left registered. +/// Why pinning could not be delivered. On failure nothing is left registered. #[derive(Debug)] pub(crate) enum PinError { - /// No CUDA runtime could be loaded. Carries one line per probed path so a - /// caller can fix their environment without reading echo's source. + /// 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 { @@ -124,13 +109,10 @@ impl std::error::Error for PinError {} // Resolving the CUDA runtime // --------------------------------------------------------------------------- -/// Versioned sonames for the current and previous CUDA major versions, then the -/// unversioned one. +/// Current and previous CUDA major versions, then the unversioned soname. /// -/// The unversioned soname is last and on its own is near-useless: the pip CUDA -/// runtime wheels ship only the versioned soname, with no unversioned symlink -/// and no ldconfig entry. Opening `libcudart.so` alone is what made this -/// feature a silent no-op for months. +/// 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. @@ -170,19 +152,17 @@ fn is_runtime_lib(path: &str) -> bool { name == "libcudart.so" || name.starts_with("libcudart.so.") } -/// Rung 1: CUDA runtimes the process already has mapped, in first-seen order. +/// Rung 1: runtimes the process already has mapped, in first-seen order. /// -/// Version- and path-agnostic, and it hits whenever the framework has already -/// initialised CUDA, which is the common case. `dlopen`-ing the absolute path -/// of an already-mapped library returns a handle to that same mapping rather -/// than loading a second copy. +/// 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. fn scan_mapped_runtimes(maps: &str) -> Vec { let mut found: Vec = Vec::new(); for line in maps.lines() { - // The pathname is the last field of a `/proc//maps` line and is - // absolute, so it starts at the first " /" — found this way rather than - // by splitting on whitespace because paths may contain spaces. An - // unlinked file gets a " (deleted)" suffix. + // 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; @@ -195,23 +175,17 @@ fn scan_mapped_runtimes(maps: &str) -> Vec { found } -/// The deepest a runtime library sits below the vendor package -/// (`nvidia/cu13/lib/libcudart.so.13` is three), plus room for a -/// reorganisation. +/// `nvidia/cu13/lib/libcudart.so.13` is three deep; one spare for a relayout. const WHEEL_SEARCH_DEPTH: usize = 4; -/// Rung 2: CUDA runtimes shipped by installed wheels, newest major first. -/// -/// `roots` are the CUDA *vendor* package directories, located through Python's -/// import machinery by the caller. Searching the whole vendor package instead -/// of a named component subpackage is load-bearing, not lazy: CUDA 13 ships one -/// consolidated wheel laid out as `nvidia/cu13/lib/`, whereas CUDA 12 ships one -/// wheel per component laid out as `nvidia/cuda_runtime/lib/`. A lookup that -/// names the component finds nothing on a CUDA 13 install. Searching beneath -/// the vendor package covers both layouts and survives the next one. +/// Rung 2: runtimes shipped by installed wheels, newest major first. Lets the +/// server be constructed before the framework has loaded CUDA. /// -/// This rung is what lets echo be constructed before the framework has loaded -/// CUDA, and it removes any need for loader-path environment variables. +/// `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. 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(); @@ -263,11 +237,8 @@ fn soname_major(name: &str) -> Option { .ok() } -/// The full resolution ladder, in the order it will be tried. -/// -/// Kept as one pure function of its inputs so that the ordering — the thing a -/// future edit could silently break, reintroducing the original bug — is -/// unit-testable without a GPU or a CUDA install. +/// 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. fn candidates(maps: &str, vendor_roots: &[PathBuf]) -> Vec { let rungs = [ (Rung::AlreadyLoaded, scan_mapped_runtimes(maps)), @@ -293,19 +264,17 @@ fn candidates(maps: &str, vendor_roots: &[PathBuf]) -> Vec { /// `dlopen` a candidate and resolve the four symbols pinning needs. /// -/// The handle is deliberately never `dlclose`d: the registrations it backs must -/// outlive it, and the runtime is process-wide state anyway. +/// 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 rather than RTLD_GLOBAL: echo should not change how any other - // library in the process resolves its symbols. + // 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 asks the loader for an already-present library first, so a - // bare soname prefers a copy the process has (under whatever path) - // over pulling in a second one from the loader path. + // 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 { @@ -395,7 +364,7 @@ fn resolve(vendor_roots: &[PathBuf]) -> Result { } /// What to report for a rung that contributed no candidate, so the error says -/// where it looked rather than staying silent. +/// 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(), @@ -417,12 +386,11 @@ fn empty_rung_note(rung: Rung, vendor_roots: &[PathBuf]) -> String { static API: OnceLock = OnceLock::new(); -/// The CUDA runtime, resolving it on first use. +/// The CUDA runtime, resolving it on first use. Only called when pinning was +/// asked for, so the default-off path loads nothing. /// -/// Only ever called when a caller asked for pinning, so the default-off path -/// loads nothing. A failure is not cached: a process that constructs a server -/// before its framework has initialised CUDA and retries later should get the -/// later, better answer. +/// Failures are not cached: a process that retries after its framework has +/// initialised CUDA should get the later, better answer. pub(crate) fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError> { if let Some(api) = API.get() { return Ok(api); @@ -439,9 +407,8 @@ pub(crate) fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError /// Page-lock every region, or leave none of them locked. /// -/// On failure the regions registered so far are unregistered before returning, -/// so a caller whose construction fails leaves nothing behind for a retry or a -/// long-lived process to accumulate. Normal teardown is `Drop`'s job. +/// 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. @@ -451,19 +418,14 @@ pub(crate) fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError /// - Each region registered here must be passed to [`unpin_all`] before its /// memory is freed. pub(crate) unsafe fn pin_all(api: &CudaApi, regions: &[Region]) -> Result<(), PinError> { - // Force runtime initialisation, best-effort: registration against an - // uninitialised runtime fails, and this is what gives the runtime a reason - // to set itself up. Freeing a null pointer frees nothing and does not - // *change* the current device — but note it is not free of consequence: it - // creates the primary CUDA context on whatever device is already current, - // which costs that context's device memory (order of 100 MB). When the - // caller follows the documented order and constructs the server after the - // framework has initialised CUDA, the context already exists and this costs - // nothing. + // 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 (~100 MB). Constructing after the framework has + // initialised CUDA — the documented order — means it already exists. // - // The result is ignored on purpose: portable registration makes the choice - // of device irrelevant, so the only failure that matters is registration's - // own, which reports the real CUDA error a few lines below. + // 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() { @@ -510,8 +472,8 @@ pub(crate) unsafe fn unpin_all(api: &CudaApi, regions: &[Region]) { } } -/// A CUDA error as its symbol (`cudaErrorInvalidValue`), so a reader can look -/// it up instead of decoding an integer. +/// 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() { @@ -523,490 +485,4 @@ fn error_name(api: &CudaApi, code: c_int) -> String { } #[cfg(test)] -mod tests { - use super::*; - use std::cell::RefCell; - - // --- rung 1: parsing the process's own mapped files --- - - /// Representative `/proc/self/maps` content: a versioned runtime mapped - /// several times (one line per segment), unrelated libraries, anonymous and - /// special mappings, a path containing a space, 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 sonames before the unversioned one: the pip CUDA - // runtime wheels ship no unversioned symlink, which is the - // original defect this ordering exists to prevent. - "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"); - } - - // --- 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); - CudaApi { - register, - unregister: stub_unregister, - error_name: stub_error_name, - free: stub_free, - } - } - - /// Dangling but never dereferenced: `pin_all` only passes them to the 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() { - // The success path's teardown, without a GPU: a real ring buffer, stubbed - // CUDA. Registrations leaked here would be invisible from Python, so this - // is checked through the injection point rather than externally. - let api = stub_api(stub_register_ok); - let mut ring = crate::ring_buf::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(crate::ring_buf::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() { - // Rollback across the real seam: 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 = crate::ring_buf::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, so rung 2 is exercisable from - /// `cargo test` (which 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*)` — the runtime's own - /// account of how an address is registered. - type HostGetFlagsFn = unsafe extern "C" fn(*mut c_uint, *mut c_void) -> c_int; - - /// Resolve the runtime, or `None` on a machine without one. Rust has no - /// native test skip, so these tests 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. - /// Confirming registration through echo's own function pointers would only - /// prove echo believes it succeeded. - /// - /// Call only after resolution has run, so the runtime is mapped and this - /// re-opens it (`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 = symbol(handle, c"cudaHostGetFlags") - .expect("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 = crate::ring_buf::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; the same address must no longer be known to the - // runtime as registered host memory. - 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 must each pin with no extra configuration. - if real_api().is_none() { - return; - } - let roots = dev_vendor_roots(); - let mut first = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); - let mut second = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); - first.pin_host_memory(&roots).expect("first ring failed"); - second.pin_host_memory(&roots).expect("second ring failed"); - } - - #[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}"); - } -} +mod tests; diff --git a/src/host_pinning/tests.rs b/src/host_pinning/tests.rs new file mode 100644 index 0000000..cd8e13c --- /dev/null +++ b/src/host_pinning/tests.rs @@ -0,0 +1,483 @@ +//! Unit tests for the private `host_pinning` module. +//! +//! In-crate rather than under `tests/`: the module is private, the stub-injection +//! tests need the internal `CudaApi`, and the resolution ladder is not public +//! surface. + +use super::*; +use std::cell::RefCell; + +// --- 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"); +} + +// --- 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); + CudaApi { + register, + unregister: stub_unregister, + error_name: stub_error_name, + free: 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 = crate::ring_buf::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(crate::ring_buf::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 = crate::ring_buf::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 = symbol(handle, c"cudaHostGetFlags") + .expect("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 = crate::ring_buf::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 = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); + let mut second = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); + first.pin_host_memory(&roots).expect("first ring failed"); + second.pin_host_memory(&roots).expect("second ring failed"); +} + +#[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}"); +} diff --git a/src/py_bindings.rs b/src/py_bindings.rs index 48656bf..28b719a 100644 --- a/src/py_bindings.rs +++ b/src/py_bindings.rs @@ -139,10 +139,9 @@ impl PyServer { Box::new(FifoRemover::new()), ); - // The page-locked guarantee: if this constructor returns, every ring - // buffer is registered. Anything short of that raises, so a caller can - // never measure a workload against a pinning path that silently did - // nothing. Nothing is loaded at all when the caller didn't ask. + // 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 { store .pin_host_memory(&cuda_vendor_roots(py)) @@ -265,15 +264,12 @@ impl PyServer { } } -/// CUDA vendor package directories, located through Python's import machinery. +/// 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. /// -/// Rung 2 of the CUDA resolution ladder searches beneath these, which is how -/// pinning works with a pip-installed CUDA runtime and without any loader-path -/// environment variable. `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; the other -/// two rungs still run, and total failure reports every path probed. +/// `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(); diff --git a/src/ring_buf.rs b/src/ring_buf.rs index 1d1f541..7359c7e 100644 --- a/src/ring_buf.rs +++ b/src/ring_buf.rs @@ -22,9 +22,8 @@ pub struct PytreeRingBuf { slot_bytes: Vec, /// Total number of slots. capacity: usize, - /// The CUDA runtime the buffers are registered with, once they are. - /// `Some` is exactly the condition for `Drop` having registrations to - /// reverse, and holding it here keeps `Drop` off any process-global. + /// 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, } @@ -54,21 +53,17 @@ impl PytreeRingBuf { } } - /// Page-lock every buffer, so a downstream host-to-device copy of a sampled - /// view is a DMA transfer rather than a chunked staging copy. + /// 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. /// - /// Either every buffer ends up locked or none does — a partial failure is - /// rolled back before the error returns. That is why this is a separate - /// fallible step on a constructed buffer rather than part of `new`: 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 for both the failure path and normal teardown. - /// - /// `cuda_vendor_roots` are the CUDA vendor package directories located - /// through Python's import machinery; see [`crate::host_pinning`]. + /// 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. + /// valid for the buffer's whole life. `cuda_vendor_roots` comes from Python's + /// import machinery; see [`crate::host_pinning`]. pub(crate) fn pin_host_memory( &mut self, cuda_vendor_roots: &[PathBuf], @@ -76,14 +71,11 @@ impl PytreeRingBuf { 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 stubbed CUDA - /// entry points, on a machine with no GPU. + /// [`Self::pin_host_memory`] against an already-resolved runtime. Split out + /// so tests can drive registration and teardown with stubs, without a GPU. pub(crate) fn pin_with(&mut self, api: CudaApi) -> Result<(), PinError> { - // SAFETY: the regions are this buffer's own allocations, which live as - // long as `self` and are never reallocated, and `Drop` unregisters them - // before they are freed. + // 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(()) @@ -180,10 +172,10 @@ impl Drop for PytreeRingBuf { let Some(api) = self.pinned_with else { return; }; - // Unregister before the backing Vecs are freed. Drop::drop runs before - // the struct's fields are dropped, so the memory is still valid. + // 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, and they are still valid here. + // 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 b11e6f0..f0aa37d 100644 --- a/src/store.rs +++ b/src/store.rs @@ -68,8 +68,8 @@ impl Store { /// Page-lock the ring buffers; see [`PytreeRingBuf::pin_host_memory`]. /// - /// `&mut self` puts this between construction and sharing the store, so a - /// caller that then fails to finish building unregisters by dropping it. + /// `&mut self` puts this before the store is shared, so a caller that fails + /// to finish building unregisters by dropping it. pub(crate) fn pin_host_memory( &mut self, cuda_vendor_roots: &[PathBuf], From eb7d4aa5e64d28c6c95d880a119fd33915497f5e Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Fri, 31 Jul 2026 15:02:22 +0200 Subject: [PATCH 06/15] refactor: split host_pinning into resolve and register MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The module did two jobs: find the CUDA runtime, and page-lock memory with it. Resolution was ~75% of the source and the part with no CUDA calls in it at all, so the two barely interact — resolve hands back a CudaApi and register uses it. host_pinning/ ├── mod.rs shared types: CudaApi, Region, PinError ├── resolve.rs the three-rung ladder ├── resolve/tests.rs ├── register.rs pin_all / unpin_all / rollback └── register/tests.rs Tests nest under the module they cover rather than sitting in one file, so each keeps its own privacy: the ladder's internals stay private to resolve, and the stub-injection tests stay next to the unsafe code they exist to cover. Only scan_mapped_runtimes widens, to pub(super), because register's GPU test needs the mapped-runtime path to look up cudaHostGetFlags independently of the pointers under test. No behaviour change; the same 20 tests pass, 10 per module. Co-Authored-By: Claude Opus 5 (1M context) --- src/host_pinning/mod.rs | 108 +++++++++ src/host_pinning/register.rs | 93 ++++++++ src/host_pinning/{ => register}/tests.rs | 208 ++---------------- .../resolve.rs} | 195 +--------------- src/host_pinning/resolve/tests.rs | 193 ++++++++++++++++ 5 files changed, 414 insertions(+), 383 deletions(-) create mode 100644 src/host_pinning/mod.rs create mode 100644 src/host_pinning/register.rs rename src/host_pinning/{ => register}/tests.rs (58%) rename src/{host_pinning.rs => host_pinning/resolve.rs} (58%) create mode 100644 src/host_pinning/resolve/tests.rs diff --git a/src/host_pinning/mod.rs b/src/host_pinning/mod.rs new file mode 100644 index 0000000..821fb3a --- /dev/null +++ b/src/host_pinning/mod.rs @@ -0,0 +1,108 @@ +//! 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. + +mod register; +mod resolve; + +pub(crate) use register::{pin_all, unpin_all}; +pub(crate) 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)` +type RegisterFn = unsafe extern "C" fn(*mut c_void, usize, c_uint) -> c_int; +/// `cudaError_t cudaHostUnregister(void*)` +type UnregisterFn = unsafe extern "C" fn(*mut c_void) -> c_int; +/// `const char* cudaGetErrorName(cudaError_t)` +type ErrorNameFn = unsafe extern "C" fn(c_int) -> *const c_char; +/// `cudaError_t cudaFree(void*)` +type FreeFn = unsafe extern "C" fn(*mut c_void) -> c_int; + +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(crate) struct CudaApi { + register: RegisterFn, + unregister: UnregisterFn, + error_name: ErrorNameFn, + free: FreeFn, +} + +/// A contiguous host allocation to page-lock. Owns nothing; the validity +/// invariants live on [`pin_all`] / [`unpin_all`]. +#[derive(Clone, Copy)] +pub(crate) struct Region { + pub ptr: *mut u8, + pub len: usize, +} + +/// Why pinning could not be delivered. On failure nothing is left registered. +#[derive(Debug)] +pub(crate) 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..8afa6f7 --- /dev/null +++ b/src/host_pinning/register.rs @@ -0,0 +1,93 @@ +//! 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. +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(crate) 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 (~100 MB). 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(crate) 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() +} + +#[cfg(test)] +mod tests; diff --git a/src/host_pinning/tests.rs b/src/host_pinning/register/tests.rs similarity index 58% rename from src/host_pinning/tests.rs rename to src/host_pinning/register/tests.rs index cd8e13c..5ed19b2 100644 --- a/src/host_pinning/tests.rs +++ b/src/host_pinning/register/tests.rs @@ -1,187 +1,18 @@ -//! Unit tests for the private `host_pinning` module. +//! Tests for registration, rollback and teardown. //! -//! In-crate rather than under `tests/`: the module is private, the stub-injection -//! tests need the internal `CudaApi`, and the resolution ladder is not public -//! surface. +//! The stubbed half runs anywhere — rollback's only visible consequence is the +//! *absence* of leaked registrations, which cannot be observed from Python, so it +//! is checked through the injected API. The rest needs a real GPU and reports a +//! skip when there isn't one. -use super::*; use std::cell::RefCell; +use std::ffi::CString; +use std::os::raw::c_char; +use std::path::PathBuf; -// --- 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"); -} +use super::*; +use crate::host_pinning::resolve::{api, scan_mapped_runtimes}; +use crate::host_pinning::RegisterFn; // --- registration and rollback, with the CUDA calls stubbed out --- @@ -397,8 +228,8 @@ fn host_get_flags() -> HostGetFlagsFn { libc::RTLD_NOW | libc::RTLD_LOCAL | libc::RTLD_NOLOAD, ); assert!(!handle.is_null(), "the mapped runtime should re-open"); - let sym = symbol(handle, c"cudaHostGetFlags") - .expect("the runtime should export cudaHostGetFlags"); + 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) } } @@ -468,16 +299,3 @@ fn pinning_twice_in_one_process_reuses_the_resolved_runtime() { first.pin_host_memory(&roots).expect("first ring failed"); second.pin_host_memory(&roots).expect("second ring failed"); } - -#[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}"); -} diff --git a/src/host_pinning.rs b/src/host_pinning/resolve.rs similarity index 58% rename from src/host_pinning.rs rename to src/host_pinning/resolve.rs index f1b69ff..fd94634 100644 --- a/src/host_pinning.rs +++ b/src/host_pinning/resolve.rs @@ -1,113 +1,15 @@ -//! Optional CUDA host-memory pinning (page-locking) of the ring buffers. +//! Finding the CUDA runtime: a three-rung ladder, tried in order, accumulating +//! every attempted path so a total failure can say what it tried. //! -//! 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. +//! The rungs exist because a single soname `dlopen` does not cover the +//! deployment shapes — that is precisely how pinning came to be a silent no-op. use std::ffi::{CStr, CString}; -use std::fmt; -use std::os::raw::{c_char, c_int, c_uint, c_void}; +use std::os::raw::c_void; use std::path::PathBuf; use std::sync::OnceLock; -/// `cudaError_t cudaHostRegister(void*, size_t, unsigned int)` -type RegisterFn = unsafe extern "C" fn(*mut c_void, usize, c_uint) -> c_int; -/// `cudaError_t cudaHostUnregister(void*)` -type UnregisterFn = unsafe extern "C" fn(*mut c_void) -> c_int; -/// `const char* cudaGetErrorName(cudaError_t)` -type ErrorNameFn = unsafe extern "C" fn(c_int) -> *const c_char; -/// `cudaError_t cudaFree(void*)` -type FreeFn = unsafe extern "C" fn(*mut c_void) -> c_int; - -const CUDA_SUCCESS: c_int = 0; - -/// `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. -const CUDA_HOST_REGISTER_PORTABLE: c_uint = 0x01; - -/// 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(crate) struct CudaApi { - register: RegisterFn, - unregister: UnregisterFn, - error_name: ErrorNameFn, - free: FreeFn, -} - -/// A contiguous host allocation to page-lock. Owns nothing; the validity -/// invariants live on [`pin_all`] / [`unpin_all`]. -#[derive(Clone, Copy)] -pub(crate) struct Region { - pub ptr: *mut u8, - pub len: usize, -} - -/// Why pinning could not be delivered. On failure nothing is left registered. -#[derive(Debug)] -pub(crate) 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 {} - -// --------------------------------------------------------------------------- -// Resolving the CUDA runtime -// --------------------------------------------------------------------------- +use super::{CudaApi, ErrorNameFn, FreeFn, PinError, RegisterFn, UnregisterFn}; /// Current and previous CUDA major versions, then the unversioned soname. /// @@ -157,7 +59,7 @@ fn is_runtime_lib(path: &str) -> bool { /// 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. -fn scan_mapped_runtimes(maps: &str) -> Vec { +pub(super) 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 @@ -401,88 +303,5 @@ pub(crate) fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError Ok(API.get_or_init(|| resolved)) } -// --------------------------------------------------------------------------- -// Registration -// --------------------------------------------------------------------------- - -/// 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(crate) 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 (~100 MB). 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(crate) 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() -} - #[cfg(test)] mod tests; diff --git a/src/host_pinning/resolve/tests.rs b/src/host_pinning/resolve/tests.rs new file mode 100644 index 0000000..3446afa --- /dev/null +++ b/src/host_pinning/resolve/tests.rs @@ -0,0 +1,193 @@ +//! Tests for the resolution ladder. No GPU and no CUDA install required: the +//! rungs are pure functions of the process's mappings and the filesystem. + +use super::*; + +// --- 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}"); +} From f2a3907d8cc1d9dc356b2a8924b3c7e3f79ee6ab Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Fri, 31 Jul 2026 15:19:41 +0200 Subject: [PATCH 07/15] refactor: move host_pinning tests to tests/, like every other module MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tests were inside src/ because host_pinning was the crate's only private module, and a private module's internals are invisible to tests/, which are separate crates. That inverted the repo's convention: every other module is pub with pub internals for precisely this reason — PytreeRingBuf::slot_mut is a pub unsafe fn so tests/ring_buf.rs can drive it. So host_pinning becomes a pub mod and tests/host_pinning.rs joins the other six, one per module. This drops ticket 01's "module stays private" checklist item, deliberately and noted there. What that constraint protected is unaffected: the crate is publish = false and ships only as a Python extension, whose entire surface for this feature is pin_host_memory=True. Two things keep the widening honest — CudaApi's fields stay private behind `unsafe fn CudaApi::new`, so no caller can fabricate one with arbitrary function pointers, and pin_all/unpin_all stay unsafe fn with their Safety contracts. Same 20 tests, still passing, GPU ones included. Co-Authored-By: Claude Opus 5 (1M context) --- src/host_pinning/mod.rs | 45 +++- src/host_pinning/register.rs | 9 +- src/host_pinning/resolve.rs | 42 ++- src/host_pinning/resolve/tests.rs | 193 -------------- src/lib.rs | 2 +- src/ring_buf.rs | 7 +- src/store.rs | 5 +- .../tests.rs => tests/host_pinning.rs | 241 ++++++++++++++++-- 8 files changed, 275 insertions(+), 269 deletions(-) delete mode 100644 src/host_pinning/resolve/tests.rs rename src/host_pinning/register/tests.rs => tests/host_pinning.rs (52%) diff --git a/src/host_pinning/mod.rs b/src/host_pinning/mod.rs index 821fb3a..524804f 100644 --- a/src/host_pinning/mod.rs +++ b/src/host_pinning/mod.rs @@ -16,49 +16,70 @@ //! 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. -mod register; -mod resolve; +pub mod register; +pub mod resolve; -pub(crate) use register::{pin_all, unpin_all}; -pub(crate) use resolve::api; +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)` -type RegisterFn = unsafe extern "C" fn(*mut c_void, usize, c_uint) -> c_int; +pub type RegisterFn = unsafe extern "C" fn(*mut c_void, usize, c_uint) -> c_int; /// `cudaError_t cudaHostUnregister(void*)` -type UnregisterFn = unsafe extern "C" fn(*mut c_void) -> c_int; +pub type UnregisterFn = unsafe extern "C" fn(*mut c_void) -> c_int; /// `const char* cudaGetErrorName(cudaError_t)` -type ErrorNameFn = unsafe extern "C" fn(c_int) -> *const c_char; +pub type ErrorNameFn = unsafe extern "C" fn(c_int) -> *const c_char; /// `cudaError_t cudaFree(void*)` -type FreeFn = unsafe extern "C" fn(*mut c_void) -> c_int; +pub type FreeFn = unsafe extern "C" fn(*mut c_void) -> c_int; -const CUDA_SUCCESS: c_int = 0; +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(crate) struct CudaApi { +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(crate) struct Region { +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(crate) enum PinError { +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 }, diff --git a/src/host_pinning/register.rs b/src/host_pinning/register.rs index 8afa6f7..8b30843 100644 --- a/src/host_pinning/register.rs +++ b/src/host_pinning/register.rs @@ -8,7 +8,7 @@ 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. -const CUDA_HOST_REGISTER_PORTABLE: c_uint = 0x01; +pub const CUDA_HOST_REGISTER_PORTABLE: c_uint = 0x01; /// Page-lock every region, or leave none of them locked. /// @@ -22,7 +22,7 @@ const CUDA_HOST_REGISTER_PORTABLE: c_uint = 0x01; /// addresses. /// - Each region registered here must be passed to [`unpin_all`] before its /// memory is freed. -pub(crate) unsafe fn pin_all(api: &CudaApi, regions: &[Region]) -> Result<(), PinError> { +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 @@ -65,7 +65,7 @@ pub(crate) unsafe fn pin_all(api: &CudaApi, regions: &[Region]) -> Result<(), Pi /// # 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(crate) unsafe fn unpin_all(api: &CudaApi, regions: &[Region]) { +pub unsafe fn unpin_all(api: &CudaApi, regions: &[Region]) { for region in regions { if region.len == 0 { continue; @@ -88,6 +88,3 @@ fn error_name(api: &CudaApi, code: c_int) -> String { .to_string_lossy() .into_owned() } - -#[cfg(test)] -mod tests; diff --git a/src/host_pinning/resolve.rs b/src/host_pinning/resolve.rs index fd94634..65be4a8 100644 --- a/src/host_pinning/resolve.rs +++ b/src/host_pinning/resolve.rs @@ -1,8 +1,5 @@ //! Finding the CUDA runtime: a three-rung ladder, tried in order, accumulating //! every attempted path so a total failure can say what it tried. -//! -//! The rungs exist because a single soname `dlopen` does not cover the -//! deployment shapes — that is precisely how pinning came to be a silent no-op. use std::ffi::{CStr, CString}; use std::os::raw::c_void; @@ -19,7 +16,7 @@ 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)] -enum Rung { +pub enum Rung { /// Rung 1: a runtime the process already has mapped. AlreadyLoaded, /// Rung 2: a runtime shipped by an installed CUDA wheel. @@ -29,7 +26,7 @@ enum Rung { } impl Rung { - fn label(self) -> &'static str { + pub fn label(self) -> &'static str { match self { Rung::AlreadyLoaded => "already-loaded scan", Rung::InstalledWheel => "installed-wheel search", @@ -40,10 +37,10 @@ impl Rung { /// One thing to hand to `dlopen`, in ladder order. #[derive(Debug, Clone, PartialEq, Eq)] -struct Candidate { +pub struct Candidate { /// An absolute path (rungs 1 and 2) or a bare soname (rung 3). - name: String, - rung: Rung, + pub name: String, + pub rung: Rung, } /// True for a CUDA runtime shared object, by file name. @@ -59,7 +56,7 @@ fn is_runtime_lib(path: &str) -> bool { /// 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(super) fn scan_mapped_runtimes(maps: &str) -> Vec { +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 @@ -88,7 +85,7 @@ const WHEEL_SEARCH_DEPTH: usize = 4; /// 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. -fn search_wheel_roots(roots: &[PathBuf]) -> Vec { +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(); @@ -141,7 +138,7 @@ fn soname_major(name: &str) -> Option { /// 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. -fn candidates(maps: &str, vendor_roots: &[PathBuf]) -> Vec { +pub fn candidates(maps: &str, vendor_roots: &[PathBuf]) -> Vec { let rungs = [ (Rung::AlreadyLoaded, scan_mapped_runtimes(maps)), (Rung::InstalledWheel, search_wheel_roots(vendor_roots)), @@ -195,21 +192,15 @@ fn open(candidate: &Candidate) -> Result { // 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 { - register: std::mem::transmute::<*mut c_void, RegisterFn>(symbol( - handle, - c"cudaHostRegister", - )?), - unregister: std::mem::transmute::<*mut c_void, UnregisterFn>(symbol( + 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", )?), - error_name: std::mem::transmute::<*mut c_void, ErrorNameFn>(symbol( - handle, - c"cudaGetErrorName", - )?), - free: std::mem::transmute::<*mut c_void, FreeFn>(symbol(handle, c"cudaFree")?), - }) + std::mem::transmute::<*mut c_void, ErrorNameFn>(symbol(handle, c"cudaGetErrorName")?), + std::mem::transmute::<*mut c_void, FreeFn>(symbol(handle, c"cudaFree")?), + )) } } @@ -293,7 +284,7 @@ static API: OnceLock = OnceLock::new(); /// /// Failures are not cached: a process that retries after its framework has /// initialised CUDA should get the later, better answer. -pub(crate) fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError> { +pub fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError> { if let Some(api) = API.get() { return Ok(api); } @@ -302,6 +293,3 @@ pub(crate) fn api(vendor_roots: &[PathBuf]) -> Result<&'static CudaApi, PinError // reference-counted, so the loser just drops an identical set of pointers. Ok(API.get_or_init(|| resolved)) } - -#[cfg(test)] -mod tests; diff --git a/src/host_pinning/resolve/tests.rs b/src/host_pinning/resolve/tests.rs deleted file mode 100644 index 3446afa..0000000 --- a/src/host_pinning/resolve/tests.rs +++ /dev/null @@ -1,193 +0,0 @@ -//! Tests for the resolution ladder. No GPU and no CUDA install required: the -//! rungs are pure functions of the process's mappings and the filesystem. - -use super::*; - -// --- 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}"); -} diff --git a/src/lib.rs b/src/lib.rs index c447663..40f6a19 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ use pyo3::prelude::*; pub mod array_spec; -mod host_pinning; +pub mod host_pinning; pub mod ingress; pub mod metrics; mod py_bindings; diff --git a/src/ring_buf.rs b/src/ring_buf.rs index 7359c7e..5715377 100644 --- a/src/ring_buf.rs +++ b/src/ring_buf.rs @@ -64,16 +64,13 @@ impl PytreeRingBuf { /// 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(crate) fn pin_host_memory( - &mut self, - cuda_vendor_roots: &[PathBuf], - ) -> Result<(), PinError> { + 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(crate) fn pin_with(&mut self, api: CudaApi) -> Result<(), PinError> { + 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())? }; diff --git a/src/store.rs b/src/store.rs index f0aa37d..4f54391 100644 --- a/src/store.rs +++ b/src/store.rs @@ -70,10 +70,7 @@ impl Store { /// /// `&mut self` puts this before the store is shared, so a caller that fails /// to finish building unregisters by dropping it. - pub(crate) fn pin_host_memory( - &mut self, - cuda_vendor_roots: &[PathBuf], - ) -> Result<(), PinError> { + pub fn pin_host_memory(&mut self, cuda_vendor_roots: &[PathBuf]) -> Result<(), PinError> { self.ring.pin_host_memory(cuda_vendor_roots) } diff --git a/src/host_pinning/register/tests.rs b/tests/host_pinning.rs similarity index 52% rename from src/host_pinning/register/tests.rs rename to tests/host_pinning.rs index 5ed19b2..472591a 100644 --- a/src/host_pinning/register/tests.rs +++ b/tests/host_pinning.rs @@ -1,18 +1,221 @@ -//! Tests for registration, rollback and teardown. +//! Tests for `host_pinning`: the CUDA-runtime resolution ladder, and +//! registration with its rollback. //! -//! The stubbed half runs anywhere — rollback's only visible consequence is the -//! *absence* of leaked registrations, which cannot be observed from Python, so it -//! is checked through the injected API. The rest needs a real GPU and reports a -//! skip when there isn't one. +//! 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; +use std::os::raw::{c_char, c_int, c_uint, c_void}; use std::path::PathBuf; -use super::*; -use crate::host_pinning::resolve::{api, scan_mapped_runtimes}; -use crate::host_pinning::RegisterFn; +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 --- @@ -58,12 +261,8 @@ 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); - CudaApi { - register, - unregister: stub_unregister, - error_name: stub_error_name, - free: stub_free, - } + // 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. @@ -126,7 +325,7 @@ 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 = crate::ring_buf::PytreeRingBuf::new(vec![64, 128], 8, 4); + let mut ring = PytreeRingBuf::new(vec![64, 128], 8, 4); ring.pin_with(api) .expect("stubbed registration should succeed"); @@ -145,7 +344,7 @@ fn dropping_a_pinned_ring_buffer_unregisters_every_buffer() { #[test] fn dropping_an_unpinned_ring_buffer_touches_no_cuda_entry_point() { let _api = stub_api(stub_register_ok); - drop(crate::ring_buf::PytreeRingBuf::new(vec![64], 8, 4)); + drop(PytreeRingBuf::new(vec![64], 8, 4)); assert!(REGISTERED.with(|c| c.borrow().is_empty())); assert!(UNREGISTERED.with(|c| c.borrow().is_empty())); } @@ -155,7 +354,7 @@ 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 = crate::ring_buf::PytreeRingBuf::new(vec![64, 64, 64, 64], 8, 4); + 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()); @@ -265,7 +464,7 @@ fn a_real_ring_buffer_pins_and_unregisters_on_drop() { return; } let roots = dev_vendor_roots(); - let mut ring = crate::ring_buf::PytreeRingBuf::new(vec![1024, 2048], 64, 8); + let mut ring = PytreeRingBuf::new(vec![1024, 2048], 64, 8); ring.pin_host_memory(&roots) .expect("registration failed on a GPU host"); @@ -294,8 +493,8 @@ fn pinning_twice_in_one_process_reuses_the_resolved_runtime() { return; } let roots = dev_vendor_roots(); - let mut first = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); - let mut second = crate::ring_buf::PytreeRingBuf::new(vec![4096], 32, 4); + 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"); } From 73e9fd6388248805415e2580281b0ce75ed758f1 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Fri, 31 Jul 2026 15:27:26 +0200 Subject: [PATCH 08/15] docs: add a host-pinning design page, reflecting the module split MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Rust-internals section had no entry for host_pinning at all — the module map in overview.md predated the feature and was never updated — and the resolution ladder's design was only described in the user-facing guide. Adds docs/src/design/host-pinning.md covering the internals: the resolve/register split and why it's clean, why resolution is a ladder and what each rung is for, why registration is a post-construction step, the function-pointer injection seam and what it exists to test, why the module is pub, and the absence of any build-time CUDA dependency. Trims ring-buffer.md's pinning section to what is actually about the ring buffer — that its buffers are never reallocated, so a registration stays valid for their life, and that Drop reverses it — with the module-level rationale moved to the new page. Adds the module map row, the nav entry, and a pointer from the guide. development.md now says where the Rust tests live, and distinguishes the Rust skip-and-pass convention from the Python `gpu` marker. Co-Authored-By: Claude Opus 5 (1M context) --- docs/src/design/host-pinning.md | 103 +++++++++++++++++++++++++ docs/src/design/overview.md | 1 + docs/src/design/ring-buffer.md | 47 +++++------ docs/src/development.md | 25 ++++-- docs/src/guides/host-memory-pinning.md | 3 + mkdocs.yml | 1 + 6 files changed, 146 insertions(+), 34 deletions(-) create mode 100644 docs/src/design/host-pinning.md diff --git a/docs/src/design/host-pinning.md b/docs/src/design/host-pinning.md new file mode 100644 index 0000000..3fcf6a8 --- /dev/null +++ b/docs/src/design/host-pinning.md @@ -0,0 +1,103 @@ +# 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. This page is about how the module is put together; for the +mechanism, the measured numbers and how to verify it engaged, see the +[host-memory pinning guide](../guides/host-memory-pinning.md). + +Two jobs, one file each: + +| 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. | + +They interact at exactly one point: `resolve` returns a `CudaApi` and `register` +takes one. `resolve` makes no CUDA calls at all beyond `dlsym`, which is why the +split is clean rather than nominal — and why resolution, three quarters of the +source, is testable with no GPU anywhere in sight. + +## Why resolution is a ladder + +The whole feature was a silent no-op for months 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. + +So resolution 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. Searching the vendor package rather than a named component + subpackage is load-bearing: CUDA 13 ships one consolidated wheel laid out as + `nvidia/cu13/lib/`, CUDA 12 one wheel per component as + `nvidia/cuda_runtime/lib/`. 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. That is deliberate: reordering +the rungs is precisely how this broke, and the ordering is the kind of thing an +unrelated edit can silently change. + +A rung that yields no candidate still reports itself, 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` is what 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, not tidiness. A constructor that returns `Err` never runs +`Drop`, so registering inside one forces a hand-written unregister loop on the +error path — the exact code most likely to be wrong and least likely to be +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. + +That exists for one reason: 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 096fc21..cf5a268 100644 --- a/docs/src/design/ring-buffer.md +++ b/docs/src/design/ring-buffer.md @@ -46,32 +46,27 @@ we'd need a scatter-gather variant. ## Optional host-memory pinning -`pin_host_memory` (see the [guide](../guides/host-memory-pinning.md) for the -mechanism and the numbers) CUDA-page-locks every backing `Vec` so a -downstream host-to-device copy of a sampled view is a DMA transfer rather than a -chunked staging copy. Two design points are worth recording here. - -**It is a separate fallible step, not part of `new`.** `PytreeRingBuf::new` and -`Store::new` stay infallible; `pin_host_memory(&mut self, ..)` runs on a -fully-constructed buffer and returns `Result`. The reason is rollback -correctness, not taste: a constructor that returns `Err` never runs `Drop`, so -registering inside the constructor would force a hand-written unregister loop on -the error path. Registering afterwards lets the existing `Drop` own rollback for -both the failure path and normal teardown. A `pinned` flag says whether `Drop` -has anything to reverse. - -Within one attempt, either every buffer ends up registered or none does — -`pin_all` unregisters what succeeded before returning the error. Registrations -are still leaked if a reference-counted `Store` outlives process shutdown, which -is accepted. - -**The registration is valid for the buffer's life** because the buffers are -allocated once in `new` and never reallocated or resized. That property is what -makes registering the whole `Vec` up front sound; a growable buffer would -invalidate the registration on its first reallocation. - -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 +`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 +belongs here is why this type can be registered at all. + +**The registration is valid for the buffer's whole life** because the buffers are +allocated once in `new` and never reallocated or resized. That is what makes +registering each `Vec` up front sound — 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. That is also why pinning is a separate step after +construction rather than part of `new`: a constructor returning `Err` never runs +`Drop`, so registering inside one would need a hand-written unregister loop on the +error path. + +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 at all on host-thread occupancy, so `Vec` stays. ## What this type does not do diff --git a/docs/src/development.md b/docs/src/development.md index 05ba0df..a5b333c 100644 --- a/docs/src/development.md +++ b/docs/src/development.md @@ -22,10 +22,19 @@ 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 -The [pinning](guides/host-memory-pinning.md) tests that need a GPU carry the -`gpu` marker and skip automatically when no device is present, so CI needs no +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. The tests that do need a +device print a skip notice and pass when there isn't one, since Rust has no native +test skip. On the Python side the equivalent tests carry a `gpu` marker (registered +in `pyproject.toml`) and `pytest` skips them automatically. Either way CI needs no per-runner configuration. To exercise the CUDA-runtime resolution ladder locally, install a runtime wheel @@ -35,13 +44,13 @@ into the checkout's venv: uv pip install nvidia-cuda-runtime # CUDA 13; use nvidia-cuda-runtime-cu12 for CUDA 12 ``` -Without it only rung 1 (already-mapped libraries) can hit, and on a machine with -no system CUDA install nothing resolves at all — the pinning tests then report -that they skipped rather than failing. The wheel is deliberately *not* a `dev` -extra: CI is CPU-only and should not download a 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 deliberately *not* a +`dev` extra: CI is CPU-only and should not download a CUDA runtime. Note that +`uv run` re-syncs the venv, so re-run the install if it disappears. -The Rust tests find the wheel by looking under `.venv/lib/python*/site-packages/` -in the checkout, since `cargo test` has no Python interpreter to ask. +`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 diff --git a/docs/src/guides/host-memory-pinning.md b/docs/src/guides/host-memory-pinning.md index eb57ba4..f5e3640 100644 --- a/docs/src/guides/host-memory-pinning.md +++ b/docs/src/guides/host-memory-pinning.md @@ -211,6 +211,9 @@ 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. + ## Verifying that pinning engaged This section is the reason this page exists. Pinning was once believed not to diff --git a/mkdocs.yml b/mkdocs.yml index 201e682..d301a7e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -94,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 From b8f2b8877cafa69b92583e7d481151e8e5da8725 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Fri, 31 Jul 2026 15:38:33 +0200 Subject: [PATCH 09/15] chore: add docs-serve-on for viewing docs from another machine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `just docs-serve` binds to loopback, so it is invisible when the checkout lives on a remote box. Adds a recipe that takes a bind address, and documents both routes in development.md — SSH port forwarding (preferred, no server-side change) and binding to a reachable interface. Also records the /echo/ path prefix that site_url imposes on the dev server, since http://127.0.0.1:8000/ on its own 404s. Co-Authored-By: Claude Opus 5 (1M context) --- docs/src/development.md | 25 +++++++++++++++++++++++-- justfile | 6 +++++- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/docs/src/development.md b/docs/src/development.md index a5b333c..6f1fd96 100644 --- a/docs/src/development.md +++ b/docs/src/development.md @@ -63,10 +63,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/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 From 6713eeb81b48e9aeb0ca1314ff72be0dfb9cdf59 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Mon, 3 Aug 2026 11:43:40 +0200 Subject: [PATCH 10/15] docs: cut rhetorical filler from the host-pinning pages The two new pages carried a lot of scaffolding that wasn't doing work: roadmap sentences ("this page covers..."), self-referential asides ("this section is the reason this page exists"), and a bolded mini-headline on every paragraph of the Measured section. Technical content is unchanged; the prose is trimmed to the voice of the surrounding docs. Substantive fixes found while editing: - The sample RuntimeError in the guide omitted the two "(none)" lines that resolve() emits for a rung that ran and found nothing, which are the whole point of reporting empty rungs. Replaced with real output. - "Two jobs, one file each" introduced a three-row table. - Dropped the CHANGELOG's Changed and Removed sections. Relative to 0.1.1 the resolution ladder is new, not changed, and ECHO_PIN_HOST_MEMORY was added and removed inside this branch, so it never shipped and users can't have depended on it. - The Added entry showed pin_host_memory=False, the default, rather than the value that turns the feature on. - Reported the primary-context cost as ~128 MB in register.rs to match the measured figure in the guide, not ~100 MB. - Replaced "~1800x" occupancy with "three orders of magnitude"; the ratio was derived from a 0.002 ms figure at timer resolution. - ring-buffer.md now links to the constructor rationale rather than restating it. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 28 ++--- docs/src/design/host-pinning.md | 64 ++++++------ docs/src/design/ring-buffer.md | 31 +++--- docs/src/development.md | 4 +- docs/src/guides/host-memory-pinning.md | 135 ++++++++++++------------- python/echo/server.py | 3 +- src/host_pinning/register.rs | 4 +- 7 files changed, 124 insertions(+), 145 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cb282ff..f442d15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,33 +9,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- `Server(..., pin_host_memory=False)` CUDA-page-locks the ring buffers, so a +- `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. 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. + 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, so a - throughput measurement can't be invalidated by the feature having done nothing. + 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 measured numbers, the unswappable footprint arithmetic, and how to verify pinning engaged (and why `VmLck` cannot). -### Changed - -- The CUDA runtime is now resolved through a three-rung ladder — already-mapped - libraries, then the installed CUDA wheels, then sonames (versioned before - unversioned). A soname symlink or `LD_LIBRARY_PATH` entry is no longer needed. - -### Removed - -- `ECHO_PIN_HOST_MEMORY`. Pinning is controlled only by the `pin_host_memory` - keyword argument, so the two can never disagree. Nothing released ever - responded to this variable — and it never worked: it opened the unversioned - `libcudart.so` soname, which the pip CUDA runtime wheels do not ship, so every - pin call was silently a no-op. - ## [0.1.1] - 2026-05-26 - `TrajectoryAccumulator` better supports single and buffered timescales diff --git a/docs/src/design/host-pinning.md b/docs/src/design/host-pinning.md index 3fcf6a8..7ebd842 100644 --- a/docs/src/design/host-pinning.md +++ b/docs/src/design/host-pinning.md @@ -2,11 +2,8 @@ `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. This page is about how the module is put together; for the -mechanism, the measured numbers and how to verify it engaged, see the -[host-memory pinning guide](../guides/host-memory-pinning.md). - -Two jobs, one file each: +staging copy. For the mechanism, the measured numbers and how to check it +engaged, see the [host-memory pinning guide](../guides/host-memory-pinning.md). | File | Role | |---|---| @@ -14,57 +11,56 @@ Two jobs, one file each: | `resolve.rs` | Find the CUDA runtime. | | `register.rs` | Page-lock memory with it, and roll back cleanly. | -They interact at exactly one point: `resolve` returns a `CudaApi` and `register` -takes one. `resolve` makes no CUDA calls at all beyond `dlsym`, which is why the -split is clean rather than nominal — and why resolution, three quarters of the -source, is testable with no GPU anywhere in sight. +`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 whole feature was a silent no-op for months because it did one thing: -`dlopen("libcudart.so")`. The pip CUDA runtime wheels ship only the *versioned* +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. -So resolution tries three rungs in order, accumulating every attempted path: +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 +2. **Installed-wheel search.** Search beneath the CUDA vendor package directories, which `py_bindings` locates through Python's import machinery and - passes down. Searching the vendor package rather than a named component - subpackage is load-bearing: CUDA 13 ships one consolidated wheel laid out as + 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/`. Naming the component finds nothing on a CUDA 13 + `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. That is deliberate: reordering -the rungs is precisely how this broke, and the ordering is the kind of thing an -unrelated edit can silently change. +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, so a failure message -distinguishes "searched and found nothing" from "never ran". +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` is what lets `cargo test` exercise it with no interpreter. +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, not tidiness. A constructor that returns `Err` never runs -`Drop`, so registering inside one forces a hand-written unregister loop on the -error path — the exact code most likely to be wrong and least likely to be -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. +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 @@ -77,12 +73,12 @@ accepted. `CudaApi` is a struct of function pointers passed explicitly to `pin_all` / `unpin_all`, rather than a global the two reach into. -That exists for one reason: 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. +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 diff --git a/docs/src/design/ring-buffer.md b/docs/src/design/ring-buffer.md index cf5a268..bc24a7e 100644 --- a/docs/src/design/ring-buffer.md +++ b/docs/src/design/ring-buffer.md @@ -48,26 +48,23 @@ we'd need a scatter-gather variant. `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 -belongs here is why this type can be registered at all. - -**The registration is valid for the buffer's whole life** because the buffers are -allocated once in `new` and never reallocated or resized. That is what makes -registering each `Vec` up front sound — 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. That is also why pinning is a separate step after -construction rather than part of `new`: a constructor returning `Err` never runs -`Drop`, so registering inside one would need a hand-written unregister loop on the -error path. +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 at all on host-thread occupancy, so `Vec` stays. +time and nothing on host-thread occupancy, so `Vec` stays. ## What this type does not do diff --git a/docs/src/development.md b/docs/src/development.md index 6f1fd96..4daa86c 100644 --- a/docs/src/development.md +++ b/docs/src/development.md @@ -45,8 +45,8 @@ uv pip install nvidia-cuda-runtime # CUDA 13; use nvidia-cuda-runtime-cu12 ``` 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 deliberately *not* a -`dev` extra: CI is CPU-only and should not download a CUDA runtime. Note that +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. Note that `uv run` re-syncs the venv, so re-run the install if it disappears. `cargo test` has no Python interpreter to ask for the wheel's location, so the diff --git a/docs/src/guides/host-memory-pinning.md b/docs/src/guides/host-memory-pinning.md index f5e3640..a6f1bf6 100644 --- a/docs/src/guides/host-memory-pinning.md +++ b/docs/src/guides/host-memory-pinning.md @@ -1,13 +1,10 @@ # 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. +`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 — the page-locked memory is not -swappable. This page covers what the mechanism actually is, what it measured, -how to size the footprint, and — the part that is easy to get wrong — **how to -confirm it engaged**. +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), @@ -25,6 +22,8 @@ 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 @@ -32,15 +31,16 @@ Install a CUDA runtime (for example the nvidia-cuda-runtime wheel), or construct the Server after your framework has initialised CUDA. ``` -That is deliberately all the API there is: no status object, no report method, no -warning. Construction succeeding *is* the assertion, and unlike an explicit check -you might add, it cannot be forgotten. If you would rather degrade than fail, -catch `RuntimeError`. +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. -## The mechanism: a pageable "async" copy is not async +## 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 @@ -50,16 +50,15 @@ cannot hand the transfer to the copy engine and walk away. Instead it: 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. Two things -follow, and the second is usually the expensive one: +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 waiting for work. + 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. @@ -72,11 +71,11 @@ allocates device memory. ## Measured -Local microbenchmark on a **NVIDIA GeForce RTX 5060 Ti (16 GB), driver +Local microbenchmark on an **NVIDIA GeForce RTX 5060 Ti (16 GB), driver 595.71.05, CUDA runtime 13.3**, sized like a large-batch learner: `batch_size=512` over four arrays (34.1 MB per batch, 102.2 MB ring), with the copy measurement isolated to a single 52.4 MB buffer. 11 repeats; median and -interquartile range: +interquartile range. Through echo's own ring buffers and sampled views (34.1 MB batch, 102.2 MB ring): @@ -94,49 +93,47 @@ memory and to time the copy call on its own: | page-locked | 3.715 ms | 3.714–3.719 | 14.1 GB/s | 0.002 ms — **0%** | | `cudaHostAlloc` (reference ceiling) | 3.658 ms | 3.657–3.658 | 14.3 GB/s | 0.002 ms | -Read the throughput and the return-latency columns separately, because they say -different things. +The throughput and return-latency columns say different things, so read them +separately. -**Bandwidth barely moves on this machine, and that is expected.** Both paths -saturate the host's link at ~14 GB/s; driver-allocated pinned memory only reaches +Bandwidth barely moves on this machine, and that is expected. Both paths saturate +the host's link at ~14 GB/s, and driver-allocated pinned memory only reaches 14.3 GB/s, so there is nothing more to win here. A host with more PCIe headroom will show a larger gap. -**Host-thread occupancy collapses by ~1800x.** That is the mechanism above, -measured: the same call goes from holding the calling thread for 3.645 ms to -0.002 ms. This is the component that scales into the driver-lock contention a +Host-thread occupancy drops by three orders of magnitude: the same call goes from +holding the calling thread for 3.645 ms to 0.002 ms. That is the mechanism above, +measured. It is the component that scales into the driver-lock contention a large-batch learner suffers, and it is the reason to turn pinning on. -**Nothing on the write side pays for it.** Host-write throughput into the ring — -what a drainer costs — is unchanged at 12.3 GB/s, and registering the 102 MB ring -took 8.9 ms once, at construction. Pinning changes only how the memory is mapped, -not any code on the ingest or sample path, so there is no per-sample or per-drain +Nothing on the write side pays for it. Host-write throughput into the ring — what +a drainer costs — is unchanged at 12.3 GB/s, and registering the 102 MB ring took +8.9 ms once, at construction. Pinning changes only how the memory is mapped, not +any code on the ingest or sample path, so there is no per-sample or per-drain cost. A batch sampled with pinning on is bit-identical to the same batch with it off. -**Only the portable flag is used.** Two variants were measured and rejected. +Only the portable flag is used; two other variants were measured and rejected. Page-aligning the ring buffers gained 1.6% of copy time and nothing at all on return latency, well under the bar set in advance, so the buffers stay plain -`Vec`. Read-only registration (`cudaHostRegisterReadOnly`) is not usable -here at all — `cudaHostRegister` returns `cudaErrorNotSupported` on this GPU, -whose `cudaDevAttrHostRegisterReadOnlySupported` is 0 — and the CUDA -documentation describes that flag as permission to register memory *mapped* -read-only rather than as a transfer optimisation, while saying nothing about host -writes to such a range. Echo's drainers write these pages continuously, so there -would be no documented basis for using it even where it is supported. +`Vec`. Read-only registration (`cudaHostRegisterReadOnly`) is not usable here +at all — `cudaHostRegister` returns `cudaErrorNotSupported` on this GPU, whose +`cudaDevAttrHostRegisterReadOnlySupported` is 0 — and the CUDA documentation +describes that flag as permission to register memory *mapped* read-only rather +than as a transfer optimisation, while saying nothing about host writes to such a +range. Echo's drainers write these pages continuously, so there would be no +documented basis for using it even where it is supported. **What this microbenchmark cannot tell you.** A single consumer GPU cannot reproduce the driver-lock contention of a real learner issuing tens of thousands -of small kernel launches concurrently with the staging copy. The table above -measures the bandwidth component and the mechanism; it does *not* size the win on -a large-batch learner, where the contention component dominates and the effect is -correspondingly larger than the 1–3% bandwidth figure here. - -So do not read 1–3% as "pinning is worth 1–3%", and do not read the occupancy -column as a step-time prediction either. Measure your own workload: turn pinning -on, confirm it engaged (below), and compare step times. - -To reproduce the table: allocate a buffer with `malloc`, `cudaMalloc` a +of small kernel launches concurrently with the staging copy. The tables above +measure the bandwidth component and the mechanism; they do not size the win on a +large-batch learner, where the contention component dominates. So the 1–3% figure +is not a prediction of what pinning is worth, and neither is the occupancy column. +Measure your own workload: turn pinning on, confirm it engaged (below), and +compare step times. + +To reproduce the tables: allocate a buffer with `malloc`, `cudaMalloc` a destination, and time `cudaMemcpyAsync` + `cudaStreamSynchronize` before and after `cudaHostRegister(ptr, size, cudaHostRegisterPortable)`. Time the `cudaMemcpyAsync` call *on its own*, without the synchronize, to see the @@ -153,13 +150,10 @@ page-locked bytes = batch_size x num_buffers x sum(leaf.nbytes for leaf in examp For a batch of 512 with `num_buffers=3` and 67 KB per sample, that is 512 x 3 x 67 KB ≈ 102 MB. -Two things to hold onto: - - **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. Size the host before the job does it - for you. + 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, so a memory-lock resource limit (`ulimit -l`) does not bind here — see @@ -177,14 +171,14 @@ initialisation best-effort, by freeing a null pointer. That frees nothing and never *changes* which device is current, so it cannot perturb your framework's device selection. -It is not entirely free, though, and this is the reason the ordering above is a -recommendation rather than a footnote: initialising the runtime creates the -primary CUDA context on whatever device is already current, and that context -costs device memory (~128 MB on the machine in the table above). Construct after -your framework has initialised CUDA and the context already exists, so echo adds -nothing. Construct before it, and echo creates the context first — which both -spends that memory early and may interact badly with a framework that -pre-allocates a fraction of *free* device memory. +It is not entirely free, though, which is why the ordering above is worth +following: initialising the runtime creates the primary CUDA context on whatever +device is already current, and that context costs device memory (~128 MB on the +machine in the tables above). Construct after your framework has initialised CUDA +and the context already exists, so echo adds nothing. Construct before it, and +echo creates the context first — which both spends that memory early and may +interact badly with a framework that pre-allocates a fraction of *free* device +memory. Echo never allocates device buffers of its own and never calls a set-device function. @@ -216,9 +210,9 @@ The ladder and why each rung exists are covered in ## Verifying that pinning engaged -This section is the reason this page exists. Pinning was once believed not to -help a workload that had in fact never executed the pinning code path, and the -knowledge of how to check lived in one person's head. +Pinning has been judged unhelpful before on a workload that turned out never to +have executed the pinning code path at all. Confirm it engaged before drawing any +conclusion from a measurement. **In-process: the constructor.** `Server(..., pin_host_memory=True)` returning *is* the assertion. There is nothing else to query. @@ -229,8 +223,8 @@ for each transfer; it must say the source is pinned/page-locked rather than pageable. This is the authoritative external signal. **From a test, without a framework.** Ask the CUDA runtime directly for the flags -on the address behind a sampled batch — this is what echo's own test suite does, -precisely so that a test cannot pass merely because echo believes it worked: +on the address behind a sampled batch. This is what echo's own test suite does, so +that a test cannot pass merely because echo believes it worked: ```python import ctypes @@ -244,13 +238,18 @@ assert code == 0 # 0 = cudaSuccess; non-zero means not registered assert flags.value & 0x01 # cudaHostRegisterPortable ``` +`CDLL` here goes through the system loader, which is exactly what a wheel-only +install does not satisfy, so pass the absolute path of the runtime if the soname +does not resolve. `python/tests/test_host_pinning.py` walks the same candidates +echo does. + !!! warning "`VmLck` is not a valid check" `VmLck` in `/proc/self/status` stays at **zero** even when pinning demonstrably works: the NVIDIA driver's page-locking does not go through - mlock accounting. Anyone who reads it as a check will conclude pinning is off - when it is on. The same goes for any tool built on mlock accounting. Use the - profile trace or `cudaHostGetFlags`. + mlock accounting. Read as a check, it reports pinning as off when it is on, + and so does any tool built on mlock accounting. Use the profile trace or + `cudaHostGetFlags`. ## What is not covered diff --git a/python/echo/server.py b/python/echo/server.py index 3e4faf3..7a7835a 100644 --- a/python/echo/server.py +++ b/python/echo/server.py @@ -37,8 +37,7 @@ class Server: Anything that would prevent that — no CUDA runtime found, no usable device, a registration rejected — raises ``RuntimeError`` here, naming every path probed and the CUDA error. There is no silent - fallback, so a throughput measurement can never be invalidated by - pinning having quietly done nothing. + fallback. Construct the server *after* your framework has initialised CUDA. The page-locked footprint is the full ring diff --git a/src/host_pinning/register.rs b/src/host_pinning/register.rs index 8b30843..0b58a72 100644 --- a/src/host_pinning/register.rs +++ b/src/host_pinning/register.rs @@ -26,8 +26,8 @@ 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 (~100 MB). Constructing after the framework has - // initialised CUDA — the documented order — means it already exists. + // 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. From feafc79f9e33202731fd42ab788f1014389713da Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Mon, 3 Aug 2026 15:26:57 +0200 Subject: [PATCH 11/15] chore: simplify docs --- README.md | 4 +- docs/src/development.md | 11 +- docs/src/guides/host-memory-pinning.md | 133 +------------------------ python/echo/server.py | 19 +--- 4 files changed, 7 insertions(+), 160 deletions(-) diff --git a/README.md b/README.md index a1f7d88..241923d 100644 --- a/README.md +++ b/README.md @@ -44,9 +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/) - of the ring buffers, so a host-to-device copy of a batch is a real DMA - transfer rather than a chunked staging copy that blocks the calling thread +- CUDA [host-memory pinning](https://instadeepai.github.io/echo/guides/host-memory-pinning/) so that even H2D doesn't require copies ## Example diff --git a/docs/src/development.md b/docs/src/development.md index 4daa86c..5a80a25 100644 --- a/docs/src/development.md +++ b/docs/src/development.md @@ -31,23 +31,18 @@ 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. The tests that do need a -device print a skip notice and pass when there isn't one, since Rust has no native -test skip. On the Python side the equivalent tests carry a `gpu` marker (registered -in `pyproject.toml`) and `pytest` skips them automatically. Either way CI needs no -per-runner configuration. +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 # CUDA 13; use nvidia-cuda-runtime-cu12 for CUDA 12 +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. Note that -`uv run` re-syncs the venv, so re-run the install if it disappears. +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. diff --git a/docs/src/guides/host-memory-pinning.md b/docs/src/guides/host-memory-pinning.md index a6f1bf6..7356347 100644 --- a/docs/src/guides/host-memory-pinning.md +++ b/docs/src/guides/host-memory-pinning.md @@ -69,76 +69,6 @@ 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. -## Measured - -Local microbenchmark on an **NVIDIA GeForce RTX 5060 Ti (16 GB), driver -595.71.05, CUDA runtime 13.3**, sized like a large-batch learner: -`batch_size=512` over four arrays (34.1 MB per batch, 102.2 MB ring), with the -copy measurement isolated to a single 52.4 MB buffer. 11 repeats; median and -interquartile range. - -Through echo's own ring buffers and sampled views (34.1 MB batch, 102.2 MB ring): - -| arm | copy time | IQR | throughput | host-write | registration | -|---|---|---|---|---|---| -| pageable | 2.441 ms | 2.440–2.442 | 14.0 GB/s | 12.32 GB/s | – | -| page-locked (portable) | 2.401 ms | 2.399–2.407 | 14.2 GB/s | 12.39 GB/s | 8.9 ms | - -And isolating one 52.4 MB buffer, to compare against driver-allocated pinned -memory and to time the copy call on its own: - -| source memory | copy time | IQR | throughput | `cudaMemcpyAsync` returns after | -|---|---|---|---|---| -| pageable | 3.751 ms | 3.748–3.752 | 14.0 GB/s | 3.645 ms — **97% of the copy** | -| page-locked | 3.715 ms | 3.714–3.719 | 14.1 GB/s | 0.002 ms — **0%** | -| `cudaHostAlloc` (reference ceiling) | 3.658 ms | 3.657–3.658 | 14.3 GB/s | 0.002 ms | - -The throughput and return-latency columns say different things, so read them -separately. - -Bandwidth barely moves on this machine, and that is expected. Both paths saturate -the host's link at ~14 GB/s, and driver-allocated pinned memory only reaches -14.3 GB/s, so there is nothing more to win here. A host with more PCIe headroom -will show a larger gap. - -Host-thread occupancy drops by three orders of magnitude: the same call goes from -holding the calling thread for 3.645 ms to 0.002 ms. That is the mechanism above, -measured. It is the component that scales into the driver-lock contention a -large-batch learner suffers, and it is the reason to turn pinning on. - -Nothing on the write side pays for it. Host-write throughput into the ring — what -a drainer costs — is unchanged at 12.3 GB/s, and registering the 102 MB ring took -8.9 ms once, at construction. Pinning changes only how the memory is mapped, not -any code on the ingest or sample path, so there is no per-sample or per-drain -cost. A batch sampled with pinning on is bit-identical to the same batch with it -off. - -Only the portable flag is used; two other variants were measured and rejected. -Page-aligning the ring buffers gained 1.6% of copy time and nothing at all on -return latency, well under the bar set in advance, so the buffers stay plain -`Vec`. Read-only registration (`cudaHostRegisterReadOnly`) is not usable here -at all — `cudaHostRegister` returns `cudaErrorNotSupported` on this GPU, whose -`cudaDevAttrHostRegisterReadOnlySupported` is 0 — and the CUDA documentation -describes that flag as permission to register memory *mapped* read-only rather -than as a transfer optimisation, while saying nothing about host writes to such a -range. Echo's drainers write these pages continuously, so there would be no -documented basis for using it even where it is supported. - -**What this microbenchmark cannot tell you.** A single consumer GPU cannot -reproduce the driver-lock contention of a real learner issuing tens of thousands -of small kernel launches concurrently with the staging copy. The tables above -measure the bandwidth component and the mechanism; they do not size the win on a -large-batch learner, where the contention component dominates. So the 1–3% figure -is not a prediction of what pinning is worth, and neither is the occupancy column. -Measure your own workload: turn pinning on, confirm it engaged (below), and -compare step times. - -To reproduce the tables: allocate a buffer with `malloc`, `cudaMalloc` a -destination, and time `cudaMemcpyAsync` + `cudaStreamSynchronize` before and -after `cudaHostRegister(ptr, size, cudaHostRegisterPortable)`. Time the -`cudaMemcpyAsync` call *on its own*, without the synchronize, to see the -occupancy column. - ## Sizing the footprint The whole ring is locked, not one batch: @@ -167,18 +97,7 @@ Registration needs an initialised CUDA runtime, so the supported order is: 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. That frees nothing and -never *changes* which device is current, so it cannot perturb your framework's -device selection. - -It is not entirely free, though, which is why the ordering above is worth -following: initialising the runtime creates the primary CUDA context on whatever -device is already current, and that context costs device memory (~128 MB on the -machine in the tables above). Construct after your framework has initialised CUDA -and the context already exists, so echo adds nothing. Construct before it, and -echo creates the context first — which both spends that memory early and may -interact badly with a framework that pre-allocates a fraction of *free* device -memory. +initialisation best-effort, by freeing a null pointer. Echo never allocates device buffers of its own and never calls a set-device function. @@ -208,57 +127,7 @@ 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. -## Verifying that pinning engaged - -Pinning has been judged unhelpful before on a workload that turned out never to -have executed the pinning code path at all. Confirm it engaged before drawing any -conclusion from a measurement. - -**In-process: the constructor.** `Server(..., pin_host_memory=True)` returning -*is* the assertion. There is nothing else to query. - -**From a profile: the copy's memory-source classification.** Profile the learner -and look at the host-to-device memcpy rows. A profiler reports the source kind -for each transfer; it must say the source is pinned/page-locked rather than -pageable. This is the authoritative external signal. - -**From a test, without a framework.** Ask the CUDA runtime directly for the flags -on the address behind a sampled batch. This is what echo's own test suite does, so -that a test cannot pass merely because echo believes it worked: - -```python -import ctypes - -cudart = ctypes.CDLL("libcudart.so.13") -flags = ctypes.c_uint(0) -code = cudart.cudaHostGetFlags( - ctypes.byref(flags), ctypes.c_void_p(batch["obs"].ctypes.data) -) -assert code == 0 # 0 = cudaSuccess; non-zero means not registered -assert flags.value & 0x01 # cudaHostRegisterPortable -``` - -`CDLL` here goes through the system loader, which is exactly what a wheel-only -install does not satisfy, so pass the absolute path of the runtime if the soname -does not resolve. `python/tests/test_host_pinning.py` walks the same candidates -echo does. - -!!! warning "`VmLck` is not a valid check" - - `VmLck` in `/proc/self/status` stays at **zero** even when pinning - demonstrably works: the NVIDIA driver's page-locking does not go through - mlock accounting. Read as a check, it reports pinning as off when it is on, - and so does any tool built on mlock accounting. Use the profile trace or - `cudaHostGetFlags`. - ## 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. -- **Anything other than the ring buffers.** Producer queues, transport staging - buffers and accumulator storage are not registered: the ring buffers are the - only source of a host-to-device copy. -- **Allocating pinned memory directly** instead of registering the existing - buffers. Steady-state transfer performance would be identical, so it would - change only construction cost while making the CUDA runtime mandatory at - allocation time. diff --git a/python/echo/server.py b/python/echo/server.py index 7a7835a..4ccedd8 100644 --- a/python/echo/server.py +++ b/python/echo/server.py @@ -28,23 +28,8 @@ 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 - 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. - - **If this constructor returns, every ring buffer is page-locked.** - Anything that would prevent that — no CUDA runtime found, no usable - device, a registration rejected — raises ``RuntimeError`` here, - naming every path probed and the CUDA error. There is no silent - fallback. - - Construct the server *after* your framework has initialised CUDA. - The page-locked footprint is the full ring - (``batch_size * num_buffers * bytes_per_sample``), is not - swappable, and multiplies by the number of servers in the process. - See the [host-memory pinning guide](../guides/host-memory-pinning.md). - + pin_host_memory: CUDA-page-lock the ring buffers, so a downstream H2D + is a real DMA transfer on the copy engine instead of a staging copy. Raises: RuntimeError: If ``pin_host_memory`` is set and the buffers could not be page-locked. From 61750a8791ad6a179a469b6c9dafb13f980ef85a Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Mon, 3 Aug 2026 15:35:48 +0200 Subject: [PATCH 12/15] test: cover the real resolution-failure message; release the GIL while pinning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The failure message is the entire interface to a resolution failure, and nothing tested it against a real failure. The Rust test hand-builds a PinError and checks Display; the Python test that would have exercised the actual accumulation was skipif(HAS_GPU), so it skipped on every machine with a GPU — i.e. everywhere this gets developed. That gap is why the guide's example output was wrong: it omitted the two "(none)" lines that resolve() emits for a rung that ran and found nothing, which are the whole reason empty rungs are reported at all. Replaced it with a subprocess test that forces every rung to fail (nothing mapped, nvidia unimportable, no LD_LIBRARY_PATH) and asserts each rung appears, in order. It runs on GPU hosts too, since it needs resolution to fail rather than a device to be absent. Verified by mutation: dropping the empty-rung note fails it, and reordering the reporting loop fails it on ordering. The soname probe runs inside the subprocess deliberately. Probing in the parent reports every soname as resolvable, because the module loads the runtime by absolute path at import and dlopen then matches the already-loaded object by soname — which would have skipped this test on any host with the wheel installed. Also release the GIL around pinning. Measured, a cold cudaFree(NULL) creating the primary context takes ~210 ms, all of it with the GIL held; detach lets other Python threads run at 82% of their unblocked rate during construction instead of stalling. detach is synchronous, so the constructor still returns only once every buffer is registered and the page-locked guarantee is unchanged. It also removes a deadlock class: holding the GIL across a CUDA call lets echo block against a thread that holds a CUDA lock and wants the GIL. Minor: note why empty_rung_note's Soname arm is unreachable, and why the wheel walk does not follow directory symlinks. Co-Authored-By: Claude Opus 5 (1M context) --- python/tests/test_host_pinning.py | 88 +++++++++++++++++++++++++++---- src/host_pinning/resolve.rs | 6 +++ src/py_bindings.rs | 11 +++- 3 files changed, 92 insertions(+), 13 deletions(-) diff --git a/python/tests/test_host_pinning.py b/python/tests/test_host_pinning.py index 00e4462..075ff23 100644 --- a/python/tests/test_host_pinning.py +++ b/python/tests/test_host_pinning.py @@ -30,6 +30,9 @@ 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), @@ -47,7 +50,7 @@ def _cudart_paths() -> list[str]: 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 += ["libcudart.so.13", "libcudart.so.12", "libcudart.so"] + candidates += list(SONAMES) return [c for c in candidates if not c.endswith(".a")] @@ -140,17 +143,80 @@ def mapped(): assert "ok" in result.stdout -@pytest.mark.skipif(HAS_GPU, reason="a usable CUDA device is present") -def test_requesting_pinning_without_cuda_raises_and_lists_probed_paths(): - """The failure a misconfigured deployment gets: loud, at startup, specific.""" - with pytest.raises(RuntimeError) as excinfo: - Server(EXAMPLE, batch_size=4, pin_host_memory=True) +@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 - message = str(excinfo.value) - assert "libcudart.so" in message, message - # Every rung reports what it tried, so the user can fix their environment - # without reading echo's source. - assert "soname load" in message, message + 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 diff --git a/src/host_pinning/resolve.rs b/src/host_pinning/resolve.rs index 65be4a8..2ae6ae3 100644 --- a/src/host_pinning/resolve.rs +++ b/src/host_pinning/resolve.rs @@ -95,6 +95,9 @@ pub fn search_wheel_roots(roots: &[PathBuf]) -> Vec { }; 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; }; @@ -272,6 +275,9 @@ fn empty_rung_note(rung: Rung, vendor_roots: &[PathBuf]) -> 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()) diff --git a/src/py_bindings.rs b/src/py_bindings.rs index 28b719a..30539e3 100644 --- a/src/py_bindings.rs +++ b/src/py_bindings.rs @@ -143,8 +143,15 @@ impl PyServer { // 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 { - store - .pin_host_memory(&cuda_vendor_roots(py)) + // Locating the vendor roots needs the GIL; page-locking does not, and + // it can take a while — a cold `cudaFree(NULL)` creates the primary + // context (~210 ms measured) when the framework has not initialised + // CUDA yet. `detach` is still synchronous: it releases the GIL, runs + // on this thread, and reacquires before returning, so the guarantee + // above holds. 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); From 2eb2f4173fd612ebcaf4b408782d74fedbd1c06c Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Mon, 3 Aug 2026 15:38:47 +0200 Subject: [PATCH 13/15] docs: repair references left dangling by the simplification pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Trimming the Measured and Verifying sections left four things pointing at content that no longer exists: - The footprint section linked to #verifying-that-pinning-engaged. Folded the part that mattered into one sentence instead: VmLck stays at zero because the driver's page-locking bypasses mlock accounting, so it is neither a binding limit nor a way to check pinning engaged. - design/host-pinning.md and the CHANGELOG both advertised the guide as covering measured numbers and how to verify pinning; they now describe what it actually covers. - server.py lost the blank line before `Raises:`, so griffe stopped parsing it as a section and rendered it as loose prose after the parameters table. Confirmed against the built HTML: the page had only a Parameters: section, and now has Parameters: and Raises:. Also: the README said pinning means "even H2D doesn't require copies". Pinning doesn't remove the copy, it makes it a DMA transfer instead of a staged one — and next to echo's zero-copy claims that reads as though H2D were free. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 4 ++-- README.md | 2 +- docs/src/design/host-pinning.md | 4 ++-- docs/src/guides/host-memory-pinning.md | 5 +++-- python/echo/server.py | 4 +++- 5 files changed, 11 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f442d15..d6683b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,8 +21,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 measured numbers, the unswappable footprint - arithmetic, and how to verify pinning engaged (and why `VmLck` cannot). + 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/README.md b/README.md index 241923d..dbc084b 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ pytree-agnostic. - GIL released while waiting for batches - FIFO sampling (with more strategies planned) - Detailed metrics exposed per batch -- CUDA [host-memory pinning](https://instadeepai.github.io/echo/guides/host-memory-pinning/) so that even H2D doesn't require copies +- 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 index 7ebd842..03f7da2 100644 --- a/docs/src/design/host-pinning.md +++ b/docs/src/design/host-pinning.md @@ -2,8 +2,8 @@ `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, the measured numbers and how to check it -engaged, see the [host-memory pinning guide](../guides/host-memory-pinning.md). +staging copy. For the mechanism and the user-facing contract, see the +[host-memory pinning guide](../guides/host-memory-pinning.md). | File | Role | |---|---| diff --git a/docs/src/guides/host-memory-pinning.md b/docs/src/guides/host-memory-pinning.md index 7356347..34ff5d6 100644 --- a/docs/src/guides/host-memory-pinning.md +++ b/docs/src/guides/host-memory-pinning.md @@ -86,8 +86,9 @@ For a batch of 512 with `num_buffers=3` and 67 KB per sample, that is 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, so a memory-lock resource limit (`ulimit -l`) does not bind here — see -[below](#verifying-that-pinning-engaged). +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 diff --git a/python/echo/server.py b/python/echo/server.py index 4ccedd8..f2cb7c6 100644 --- a/python/echo/server.py +++ b/python/echo/server.py @@ -29,7 +29,9 @@ class Server: 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 - is a real DMA transfer on the copy engine instead of a staging copy. + copy is a real DMA transfer on the copy engine instead of a staging + copy. + Raises: RuntimeError: If ``pin_host_memory`` is set and the buffers could not be page-locked. From efbf94b90c614d73810910f477785f09f42064f3 Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Mon, 3 Aug 2026 15:44:57 +0200 Subject: [PATCH 14/15] chore: simplify docs --- python/echo/server.py | 3 +-- src/py_bindings.rs | 9 ++------- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/python/echo/server.py b/python/echo/server.py index f2cb7c6..f9cea91 100644 --- a/python/echo/server.py +++ b/python/echo/server.py @@ -29,8 +29,7 @@ class Server: 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 on the copy engine instead of a staging - copy. + copy is a real DMA transfer instead of a staging copy. Raises: RuntimeError: If ``pin_host_memory`` is set and the buffers could not diff --git a/src/py_bindings.rs b/src/py_bindings.rs index 30539e3..1470638 100644 --- a/src/py_bindings.rs +++ b/src/py_bindings.rs @@ -143,13 +143,8 @@ impl PyServer { // 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 { - // Locating the vendor roots needs the GIL; page-locking does not, and - // it can take a while — a cold `cudaFree(NULL)` creates the primary - // context (~210 ms measured) when the framework has not initialised - // CUDA yet. `detach` is still synchronous: it releases the GIL, runs - // on this thread, and reacquires before returning, so the guarantee - // above holds. Holding the GIL across a CUDA call would also let echo - // deadlock against a thread that holds a CUDA lock and wants the GIL. + // 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()))?; From e12c75451463d4930bcb52244b2b309d8c3b2a8e Mon Sep 17 00:00:00 2001 From: Sasha Abramowitz Date: Mon, 3 Aug 2026 15:50:26 +0200 Subject: [PATCH 15/15] chore: date the 0.2.0 changelog entry for release Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d6683b2..1e36d39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ 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). -## [0.2.0] - unreleased +## [0.2.0] - 2026-08-03 ### Added