[DRAFT] refactor(sc): interface-driven staleness samplers (design demo for #3220)

[terrykong]

⚠️ DRAFT — do not merge

Design demonstration for the sampler-refactor discussion on #3220 (the StalenessSampler interface-driven split). This is a strawman to look at, built on top of #3220's head so the diff is just the refactor. Not tested locally (macOS/no-GPU, Linux-only lockfile) — treat every behavior claim as static-only.

Posting this instead of a wall of review text so the shape is concrete. Below is before/after for each touched section; permalinks point at the surrounding code (before = base d3188b24b, after = this branch 1398b010).

---

1. Sampler: one flag-driven class → an interface + one policy per behavior

Before — staleness_sampler.py:21-53: StalenessSampler, five mode-encoding flags whose valid combinations are policed at runtime; select is three disjoint branches; the shared evict only fits two of the three modes.

class StalenessSampler:
    def __init__(
        self,
        buffer: TQReplayBuffer,
        max_staleness_versions: int,
        sample_freshest_first: bool = False,
        strict_weight_fifo: bool = False,
        force_in_order: bool = False,
    ) -> None:
        if strict_weight_fifo and sample_freshest_first:
            raise ValueError("strict_weight_fifo and sample_freshest_first are mutually exclusive")
        ...
    async def select(...):   # if force_in_order: ... elif strict_weight_fifo: ... else: ...
    async def evict(...):     # weight-window only; wrong for force_in_order

After — PromptGroupSampler interface + one named policy each (WindowedSampler, WeightFifoSampler, InOrderSampler). The three new members (admit, and the is_on_policy/required_buffer_capacity derived facts) are the seam that lets the pumps and validation stop reading raw knobs.

class PromptGroupSampler(Protocol):
    async def admit(self, *, trainer_version_fn) -> Optional[int]: ...   # NEW: rollout-pump gate + target_step stamp
    async def select(self, *, current_train_weight, min_prompt_groups, max_prompt_groups): ...
    async def evict(self, *, current_train_weight) -> int: ...
    @property
    def is_on_policy(self) -> bool: ...                        # NEW: weight-sync (invalidate_kv_cache, wired in #3266)
    def required_buffer_capacity(self, groups_per_step): ...   # NEW: replaces the scattered capacity assertion

Constructor contract (methods only — a runtime_checkable Protocol can't express __init__): create_sampler builds every policy as Sampler(buffer, **extra) — the shared TQReplayBuffer is the first positional arg. This is documented on the Protocol, and the custom-FQN branch wraps construction so a sampler that doesn't take buffer-first fails with a clear message rather than an opaque TypeError; subclassing BaseSampler satisfies it.

Two review findings get fixed by construction: the evict/select asymmetry (r3617016004 — InOrderSampler.evict keys on target_step, matching its select) and the dead sample_freshest_first (r3617016008 — it exists only on WindowedSampler).

---

2. SingleControllerConfig: four indirect knobs → one discriminated sampler

Before — single_controller.py:107-125

    max_weight_staleness_versions: int = 1
    ...
    batch_selection_strategy: Literal["strict_on_policy", "staleness_window"] = "strict_on_policy"
    ...
    # over_sampling=False gates each batch on max_rollout_version vs trainer_version;
    # force_in_order stamps target_step on each dispatch ...
    over_sampling: bool = False
    force_in_order: bool = False

After — single_controller.py:119-126

    # 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. ``custom`` takes a
    # ``module:ClassName`` target to load a PromptGroupSampler from outside this repo.
    sampler: SamplerConfig = Field(default_factory=InOrderSamplerConfig)

SamplerConfig is a pydantic discriminated union; name: custom + target: "module:ClassName" loads an out-of-repo sampler, mirroring the existing create_checkpoint_engine convention (nemo_rl/utils/checkpoint_engines/base.py). That's the concrete shape for the custom-sampler support tracked in #2625 item 3.

---

3. SingleControllerActor.__init__: runtime knob-policing → derived-fact check

Before — single_controller.py:271-288 (coerce one mode, raise on two invalid knob combos) then :329-336 (warn, then translate knobs into the sampler):

        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 ...")
        if cfg.max_weight_staleness_versions == 0 and cfg.over_sampling:
            raise ValueError("max_weight_staleness_versions=0 requires over_sampling=False: ...")
        if cfg.force_in_order and cfg.over_sampling:
            raise ValueError("force_in_order=True requires over_sampling=False ...")
        ...
        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,
        )

After — single_controller.py:260-276: the union makes those states unrepresentable, so the coercion/raises/warn_if_... are gone; the one remaining check reads a fact the sampler owns, so it can't drift from the admission logic:

        # (strict_on_policy coercion + over_sampling/force_in_order raises + warn fn removed)
        self._sampler = create_sampler(self._buffer, cfg.sampler)

        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 the "
                f"{type(self._sampler).__name__}'s required capacity ({required_capacity}) ..."
            )

warn_if_staleness_window_below_minibatches (~50 lines) is deleted outright.

---

4. _rollout_pump: inline gate + actor-owned counter → sampler.admit()

Before — single_controller.py:499-510: the dispatch gate, the target_step derivation, and the _max_rollout_version counter all live in the pump (a second copy of the sampler's staleness logic):

                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 = self._max_rollout_version if force_in_order else None

After — single_controller.py:445-449: the sampler owns admission and the dispatch counter; the pump follows whichever policy is selected with no gating logic of its own (self._max_rollout_version is deleted from the actor):

                for prompt_batch in self._dataloader:
                    target_step = await self._sampler.admit(
                        trainer_version_fn=lambda: self._trainer_version
                    )

---

No legacy path

Per the discussion, this drops the legacy knobs entirely rather than keeping a shim — the sampler config's name is the only switch. Removed: the four config knobs above; the strict_on_policy coercion + two raises; warn_if_staleness_window_below_minibatches; _max_rollout_version on the actor; and the monolithic StalenessSampler + its 526-line test (replaced by test_sampler_interface.py).

Prior art

• verl (fully-async): a single async_training.staleness_threshold, interpreted by the producer only — same "one owner for staleness" lesson, but a scalar can't express in-order exact matching, and there's no pluggable sampler.
• miles/slime: buffer_filter = load_function(args.buffer_filter_path) — direct precedent for the FQN hook, but its staleness control is disconnected from the filter, so a custom filter can't change dispatch. The admit()-owning policy here is strictly stronger.

Not done (draft)

• Not run: no compatible local env. test_sampler_interface.py (and the updated pump tests) are written to the intended contract but unverified.
• is_on_policy has no consumer at feat(sc): train path — StalenessSampler + _train_pump rewrite #3220's head (the invalidate_kv_cache call lands in feat(sc): setup + entrypoint #3266); included to show the derived-fact seam.
• Exemplar/recipe YAMLs still set the old knobs — they'd move to an async_rl.sampler: block (not done here).

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.