`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.
`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 built-in sampler’s required capacity without constructing it.
`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.

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; drop them locally.

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).

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

set_dispatch_index(resume_from_step: int) → None#: Seed the dispatch cursor when resuming from a checkpoint.

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

set_dispatch_index(resume_from_step: int) → None#

Seed the dispatch cursor when resuming from a checkpoint.

Parameters:: resume_from_step – Trainer step this run starts from — 0 for a fresh run, the restored current_step when resuming. Sets the cursor to resume_from_step - 1 so gated admit and InOrderSampler’s target_step stamps line up with the restored trainer version exactly as at step 0 of a fresh run. 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, )#

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

_eviction_window() → int#

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

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

_eviction_window() → int#

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

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

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

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, )#

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.

Initialization

_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

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, ) → Optional[int]#: Return a built-in sampler’s required capacity without constructing it.

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, ) → nemo_rl.algorithms.async_utils.staleness_sampler.PromptGroupSampler#: Build a sampler from its config (or import one by FQN).

nemo_rl.algorithms.async_utils.staleness_sampler#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.async_utils.staleness_sampler`#