`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. With policy.generation.top_logprobs the worker adds two more jagged columns to the same staging put through the opaque attachment mapping: the sampler’s top-k ids and log-probs of the call’s generated tokens. They are not digest-covered; the finalizer checks their shape against the generated length.

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`.
`StagedPrefix`	One resolved `staging_chain`: its tokens plus how much media it staged.
`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_item_sizes`	Group `imgs_sizes` rows into media items.
`_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.
`resolve_admission_prefix_chains`	Resolve a `CaptureAdmission` to its prefix and staged media count.
`_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.
`_row_media_item_sizes`	The `(height, width)` rows of each item one media-enabled row staged.

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

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#

property item_sizes: nemo_rl.data_plane.tq_token_sink.MediaItemSizes#: The (height, width) rows of each media item, see media_item_sizes.

nemo_rl.data_plane.tq_token_sink.MediaItemSizes#: None

nemo_rl.data_plane.tq_token_sink.media_item_sizes( imgs_sizes: torch.Tensor, num_frames: torch.Tensor | None, ) → nemo_rl.data_plane.tq_token_sink.MediaItemSizes#

Group imgs_sizes rows into media items.

The one definition of an item boundary: a video is one item spanning its num_frames rows of imgs_sizes, a still image one item of one row. The sink’s prefix reader (media_count / media_sizes) and the Megatron stager’s media slicing both use it, so the count the preparer carries forward and the slice point the stager cuts at cannot disagree.

Raises:: ValueError – num_frames does not tile the imgs_sizes rows.

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. sampler_topk is the call’s staged sampler top-k (ids, logprobs), each [generated_len, k], only when the fetch requested it.

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

sampler_topk: tuple[torch.Tensor, torch.Tensor] | 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.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 media attachments. Route attachments are supported on both kinds of partition.

Initialization

property media_pixel_dtype: torch.dtype | None#: The pixel dtype this partition’s media column is pinned to (None for text).

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 tokens, routes, media and sampler 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.StagedPrefix#

One resolved staging_chain: its tokens plus how much media it staged.

expanded is the concatenated token_ids_delta chain: what the engine prompt must start with, what the Megatron preparer splices in, and what Gym’s capture core verifies. media_count is how many media items (images, or videos) the chain’s rows staged, so the next call can stage only the media new to it and the engine expands only the media after it. media_sizes gives each of those items’ (height, width) rows (one for a still, one per frame for a video) so the next call can check that the engine saw the retained media at the staged geometry, not just as many items.

expanded: list[int]#: None

media_count: int#: 0

media_sizes: nemo_rl.data_plane.tq_token_sink.MediaItemSizes#: ()

__add__( other: nemo_rl.data_plane.tq_token_sink.StagedPrefix, ) → nemo_rl.data_plane.tq_token_sink.StagedPrefix#

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

Worker-local cache of resolved staging_chain prefixes.

Initialization

install( source: nemo_rl.data_plane.tq_token_sink.TQTokenSource, ) → None#: Attach (or replace) the TQTokenSource and drop cached chains.

fetch(staging_chain: list[str]) → list[int]#

Assemble the flat (expanded) prefix token ids from staging_chain.

The vLLM worker’s path. It goes through fetch_chains so both entry points fill the one cache with the same entry: the media count is never cached as 0 for a chain that staged media.

fetch_chains( staging_chain: list[str], ) → nemo_rl.data_plane.tq_token_sink.StagedPrefix#

Assemble the prefix and its media count from staging_chain.

Resolves through 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.resolve_admission_prefix_chains( admission: Any, chain_prefix: nemo_rl.data_plane.tq_token_sink.ChainPrefixCache, ) → nemo_rl.data_plane.tq_token_sink.StagedPrefix#

Resolve a CaptureAdmission to its prefix and staged media count.

An inline required_prefix_token_ids prefix carries no media: Gym only inlines prefixes for text chains.

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_chains( staging_keys: list[str], ) → nemo_rl.data_plane.tq_token_sink.StagedPrefix#

Bulk-fetch the ordered delta chain and describe the media it staged.

media_count and media_sizes are read off the small media columns (never the pixels) and are only computed when this source was built with capture_media=True; otherwise it is 0 regardless of what the rows staged. The Megatron preparer’s media_prev_count therefore depends on the source’s capture_media matching the sink’s – both are set from setup’s capture_media in megatron_worker.setup_token_capture.

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, include_sampler_topk: 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. include_sampler_topk adds the sampler top-k columns (policy.generation.top_logprobs); a row without them fails the fetch like a missing row.

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._row_media_item_sizes( row: Any, ) → nemo_rl.data_plane.tq_token_sink.MediaItemSizes#

The (height, width) rows of each item one media-enabled row staged.

Item boundaries come from media_item_sizes, shared with the Megatron stager’s slicing.

Raises:: ValueError – num_frames does not tile the imgs_sizes rows.

nemo_rl.data_plane.tq_token_sink#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.data_plane.tq_token_sink`#