`nemo_rl.data_plane.adapters.transfer_queue`#

Adapter wiring :class:DataPlaneClient onto the transfer_queue package.

Pure plumbing — it owns the TQ controller / client handle and translates

class:: KVBatchMeta ↔ TQ’s own BatchMeta / KVBatchMeta. No business logic. Backend init is lifted from rl-arena/arena/backends.py; the call shapes are lifted from rl-arena/arena/dataplane_client.py.

Module Contents#

Classes#

`_StagingPool`	RDMA-registered host buffers, owned by one mooncake client.
`_StagingPoolRegistry`	Owns each client’s staging pool, keyed weakly so it dies with the client.
`_MooncakeMasterArgv`	Stand-in for the mooncake bootstrap module’s `subprocess` reference.
`TQDataPlaneClient`	Adapter façade — maps NeMo-RL calls onto TransferQueue’s public API.

Functions#

`_get_local_node_ip`	Return THIS process’s host IP, not the cluster head’s.
`rdma_devices`	Return this host’s RDMA devices as mooncake’s comma-separated list.
`_mooncake_transport_config`
`_memlock_limit`	Return this process’s RLIMIT_MEMLOCK soft limit, for error messages.
`_register_checked`	`store.register_buffer` with its status actually checked.
`_tq_shape_drift_error`	Build the error shared by the monkey-patch shape guards below.
`_patch_mooncake_register_check`	Make a failed RDMA registration fail at the registration.
`_assert_tq_stores_scalar_rows_0d`	Confirm a dense 1-D field really is stored as 0-d rows.
`_patch_scalar_field_schema`	Report the true `()` sample shape for dense 1-D fields.
`_patch_mooncake_staging_buffers`	Reuse RDMA-registered host buffers for mooncake tensor GETs and PUTs.
`_patch_mooncake_master_metrics_port`	Move mooncake_master’s metrics server onto a reserved port.
`_connect_existing`	Worker-process path: connect this process’s client to the Ray cluster.
`_init_tq`	Driver-process path: bootstrap the TQ controller for the chosen backend.
`_assert_no_key_loss`	Guard against silent leaf drops through TensorDict constructor rebuild.
`_from_wire`	Densify uniform nested tensors coming back from TQ.

Data#

`LOGGER`
`_STAGING_SLOT_TIMEOUT_S`
`_METRICS_PORT_DRIFT_CONSEQUENCE`

API#

nemo_rl.data_plane.adapters.transfer_queue.LOGGER#: ‘getLogger(…)’

nemo_rl.data_plane.adapters.transfer_queue._get_local_node_ip() → str#

Return THIS process’s host IP, not the cluster head’s.

Each Ray actor process must use its own node’s IP so Mooncake’s announce address (MC_TCP_BIND_ADDRESS → desc.ip_or_host_name in transfer_engine_impl.cpp) is routable cross-node. Non-routable addresses are rejected:

Link-local (169.254/16, fe80::/10) — gethostbyname can resolve to APIPA on hosts where avahi-autoipd is active.
Loopback (127.0.0.0/8, ::1) — hosts whose /etc/hosts maps the hostname to 127.0.0.1 would otherwise announce an unroutable address to Mooncake peers, causing cross-node connection refused.

nemo_rl.data_plane.adapters.transfer_queue.rdma_devices() → str#

Return this host’s RDMA devices as mooncake’s comma-separated list.

MC_MOONCAKE_DEVICE wins and is passed through verbatim (one device or a list). Otherwise every rail is offered: the NICs are split across NUMA domains, so naming only one device makes the other domain’s ranks cross the socket on every transfer.

Offering every rail is only safe because MC_ENABLE_DEST_DEVICE_AFFINITY pins each transfer’s peer rail to the local one by name, so a cross-rail pair is never formed — see

Mod:: nemo_rl.data_plane.adapters.transfer_queue_env. Without it, on a fabric where each rail is its own subnet, a cross-rail draw has no route and dies with “transport retry counter exceeded”.

IB and RoCE are never mixed; InfiniBand is preferred when present.

Also the skip predicate for the mooncake tests — mooncake_cpu is RDMA-only, so they cannot run without a device.

nemo_rl.data_plane.adapters.transfer_queue._mooncake_transport_config() → dict#

nemo_rl.data_plane.adapters.transfer_queue._STAGING_SLOT_TIMEOUT_S#: 600.0

nemo_rl.data_plane.adapters.transfer_queue._memlock_limit() → str#: Return this process’s RLIMIT_MEMLOCK soft limit, for error messages.

