nemo_rl.data_plane.driver_mixin#

Driver-side TransferQueue helpers shared by TQPolicy and TQValue.

Module Contents#

Classes#

TQDriverMixin

Pad-target minting and column read/write against the data plane.

API#

class nemo_rl.data_plane.driver_mixin.TQDriverMixin#

Pad-target minting and column read/write against the data plane.

Hosts must provide cfg, dp_client, and the use_dynamic_batches / use_sequence_packing attribute pairs that Policy and Value both set.

_packing_args(
mb_tokens_key: str,
) → tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]#

Resolve (sequence_packing_args, dynamic_batching_args) for a given stage.

The stage is identified by mb_tokens_key ("logprob_mb_tokens" or "train_mb_tokens").

_stamp_pad_seqlen(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
) → None#

Mint GLOBAL_FORWARD_PAD_SEQLEN onto meta.extra_info (idempotent).

Cross-DP forward pad target. Preshard shards inherit it via dict(meta.extra_info) propagation.

_isolated_meta(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
*,
fields: list[str],
task_name: str,
) → nemo_rl.data_plane.interfaces.KVBatchMeta#

Narrow meta for one model’s dispatch and mint it a fresh pad target.

The mint is idempotent, so sharing or inheriting the target would let whichever model dispatches first decide the forward pad for the rest – and with ppo_epochs > 1 the caller’s meta is already stamped when the critic dispatches again.

read_from_dataplane(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
*,
select_fields: list[str],
pad_value_dict: Optional[dict[str, Any]] = None,
) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]#

Fetch + materialize columns from the data plane (TQ).

read_columns pads to meta.extra_info[GLOBAL_FORWARD_PAD_SEQLEN] — the same value workers pad to in their forward pass. Driver and workers thus return columns at one identical seq dim, with no driver-side knowledge of sequence_length_round.

write_to_dataplane(
meta: nemo_rl.data_plane.interfaces.KVBatchMeta,
fields: dict[str, Any],
) → None#

Write driver-computed columns to the data plane (TQ).