`nemo_rl.data_plane.tq_token_sink`#

TransferQueue implementations of NeMo-Gym’s token staging protocols.

TQTokenSink/TQTokenSource are NeMo-RL’s providers for the ledger-authoritative capture design: the sink is the worker-side write of one model call’s token delta to the rollout_staging partition — the design’s only heavy token hop — and the source is the finalizer’s read-back of those rows by staging key. This module is the only hot-path file that knows tokens live in TQ; Gym sees opaque staging keys.

Each staged row carries three jagged columns (token_ids_delta, token_mask_delta, generation_logprobs_delta), the complete receipt identity/lineage metadata, and all digest inputs so it round-trips to a normally validated StagedCallBaseSnapshot. Masks/logprobs are float32 on the wire, matching compute_staging_digest’s float32-bit-pattern scheme, so digest recomputation over fetched values is byte-exact. Route payloads never ride inside snapshots: the source returns them as separate RouteFragment values keyed by staging key, digest-verified by the plan executor at point of use.

Module Contents#

Classes#

`FetchedStagedCall`	One explicitly identified small-column finalization fetch result.
`TQTokenSink`	Gym `StagingSink` over `DataPlaneClient.put_samples`.
`TQTokenSource`	Gym `StagingSource` over `DataPlaneClient.get_samples`.

Functions#

`_bytes_tensor`	Encode non-empty bytes as one jagged TQ row.
`_optional_digest_fields`
`_call_dp`	Call a DataPlaneClient method on a local client or a Ray actor handle.
`_select_row`	Slice one row out of a batched fetch, restoring single-row shapes.
`_row_leaf`
`_row_text`
`_row_to_base_snapshot`	Rebuild one normally validated base snapshot; route bytes never enter it.
`_row_to_route_fragment`	Extract one staged route payload beside (never inside) the snapshot.
`_row_scalar_int`	Read one required scalar from a single-row TQ result.

Data#

`STAGING_FIELDS`
`_MODE_TO_CODE`
`_CODE_TO_MODE`

API#

nemo_rl.data_plane.tq_token_sink.STAGING_FIELDS#: None

nemo_rl.data_plane.tq_token_sink._MODE_TO_CODE#: None

nemo_rl.data_plane.tq_token_sink._CODE_TO_MODE#: None

nemo_rl.data_plane.tq_token_sink._bytes_tensor(value: bytes) → torch.Tensor#: Encode non-empty bytes as one jagged TQ row.

nemo_rl.data_plane.tq_token_sink._optional_digest_fields( value: str | None, ) → tuple[torch.Tensor, torch.Tensor]#

class nemo_rl.data_plane.tq_token_sink.FetchedStagedCall#

One explicitly identified small-column finalization fetch result.

fragment is populated only when the fetch requested route payloads (direct mode); deferred finalization leaves route bytes in TQ and carries only routed_len transport metadata.

staging_key: str#: None

snapshot: nemo_gym.token_id_capture.staging.records.StagedCallBaseSnapshot#: None

routed_len: int#: None

fragment: nemo_rl.experience.route_assembly.RouteFragment | None#: None

nemo_rl.data_plane.tq_token_sink._call_dp(

dp_client: Any,

method_name: str,

**kwargs: Any,

) → Any#: Call a DataPlaneClient method on a local client or a Ray actor handle.

class nemo_rl.data_plane.tq_token_sink.TQTokenSink(dp_client: Any, *, staging_partition: str)#

Gym StagingSink over DataPlaneClient.put_samples.

stage is synchronous and returns only after TQ acknowledged the write, so the capture layer’s fail-closed ordering (bytes durable before the model call is acked) holds by construction. Failures are reported in the StageResult; the finalizer turns a poisoned rollout into a placeholder row (see RolloutReassembler.finalize_group).

stage is thread-safe per the StagingSink contract: it holds no per-call mutable state, so the capture host may run writes for unrelated calls concurrently.

Initialization

stage( record: nemo_gym.token_id_capture.staging.records.StagedCallRecord, ) → nemo_gym.token_id_capture.staging.records.StageResult#

clear(staging_keys: list[str]) → None#: Drop staged rows (finalizer / eviction cleanup).

class nemo_rl.data_plane.tq_token_sink.TQTokenSource(dp_client: Any, *, staging_partition: str)#

Gym StagingSource over DataPlaneClient.get_samples.

All requested rows are fetched in a single batched get_samples call (TQ returns jagged delta columns as nested tensors; _from_wire preserves the raggedness), in the order requested. A missing or unreadable row raises KeyError per the protocol — the finalizer maps that to a placeholder, never a silent skip. TQ’s field-readiness check is all-or-nothing across a batch, so the extras fallback is batch-level: extras-free runs land in the base schema exactly like the old per-key probe, but a batch with mixed extras presence degrades every row to the base schema (worker feature-gating makes presence uniform per run).

Initialization

fetch( staging_keys: list[str], ) → list[nemo_gym.token_id_capture.staging.records.StagedCallBaseSnapshot]#: Gym StagingSource conformance: base snapshots only, in order.

fetch_prefix_token_ids(staging_keys: list[str]) → list[int]#: Bulk-fetch ordered delta chain and concatenate token_ids_delta into a prefix.

fetch_for_finalization( staging_keys: list[str], *, include_route_fragments: bool = False, ) → list[nemo_rl.data_plane.tq_token_sink.FetchedStagedCall]#

Fetch digest-covered base columns, plus route payloads when requested.

Deferred mode (the default) never selects routed_experts — route bytes stay in TQ for the policy worker. Direct mode passes include_route_fragments=True to pull the payloads in the same batched read and receives them as RouteFragment values beside the base snapshots, never inside them.

nemo_rl.data_plane.tq_token_sink._select_row( rows: tensordict.TensorDict, index: int, ) → dict[str, torch.Tensor]#

Slice one row out of a batched fetch, restoring single-row shapes.

_row_to_base_snapshot predates batching and expects each field with a leading batch dim of 1 (the shape a single-key get_samples returns), so re-add it after indexing. Indexing a nested tensor yields that row’s dense component, which is exactly the jagged-row payload.

nemo_rl.data_plane.tq_token_sink._row_leaf(row: Any, name: str) → torch.Tensor#

nemo_rl.data_plane.tq_token_sink._row_text(row: Any, name: str) → str#

nemo_rl.data_plane.tq_token_sink._row_to_base_snapshot( row: Any, ) → nemo_gym.token_id_capture.staging.records.StagedCallBaseSnapshot#: Rebuild one normally validated base snapshot; route bytes never enter it.

nemo_rl.data_plane.tq_token_sink._row_to_route_fragment( row: Any, ) → nemo_rl.experience.route_assembly.RouteFragment | None#: Extract one staged route payload beside (never inside) the snapshot.

nemo_rl.data_plane.tq_token_sink._row_scalar_int(row: Any, field_name: str) → int#: Read one required scalar from a single-row TQ result.

nemo_rl.data_plane.tq_token_sink#