nemo_rl.data_plane.adapters.transfer_queue._register_checked(store: Any, ptr: int, nbytes: int) → None#

store.register_buffer with its status actually checked.

Mooncake returns a status int here, and TQ drops it at every call site (mooncake_client.py’s _register_all_buffers). A registration that fails is then invisible: the transfer into that unmapped region comes back as the generic TRANSFER_FAIL (-800), which carries no root cause, and burns its three retries against the same unmapped memory. Registration pins pages with ibv_reg_mr once per RDMA rail, so it is exactly the call that a memlock rlimit or a missing IPC_LOCK breaks.

None counts as success — the binding’s return type has varied across mooncake wheels, so only an explicit non-zero status is a failure.

class nemo_rl.data_plane.adapters.transfer_queue._StagingPool(store: Any, n_slots: int, max_bytes: int)#

RDMA-registered host buffers, owned by one mooncake client.

Not thread-local: the ThreadPoolExecutor is rebuilt inside each get/put, so thread-local buffers would be discarded every call. Sized to the executor width so no worker normally waits for a slot.

A slot’s buffer is registered for as long as the pool holds it. The invariant that matters is the converse: no buffer is ever freed while still registered, because mooncake would keep a mapping over an address the allocator immediately hands to the next caller.

Initialization

buffer(nbytes: int)#

class nemo_rl.data_plane.adapters.transfer_queue._StagingPoolRegistry(n_slots: int, max_bytes: int)#

Owns each client’s staging pool, keyed weakly so it dies with the client.

Weak keys because the registry is reachable from the patched class for the process lifetime; a strong table would pin every client’s registered buffers for that long.

Initialization

pool_for( client: Any, ) → nemo_rl.data_plane.adapters.transfer_queue._StagingPool#

Return client’s pool, building it at most once across threads.

Locked because put/get drive the thread workers from a ThreadPoolExecutor, so two of them reach a cold client at once whenever a call splits into more than one BATCH_SIZE_LIMIT batch. Unsynchronized, the loser’s pool is dropped on the floor and its buffers are freed while still registered — see :func:_register_checked for why that surfaces as a bare TRANSFER_FAIL. The lock is taken on every lookup rather than double-checked: it is uncontended after the first transfer, and nanoseconds against a millisecond RDMA transfer is not worth reasoning about visibility.

nemo_rl.data_plane.adapters.transfer_queue._tq_shape_drift_error( missing: str, consequence: str, target: str, *, opt_out: bool = False, ) → RuntimeError#

Build the error shared by the monkey-patch shape guards below.

All three guards fire for the same reason — the pinned transfer_queue revision no longer has the internals a patch depends on — so they share this message shape rather than each hand-rolling it.

nemo_rl.data_plane.adapters.transfer_queue._patch_mooncake_register_check() → None#

Make a failed RDMA registration fail at the registration.

Upstream’s _register_all_buffers ignores register_buffer’s status, so every worker that uses it — including the two bytes workers

Func:: _patch_mooncake_staging_buffers leaves alone — transfers into memory the NIC may never have mapped and reports only TRANSFER_FAIL (-800). Applied for every mooncake_cpu client, independent of reuse_registered_buffers, so the check survives disabling the pool.

Raises if _register_all_buffers is missing, rather than returning early: this check has no reuse_registered_buffers-style opt-out, so silently skipping it would put a failed registration back to surfacing only as a bare TRANSFER_FAIL, which is exactly the diagnosability this patch exists to add.

nemo_rl.data_plane.adapters.transfer_queue._assert_tq_stores_scalar_rows_0d() → None#

Confirm a dense 1-D field really is stored as 0-d rows.

Func:: _patch_scalar_field_schema rewrites the reported sample shape to () on that premise, and nothing reshapes the payload to compensate any more. If a TQ revision started storing 1-D fields as (1,) rows instead — fixing the same bug from the other side — the rewrite would turn a correct schema into a wrong one, and the symptom would be corrupt reads rather than an import error.

So ask TQ directly rather than trusting the pin.

Raises rather than skipping when the storage module is gone: the caller reached here only after importing transfer_queue.metadata, so “TQ isn’t installed” is no longer a live explanation — a missing module means the layout moved, which is exactly what this guard exists to catch.

nemo_rl.data_plane.adapters.transfer_queue._patch_scalar_field_schema() → None#

Report the true () sample shape for dense 1-D fields.

Upstream transfer_queue.metadata.extract_field_schema rebinds a local for 1-D inputs::

if len(value.shape) == 1:
    value = value.unsqueeze(-1)     # local only
