`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#

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.
`_mooncake_transport_config`
`_connect_existing`	Worker-process path: connect this process’s client to the Ray cluster.
`_resolve_tq_pin`	Return the `TransferQueue` requirement string from nemo-rl metadata.
`_patch_tq_actor_runtime_env`	Inject a per-actor `runtime_env` pin into TQ’s actor `.options()`.
`_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.
`_promote_1d_leaves`	Promote declared scalar leaves to `(N, 1)` for Mooncake.
`_from_wire`	Normalize TQ reads and invert :func:`_promote_1d_leaves` when needed.

Data#

_TQ_RUNTIME_ENV_PATCHED

API#

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._mooncake_transport_config() → dict#

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._TQ_RUNTIME_ENV_PATCHED#: False

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

Return the TransferQueue requirement string from nemo-rl metadata.

Single source of truth is pyproject.toml — we read it back via importlib.metadata.requires so the runtime_env injection cannot drift from the dependency declaration.

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

Inject a per-actor runtime_env pin into TQ’s actor .options().

TQ spawns SimpleStorageUnit and TransferQueueController via Cls.options(...).remote(...) without a runtime_env, so they inherit the job-level env. In a multi-node container deployment where each node has its own /opt/nemo_rl_venv, the driver’s uv sync only updates ray-head’s venv and a worker-node actor fails with ModuleNotFoundError. This monkey-patch makes Ray pip-install TQ into a per-actor runtime_env on first spawn (cached per-node by Ray afterwards). Idempotent. Couples us to TQ’s internal class layout — if TQ restructures, this becomes a no-op with a logged warning and we fall back to per-node uv sync.

The pin is sourced from nemo-rl’s installed metadata via

Func:: _resolve_tq_pin so it cannot drift from pyproject.toml.

TODO(zhiyul): remove this patch once the nightly container image is published with TransferQueue baked in via pyproject.toml. When every node starts from that image, the base env already has TQ and Ray actors inherit it — this injection then becomes pure overhead (Ray builds a redundant per-actor pip env on top of the container’s existing TQ install). Drop the call from TQDataPlaneClient.__init__ and delete this function.

nemo_rl.data_plane.adapters.transfer_queue._init_tq(cfg: nemo_rl.data_plane.interfaces.DataPlaneConfig) → 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._promote_1d_leaves(td: tensordict.TensorDict) → tensordict.TensorDict#

Promote declared scalar leaves to (N, 1) for Mooncake.

The authoritative field list lives in

Data:: nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS. Declared fields must arrive as dense (N,) tensors. Any other dense 1D tensor is rejected so it cannot silently encounter TQ v0.1.9’s schema/data mismatch. NonTensorStack and NonTensorData leaves pass through.
Parameters:: td – TensorDict to validate and encode for the Mooncake wire format.
Returns:: TensorDict with declared scalar leaves promoted to (N, 1).
Raises:: ValueError – If a declared field is not a dense 1D tensor, or an undeclared field is a dense 1D tensor.

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

Normalize TQ reads and invert :func:_promote_1d_leaves when needed.

Both TQ v0.1.9 storage managers reconstruct every non-scalar field as a nested tensor, including fields whose rows all have the same shape. Densify those uniform nested tensors first so regular batched inputs retain their dense representation. Truly ragged fields remain nested. Finally, squeeze only singleton dimensions declared in

Data:: nemo_rl.data_plane.schema.PROMOTE_1D_FIELDS.

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

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).

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#

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

close() → None#

nemo_rl.data_plane.adapters.transfer_queue#