`nemo_rl.data_plane.nixl_storage_unit`#

NixlStorageUnit: a CPU Ray actor that owns one NIXL-registered DRAM slab.

Passive on the data path: clients WRITE and READ one-sided. The unit only serves small control RPCs (alloc, commit, resolve_keys, release_rows) and tracks pins and refcounts so a region is never reused under an in-flight read.

Index-addressed, like TransferQueue’s SimpleStorage: a row (TQ global index) belongs to unit unit_ids[gidx % N], and that unit keeps the row’s field -> (blob, entry) index. The TQ controller stores no locations.

Started by :func:nemo_rl.data_plane.adapters.tq_nixl.initialize_nixl_storage; the store-side pieces it uses live in :mod:nemo_rl.data_plane.nixl.

Module Contents#

Classes#

`BlobRec`
`NixlStorageUnit`	Owns one registered DRAM slab.

Functions#

`split_key`	`"<global_index>@<field>"` (TQ’s KV key) -> `(global_index, field)`.
`_free_at`
`_close`
`_write_all`
`_copy_out`	Append the first `nbytes` of file `fd` to open file `f`.
`_locked`	Serialise a unit method under the unit lock.

Data#

`log`
`_MAX_IO_JOBS`
`RowEntry`
`_COPY_CHUNK`

API#

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

nemo_rl.data_plane.nixl_storage_unit._MAX_IO_JOBS#: 256

nemo_rl.data_plane.nixl_storage_unit.RowEntry#: None

nemo_rl.data_plane.nixl_storage_unit.split_key(key: str) → tuple[int, str]#: "<global_index>@<field>" (TQ’s KV key) -> (global_index, field).

class nemo_rl.data_plane.nixl_storage_unit.BlobRec#

off: int#: None

nbytes: int#: None

refs: int#: None

created_at: float#: None

published: bool#: False

entries: list[nemo_rl.data_plane.nixl.blob_format.Entry] | None#: None

last_resolve_at: float#: 0.0

pins: dict[str, float]#: ‘field(…)’

holds: int#: 0

in_slab: bool#: True

disk: int | None#: None

staged_off: int | None#: None

staging: bool#: False

stage_failed: bool#: False

staged_at: float#: 0.0

picked_up: bool#: False

wait_since: float#: 0.0

pin_expiry(now: float) → float#

nemo_rl.data_plane.nixl_storage_unit._free_at( a: nemo_rl.data_plane.nixl.allocator.SlabAllocator, off: int, nbytes: int, at: float, now: float, ) → None#

nemo_rl.data_plane.nixl_storage_unit._close(fd: int) → None#

nemo_rl.data_plane.nixl_storage_unit._write_all(fd: int, view: memoryview) → None#

nemo_rl.data_plane.nixl_storage_unit._COPY_CHUNK#: None

nemo_rl.data_plane.nixl_storage_unit._copy_out(fd: int, f, nbytes: int) → None#: Append the first nbytes of file fd to open file f.

nemo_rl.data_plane.nixl_storage_unit._locked(fn)#

Serialise a unit method under the unit lock.

Unit state is touched by the ZMQ server thread and by (rare) Ray-called methods; one re-entrant lock serialises both.

class nemo_rl.data_plane.nixl_storage_unit.NixlStorageUnit( unit_id: int, slab_bytes: int, *, read_pin_s: float, nixl_backend: str = 'UCX', nixl_init_params: dict[str, Any] | None = None, directory: Any = None, numa_mode: str = 'auto', numa_index: int = 0, ssd_dir: str | None = None, staging_bytes: int = 0, pickup_s: float = 10.0, alloc_lease_s: float | None = None, )#

Owns one registered DRAM slab.

Clients reach it over ZMQ (info()["zmq"]) for alloc / commit / resolve_keys / release_rows; Ray is only its lifecycle and admin path.

Initialization

ssd_dir: str | None#: None

staging_: nemo_rl.data_plane.nixl.allocator.SlabAllocator | None#: None

staging_base: int#: None

pickup_s: float#: None

ssd_stats: dict[str, float]#: None

_io_jobs: int#: None

_max_io_jobs: int#: None

_staged: dict[str, nemo_rl.data_plane.nixl_storage_unit.BlobRec]#: None

alloc_: nemo_rl.data_plane.nixl.allocator.SlabAllocator#: None

_worker: concurrent.futures.ThreadPoolExecutor#: None

_init_ssd( slab_bytes: int, ssd_dir: str | None, staging_bytes: int, pickup_s: float, ) → None#

_start_ticker() → None#: Start the maintenance ticker; call once the unit is fully built.

_tick(every: float = 1.0, report_s: float = 30.0) → None#

Ticker thread: free what aged out and what dead clients left.

Also drops expired pins and stale allocations, and (SSD on) prints stats when they change.

_ssd_stats() → dict[str, Any]#

_reclaim(now: float) → None#

_sweep(now: float) → None#

Forget what dead clients left behind (caller holds _mu).

Pins of readers that never unpinned are dropped once expired. An allocation never committed within alloc_lease_s is retired, not freed: its writer may have a WRITE still in flight (as abort(in_flight=True)).

_alloc_io(nbytes: int, now: float) → int#

IO-area region for a put the put area cannot take.

Busy if something will free room (jobs, staged copies, pins, quarantine), UnitFull if nothing will.