first_item = value[0]               # -> shape (1,)

but the value that reaches storage is the original (N,) tensor, which KVStorageManager._generate_values iterates into N 0-d rows. So the schema claims a per-sample shape of (1,) while the stored rows are ().

Only the KV path notices. BatchMeta.get_shapes repeats the uniform shape per sample for non-nested fields, and KVStorageManager hands that list to the client, which reshapes raw bytes with it — so a scalar column reconstructs as (1,) rows and _merge_tensors_to_tensordict then re-nests it instead of taking its all(dim() == 0) -> torch.stack branch. SimpleStorage fetches stored objects by (index, field) and never consults the schema, which is why the symptom is mooncake_cpu-only.

Byte counts are unaffected either way (prod(()) == prod((1,)) == 1); this is a reshape/dtype-of-container bug, not a sizing one.

Applied on every backend so one partition’s schema cannot disagree with itself across processes. There is no payload-side fallback, so the premise is verified against TQ itself before the patch is installed — see :func:_assert_tq_stores_scalar_rows_0d.

nemo_rl.data_plane.adapters.transfer_queue._patch_mooncake_staging_buffers(max_bytes: int) → None#

Reuse RDMA-registered host buffers for mooncake tensor GETs and PUTs.

Upstream’s thread workers allocate a fresh destination per call and register/unregister it on the critical path. Pinning pages for DMA costs several times the wire time for the same bytes, and because the buffers are freed each call the pointers are always new, so nothing can be cached. This keeps a small pool of registered buffers alive instead.

Monkey-patched because TransferQueue is pinned by git SHA in pyproject.toml. Raises if the internals it drives are not shaped as expected, rather than returning early: a silent return would leave reuse_registered_buffers: true reading as on while the pool is never built, with no symptom besides lost throughput.

class nemo_rl.data_plane.adapters.transfer_queue._MooncakeMasterArgv(wrapped: Any, metrics_port: int)#

Stand-in for the mooncake bootstrap module’s subprocess reference.

Overrides Popen, and only for mooncake_master’s argv; everything else the bootstrap reaches for (STDOUT, the offload client’s launch) delegates to the real module untouched.

Initialization

__getattr__(name: str) → Any#

Popen(

args: Any,

*rest: Any,

**kwargs: Any,

) → Any#: Append --metrics_port when this is the master being launched.

nemo_rl.data_plane.adapters.transfer_queue._METRICS_PORT_DRIFT_CONSEQUENCE#: “the –metrics_port TQ omits cannot be applied, leaving mooncake_master’s metrics server on its 9003 …”

nemo_rl.data_plane.adapters.transfer_queue._patch_mooncake_master_metrics_port(port: int) → None#

Move mooncake_master’s metrics server onto a reserved port.

MasterAdminServer::Start binds the metrics socket before it consults enable_metric_reporting, and the master exits non-zero if that bind fails, so the metrics_port gflag default (9003) is a port the job depends on whether or not anything scrapes it — and it sits inside the ephemeral range these nodes hand out as source ports. TQ forwards no --metrics_port, nor a --config_path file that could carry one, so without this the metrics server is the one data-plane port that cannot move into ray.sub’s band and the master can still lose a startup race it has no reason to be in.

TQ does build the master’s argv in this process, though: tq.init -> _maybe_create_tq_storage -> initialize_mooncake_storage all run on the driver, so appending the flag to the subprocess.Popen the bootstrap calls is enough. gflags takes the last occurrence of a repeated flag, so this stays correct if a future TQ revision starts passing its own.

The provider registry holds a functools.wraps wrapper closed over the original bootstrap function, so rebinding the module attribute alone would never be called — the same trap extract_field_schema has. Re-registering is also where the drift check lives: if the bootstrap ever launches the master by some other route the flag stops landing silently, putting the metrics server back on 9003, so the bootstrap is required to have gone through the wrapped Popen.

nemo_rl.data_plane.adapters.transfer_queue._connect_existing() → None#

Worker-process path: connect this process’s client to the Ray cluster.

Connects to the already-running named controller actor. Mirrors rl-arena/arena/dataplane_client.py’s tq.init() (no args) call.

nemo_rl.data_plane.adapters.transfer_queue._init_tq( cfg: nemo_rl.data_plane.interfaces.DataPlaneConfig, *, checkpointing: bool = False, ) → None#: Driver-process path: bootstrap the TQ controller for the chosen backend.

nemo_rl.data_plane.adapters.transfer_queue._assert_no_key_loss( src_dict: dict, new_td: tensordict.TensorDict, fn: str, ) → None#

