nemo_rl.data_plane.nixl.bufpool#

Registered host buffers for one client process, reused across operations.

Registration is the expensive part of moving bytes with NIXL: a fresh 1 GiB buffer costs ~70 ms to register and deregister on GB300, against ~10 ms to copy it with torch’s parallel copy_ (scripts/bench_register_vs_copy.py). So the pool keeps buffers registered:

  • base_count buffers of base_bytes, registered at start. An op that fits takes one; with several, independent ops proceed concurrently.

  • Larger ops get a buffer rounded up to a power of two, registered on first use and kept (LRU, free buffers only) until max_oversize_bytes of oversize buffers are registered. Only then is anything deregistered.

A buffer that a failed transfer may still use (ep.touches_unfinished: a READ writing into it, or a WRITE reading from it) is never handed out again. It stays referenced and registered in retired until the process exits; a lost base buffer is replaced with a fresh one. If that replacement cannot be registered the pool shrinks; once no base buffer is left, acquire raises :class:PoolExhausted instead of waiting forever.

The endpoint is not thread-safe, so register / deregister go through ep_lock; the pool’s own bookkeeping has its own lock.

Module Contents#

Classes#

BufferPool

Functions#

_round_pow2

Data#

log

API#

nemo_rl.data_plane.nixl.bufpool.log#

‘getLogger(…)’

exception nemo_rl.data_plane.nixl.bufpool.PoolExhausted#

Bases: RuntimeError

No base buffer is left, or none came free before the caller’s timeout.

Initialization

Initialize self. See help(type(self)) for accurate signature.

exception nemo_rl.data_plane.nixl.bufpool.PoolTimeout#

Bases: nemo_rl.data_plane.nixl.bufpool.PoolExhausted

No base buffer came free before the caller’s timeout (retryable).

Initialization

Initialize self. See help(type(self)) for accurate signature.

nemo_rl.data_plane.nixl.bufpool._round_pow2(n: int) → int#
class nemo_rl.data_plane.nixl.bufpool.BufferPool(
ep: Any,
ep_lock: Any,
*,
base_bytes: int,
base_count: int = 2,
max_oversize_bytes: int = 8 << 30,
)#

Initialization

_register(buf: numpy.ndarray) → tuple[int, numpy.ndarray]#
_deregister(addr: int) → None#
acquire(
nbytes: int,
timeout_s: float | None = None,
) → Iterator[tuple[int, numpy.ndarray]]#

Yield (registered_addr, uint8 array) with at least nbytes.

Waits at most timeout_s (forever if None) for a base buffer.

_replace_base() → tuple[int, numpy.ndarray] | None#

A fresh registered base buffer, or None if registration fails.

Never raises: it runs in acquire’s finally and must not mask the caller’s own exception.

_retire(item: tuple[int, numpy.ndarray]) → bool#

Keep item out of the pool if a failed transfer may still use it.

_evict_for(incoming: int) → None#

Deregister least-recently-used free oversize buffers to stay in budget.

close() → None#