_submit_io( fn: Any, b: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → None#

Queue an IO job that owns rec.staged_off (caller holds _mu).

Raises before any state changes, so a failed submit leaves nothing to undo.

_spill( b: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → None#: IO worker: write a put’s IO-area region to a file, then publish it as a disk blob.

_new_file() → int#

Open a fresh file in the SSD dir and unlink it at once; the fd keeps it.

Nothing is left on the node’s disk when the unit dies, however it dies. Named by a fresh uuid, not the client-supplied blob id, so no id can steer the path.

_start_stage( b: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, now: float, ) → None#: Stage a disk blob into the IO area, if room can be found.

_evict_staged(nbytes: int, now: float) → None#

Drop unpinned staged copies, oldest first.

A copy no reader has pinned yet is kept for pickup_s after its stage, so the reader polling for it finds it (design §4.2).

_drop_staged( b: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, at: float, now: float, ) → None#: Free a blob’s IO-area region at at (caller holds _mu).

_stage( b: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → None#

IO worker: read a disk blob’s file into its IO-area region.

While a stage runs, the fd is closed only here (_dispose skips it).

info() → dict[str, Any]#

ping() → int#

alloc( blob_id: str, nbytes: int, n_entries: int, spill: bool = True, ) → int#

commit(blob_id: str) → str#

Publish a blob whose WRITE is DONE: "done" once its rows are indexed.

The blob’s own embedded index names its gidx@field keys; they go into the row index, and a key that already pointed elsewhere drops its old blob’s reference (an overwrite). Idempotent: a published blob answers "done" again, so a client whose reply was lost can ask again. A spilled blob is "pending" while its disk write runs; "lost" if the blob is gone (aborted or its write failed), and the put then fails.

_blob_view( rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → memoryview#: The blob’s bytes in the put area or, for a spill in progress, the IO area.

_io_view( rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → memoryview#: The blob’s IO-area region (staged_off is set while it holds one).

_publish( blob_id: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → None#: Install a committed blob’s entries in the row index (caller holds _mu).

abort(blob_id: str, in_flight: bool = True) → None#

Drop a blob whose WRITE failed; nothing was published.

If the WRITE may still be running (in_flight), the region is never freed: a late WRITE would land in whatever blob reused it. The leak is counted in retired_bytes. A published blob is not aborted: its rows are live (the commit reply was lost), and only clears release them.

release_rows(gidxs: list[int]) → int#

Drop every field of these rows; returns how many entries were dropped.

Row-wide, like SimpleStorage’s clear: TQ’s KV clear only names the fields all cleared rows share. Unknown rows are ignored (idempotent).

_unref(blob_id: str) → None#: One row-index entry stopped pointing at blob_id (caller holds _mu).

_drop( blob_id: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → None#: Free a gone blob’s copies, or after the running copy of it ends.

_dispose( blob_id: str, rec: nemo_rl.data_plane.nixl_storage_unit.BlobRec, ) → None#

Free every copy of a released blob (caller holds _mu).

A copy is free once its pins end: every read pins what it reads.

resolve(blob_ids: list[str]) → dict[str, Any]#

Return current slab offsets for blobs and pin them for read_pin_s.

Unknown blobs are listed under missing rather than raising, so the client can report exactly which keys are lost.

resolve_keys(keys: list[str]) → dict[str, Any]#

Locate gidx@field keys in the row index and pin their blobs.

Like :meth:resolve, but per key: entries maps a key to [slab_off, len, kind, dtype, shape] (slab_off already includes the entry’s offset in its blob); missing / pending / failed list keys, with :meth:resolve’s meanings.

unpin(pin_id: str) → None#

scan() → list[dict[str, Any]]#

stats() → dict[str, Any]#

read_bytes(off: int, nbytes: int) → bytes#: Test helper: copy bytes out of the slab through the actor.

save_shard(path: str) → dict[str, Any]#

Dump every live blob back to back into path, its live keys into path + ".rows".

Returns {"blobs": {blob_id: [off_in_file, nbytes, refs]}, "nbytes": total, "clears_during_save": n}. The rows file maps each blob to the gidx@field keys the row index points at it with; a blob’s own embedded index may still list overwritten or cleared keys, which a restore must not bring back.

The snapshot is taken under the unit lock; the file I/O runs without it (blocking syscalls, design §2.5), so puts, reads and clears keep being served. Each snapshotted blob carries a hold: one released meanwhile is freed only after it is copied. Every exit, including an exception at any step, drops exactly the holds taken (one outer try).

SSD on: a disk blob is copied from its file, through a dup of its fd taken just before the copy. A disk blob released meanwhile may already have its file closed; the save then raises (a clear during a checkpoint is outside TQ’s contract). A spill not acknowledged yet is not published, so it is skipped.

load_shard( path: str, blobs: dict[str, list[int]], spill: bool = False, ) → dict[str, list[str]]#

Read blobs (blob_id -> [off_in_file, nbytes, refs]) from path into fresh regions.

Blobs that do not fit are returned under failed so the caller can ask this unit again with spill; already-present blobs count as loaded. With spill and the SSD tier on, a blob the put area cannot take goes to a file instead, as a spilled put would (the caller asks for that only after a RAM-only pass on every unit, _load_in_place).

Regions are allocated under the unit lock but filled without it; a blob is published only once its bytes are in place, and only with the keys path + ".rows" lists as live (see :meth:save_shard).

_restore_rows( b: str, entries: list[nemo_rl.data_plane.nixl.blob_format.Entry], keep: set[str], ) → None#: Publish a restored blob with only its live keys (caller holds _mu).

_load_to_disk(f, b: str, off: int, nbytes: int) → bool#

Restore one blob from checkpoint file f into a fresh SSD file.

False (the caller tries another unit) with the SSD tier off, if the blob could never be staged back for a read, or if the disk write fails.

nemo_rl.data_plane.nixl_storage_unit#