`nemo_rl.algorithms.async_utils.staleness_sampler`#

Prompt-group staleness policies over a TQReplayBuffer.

A staleness policy owns the whole off-policyness contract for the SC async loop in one object, injected into both pumps:

admit (rollout-pump side): block until the next prompt batch may dispatch, then return the target_step stamp for that batch (None when the policy doesn’t stamp). Owning admission here is what lets the rollout pump follow whichever sampling algorithm is selected without a second, hand-kept copy of the gating logic.
select / evict (train-pump side): pick / drop prompt groups.
is_on_policy / required_buffer_capacity (derived facts): what the weight-sync and capacity-validation paths need without re-reading raw knobs, so those consumers can’t drift out of sync with the sampler.

PromptGroupSampler is the interface; WindowedSampler / ReadyFirstSampler / WeightFifoSampler / InOrderSampler are the built-in policies, one per behavior, each owning only the args that apply to it. create_sampler builds one from a discriminated-union config (or a module:ClassName FQN for a policy defined outside this repo) — the config’s name is the single source of truth for which behavior runs, so there are no cross-field knob combinations to validate.

Module Contents#

Classes#

`PromptGroupSampler`	Staleness policy shared by the SC rollout and train pumps.
`TransactionalAdmissionSampler`	Sampler whose blocking wait is separate from its cursor mutation.
`BaseSampler`	Shared machinery for the built-in policies.
`WindowedSampler`	Over-sampled windowed selection.
`_GatedSampler`	Base for policies that admit exactly one dispatch batch per trainer step.
`ReadyFirstSampler`	Gated admission with ready-first, mixed-version selection.
`WeightFifoSampler`	Gated, strict weight-version FIFO.
`InOrderSampler`	Gated, exact batch->step matching.
`WindowedSamplerConfig`
`ReadyFirstSamplerConfig`
`WeightFifoSamplerConfig`
`InOrderSamplerConfig`
`CustomSamplerConfig`

Functions#

`_gated_required_buffer_capacity`	Return capacity for one live batch plus each lookahead batch.
`required_buffer_capacity_for_config`	Return a sampler config’s required capacity without constructing it.
`_custom_sampler_class`	Import and return a custom sampler class without constructing it.
`_sampler_class_for_config`	Return the sampler class selected by a built-in or custom config.
`sampler_supports_buffer_checkpoint`	Return a sampler class’s static replay-checkpoint capability.
`sampler_supports_training_claims`	Return whether selection transfers rows into training ownership.
`create_sampler`	Build a sampler from its config (or import one by FQN).

Data#

`_GATE_POLL_SECONDS`
`SamplerConfig`

API#

nemo_rl.algorithms.async_utils.staleness_sampler._GATE_POLL_SECONDS#: 0.005

class nemo_rl.algorithms.async_utils.staleness_sampler.PromptGroupSampler#

Bases: typing.Protocol

Staleness policy shared by the SC rollout and train pumps.

Implement this (or subclass BaseSampler) to add a custom sampling algorithm; point async_rl.sampler at module:ClassName to load it. A custom sampler that supports replay recovery must explicitly declare supports_buffer_checkpoint = True. It must additionally declare supports_training_claims = True before periodic rollout snapshots may be enabled; omitting that optional capability preserves the legacy remove-on-selection behavior.

async admit( *, trainer_version_fn: Callable[[], int], ) → Optional[int]#

Block until the next prompt batch may dispatch.

Parameters:: trainer_version_fn – Zero-arg accessor for the live trainer version (polled, so a blocking policy sees updates while it waits).
Returns:: The target_step to stamp on this batch’s slots, or None when the policy does not stamp target steps.

async select( *, current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, ) → tuple[Optional[nemo_rl.data_plane.KVBatchMeta], int]#

Pick up to max_prompt_groups eligible groups for training.

Claim-aware samplers transfer the groups from ordinary replay-buffer selection into training ownership until the controller releases them. Legacy custom samplers may still remove selected groups immediately.

async evict(*, current_train_weight: int) → int#: Drop groups that can no longer be selected; clear their DP rows.

should_abort_inflight( *, start_weight_version: int, current_train_weight: int, ) → bool#: Return whether an unfinished rollout can no longer be selected.

