From 4e064123389254242925c6221b2456211c236ca3 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Wed, 22 Jul 2026 17:03:49 -0700 Subject: [PATCH] refactor(sc): interface-driven staleness samplers (DRAFT design demo) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the monolithic StalenessSampler (five mode-encoding flags, three disjoint select branches, a shared evict fitting only two of the modes) with a PromptGroupSampler interface — admit/select/evict + is_on_policy / required_buffer_capacity — and one named policy per behavior: WindowedSampler / WeightFifoSampler / InOrderSampler. The sampler is selected by a pydantic discriminated-union config (name = windowed|weight_fifo|in_order|custom) or a module:ClassName FQN, so invalid knob combinations are unrepresentable and there are no cross-field validations in __init__. admit() moves the rollout-pump dispatch gate + target_step stamping into the sampler (the pump follows whichever policy is selected); InOrderSampler keys evict on target_step so evict/select agree by construction. No legacy-knob path: over_sampling / force_in_order / batch_selection_strategy / max_weight_staleness_versions and warn_if_staleness_window_below_minibatches are removed; the actor no longer owns _max_rollout_version. test_staleness_sampler (monolith) is replaced by test_sampler_interface. DRAFT / design demonstration for the #3220 sampler discussion. Not merged; not run locally (no compatible env). Signed-off-by: Terry Kong --- .../async_utils/staleness_sampler.py | 487 ++++++++++++---- nemo_rl/algorithms/single_controller.py | 173 ++---- .../single_controller/test_rollout_pump.py | 39 +- .../test_sampler_interface.py | 304 ++++++++++ .../test_staleness_sampler.py | 526 ------------------ .../unit/single_controller/test_train_pump.py | 5 +- 6 files changed, 773 insertions(+), 761 deletions(-) create mode 100644 tests/unit/single_controller/test_sampler_interface.py delete mode 100644 tests/unit/single_controller/test_staleness_sampler.py diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 86251dbfba..c8d699343b 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -12,71 +12,162 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prompt-group selection over a TQReplayBuffer.""" +"""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`` / +``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. +""" + +from __future__ import annotations + +import abc +import asyncio +import importlib +from typing import ( + Annotated, + Callable, + Literal, + Optional, + Protocol, + Union, + runtime_checkable, +) + +from pydantic import BaseModel, Field from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer from nemo_rl.data_plane import KVBatchMeta +# Poll interval for the rollout-pump admission gate. +_GATE_POLL_SECONDS = 0.005 -class StalenessSampler: - """Pick complete prompt groups from a TQReplayBuffer. - Args: - buffer: Shared TQReplayBuffer holding the candidate slots. - max_staleness_versions: Max weight-version gap a sample may have from the trainer. - sample_freshest_first: Prefer smallest lag when picking from the in-window set. - strict_weight_fifo: Consume in strict weight_version FIFO — take only from the oldest in-window weight_version and wait for its batch to fill. - force_in_order: Match each slot's target_step against current_train_weight, ignoring the window; mirrors legacy async_grpo target_weight semantics. +@runtime_checkable +class PromptGroupSampler(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. """ - def __init__( + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + """Block until the next prompt batch may dispatch. + + Args: + 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 def select( self, - buffer: TQReplayBuffer, - max_staleness_versions: int, - sample_freshest_first: bool = False, - strict_weight_fifo: bool = False, - force_in_order: bool = False, - ) -> None: - if max_staleness_versions < 0: - raise ValueError( - f"max_staleness_versions must be non-negative, got " - f"{max_staleness_versions}" - ) - if strict_weight_fifo and sample_freshest_first: - raise ValueError( - "strict_weight_fifo and sample_freshest_first are mutually exclusive" - ) + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + """Pick up to ``max_prompt_groups`` eligible groups; drop them locally.""" + ... + + async def evict(self, *, current_train_weight: int) -> int: + """Drop groups that can no longer be selected; clear their DP rows.""" + ... + + @property + def is_on_policy(self) -> bool: + """True when the policy admits zero staleness (sync mode).""" + ... + + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + """Buffer-capacity the policy needs, or ``None`` if unconstrained.""" + ... + + +class BaseSampler(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. + """ + + def __init__(self, buffer: TQReplayBuffer) -> None: self._buffer = buffer - self.max_staleness_versions = max_staleness_versions - self.sample_freshest_first = sample_freshest_first - self.strict_weight_fifo = strict_weight_fifo - self.force_in_order = force_in_order + # Pre-incremented before each admitted batch, so -1 lets the first + # batch through a zero-staleness gate. + self._dispatch_index: int = -1 + + # ── rollout-pump side ──────────────────────────────────────────────── + @abc.abstractmethod + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + ... + # ── train-pump side ────────────────────────────────────────────────── + @abc.abstractmethod async def select( self, *, current_train_weight: int, min_prompt_groups: int, max_prompt_groups: int, - ) -> tuple[KVBatchMeta | None, int]: - """Concat up to max_prompt_groups eligible groups and drop them from the buffer. + ) -> tuple[Optional[KVBatchMeta], int]: + ... - Eligibility = ready and weight in - [current_train_weight - max_staleness_versions, current_train_weight]. - DataPlane rows survive the local drop; caller clears them at step boundary. - - Args: - current_train_weight: Current trainer weight version. - min_prompt_groups: Minimum groups required; returns (None, 0) below this. - max_prompt_groups: Cap on groups returned when the threshold is met. - Selection is greedy without waiting: it returns all currently - eligible groups up to this cap (never deliberately fewer, and it - does not wait for the cap to fill). + async def evict(self, *, current_train_weight: int) -> int: + """Default: drop *ready* groups below the weight window. - Returns: - meta: Concatenated KVBatchMeta, or None if not enough groups. - num_groups: Number of prompt groups in meta; 0 when meta is None. + 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. """ + min_valid_version = max(0, current_train_weight - self._eviction_window()) + stale_idxs = [ + i + for i, weight in enumerate(self._buffer.start_weight_list) + if weight < min_valid_version and self._buffer.ready_list[i] + ] + if not stale_idxs: + return 0 + return await self._buffer.remove(stale_idxs, remove_in_dp=True) + + # ── derived facts ──────────────────────────────────────────────────── + @property + def is_on_policy(self) -> bool: + return self._eviction_window() == 0 + + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + return None + + # ── shared helpers ─────────────────────────────────────────────────── + def _eviction_window(self) -> int: + """Weight-version span kept selectable; drives the default ``evict``.""" + return 0 + + @staticmethod + def _validate_group_bounds(min_prompt_groups: int, max_prompt_groups: int) -> None: if min_prompt_groups < 1: raise ValueError(f"min_prompt_groups must be >= 1, got {min_prompt_groups}") if max_prompt_groups < min_prompt_groups: @@ -85,42 +176,77 @@ async def select( f"min_prompt_groups ({min_prompt_groups})" ) - if self.force_in_order: - # target_step exact match; staleness window ignored. - valid_idxs = [ - i - for i, target in enumerate(self._buffer.target_step_list) - if target == current_train_weight and self._buffer.ready_list[i] - ] - else: - min_valid_version = max( - 0, current_train_weight - self.max_staleness_versions - ) - if self.strict_weight_fifo: - in_window = [ - weight - for weight in self._buffer.start_weight_list - if min_valid_version <= weight <= current_train_weight - ] - if not in_window: - return None, 0 - target_version = min(in_window) - valid_idxs = [ - i - for i, weight in enumerate(self._buffer.start_weight_list) - if weight == target_version and self._buffer.ready_list[i] - ] - else: - valid_idxs = [ - i - for i, weight in enumerate(self._buffer.start_weight_list) - if min_valid_version <= weight <= current_train_weight - and self._buffer.ready_list[i] - ] + async def _finalize_selection( + self, + valid_idxs: list[int], + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[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``. + """ if len(valid_idxs) < min_prompt_groups: return None, 0 + requested_groups = min(len(valid_idxs), max_prompt_groups) + selected_idxs = valid_idxs[:requested_groups] + selected_metas = [self._buffer.meta_list[i] for i in selected_idxs] + await self._buffer.remove(selected_idxs, remove_in_dp=False) + return ( + selected_metas[0].concat(*selected_metas[1:]), # type: ignore + len(selected_idxs), + ) + + +class WindowedSampler(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. + """ + + def __init__( + self, + buffer: TQReplayBuffer, + *, + max_staleness_versions: int, + sample_freshest_first: bool = False, + ) -> None: + super().__init__(buffer) + if max_staleness_versions < 0: + raise ValueError( + f"max_staleness_versions must be non-negative, got " + f"{max_staleness_versions}" + ) + self.max_staleness_versions = max_staleness_versions + self.sample_freshest_first = sample_freshest_first + + def _eviction_window(self) -> int: + return self.max_staleness_versions + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + # Over-sampled: dispatch is bounded by buffer capacity, not by version. + return None + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + self._validate_group_bounds(min_prompt_groups, max_prompt_groups) + min_valid_version = max(0, current_train_weight - self.max_staleness_versions) + valid_idxs = [ + i + for i, weight in enumerate(self._buffer.start_weight_list) + if min_valid_version <= weight <= current_train_weight + and self._buffer.ready_list[i] + ] if self.sample_freshest_first: valid_idxs.sort( key=lambda i: ( @@ -128,36 +254,205 @@ async def select( i, ) ) + return await self._finalize_selection( + valid_idxs, min_prompt_groups, max_prompt_groups + ) - requested_groups = min(len(valid_idxs), max_prompt_groups) - selected_idxs = valid_idxs[:requested_groups] - selected_metas = [self._buffer.meta_list[i] for i in selected_idxs] - await self._buffer.remove(selected_idxs, remove_in_dp=False) +class _GatedSampler(BaseSampler): + """Base for policies that admit exactly one dispatch batch per trainer step. - return ( - selected_metas[0].concat(*selected_metas[1:]), # type: ignore - len(selected_idxs), - ) + The gate bounds how far generation may run ahead of the trainer + (``gate_window`` versions of lookahead). + """ - async def evict(self, *, current_train_weight: int) -> int: - """Drop groups whose weight falls below the staleness window. + def __init__(self, buffer: TQReplayBuffer, *, gate_window: int) -> None: + super().__init__(buffer) + if gate_window < 0: + raise ValueError(f"gate_window must be non-negative, got {gate_window}") + self._gate_window = gate_window - Future entries (weight > current_train_weight) are left alone. + def _eviction_window(self) -> int: + return self._gate_window - Args: - current_train_weight: Current trainer weight version; groups with - weight < current_train_weight - max_staleness_versions are dropped. + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + # One batch of lookahead per version in the window, plus the live batch. + return groups_per_step * (self._gate_window + 1) - Returns: - Number of group entries removed from the buffer. - """ + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + while self._dispatch_index >= trainer_version_fn() + self._gate_window: + await asyncio.sleep(_GATE_POLL_SECONDS) + self._dispatch_index += 1 + return self._stamp() + + def _stamp(self) -> Optional[int]: + return None + + +class WeightFifoSampler(_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). + """ + + def __init__(self, buffer: TQReplayBuffer, *, max_staleness_versions: int) -> None: + super().__init__(buffer, gate_window=max_staleness_versions) + self.max_staleness_versions = max_staleness_versions + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + self._validate_group_bounds(min_prompt_groups, max_prompt_groups) min_valid_version = max(0, current_train_weight - self.max_staleness_versions) - stale_idxs = [ + in_window = [ + weight + for weight in self._buffer.start_weight_list + if min_valid_version <= weight <= current_train_weight + ] + if not in_window: + return None, 0 + target_version = min(in_window) + valid_idxs = [ i for i, weight in enumerate(self._buffer.start_weight_list) - if weight < min_valid_version + if weight == target_version and self._buffer.ready_list[i] + ] + return await self._finalize_selection( + valid_idxs, min_prompt_groups, max_prompt_groups + ) + + +class InOrderSampler(_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. + """ + + def __init__(self, buffer: TQReplayBuffer, *, max_lookahead_versions: int) -> None: + super().__init__(buffer, gate_window=max_lookahead_versions) + self.max_lookahead_versions = max_lookahead_versions + + def _stamp(self) -> Optional[int]: + return self._dispatch_index + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + self._validate_group_bounds(min_prompt_groups, max_prompt_groups) + valid_idxs = [ + i + for i, target in enumerate(self._buffer.target_step_list) + if target == current_train_weight and self._buffer.ready_list[i] + ] + return await self._finalize_selection( + valid_idxs, min_prompt_groups, max_prompt_groups + ) + + async def evict(self, *, current_train_weight: int) -> int: + # Keyed on target_step (matches `select`): a ready group whose target + # step has already passed can never be selected, so it is stale. Unready + # slots are skipped to avoid racing a concurrent commit. + stale_idxs = [ + i + for i, target in enumerate(self._buffer.target_step_list) + if target is not None + and target < current_train_weight + and self._buffer.ready_list[i] ] if not stale_idxs: return 0 return await self._buffer.remove(stale_idxs, remove_in_dp=True) + + +# ── config + factory ──────────────────────────────────────────────────────── + + +class WindowedSamplerConfig(BaseModel, extra="allow"): + name: Literal["windowed"] = "windowed" + # Max weight-version gap a selected group may have from the trainer. + max_staleness_versions: int = 1 + # Prefer smallest lag when picking from the in-window set. + sample_freshest_first: bool = False + + +class WeightFifoSamplerConfig(BaseModel, extra="allow"): + name: Literal["weight_fifo"] = "weight_fifo" + # Lookahead + selectable weight window, in trainer versions. + max_staleness_versions: int = 1 + + +class InOrderSamplerConfig(BaseModel, extra="allow"): + name: Literal["in_order"] = "in_order" + # How far generation may run ahead of the trainer, in dispatch batches. + max_lookahead_versions: int = 1 + + +class CustomSamplerConfig(BaseModel, extra="allow"): + name: Literal["custom"] = "custom" + # "module:ClassName" of a PromptGroupSampler defined outside this repo. + # Extra keys are forwarded to the constructor (after ``buffer``). + target: str + + +# Discriminated on ``name`` so each variant carries only its own typed args and +# pydantic narrows the type at construction — invalid arg combinations are +# unrepresentable rather than caught by a runtime assert. +SamplerConfig = Annotated[ + Union[ + WindowedSamplerConfig, + WeightFifoSamplerConfig, + InOrderSamplerConfig, + CustomSamplerConfig, + ], + Field(discriminator="name"), +] + + +def create_sampler( + buffer: TQReplayBuffer, + cfg: SamplerConfig, +) -> PromptGroupSampler: + """Build a sampler from its config (or import one by FQN).""" + if isinstance(cfg, WindowedSamplerConfig): + return WindowedSampler( + buffer, + max_staleness_versions=cfg.max_staleness_versions, + sample_freshest_first=cfg.sample_freshest_first, + ) + if isinstance(cfg, WeightFifoSamplerConfig): + return WeightFifoSampler( + buffer, max_staleness_versions=cfg.max_staleness_versions + ) + if isinstance(cfg, InOrderSamplerConfig): + return InOrderSampler( + buffer, max_lookahead_versions=cfg.max_lookahead_versions + ) + if isinstance(cfg, CustomSamplerConfig): + module_name, sep, class_name = cfg.target.partition(":") + if not sep: + raise ValueError( + f"custom sampler target must be 'module:ClassName', got " + f"{cfg.target!r}" + ) + sampler_cls = getattr(importlib.import_module(module_name), class_name) + sampler = sampler_cls(buffer, **(cfg.model_extra or {})) + if not isinstance(sampler, PromptGroupSampler): + raise TypeError( + f"{cfg.target} does not implement the PromptGroupSampler " + f"interface (needs admit/select/evict)" + ) + return sampler + raise ValueError(f"unknown sampler config {type(cfg).__name__}") diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 2047e70c9f..5246a7ec11 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -50,7 +50,11 @@ import torch from pydantic import BaseModel, Field -from nemo_rl.algorithms.async_utils.staleness_sampler import StalenessSampler +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + SamplerConfig, + create_sampler, +) from nemo_rl.algorithms.grpo import GRPOLoggerConfig from nemo_rl.algorithms.single_controller_utils.utils import ( aggregate_step_metrics, @@ -100,29 +104,26 @@ class SingleControllerConfig(BaseModel, extra="allow"): are validated at construction by pydantic — no runtime assert needed. """ - # Staleness. A "group" is the atomic training unit: group_size + # Group geometry. A "group" is the atomic training unit: group_size # samples sharing one source prompt. GRPO consumers set group_size = # num generations per prompt; SFT/OPD-style consumers degenerate # cleanly to group_size=1 (every sample its own group). - max_weight_staleness_versions: int = 1 min_groups_per_batch: int = 2 target_groups_per_step: Optional[int] = None group_size: int = 4 - batch_selection_strategy: Literal[ - "strict_on_policy", - "staleness_window", - ] = "strict_on_policy" # Concurrency limits max_inflight_prompts: int = 8 max_buffered_rollouts: int = 8 # _buffer_capacity semaphore size - # Rollout dispatch gating (read by _rollout_pump). over_sampling=False - # gates each batch on max_rollout_version vs trainer_version; force_in_order - # stamps target_step on each dispatch so downstream consumers can match - # rollout batches to trainer steps exactly. - over_sampling: bool = False - force_in_order: bool = False + # Staleness policy — the single source of truth for how rollouts are gated + # (rollout pump) and selected/evicted (train pump). Discriminated on + # ``name`` (windowed | weight_fifo | in_order | custom); each variant carries + # only its own args, so invalid combinations are unrepresentable and there + # are no cross-field knobs to validate at runtime. ``custom`` takes a + # ``module:ClassName`` target to load a PromptGroupSampler from outside this + # repo. Default mirrors the exemplar (exact batch->step matching). + sampler: SamplerConfig = Field(default_factory=InOrderSamplerConfig) # Training max_train_steps: int = 10 @@ -164,57 +165,6 @@ class SingleControllerConfig(BaseModel, extra="allow"): logger: GRPOLoggerConfig -def warn_if_staleness_window_below_minibatches( - cfg: SingleControllerConfig, -) -> None: - """Warn when the staleness window is too small to cover one outer step. - - ``_train_pump`` runs ``num_minibatches`` begin/finish cycles per outer - step, bumping ``_trainer_version`` after each ``finish_train_step`` - (one opt.step) but refreshing generation weights (``_sync_weights``) - only once, at the end of the outer step. A group produced at the version - the step started on therefore ages by up to ``num_minibatches - 1`` - versions before the next sync. When the effective staleness window is - smaller than that, such groups become unselectable mid-step — the - sampler skips them and ``sampler.evict`` may drop them as stale — - so the pump can spin or thrash instead of consuming them. - - Warns rather than raises: a run may intentionally accept the resulting - eviction churn. Self-contained (does not assume ``cfg`` has been coerced - by ``__init__``) so it is unit-testable without constructing the actor. - - Args: - cfg: SingleController config. ``strict_on_policy`` pins the effective - window to 0 regardless of ``max_weight_staleness_versions``. - """ - effective_window = ( - 0 - if cfg.batch_selection_strategy == "strict_on_policy" - else cfg.max_weight_staleness_versions - ) - target_groups = cfg.target_groups_per_step or cfg.min_groups_per_batch - samples_per_step = target_groups * cfg.group_size - train_global_batch_size = cfg.train_global_batch_size or samples_per_step - groups_per_minibatch = train_global_batch_size // cfg.group_size - if groups_per_minibatch <= 0: - return - num_minibatches = target_groups // groups_per_minibatch - if num_minibatches > 1 and effective_window < num_minibatches - 1: - print( - f"WARNING: max_weight_staleness_versions (effective window " - f"{effective_window}) is smaller than num_minibatches - 1 " - f"({num_minibatches - 1}): each outer step runs {num_minibatches} " - f"optimizer steps but syncs generation weights only once, at the " - f"end, so groups produced at the version the step began on age by " - f"up to {num_minibatches - 1} versions before the next sync and " - f"become unselectable (skipped, then evicted as stale) mid-step " - f"— the train pump may spin or thrash. Remedy: raise " - f"max_weight_staleness_versions to at least {num_minibatches - 1}, " - f"or lower target_groups_per_step * group_size / " - f"train_global_batch_size to reduce num_minibatches.", - flush=True, - ) - @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: @@ -266,27 +216,10 @@ def __init__( "advantage_enabled=True requires an advantage_estimator instance" ) - # batch_selection_strategy is Literal-typed; pydantic rejects unknown - # values at config construction, so no runtime assert is needed here. - if cfg.batch_selection_strategy == "strict_on_policy": - cfg.max_weight_staleness_versions = 0 - cfg.over_sampling = False - print( - "Using strict_on_policy, auto setting max_weight_staleness_versions " - "to 0 and over_sampling to False.", - flush=True, - ) - if cfg.max_weight_staleness_versions == 0 and cfg.over_sampling: - raise ValueError( - "max_weight_staleness_versions=0 requires over_sampling=False: " - "with zero staleness the dispatch gate needs to advance one batch " - "per trainer_version, which over_sampling=True bypasses." - ) - if cfg.force_in_order and cfg.over_sampling: - raise ValueError( - "force_in_order=True requires over_sampling=False so that each " - "dispatched batch corresponds to exactly one target training step." - ) + # Staleness-mode validation that used to live here (strict_on_policy + # coercion, over_sampling/force_in_order mutual-exclusion raises) is gone: + # the sampler config's discriminated union makes those states + # unrepresentable — see ``cfg.sampler`` / ``create_sampler``. if cfg.target_groups_per_step is None: cfg.target_groups_per_step = cfg.min_groups_per_batch @@ -323,17 +256,28 @@ def __init__( f"({cfg.group_size}) so a mini-batch contains " f"whole prompt groups" ) - # num_minibatches > 1 runs multiple opt.steps (each bumps - # trainer_version) between weight syncs; warn if the staleness window - # cannot keep same-step groups selectable across those bumps. - warn_if_staleness_window_below_minibatches(cfg) - - self._sampler = StalenessSampler( - self._buffer, - max_staleness_versions=cfg.max_weight_staleness_versions, - strict_weight_fifo=not cfg.over_sampling, - force_in_order=cfg.force_in_order, + # Staleness policy owns admit (rollout side) + select/evict (train side) + # + is_on_policy / required_buffer_capacity (derived facts). ``name`` in + # the config is the single switch for which behavior runs. + self._sampler = create_sampler(self._buffer, cfg.sampler) + + # Capacity requirement is a derived fact the sampler owns, so this check + # can't drift from the admission logic. None means the policy self-bounds + # via backpressure (e.g. the over-sampled windowed policy). + required_capacity = self._sampler.required_buffer_capacity( + cfg.target_groups_per_step ) + if ( + required_capacity is not None + and cfg.max_buffered_rollouts < required_capacity + ): + raise ValueError( + f"max_buffered_rollouts ({cfg.max_buffered_rollouts}) is below " + f"the {type(self._sampler).__name__}'s required capacity " + f"({required_capacity} = target_groups_per_step " + f"{cfg.target_groups_per_step} * (lookahead + 1)); the rollout " + f"pump would deadlock waiting for buffer slots." + ) # ── asyncio state ────────────────────────────────────────────────── # Gate: cleared during _sync_weights, set when generation may proceed @@ -343,9 +287,9 @@ def __init__( # Count of in-flight generate_and_push calls self._inflight_rollouts: int = 0 - # Rollout batch counter — pre-incremented before each dispatch, so start - # at -1 to allow the first batch through the strict_on_policy gate. - self._max_rollout_version: int = -1 + # The rollout batch counter (dispatch index) is now owned by the sampler + # (``BaseSampler._dispatch_index``), which gates admission and stamps + # target_step — see ``_rollout_pump`` / ``PromptGroupSampler.admit``. # Active rollout tasks used by downstream synchronization/drain paths. # TaskGroup remains responsible for task ownership and cancellation. @@ -370,8 +314,7 @@ def __init__( self._current_epoch: int = 0 print( - f"SingleControllerActor: staleness_cap=" - f"{cfg.max_weight_staleness_versions} " + f"SingleControllerActor: sampler={cfg.sampler.name} " f"buffer={cfg.max_buffered_rollouts} " f"inflight={cfg.max_inflight_prompts} " f"transport={cfg.refit_cfg.impl}", @@ -437,9 +380,10 @@ async def _call_dp(self, method_name: str, **kwargs) -> Any: async def _rollout_pump(self) -> None: """Continuously dispatch rollout tasks until cancellation. - Per batch (over_sampling=False): - 0. Wait while _max_rollout_version >= trainer_version + max_staleness, - then claim the next step by incrementing _max_rollout_version. + Per batch: + 0. await sampler.admit(...) — blocks until the batch may dispatch and + returns the target_step stamp (the staleness gate lives in the + sampler, not here). Per prompt: 1. Acquire _buffer_capacity slot (backpressure) @@ -451,9 +395,6 @@ async def _rollout_pump(self) -> None: 5. Decrement _inflight_rollouts """ sem = asyncio.Semaphore(self._cfg.max_inflight_prompts) - over_sampling = self._cfg.over_sampling - max_staleness = self._cfg.max_weight_staleness_versions - force_in_order = self._cfg.force_in_order print("rollout_pump: starting", flush=True) async def _dispatch_one_prompt( @@ -497,17 +438,15 @@ def _release_permits_if_task_not_started( async with asyncio.TaskGroup() as rollout_tasks: while max_epochs is None or self._current_epoch < max_epochs: for prompt_batch in self._dataloader: - # over_sampling=False: batch-level gate on max_rollout_version. - if not over_sampling: - while ( - self._max_rollout_version - >= self._trainer_version + max_staleness - ): - await asyncio.sleep(0.005) - self._max_rollout_version += 1 - - # target_step = batch dispatch index when force_in_order is on. - target_step = self._max_rollout_version if force_in_order else None + # The sampler owns admission: it blocks until this batch may + # dispatch (per its own staleness algorithm) and returns the + # target_step to stamp (None when the policy doesn't stamp). + # Keeping the gate here means the rollout pump follows + # whichever sampler is selected without a second copy of the + # gating logic to keep in sync. + target_step = await self._sampler.admit( + trainer_version_fn=lambda: self._trainer_version + ) for prompt_idx in range(prompt_batch.size): prompt: DatumSpec = { # type: ignore diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 372072403d..9b32628b23 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -26,6 +26,12 @@ import torch from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSampler, + WeightFifoSampler, + WindowedSampler, + WindowedSamplerConfig, +) from nemo_rl.algorithms.single_controller import ( SingleControllerActor, SingleControllerConfig, @@ -53,14 +59,16 @@ @pytest.mark.parametrize( - ("force_in_order", "expected_target_steps"), + ("make_sampler", "expected_target_steps"), [ - (False, [None, None]), - (True, [0, 1]), + # weight_fifo gates but does not stamp target_step. + (lambda buf: WeightFifoSampler(buf, max_staleness_versions=1), [None, None]), + # in_order stamps the dispatch index as target_step. + (lambda buf: InOrderSampler(buf, max_lookahead_versions=1), [0, 1]), ], ) def test_rollout_pump_stamps_target_steps( - force_in_order: bool, + make_sampler, expected_target_steps: list[int | None], ) -> None: class _RecordingBuffer: @@ -85,13 +93,13 @@ async def generate_and_push( ctrl = object.__new__(controller_cls) ctrl._cfg = SimpleNamespace( max_inflight_prompts=2, - over_sampling=False, - max_weight_staleness_versions=1, - force_in_order=force_in_order, diagnostics=False, max_num_epochs=1, ) ctrl._rollout_manager = _RecordingRolloutManager(buffer) + # The sampler owns admission + target_step stamping (the dispatch counter + # lives on the sampler, not the actor). + ctrl._sampler = make_sampler(buffer) prompt_batch = BatchedDataDict( {"message_log": [[{"role": "user", "content": "prompt"}]]} ) @@ -101,7 +109,6 @@ async def generate_and_push( ctrl._buffer_capacity = asyncio.Semaphore(2) ctrl._inflight_rollouts = 0 ctrl._dispatched_rollouts = set() - ctrl._max_rollout_version = -1 ctrl._trainer_version = 0 ctrl._current_epoch = 0 @@ -142,13 +149,12 @@ async def _main() -> None: ctrl = object.__new__(controller_cls) ctrl._cfg = SimpleNamespace( max_inflight_prompts=2, - over_sampling=True, - max_weight_staleness_versions=1, - force_in_order=False, diagnostics=False, max_num_epochs=1, ) ctrl._rollout_manager = manager + # Over-sampled windowed policy: admit never gates (buffer unused here). + ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ BatchedDataDict( { @@ -164,7 +170,6 @@ async def _main() -> None: ctrl._buffer_capacity = asyncio.Semaphore(2) ctrl._inflight_rollouts = 0 ctrl._dispatched_rollouts = set() - ctrl._max_rollout_version = -1 ctrl._trainer_version = 0 ctrl._current_epoch = 0 @@ -221,13 +226,12 @@ async def _main() -> None: ctrl = object.__new__(controller_cls) ctrl._cfg = SimpleNamespace( max_inflight_prompts=1, - over_sampling=True, - max_weight_staleness_versions=1, - force_in_order=False, diagnostics=False, max_num_epochs=1, ) ctrl._rollout_manager = _NeverCalledRolloutManager() + # Over-sampled windowed policy: admit never gates (buffer unused here). + ctrl._sampler = WindowedSampler(None, max_staleness_versions=1) ctrl._dataloader = [ BatchedDataDict({"message_log": [[{"role": "user", "content": "prompt"}]]}) ] @@ -236,7 +240,6 @@ async def _main() -> None: ctrl._buffer_capacity = real_semaphore(1) ctrl._inflight_rollouts = 0 ctrl._dispatched_rollouts = set() - ctrl._max_rollout_version = -1 ctrl._trainer_version = 0 ctrl._current_epoch = 0 @@ -277,15 +280,13 @@ def test_rollout_pump_writes_expected_tq_data( dp_adapter = _SyncDPAdapter(tq_actor) cfg = SingleControllerConfig.model_construct( - batch_selection_strategy="staleness_window", - max_weight_staleness_versions=1, + sampler=WindowedSamplerConfig(max_staleness_versions=1), min_groups_per_batch=1, group_size=num_generations, max_inflight_prompts=num_prompts, max_buffered_rollouts=num_prompts, max_train_steps=1, max_num_epochs=None, - over_sampling=True, partition_id=_PARTITION_ID, logger={ "log_dir": str(tmp_path / "logs"), diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py new file mode 100644 index 0000000000..386798386e --- /dev/null +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -0,0 +1,304 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Contract tests for the split PromptGroupSampler policies + factory. + +CPU-only; mirrors the FakeBuffer surface used by test_staleness_sampler.py. +Covers what the split buys over the monolithic sampler: per-policy admission, +InOrderSampler's target_step-keyed evict (evict/select agree by construction), +and FQN loading of an out-of-repo sampler. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSampler, + InOrderSamplerConfig, + PromptGroupSampler, + WeightFifoSampler, + WindowedSampler, + WindowedSamplerConfig, + create_sampler, +) +from nemo_rl.data_plane import KVBatchMeta + + +class FakeBuffer: + """Minimal TQReplayBuffer surface the samplers read/mutate.""" + + def __init__(self, partition_id: str = "rollout_data") -> None: + self._partition_id = partition_id + self.meta_list: list[KVBatchMeta | None] = [] + self.start_weight_list: list[int] = [] + self.end_weight_list: list[int] = [] + self.target_step_list: list[int | None] = [] + self.ready_list: list[bool] = [] + self.remove_calls: list[tuple[list[int], bool]] = [] + + def add( + self, + group_id: str, + weight: int, + *, + ready: bool = True, + target_step: int | None = None, + ) -> None: + meta = KVBatchMeta( + partition_id=self._partition_id, + task_name=None, + sample_ids=[f"{group_id}_g0"], + tags=[{"weight_version": weight, "group_id": group_id}], + ) + self.meta_list.append(meta if ready else None) + self.start_weight_list.append(weight) + self.end_weight_list.append(weight) + self.target_step_list.append(target_step) + self.ready_list.append(ready) + + async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: + self.remove_calls.append((list(idxs), remove_in_dp)) + for i in sorted(idxs, reverse=True): + del self.meta_list[i] + del self.start_weight_list[i] + del self.end_weight_list[i] + del self.target_step_list[i] + del self.ready_list[i] + return len(idxs) + + +def _run(coro): + return asyncio.run(coro) + + +class TestBuiltinsImplementInterface: + @pytest.mark.parametrize( + "sampler", + [ + WindowedSampler(FakeBuffer(), max_staleness_versions=1), + WeightFifoSampler(FakeBuffer(), max_staleness_versions=1), + InOrderSampler(FakeBuffer(), max_lookahead_versions=1), + ], + ) + def test_isinstance_protocol(self, sampler): + assert isinstance(sampler, PromptGroupSampler) + + +class TestAdmission: + def test_windowed_never_gates_and_never_stamps(self): + s = WindowedSampler(FakeBuffer(), max_staleness_versions=2) + # trainer stuck at 0, but over-sampled admission returns immediately. + assert _run(s.admit(trainer_version_fn=lambda: 0)) is None + assert _run(s.admit(trainer_version_fn=lambda: 0)) is None + + def test_in_order_stamps_monotonic_dispatch_index(self): + s = InOrderSampler(FakeBuffer(), max_lookahead_versions=5) + assert _run(s.admit(trainer_version_fn=lambda: 10)) == 0 + assert _run(s.admit(trainer_version_fn=lambda: 10)) == 1 + + def test_weight_fifo_gates_on_lookahead_and_does_not_stamp(self): + # dispatch_index starts at -1; window 0 => admits exactly one batch + # ahead of the trainer, then blocks. Assert the second admit would block + # by giving it a trainer_version that keeps the gate closed. + s = WeightFifoSampler(FakeBuffer(), max_staleness_versions=0) + assert _run(s.admit(trainer_version_fn=lambda: 0)) is None # -1 -> 0 + # Now dispatch_index=0, trainer=0, window=0 -> 0 >= 0 blocks forever. + with pytest.raises(asyncio.TimeoutError): + _run(asyncio.wait_for(s.admit(trainer_version_fn=lambda: 0), timeout=0.05)) + + +class TestInOrderEvictMatchesSelect: + """The bug the split fixes: monolithic evict keyed on weight could drop a + slot whose target_step was still upcoming. InOrderSampler keys evict on + target_step, so it never drops a slot select would later match.""" + + def test_future_target_not_evicted_even_if_weight_out_of_window(self): + buf = FakeBuffer() + # weight far below the window, but target_step is still upcoming. + buf.add("g", weight=0, ready=True, target_step=2) + s = InOrderSampler(buf, max_lookahead_versions=1) + removed = _run(s.evict(current_train_weight=2)) + assert removed == 0 # target_step 2 == current, not past -> kept + assert len(buf.target_step_list) == 1 + + def test_past_target_ready_slot_is_evicted(self): + buf = FakeBuffer() + buf.add("g", weight=0, ready=True, target_step=1) + s = InOrderSampler(buf, max_lookahead_versions=1) + removed = _run(s.evict(current_train_weight=3)) # target 1 < 3 -> stale + assert removed == 1 + + def test_unready_slot_is_never_evicted(self): + buf = FakeBuffer() + buf.add("g", weight=0, ready=False, target_step=1) + s = InOrderSampler(buf, max_lookahead_versions=1) + # past target, but unready -> skipped to avoid the commit race. + assert _run(s.evict(current_train_weight=5)) == 0 + + +class TestFactory: + def test_windowed_config_builds_windowed(self): + s = create_sampler(FakeBuffer(), WindowedSamplerConfig(max_staleness_versions=3)) + assert isinstance(s, WindowedSampler) + assert s.max_staleness_versions == 3 + + def test_in_order_config_builds_in_order(self): + s = create_sampler(FakeBuffer(), InOrderSamplerConfig(max_lookahead_versions=2)) + assert isinstance(s, InOrderSampler) + assert s.max_lookahead_versions == 2 + + def test_weight_fifo_config_builds_weight_fifo(self): + from nemo_rl.algorithms.async_utils.staleness_sampler import ( + WeightFifoSamplerConfig, + ) + + s = create_sampler(FakeBuffer(), WeightFifoSamplerConfig(max_staleness_versions=4)) + assert isinstance(s, WeightFifoSampler) + assert s.max_staleness_versions == 4 + + +class TestCustomFqnSampler: + def test_custom_target_loads_out_of_repo_sampler(self): + # A user sampler defined anywhere importable; here, this test module. + from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, + ) + + s = create_sampler( + FakeBuffer(), + CustomSamplerConfig( + target=f"{__name__}:EchoSampler", max_lookahead_versions=1 + ), + ) + assert isinstance(s, EchoSampler) + assert isinstance(s, PromptGroupSampler) + assert s.max_lookahead_versions == 1 + + +class TestWindowedSelect: + def test_selects_ready_groups_in_window(self): + buf = FakeBuffer() + buf.add("a", weight=3) + buf.add("b", weight=5) # current + buf.add("c", weight=1) # below window (5-2=3) + s = WindowedSampler(buf, max_staleness_versions=2) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) + assert n == 2 # a(3) and b(5); c(1) excluded + assert len(buf.start_weight_list) == 1 # only c remains + + def test_below_min_returns_none(self): + buf = FakeBuffer() + buf.add("a", weight=5) + s = WindowedSampler(buf, max_staleness_versions=2) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=2, max_prompt_groups=8) + ) == (None, 0) + + def test_unready_excluded(self): + buf = FakeBuffer() + buf.add("a", weight=5, ready=False) + s = WindowedSampler(buf, max_staleness_versions=2) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) == (None, 0) + + def test_freshest_first_orders_by_lag(self): + buf = FakeBuffer() + buf.add("old", weight=1) + buf.add("new", weight=5) + s = WindowedSampler(buf, max_staleness_versions=10, sample_freshest_first=True) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1) + ) + # freshest (weight 5) picked first -> "old" (weight 1) remains. + assert n == 1 + assert buf.start_weight_list == [1] + + +class TestWeightFifoSelect: + def test_drains_oldest_in_window_weight_first(self): + buf = FakeBuffer() + buf.add("old1", weight=3) + buf.add("new", weight=5) + buf.add("old2", weight=3) + s = WeightFifoSampler(buf, max_staleness_versions=5) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) + assert n == 2 # both weight-3 groups; weight-5 waits its turn + assert buf.start_weight_list == [5] + + def test_waits_for_partial_oldest_batch(self): + buf = FakeBuffer() + buf.add("old", weight=3) + s = WeightFifoSampler(buf, max_staleness_versions=5) + # oldest weight has only 1 group but min is 2 -> wait (None), don't skip + # ahead to a newer weight. + assert _run( + s.select(current_train_weight=5, min_prompt_groups=2, max_prompt_groups=8) + ) == (None, 0) + + def test_empty_window_returns_none(self): + buf = FakeBuffer() + buf.add("future", weight=9) + s = WeightFifoSampler(buf, max_staleness_versions=2) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) == (None, 0) + + +class TestInOrderSelect: + def test_matches_target_step_ignoring_weight_window(self): + buf = FakeBuffer() + # weight far outside any window, but target_step == trainer version. + buf.add("g", weight=100, target_step=5) + s = InOrderSampler(buf, max_lookahead_versions=1) + meta, n = _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) + assert n == 1 + + def test_non_matching_target_not_selected(self): + buf = FakeBuffer() + buf.add("g", weight=5, target_step=6) + s = InOrderSampler(buf, max_lookahead_versions=1) + assert _run( + s.select(current_train_weight=5, min_prompt_groups=1, max_prompt_groups=8) + ) == (None, 0) + + +class TestDefaultEvictSkipsUnready: + def test_windowed_evict_drops_ready_below_window(self): + buf = FakeBuffer() + buf.add("stale", weight=0, ready=True) + buf.add("fresh", weight=5, ready=True) + s = WindowedSampler(buf, max_staleness_versions=1) + removed = _run(s.evict(current_train_weight=5)) # min_valid = 4 + assert removed == 1 + assert buf.start_weight_list == [5] + + def test_windowed_evict_skips_unready_stale(self): + buf = FakeBuffer() + buf.add("stale_unready", weight=0, ready=False) + s = WindowedSampler(buf, max_staleness_versions=1) + assert _run(s.evict(current_train_weight=5)) == 0 + + +class EchoSampler(InOrderSampler): + """Stand-in for a user-defined sampler loaded by FQN.""" diff --git a/tests/unit/single_controller/test_staleness_sampler.py b/tests/unit/single_controller/test_staleness_sampler.py deleted file mode 100644 index 69f580caf7..0000000000 --- a/tests/unit/single_controller/test_staleness_sampler.py +++ /dev/null @@ -1,526 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for StalenessSampler (pure filter over TQReplayBuffer state).""" - -from __future__ import annotations - -import asyncio - -import pytest - -from nemo_rl.algorithms.async_utils.staleness_sampler import StalenessSampler -from nemo_rl.data_plane import KVBatchMeta - - -class FakeBuffer: - """Minimal TQReplayBuffer surface used by StalenessSampler tests.""" - - def __init__(self, partition_id: str = "rollout_data") -> None: - self._partition_id = partition_id - self.meta_list: list[KVBatchMeta | None] = [] - self.start_weight_list: list[int] = [] - self.end_weight_list: list[int] = [] - self.target_step_list: list[int | None] = [] - self.ready_list: list[bool] = [] - self.remove_calls: list[tuple[list[int], bool]] = [] - - def add( - self, - group_id: str, - weight: int, - group_size: int = 1, - ready: bool = True, - end_weight: int | None = None, - target_step: int | None = None, - ) -> KVBatchMeta: - sample_ids = [f"{group_id}_g{i}" for i in range(group_size)] - meta = KVBatchMeta( - partition_id=self._partition_id, - task_name=None, - sample_ids=sample_ids, - tags=[{"weight_version": weight, "group_id": group_id}] * group_size, - ) - self.meta_list.append(meta if ready else None) - self.start_weight_list.append(weight) - self.end_weight_list.append(weight if end_weight is None else end_weight) - self.target_step_list.append(target_step) - self.ready_list.append(ready) - return meta - - async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: - self.remove_calls.append((list(idxs), remove_in_dp)) - for i in sorted(idxs, reverse=True): - del self.meta_list[i] - del self.start_weight_list[i] - del self.end_weight_list[i] - del self.target_step_list[i] - del self.ready_list[i] - return len(idxs) - - -def _run(coro): - return asyncio.run(coro) - - -class TestStalenessSamplerSelect: - def test_select_returns_none_when_insufficient(self): - buf = FakeBuffer() - buf.add("g0", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=2) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - - def test_select_filters_by_staleness_window(self): - buf = FakeBuffer() - # Weights 3, 4, 5, 2, 6 against trainer=5, max_staleness=2: - # lags = 2, 1, 0, 3 (stale), -1 (future) - for i, w in enumerate([3, 4, 5, 2, 6]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler( - buf, max_staleness_versions=2, sample_freshest_first=True - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - # Freshest first → g2 (lag 0), g1 (lag 1) - assert selected.sample_ids == ["g2_g0", "g1_g0"] - assert num_groups == 2 - - def test_select_fifo_orders_by_insertion(self): - buf = FakeBuffer() - for w in [3, 4, 5]: - buf.add(f"v{w}", weight=w) - sampler = StalenessSampler( - buf, max_staleness_versions=5, sample_freshest_first=False - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=6, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert selected is not None - assert selected.sample_ids == ["v3_g0", "v4_g0"] - assert num_groups == 2 - - def test_select_skips_future_weight(self): - buf = FakeBuffer() - buf.add("now", weight=5) - buf.add("future", weight=7) - sampler = StalenessSampler(buf, max_staleness_versions=10) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["now_g0"] - assert num_groups == 1 - - def test_select_concats_groups(self): - buf = FakeBuffer() - buf.add("g0", weight=5, group_size=2) - buf.add("g1", weight=5, group_size=2) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - assert selected.sample_ids == [ - "g0_g0", - "g0_g1", - "g1_g0", - "g1_g1", - ] - # Two groups concatenated, each of size 2 → 4 sample_ids total. - assert num_groups == 2 - - def test_select_strict_on_policy_requires_exact_version(self): - buf = FakeBuffer() - for i, w in enumerate([4, 5, 5, 6]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - # 3 eligible (need weight=5), only have 2 - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=3, max_prompt_groups=3 - ) - ) - assert result == (None, 0) - - # Buffer still intact: select with min=3 returned None without dropping anything. - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert selected is not None - assert selected.sample_ids == ["g1_g0", "g2_g0"] - assert num_groups == 2 - - def test_select_drops_returned_entries_from_buffer(self): - buf = FakeBuffer() - for i, w in enumerate([5, 5, 5]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - first_meta, first_num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - assert first_meta is not None - assert first_meta.sample_ids == ["g0_g0"] - assert first_num_groups == 1 - assert buf.start_weight_list == [5, 5] - # remove_in_dp=False; DP rows kept for trainer. - assert buf.remove_calls[-1][1] is False - - second_meta, second_num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - assert second_meta is not None - assert second_meta.sample_ids == ["g1_g0"] - assert second_num_groups == 1 - - def test_select_rejects_zero_min_prompt_groups(self): - buf = FakeBuffer() - sampler = StalenessSampler(buf, max_staleness_versions=0) - with pytest.raises(ValueError): - _run( - sampler.select( - current_train_weight=0, min_prompt_groups=0, max_prompt_groups=0 - ) - ) - - def test_select_rejects_max_less_than_min(self): - buf = FakeBuffer() - for i in range(3): - buf.add(f"g{i}", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - with pytest.raises(ValueError): - _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=1 - ) - ) - - def test_select_caps_at_max_prompt_groups(self): - buf = FakeBuffer() - for i in range(5): - buf.add(f"g{i}", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=3 - ) - ) - - assert selected is not None - # FIFO order; capped at max=3 even though 5 are eligible. - assert selected.sample_ids == ["g0_g0", "g1_g0", "g2_g0"] - assert num_groups == 3 - # The remaining two stay in the buffer. - assert buf.start_weight_list == [5, 5] - - def test_select_takes_all_available_when_between_min_and_max(self): - buf = FakeBuffer() - for i in range(3): - buf.add(f"g{i}", weight=5) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=8 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g0_g0", "g1_g0", "g2_g0"] - assert num_groups == 3 - assert buf.start_weight_list == [] - - -class TestStalenessSamplerEvict: - def test_evict_removes_stale_groups(self): - buf = FakeBuffer() - # trainer=5, max_staleness=1 → lag >1 means stale (weights 0, 1, 2 stale; 4, 5 fresh) - for i, w in enumerate([0, 1, 4, 5, 2]): - buf.add(f"g{i}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=1) - - dropped = _run(sampler.evict(current_train_weight=5)) - - assert dropped == 3 - assert buf.start_weight_list == [4, 5] - # Survivors' sample_ids - assert [m.sample_ids[0] for m in buf.meta_list] == ["g2_g0", "g3_g0"] - - def test_evict_returns_zero_when_nothing_stale(self): - buf = FakeBuffer() - for w in [4, 5]: - buf.add(f"v{w}", weight=w) - sampler = StalenessSampler(buf, max_staleness_versions=1) - - assert _run(sampler.evict(current_train_weight=5)) == 0 - assert buf.remove_calls == [] - - def test_evict_keeps_future_groups(self): - buf = FakeBuffer() - buf.add("future", weight=7) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - assert _run(sampler.evict(current_train_weight=5)) == 0 - assert buf.start_weight_list == [7] - - def test_evict_drops_whole_group(self): - buf = FakeBuffer() - buf.add("stale", weight=1, group_size=4) - buf.add("fresh", weight=5, group_size=4) - sampler = StalenessSampler(buf, max_staleness_versions=1) - - dropped = _run(sampler.evict(current_train_weight=5)) - - assert dropped == 1 - assert buf.remove_calls == [([0], True)] - assert buf.start_weight_list == [5] - assert [m.sample_ids[0] for m in buf.meta_list] == ["fresh_g0"] - - -class TestStalenessSamplerInit: - def test_rejects_negative_max_staleness(self): - buf = FakeBuffer() - with pytest.raises(ValueError): - StalenessSampler(buf, max_staleness_versions=-1) - - def test_rejects_strict_weight_fifo_with_freshest_first(self): - buf = FakeBuffer() - with pytest.raises(ValueError): - StalenessSampler( - buf, - max_staleness_versions=0, - sample_freshest_first=True, - strict_weight_fifo=True, - ) - - -class TestStalenessSamplerReady: - def test_default_mode_skips_unready_slots(self): - buf = FakeBuffer() - buf.add("g0", weight=5, ready=False) - buf.add("g1", weight=5, ready=True) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=1 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g1_g0"] - assert num_groups == 1 - - def test_default_mode_waits_when_too_few_ready(self): - buf = FakeBuffer() - buf.add("g0", weight=5, ready=False) - buf.add("g1", weight=5, ready=True) - sampler = StalenessSampler(buf, max_staleness_versions=0) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - - -class TestStalenessSamplerStrictWeightFifo: - def test_consumes_oldest_batch_first(self): - buf = FakeBuffer() - # Two complete batches: v=4 then v=5; strict_weight_fifo must take v=4 first. - for i, w in enumerate((4, 4, 5, 5)): - buf.add(f"v{w}_{i}", weight=w) - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - # Insertion-order FIFO inside the oldest batch. - assert selected.sample_ids == ["v4_0_g0", "v4_1_g0"] - assert num_groups == 2 - assert buf.start_weight_list == [5, 5] - - def test_waits_when_oldest_batch_partially_ready(self): - buf = FakeBuffer() - # Oldest batch v=4 has 1 ready + 1 unready; v=5 batch is fully ready. - # strict_weight_fifo must NOT skip ahead to v=5. - buf.add("v4_a", weight=4, ready=True) - buf.add("v4_b", weight=4, ready=False) - buf.add("v5_a", weight=5, ready=True) - buf.add("v5_b", weight=5, ready=True) - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - # Buffer untouched: nothing removed. - assert buf.start_weight_list == [4, 4, 5, 5] - assert buf.ready_list == [True, False, True, True] - - def test_returns_none_when_oldest_batch_not_filled(self): - buf = FakeBuffer() - buf.add("v4_a", weight=4, ready=True) - # Only 1 ready in oldest batch; need 2. - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - assert result == (None, 0) - - def test_ignores_future_versions_when_picking_target(self): - buf = FakeBuffer() - # Trainer at 5, staleness 1: window is [4, 5]; v=7 (future) must not - # become the oldest target. - buf.add("v7", weight=7, ready=True) - buf.add("v5_a", weight=5, ready=True) - buf.add("v5_b", weight=5, ready=True) - sampler = StalenessSampler( - buf, max_staleness_versions=1, strict_weight_fifo=True - ) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=2, max_prompt_groups=2 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["v5_a_g0", "v5_b_g0"] - assert num_groups == 2 - assert buf.start_weight_list == [7] - - -class TestStalenessSamplerForceInOrder: - def test_selects_only_groups_whose_target_step_matches_current(self): - buf = FakeBuffer() - # Two matching-target groups and one future-target group. - buf.add("g_match_a", weight=5, target_step=5) - buf.add("g_match_b", weight=5, target_step=5) - buf.add("g_future", weight=5, target_step=7) - sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g_match_a_g0", "g_match_b_g0"] - assert num_groups == 2 - # The future-target group survives; only the two matches were removed. - assert buf.target_step_list == [7] - assert buf.start_weight_list == [5] - - def test_ignores_stale_target_step(self): - buf = FakeBuffer() - # Trainer at 5; a slot with target_step=4 must NOT satisfy force_in_order. - buf.add("g_stale", weight=4, target_step=4) - sampler = StalenessSampler( - buf, - # Wide staleness window: force_in_order must ignore it entirely. - max_staleness_versions=10, - force_in_order=True, - ) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert result == (None, 0) - # Buffer untouched. - assert buf.target_step_list == [4] - - def test_selects_when_target_matches_even_if_weight_outside_window(self): - buf = FakeBuffer() - # start_weight=100 is well outside any realistic staleness window, - # but target_step matches current — force_in_order must ignore the window. - buf.add("g_match_stale_weight", weight=100, target_step=5) - sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) - - selected, num_groups = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert selected is not None - assert selected.sample_ids == ["g_match_stale_weight_g0"] - assert num_groups == 1 - - def test_ignores_unready_group_with_matching_target(self): - buf = FakeBuffer() - buf.add("g_unready", weight=5, target_step=5, ready=False) - sampler = StalenessSampler(buf, max_staleness_versions=0, force_in_order=True) - - result = _run( - sampler.select( - current_train_weight=5, min_prompt_groups=1, max_prompt_groups=4 - ) - ) - - assert result == (None, 0) - # Buffer untouched. - assert buf.ready_list == [False] diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index bc2b64cf23..b6134df948 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -27,6 +27,7 @@ from tensordict import TensorDict from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig from nemo_rl.algorithms.single_controller import ( SingleControllerActor, SingleControllerConfig, @@ -298,16 +299,14 @@ def test_train_pump_drives_mcore_training_step( rollout_manager = SimpleNamespace(_tq_buffer=None) cfg = SingleControllerConfig.model_construct( - max_weight_staleness_versions=1, + sampler=WindowedSamplerConfig(max_staleness_versions=1), min_groups_per_batch=num_prompts, target_groups_per_step=num_prompts, group_size=num_generations, - batch_selection_strategy="staleness_window", max_inflight_prompts=num_prompts, max_buffered_rollouts=num_prompts, max_train_steps=train_steps, max_num_epochs=None, - over_sampling=True, train_global_batch_size=train_gbs, partition_id=_PARTITION_ID, advantage_enabled=True,