`nemo_rl.data_plane.nixl.nixl_io`#

NIXL endpoint for the data plane: registration, metadata, blocking transfers.

The agent is created by :func:nemo_rl.utils.checkpoint_engines.nixl._create_nixl_agent, the same path the checkpoint-engine refit uses.

One NixlEndpoint per process. The agent is not thread-safe; callers serialise access (the KV client holds a lock, the unit is a Ray actor).

Memory type: DRAM host buffers over UCX.

Module Contents#

Classes#

`Graveyard`	Every NIXL request in this process that failed short of DONE.
`NixlEndpoint`

Functions#

`check_rdma_only`	Raise :class:`TransportPolicyError` unless UCX can only use RDMA.
`nixl_available`
`addr_of`

Data#

`log`
`_RDMA_TLS_PREFIXES`
`GRAVEYARD`

API#

nemo_rl.data_plane.nixl.nixl_io.log#: ‘getLogger(…)’

nemo_rl.data_plane.nixl.nixl_io._RDMA_TLS_PREFIXES#: (‘rc’, ‘dc’, ‘ud’, ‘ib’)

nemo_rl.data_plane.nixl.nixl_io.check_rdma_only(env: Mapping[str, str] | None = None) → str#

Raise :class:TransportPolicyError unless UCX can only use RDMA.

Allowed: an explicit UCX_TLS list with at least one RDMA transport (rc, rc_x, rc_mlx5, dc, ud, ib …) and no tcp or all, plus same-host helpers (self, sm, shm, cma …); or an exclusion list (^...) that names tcp. UCX_NET_DEVICES, if set, must name RDMA devices (<dev>:<port>), never a netdev such as eth0. Unset UCX_TLS is rejected: UCX would then fall back to TCP whenever RDMA is unavailable. Returns a one-line description.

class nemo_rl.data_plane.nixl.nixl_io.Graveyard(max_bytes: int = 4 << 30)#

Every NIXL request in this process that failed short of DONE.

One per process (:data:GRAVEYARD), so dropping an endpoint or adapter can’t free an entry through garbage collection. An entry holds the request’s handle (never released), its descriptor lists, the agent that owns it, and its local ranges; callers keep the memory behind those ranges alive (the pool retires it, the adapter keeps zero-copy sources).

Past max_bytes of retained memory (local ranges, plus every buffer the pool retires whole) the process is POISONED: every later transfer raises

Class:: DataPlaneFatal before touching NIXL.

Initialization

retain( agent: Any, h: Any, ldesc: Any, rdesc: Any, local, ) → None#

retain_buffer(nbytes: int) → None#

Charge a pool buffer retired whole because a failed transfer touches it.

Its transferred ranges were charged by :meth:retain, so they count twice: the limit errs on the side of poisoning.

_charge(nbytes: int) → None#

check() → None#

nemo_rl.data_plane.nixl.nixl_io.GRAVEYARD#: ‘Graveyard(…)’

nemo_rl.data_plane.nixl.nixl_io.nixl_available() → bool#

nemo_rl.data_plane.nixl.nixl_io.addr_of(arr: numpy.ndarray) → int#

class nemo_rl.data_plane.nixl.nixl_io.NixlEndpoint( name: str, *, backend: str = NIXL_DEFAULT_BACKEND_NAME, init_params: dict[str, Any] | None = None, prog_thread: bool = False, )#

Initialization

touches_unfinished(addr: int, nbytes: int) → bool#: True if [addr, addr+nbytes) overlaps a transfer that may still run.

register(arr: numpy.ndarray) → int#: Register a contiguous uint8 host array; returns its base address.

deregister(addr: int) → None#

metadata() → bytes#

add_remote(md: bytes) → str#

has_remote(name: str) → bool#

remove_remote(name: str) → None#

transfer( op: str, local: Sequence[tuple[int, int]], remote: Sequence[tuple[int, int]], remote_agent: str, *, timeout_s: float = 120.0, ) → None#

Blocking one-sided WRITE (local→remote) or READ (remote→local).

local and remote are parallel lists of (addr, nbytes).

transfer_many( ops: Sequence[tuple[str, Sequence[tuple[int, int]], Sequence[tuple[int, int]], str]], *, timeout_s: float = 120.0, ) → list[nemo_rl.data_plane.nixl.errors.TransferError | None]#

Post every (op, local, remote, remote_agent), then wait for all of them.

The transfers run at the same time, so N units cost about the slowest one, not the sum. One timeout_s bounds the wait for the whole batch. Returns one entry per op: None once it is DONE, else its

Class:: TransferError. Ops are independent: a failure of one does not stop the others. Invalid arguments raise ValueError before anything is posted.

Every posted transfer ends DONE (handle released) or retained (see

Meth:: _retain), on every exit, including an exception.

_post( i, op, local, remote, remote_agent, live, ) → nemo_rl.data_plane.nixl.errors.TransferError | None#: Post one transfer; live[i] is set while it is still running.

_fail( i, op_args, live, why, cause=None, ) → nemo_rl.data_plane.nixl.errors.TransferError#: Retain transfer i (it may still run) and return its error.

_retain( xfer: tuple[Any, Any, Any], local, ) → None#

Keep a transfer that may still run, for good.

The handle is never released: POSIX hands the request object to its I/O callbacks and deleting it under a pending I/O is a use-after-free.

_done(xfer: tuple[Any, Any, Any]) → None#

close() → None#

nemo_rl.data_plane.nixl.nixl_io#