Guard against silent leaf drops through TensorDict constructor rebuild.

tensordict’s constructor has historically dropped NonTensorStack / NonTensorData leaves when built from a plain dict. Compare the source dict’s keys against the rebuilt TD’s top-level keys.

nemo_rl.data_plane.adapters.transfer_queue._from_wire(td: tensordict.TensorDict) → tensordict.TensorDict#

Densify uniform nested tensors coming back from TQ.

Both storage managers reconstruct every non-scalar field as a nested tensor, including fields whose rows all share a shape. Densify those so regular batched inputs retain their dense representation; truly ragged fields stay nested.

Per-sample scalar columns need no handling here: with

Func:: _patch_scalar_field_schema applied they are stored and reported as 0-d rows, which _merge_tensors_to_tensordict stacks into a dense (N,) column before it ever reaches this function.

Packed multimodal fields are excluded: their rows are per-sample media, not a padded sequence, and “all rows share a shape” is a data-dependent accident (every sample happening to carry one image). Stacking them discards the row boundaries that PackedTensor.from_wire needs, and the dense value then fails the is_nested check in codec.materialize and reaches get_multimodal_dict unreassembled. codec.materialize applies the same exclusion.

class nemo_rl.data_plane.adapters.transfer_queue.TQDataPlaneClient( cfg: nemo_rl.data_plane.interfaces.DataPlaneConfig, *, bootstrap: bool = True, checkpointing: bool = False, )#

Bases: nemo_rl.data_plane.interfaces.DataPlaneClient

Adapter façade — maps NeMo-RL calls onto TransferQueue’s public API.

Initialization

Construct a TQ-backed client.

Parameters:

cfg – data-plane config (backend selection, poll cadence, …).
bootstrap – True (driver) bootstraps the TQ controller using cfg. False (worker) connects this process to an already-running named controller actor in the Ray cluster — cfg is then only consulted for client-side knobs (poll interval).
checkpointing – Whether the caller will save or restore data-plane state. Used only at bootstrap; workers inherit the mode from TQ.

_gdr_requested: bool#: False

_gdr_put_confirmed: bool#: False

__getstate__() → dict[str, Any]#: Serialize the config needed to rebuild a process-local TQ client.

__setstate__(state: dict[str, Any]) → None#: Rebuild process-local TQ state after Ray deserialization.

static _read_complete_checkpoint_metadata( checkpoint_dir: str | pathlib.Path, ) → dict[str, Any]#: Read TQ metadata and require a complete storage payload.

_require_checkpointing_support() → None#: Reject backends that cannot round-trip all data-plane state.

_mark_data_operation_started() → None#: Make a later checkpoint load fail instead of mixing TQ states.

_require_clean_for_load() → None#: Reject restore after this client has performed a data operation.

register_partition( partition_id: str, fields: list[str], num_samples: int, consumer_tasks: list[str], grpo_group_size: int | None = None, enums: dict[str, list[str]] | None = None, ) → None#

claim_meta( partition_id: str, task_name: str, required_fields: list[str], batch_size: int, dp_rank: int | None = None, blocking: bool = True, timeout_s: float = 60.0, ) → nemo_rl.data_plane.interfaces.KVBatchMeta#

get_data( meta: nemo_rl.data_plane.interfaces.KVBatchMeta, select_fields: list[str] | None = None, ) → tensordict.TensorDict#

check_consumption_status( partition_id: str, task_names: list[str], ) → bool#

put_samples( sample_ids: list[str], partition_id: str, fields: tensordict.TensorDict | None = None, tags: list[dict[str, Any]] | None = None, ) → nemo_rl.data_plane.interfaces.KVBatchMeta#

get_samples( sample_ids: list[str], partition_id: str, select_fields: list[str], ) → tensordict.TensorDict#

list_sample_ids(partition_id: str) → list[str]#: List TQ keys in partition_id without fetching tensor payloads.

clear_samples(sample_ids: list[str] | None, partition_id: str) → None#

save_checkpoint( checkpoint_dir: str | pathlib.Path, *, metadata: dict[str, Any] | None = None, ) → None#: Save TQ controller metadata and storage data.

load_checkpoint( checkpoint_dir: str | pathlib.Path, ) → dict[str, Any]#

Restore TQ state after initialization and before data operations.

The local lifecycle guard cannot observe operations issued by another TQ client, so the recovery coordinator must also guarantee globally clean setup ordering.

close() → None#

nemo_rl.data_plane.adapters.transfer_queue#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.data_plane.adapters.transfer_queue`#