property is_on_policy: bool#: True when the policy admits zero staleness (sync mode).

supports_buffer_checkpoint: ClassVar[bool]#

None

Whether completed buffered groups can be restored safely.

supports_training_claims: ClassVar[bool]#

None

Whether selected groups remain owned until the train step commits.

required_buffer_capacity(groups_per_step: int) → Optional[int]#: Buffer-capacity the policy needs, or None if unconstrained.

property dispatch_index: int#: Last admitted dispatch batch index.

set_dispatch_index(resume_from_trainer_version: int) → None#: Seed the cursor for checkpoints that predate exact sampler state.

restore_dispatch_index(dispatch_index: int) → None#: Restore the exact dispatch cursor from controller state.

class nemo_rl.algorithms.async_utils.staleness_sampler.TransactionalAdmissionSampler#

Bases: typing.Protocol

Sampler whose blocking wait is separate from its cursor mutation.

async wait_until_admissible( *, trainer_version_fn: Callable[[], int], ) → None#: Wait until one admission can commit without mutating sampler state.

commit_admission( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, ) → Optional[int]#: Advance the cursor under a live data-plane mutation cut.

class nemo_rl.algorithms.async_utils.staleness_sampler.BaseSampler( buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer, )#

Bases: abc.ABC

Shared machinery for the built-in policies.

Owns the monotonic dispatch counter (the batch index formerly tracked as SingleControllerActor._max_rollout_version) and the common select-finalize / weight-window-evict helpers.

Initialization

supports_buffer_checkpoint: ClassVar[bool]#: False

supports_training_claims: ClassVar[bool]#: True

property dispatch_index: int#: Return the last admitted dispatch batch index.

set_dispatch_index(resume_from_trainer_version: int) → None#

Seed the cursor for checkpoints that predate exact sampler state.

Parameters:: resume_from_trainer_version – Trainer version from which the run resumes. The next admitted batch receives that version.

restore_dispatch_index(dispatch_index: int) → None#

Restore the exact dispatch cursor.

Parameters:: dispatch_index – Last admitted batch index, or -1 when no batch has been admitted. Call before the first admit.

abstractmethod async admit( *, trainer_version_fn: Callable[[], int], ) → Optional[int]#

abstractmethod async select( *, current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, ) → tuple[Optional[nemo_rl.data_plane.KVBatchMeta], int]#

async evict(*, current_train_weight: int) → int#

Default: drop ready groups below the weight window.

Skips unready (reserved-but-uncommitted) slots so eviction can’t race a concurrent commit that re-looks-up the slot after its await. Policies whose select key isn’t the start weight (e.g. InOrderSampler) override this so evict and select agree.

should_abort_inflight( *, start_weight_version: int, current_train_weight: int, ) → bool#

property is_on_policy: bool#

required_buffer_capacity(groups_per_step: int) → Optional[int]#

_eviction_window() → int#: Weight-version span kept selectable; drives the default evict.

static _validate_group_bounds( min_prompt_groups: int, max_prompt_groups: int, ) → None#

async _finalize_selection( valid_idxs: list[int], min_prompt_groups: int, max_prompt_groups: int, ) → tuple[Optional[nemo_rl.data_plane.KVBatchMeta], int]#

Cap, drop from the buffer, and concat the chosen groups.

Greedy without waiting: returns all currently-eligible groups up to max_prompt_groups (never fewer on purpose, never waits to fill it), or (None, 0) below min_prompt_groups.

class nemo_rl.algorithms.async_utils.staleness_sampler.WindowedSampler( buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer, *, max_staleness_versions: int, sample_freshest_first: bool = False, min_groups_for_streaming_train: Optional[int] = None, )#

Bases: nemo_rl.algorithms.async_utils.staleness_sampler.BaseSampler

Over-sampled windowed selection.

Rollout never gates on the trainer version — the pump keeps producing and samples aged past the window are evicted. select takes any ready group within [train_weight - max_staleness_versions, train_weight], optionally freshest-first.

Initialization

supports_buffer_checkpoint: ClassVar[bool]#: True

_eviction_window() → int#

required_buffer_capacity(groups_per_step: int) → Optional[int]#

