`nemo_rl.data_plane.tq_token_sink`#

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

TQStagingStore is NeMo RL’s single keyed-row transport for token custody. Both vLLM and MInf write canonical Gym call deltas through TQTokenSink. 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. On media-enabled partitions the same put also carries the processed media the engine ran on (see MEDIA_STAGING_FIELDS), so staged coordinates acknowledge tokens and pixels together. 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#

`StagedMediaTensors`	One staged call’s validated media bundle, in the engine’s native shapes.
`FetchedStagedCall`	One explicitly identified small-column finalization fetch result.
`TQStagingStore`	Shared keyed-row transport for all token-capture TQ codecs.
`TQTokenSink`	Gym `StagingSink` over `DataPlaneClient.put_samples`.
`ChainPrefixCache`	Worker-local cache of resolved `staging_chain` prefixes.
`TQTokenSource`	Gym `StagingSource` over `DataPlaneClient.get_samples`.

Functions#

`_bytes_tensor`	Encode non-empty bytes as one jagged TQ row.
`_optional_digest_fields`
`_media_sentinels`	Placeholders for absent media, one per tensor column, in that column’s dtype.
`validate_media_tensors`	Check a media bundle against the Omni capture contract.
`_check_int32_positive`
`_call_dp`	Call a DataPlaneClient method on a local client or a Ray actor handle.
`resolve_admission_prefix`	Resolve a `CaptureAdmission` to the flat prefix the engine prompt starts with.
`_media_columns`	Encode one row’s media columns for a media-enabled partition.
`_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_extras`	Verify media-partition metadata before exposing the decoded extras.
`_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.
`_row_scalar_bool`	Read one required bool flag from a single-row TQ result.

Data#

`MEDIA_PRESENT_FIELD`
`MEDIA_HAS_FRAMES_FIELD`
`MEDIA_IMGS_FIELD`
`MEDIA_IMGS_SIZES_FIELD`
`MEDIA_NUM_FRAMES_FIELD`
`MEDIA_METADATA_DIGEST_FIELD`
`MEDIA_TENSOR_COLUMNS`
`MEDIA_FLAG_FIELDS`
`MEDIA_METADATA_FIELDS`
`MEDIA_STAGING_FIELDS`
`_MEDIA_REQUIRED`
`_MEDIA_PIXEL_DTYPES`
`_MEDIA_INDEX_DTYPES`
`STAGING_FIELDS`
`_MODE_TO_CODE`
`_CODE_TO_MODE`

API#

nemo_rl.data_plane.tq_token_sink.MEDIA_PRESENT_FIELD#: ‘media_present’

nemo_rl.data_plane.tq_token_sink.MEDIA_HAS_FRAMES_FIELD#: ‘media_has_frames’

nemo_rl.data_plane.tq_token_sink.MEDIA_IMGS_FIELD#: ‘media_imgs’

nemo_rl.data_plane.tq_token_sink.MEDIA_IMGS_SIZES_FIELD#: ‘media_imgs_sizes’

nemo_rl.data_plane.tq_token_sink.MEDIA_NUM_FRAMES_FIELD#: ‘media_num_frames’

nemo_rl.data_plane.tq_token_sink.MEDIA_METADATA_DIGEST_FIELD#: ‘media_metadata_digest’

nemo_rl.data_plane.tq_token_sink.MEDIA_TENSOR_COLUMNS: dict[str, str]#: None

nemo_rl.data_plane.tq_token_sink.MEDIA_FLAG_FIELDS#: None

nemo_rl.data_plane.tq_token_sink.MEDIA_METADATA_FIELDS#: None

nemo_rl.data_plane.tq_token_sink.MEDIA_STAGING_FIELDS#: None

nemo_rl.data_plane.tq_token_sink._MEDIA_REQUIRED#: (‘imgs’, ‘imgs_sizes’)

nemo_rl.data_plane.tq_token_sink._MEDIA_PIXEL_DTYPES#: ()

nemo_rl.data_plane.tq_token_sink._MEDIA_INDEX_DTYPES#: ()

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

exception nemo_rl.data_plane.tq_token_sink.MediaMetadataIntegrityError#

Bases: ValueError

Stored media metadata does not match its framework-owned checksum.

Initialization

Initialize self. See help(type(self)) for accurate signature.

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.StagedMediaTensors#

One staged call’s validated media bundle, in the engine’s native shapes.

imgs is [1, total_patches, 3*P*P] packed patches, imgs_sizes [N, 2] per-frame [height, width], num_frames [N_videos] frame counts partitioning imgs_sizes (None for still images).

imgs: torch.Tensor#: None

imgs_sizes: torch.Tensor#: None

num_frames: torch.Tensor | None#: None

property patch_size: int#

nemo_rl.data_plane.tq_token_sink._media_sentinels(pixel_dtype: torch.dtype) → dict[str, torch.Tensor]#

Placeholders for absent media, one per tensor column, in that column’s dtype.

TQ rows cannot be empty, and TQ keeps one dtype per field across all live rows: a later put with a different dtype is logged by the controller and silently loses its shape metadata, which the KV (Mooncake) backend needs to reconstruct the row. So a sentinel must never introduce a second dtype into a column; each keeps its real column’s dtype and rank.

nemo_rl.data_plane.tq_token_sink.validate_media_tensors( attachments: collections.abc.Mapping[str, Any] | None, ) → nemo_rl.data_plane.tq_token_sink.StagedMediaTensors | None#

Check a media bundle against the Omni capture contract.

Shared by the sink (before any write) and the source (after every read), so malformed geometry is rejected at both ends instead of surfacing as a reshape error in the learner. None (no attachments) returns None; anything else must be a mapping with imgs and imgs_sizes and an optional num_frames. Raises TypeError / ValueError; supported dtypes are preserved, never cast.

nemo_rl.data_plane.tq_token_sink._check_int32_positive(tensor: torch.Tensor, name: str) → None#

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

extras: dict[str, Any] | None#: None

media_present: bool#: False

media_has_frames: bool#: False

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.TQStagingStore(dp_client: Any, *, staging_partition: str)#

Shared keyed-row transport for all token-capture TQ codecs.

Initialization

put( key: str, field_dict: dict[str, torch.Tensor], *, tags: dict[str, Any] | None = None, ) → None#

get( keys: list[str], *, select_fields: list[str], ) → tensordict.TensorDict#

clear(keys: list[str]) → None#

class nemo_rl.data_plane.tq_token_sink.TQTokenSink( dp_client: Any, *, staging_partition: str, capture_media: bool = False, media_pixel_dtype: torch.dtype | None = None, )#

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.

capture_media mirrors the staging partition’s schema: a media-enabled partition registers MEDIA_STAGING_FIELDS and every row written here carries them (flags False + sentinels for text calls); a text-only partition rejects attachments outright.

Initialization

stage( record: nemo_gym.token_id_capture.staging.records.StagedCallRecord, *, attachments: collections.abc.Mapping[str, Any] | None = None, ) → nemo_gym.token_id_capture.staging.records.StageResult#

Write the token row and its media attachments in one put.

Success is returned only after the combined write was acknowledged, so staged coordinates vouch for tokens and pixels together. TQ has no transactional rollback: if the write raises, the attempted key is discarded best-effort before the failure is reported (see _discard_failed_write).

_discard_failed_write(key: str) → None#

Reclaim whatever a failed combined write may have left behind.

TQ writes a row field by field and only then publishes readiness, so a raise mid-write can leave partial field keys with no logical row. Gym’s capture_failed coordinates carry no staging key, so the finalizer never learns about this key; the sink is the only owner able to clean it. The discard is idempotent (clearing an unknown key is a no-op). If it fails too the storage state is uncertain and is logged at ERROR for the operator; the call is still reported failed.

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

class nemo_rl.data_plane.tq_token_sink.ChainPrefixCache(source: Any | None = None)#

Worker-local cache of resolved staging_chain prefixes.

Initialization

install(source: Any) → None#: Attach (or replace) the TQTokenSource and drop cached chains.

fetch(staging_chain: list[str]) → list[int]#: Assemble prefix token ids from staging_chain, with a worker-local FIFO (256-entry) cache.

nemo_rl.data_plane.tq_token_sink.resolve_admission_prefix( admission: Any, chain_prefix: nemo_rl.data_plane.tq_token_sink.ChainPrefixCache, ) → list[int]#: Resolve a CaptureAdmission to the flat prefix the engine prompt starts with.

nemo_rl.data_plane.tq_token_sink._media_columns( media: nemo_rl.data_plane.tq_token_sink.StagedMediaTensors | None, sentinels: dict[str, torch.Tensor], ) → dict[str, torch.Tensor]#

Encode one row’s media columns for a media-enabled partition.

Tensors are written with TQ’s row dimension prepended and otherwise native: imgs drops its own leading 1 so the row is [total_patches, F] and the patch dim is the row’s leading (ragged) dim, exactly like token_ids_delta / routed_experts, which is what a batched nested read requires. fetch_media restores [1, total_patches, F]. Integer geometry is written as int32 (validated to fit) so each column has one dtype across real rows and sentinels.

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

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_media( items: list[nemo_rl.data_plane.tq_token_sink.FetchedStagedCall], ) → list[nemo_rl.data_plane.tq_token_sink.StagedMediaTensors]#

One batched read of the media tensor columns for rows known to carry media.

items come from fetch_for_finalization (their flags were read with the base columns) and must all have media_present=True and the same media_has_frames: a sentinel num_frames row is all zeros and would fail validation beside real frame counts. Results are returned in request order. A transport miss raises KeyError; malformed columns raise TypeError / ValueError.

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_extras( row: Any, *, verify_media: bool, ) → dict[str, Any] | None#: Verify media-partition metadata before exposing the decoded extras.

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._row_scalar_bool(row: Any, field_name: str) → bool#: Read one required bool flag from a single-row TQ result.

nemo_rl.data_plane.tq_token_sink#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.data_plane.tq_token_sink`#