should_abort_inflight( *, start_weight_version: int, current_train_weight: int, ) → bool#

async wait_until_admissible( *, trainer_version_fn: Callable[[], int], ) → None#: Return immediately because buffer capacity is this policy’s gate.

commit_admission( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, ) → Optional[int]#: Return the unstamped admission result without changing a cursor.

async admit( *, trainer_version_fn: Callable[[], int], ) → Optional[int]#

async select( *, current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, ) → tuple[Optional[nemo_rl.data_plane.KVBatchMeta], int]#

nemo_rl.algorithms.async_utils.staleness_sampler._gated_required_buffer_capacity( groups_per_step: int, *, gate_window: int, ) → int#: Return capacity for one live batch plus each lookahead batch.

class nemo_rl.algorithms.async_utils.staleness_sampler._GatedSampler( buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer, *, gate_window: int, )#

Bases: nemo_rl.algorithms.async_utils.staleness_sampler.BaseSampler

Base for policies that admit exactly one dispatch batch per trainer step.

The gate bounds how far generation may run ahead of the trainer (gate_window versions of lookahead).

Initialization

set_gate_window(gate_window: int) → None#

_eviction_window() → int#

required_buffer_capacity(groups_per_step: int) → Optional[int]#

async wait_until_admissible( *, trainer_version_fn: Callable[[], int], ) → None#: Wait for the gate without advancing the durable dispatch cursor.

commit_admission( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, ) → Optional[int]#: Advance the cursor after the controller enters its mutation cut.

_commit_admission() → Optional[int]#: Advance admission for the legacy monolithic API.

async admit( *, trainer_version_fn: Callable[[], int], ) → Optional[int]#

_stamp() → Optional[int]#

class nemo_rl.algorithms.async_utils.staleness_sampler.ReadyFirstSampler( buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer, *, max_staleness_versions: int, )#

Bases: nemo_rl.algorithms.async_utils.staleness_sampler._GatedSampler

Gated admission with ready-first, mixed-version selection.

Admission limits generation to max_staleness_versions dispatch batches ahead of the current trainer version. Selection remains ready-first across weight versions instead of draining one version at a time: every ready group generated by a policy version no newer than the trainer stays selectable, including late stragglers outside the admission window, so no rollout is ever discarded.

Initialization

supports_buffer_checkpoint: ClassVar[bool]#: True

async select( *, current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, ) → tuple[Optional[nemo_rl.data_plane.KVBatchMeta], int]#

async evict(*, current_train_weight: int) → int#

class nemo_rl.algorithms.async_utils.staleness_sampler.WeightFifoSampler( buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer, *, max_staleness_versions: int, )#

Bases: nemo_rl.algorithms.async_utils.staleness_sampler._GatedSampler

Gated, strict weight-version FIFO.

select drains the oldest in-window start_weight first and waits for that weight’s batch to fill. Evict uses the weight window (default).

Initialization

supports_buffer_checkpoint: ClassVar[bool]#: True

async select( *, current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, ) → tuple[Optional[nemo_rl.data_plane.KVBatchMeta], int]#

class nemo_rl.algorithms.async_utils.staleness_sampler.InOrderSampler( buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer, *, max_lookahead_versions: int, warmup_lookahead_versions: Optional[int] = None, )#

Bases: nemo_rl.algorithms.async_utils.staleness_sampler._GatedSampler

Gated, exact batch->step matching.

Each dispatched batch is stamped with its dispatch index as target_step; select takes the batch whose target_step equals the trainer version (the staleness window is not used for selection). evict is keyed on target_step — not the start weight — so a slot whose target step is still upcoming is never dropped early, and evict/select can’t disagree.

warmup_lookahead_versions widens the gate while the PPO policy is frozen, so more batches stay in flight during critic warmup. The driver retunes the window every step and shrinks it back to max_lookahead_versions once the policy starts training, so the widened lookahead does not turn into permanent extra staleness. Buffer capacity is sized for the peak of the two.

Initialization

supports_buffer_checkpoint: ClassVar[bool]#: True

required_buffer_capacity(groups_per_step: int) → Optional[int]#

_stamp() → Optional[int]#

async select( *, current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, ) → tuple[Optional[nemo_rl.data_plane.KVBatchMeta], int]#

async evict(*, current_train_weight: int) → int#

class nemo_rl.algorithms.async_utils.staleness_sampler.WindowedSamplerConfig#

Bases: pydantic.BaseModel

name: Literal[windowed]#: ‘windowed’

max_staleness_versions: pydantic.NonNegativeInt#: 1

sample_freshest_first: bool#: False

class nemo_rl.algorithms.async_utils.staleness_sampler.ReadyFirstSamplerConfig#

Bases: pydantic.BaseModel

name: Literal[ready_first]#: ‘ready_first’

max_staleness_versions: pydantic.NonNegativeInt#: 1

class nemo_rl.algorithms.async_utils.staleness_sampler.WeightFifoSamplerConfig#

Bases: pydantic.BaseModel

name: Literal[weight_fifo]#: ‘weight_fifo’

max_staleness_versions: pydantic.NonNegativeInt#: 1

class nemo_rl.algorithms.async_utils.staleness_sampler.InOrderSamplerConfig#

Bases: pydantic.BaseModel

name: Literal[in_order]#: ‘in_order’

max_lookahead_versions: pydantic.NonNegativeInt#: 1

warmup_lookahead_versions: Optional[pydantic.NonNegativeInt]#: None

validate_warmup_lookahead() → nemo_rl.algorithms.async_utils.staleness_sampler.InOrderSamplerConfig#

property peak_lookahead_versions: int#: Widest window the run can reach; what buffer capacity must cover.

class nemo_rl.algorithms.async_utils.staleness_sampler.CustomSamplerConfig#

Bases: pydantic.BaseModel

name: Literal[custom]#: ‘custom’

target: str#: None

nemo_rl.algorithms.async_utils.staleness_sampler.SamplerConfig#: None

nemo_rl.algorithms.async_utils.staleness_sampler.required_buffer_capacity_for_config( cfg: nemo_rl.algorithms.async_utils.staleness_sampler.SamplerConfig, groups_per_step: int, *, min_groups_for_streaming_train: int, ) → Optional[int]#: Return a sampler config’s required capacity without constructing it.

nemo_rl.algorithms.async_utils.staleness_sampler._custom_sampler_class( cfg: nemo_rl.algorithms.async_utils.staleness_sampler.CustomSamplerConfig, ) → type#: Import and return a custom sampler class without constructing it.

nemo_rl.algorithms.async_utils.staleness_sampler._sampler_class_for_config( cfg: nemo_rl.algorithms.async_utils.staleness_sampler.SamplerConfig, ) → type#: Return the sampler class selected by a built-in or custom config.

nemo_rl.algorithms.async_utils.staleness_sampler.sampler_supports_buffer_checkpoint( cfg: nemo_rl.algorithms.async_utils.staleness_sampler.SamplerConfig, ) → bool#

Return a sampler class’s static replay-checkpoint capability.

Custom classes are imported but not instantiated, allowing setup to fail before allocating cluster resources or triggering constructor side effects.

nemo_rl.algorithms.async_utils.staleness_sampler.sampler_supports_training_claims( cfg: nemo_rl.algorithms.async_utils.staleness_sampler.SamplerConfig, ) → bool#

Return whether selection transfers rows into training ownership.

Built-in samplers use :meth:TQReplayBuffer.claim_for_training. Custom samplers retain the legacy local-removal contract unless they explicitly opt in, so enabling periodic snapshots cannot silently assume ownership metadata that the sampler never created.

nemo_rl.algorithms.async_utils.staleness_sampler.create_sampler( buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer, cfg: nemo_rl.algorithms.async_utils.staleness_sampler.SamplerConfig, *, min_groups_for_streaming_train: Optional[int] = None, ) → nemo_rl.algorithms.async_utils.staleness_sampler.PromptGroupSampler#

Build a sampler from its config (or import one by FQN).

Parameters:

buffer – Shared TQReplayBuffer holding the candidate slots.
cfg – Discriminated sampler config selecting the policy.
min_groups_for_streaming_train – Streaming floor used by the windoweds sampler.

nemo_rl.algorithms.async_utils.staleness_sampler#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.async_utils.staleness_sampler`#