dew.objectives.rl.scheduler

Keep agent rollouts in flight and admit complete groups within a staleness bound.

RolloutScheduler is the trainer’s Rollout for any SessionSource: an in-process environment, a single-turn prompt set, or a harness behind a recording gateway. Wrap the task dataset with scheduler.tasks(dataset): the wrapped stream registers each task batch as the trainer’s prefetch reads it, so when the trainer hands over batch i the scheduler submits batches i + 1 ... i + ahead under the version the engines serve now. Batch i’s own rollouts were submitted ahead calls earlier and have been running, across weight pushes, since. Nothing is read ahead of the trainer’s own prefetch, so the checkpointed data position stays the trainer’s: a resumed run re-reads and resubmits whatever was in flight, and reopening the stream cancels what the old one left running.

Each task becomes one group of groups rollouts, relabelled with the scheduler’s own group id, sample index and attempt, so a resubmitted sample rejoins its group. Admission is per rollout, by status:

• COMPLETED and AGENT_ERROR are admitted; the verifier scored them.
• TRUNCATED is admitted and trains as truncation says: mask (the default, for agentic context, turn and wall-clock limits), score on its verifier reward (single-turn RLVR), or zero (a length penalty); see dew.objectives.rl.sessions. Under score, a truncation without a reward is resubmitted as a failed attempt (cause unscored), like an INFRA_ERROR: its verifier never ran, so masking it would drop a sample for a reason unrelated to the policy, and training it has no reward.
• INFRA_ERROR and CANCELLED are never scored. The sample is submitted again under the served weights, up to max_attempts failures per sample, after which the group is abandoned rather than trained incomplete.
• A rollout whose oldest call is more than max_lag updates behind is discarded and resubmitted. One still running whose submission is already past the bound is cancelled before anyone waits on it: its first call was made under that version.
• A rollout still running timeout seconds after its submission is cancelled and resubmitted as a failed attempt, like an INFRA_ERROR. Cancelling asks the source to stop; a thread stuck inside an environment step cannot be reclaimed, so environments must bound their own step time.

A source that raises instead of returning a status is broken; the exception propagates after the batch’s work is cancelled.

The long tail is cut two ways, both keeping the batch shape fixed. oversample extra samples run per group and a group is admitted when its first groups rollouts finish; the rest are cancelled (APRIL’s active partial rollouts, arXiv:2509.18521, without the carry-over). admit below the task batch size admits the first admit groups to complete and cancels the others (slime’s over-sampling batch). Both select by completion time, which favors short rollouts; that bias is the price of not waiting on the tail. Rollouts running when a weight push lands keep running: their later calls carry the new version and the rollout’s staleness is its oldest call’s (Kimi K2’s partial rollouts, arXiv:2507.20534 section 3.3.4).

Weights are pushed through weights when the served version falls sync_every updates behind, so a first submission is at most ahead + sync_every - 1 updates stale at consumption, and construction refuses a max_lag below that. Admitted groups are packed by pack into fixed [rows, width] rows; pack computes the per-rollout advantages and masks. A complete group whose chains do not fit rows beside the groups admitted before it is cut, never packed into a failing step. The proximal policy is the trainer’s current weights, rescored over the packed rows with GRPOObjective.packed_log_probs (decoupled PPO, AReaL arXiv:2505.24298): sources report behavior likelihoods only, and the objective’s behavior_importance weights each token by proximal over behavior.

Every process of a multi-process trainer runs its own scheduler, over the task rows its data stream reads and on its own source, and packs its own rows: the step’s batch is every process’s rows together, sharded as the trainer shards any batch, and no rollout crosses a process. The pool meets twice a call. The weight push is one call every process makes, so the publisher has to agree on it across the pool. The proximal rescoring runs once over the pool’s batch, and each process reads its own rows back. A process whose admission fails raises on every process at the agreement point, rather than leave the others waiting in the rescoring. Where several processes read one share, as a tensor or sequence axis across processes makes them, the share’s first reader alone samples, and the others train on the rows it packed (first_reader_batch): their devices hold the same rows, which independent draws would not give them.

Name	Summary
Publisher	Where the trainer’s weights go: load serves them under version.
SchedulerRecord	What one trainer call consumed and what it cost.
task_ids	One task per integer task_id row, named by its decimal id.
RolloutScheduler	Train on complete rollout groups from source, ahead task batches early.

Publisher

class source

class Publisher(Protocol)

Where the trainer’s weights go: load serves them under version.

A RolloutServer is one; so is an engine fleet’s publish sequence.

Publisher.load

def load(variables: Variables, version: int) -> None

SchedulerRecord

dataclass source

class SchedulerRecord(
    updates: int,
    version: int,
    lag: int,
    groups: int,
    resubmitted: Mapping[str, int],
    cancelled: int,
    abandoned: int,
    cut: int,
    waited: float,
    metrics: Mapping[str, float],
)

What one trainer call consumed and what it cost.

version and lag are the oldest admitted call’s; groups counts admitted groups; resubmitted counts resubmissions by cause (infra_error, cancelled, stale, timeout, unscored); cancelled counts in-flight rollouts cancelled as surplus, stale or abandoned; abandoned counts groups given up after max_attempts; cut counts complete groups left out because their chains did not fit the batch’s rows beside the groups admitted before them; waited is the seconds the trainer waited. metrics is session_metrics over the admitted rollouts and their packed batch: merge ratio, status shares and masked shares, mean reward and reward components, submission-to-finish latency tail and token lag. The loss reports the trainer-engine mismatch.

task_ids

function source

def task_ids(batch: Batch) -> list[Task]

One task per integer task_id row, named by its decimal id.

RolloutScheduler

class source

class RolloutScheduler(
    objective: GRPOObjective,
    source: SessionSource,
    weights: Publisher,
    *,
    width: int,
    rows: int,
    tasks: Callable[[Batch], Sequence[Task]] = task_ids,
    groups: int = 4,
    oversample: int = 0,
    admit: int | None = None,
    max_lag: int = 1,
    ahead: int = 1,
    sync_every: int = 1,
    max_attempts: int = 3,
    timeout: float | None = None,
    estimator: str = 'group',
    truncation: str = 'mask',
    support_capacity: int | None = None,
    log: Callable[[SchedulerRecord], None] | None = None,
)

Train on complete rollout groups from source, ahead task batches early.

tasks turns one registered batch into its tasks (task_ids for integer task_id rows). timeout is each rollout’s deadline in seconds from its submission; a rollout past it is cancelled and resubmitted as a failed attempt. width and rows fix the packed batch shape, rows for this process’s share of a multi-process trainer’s batch; estimator and truncation are pack’s advantage family and truncation policy, and support_capacity its per-row support length, which a filtered-sampling source requires. log, when given, receives a SchedulerRecord per call, of this process’s rollouts; a share’s later readers sample none and log nothing.

RolloutScheduler.tasks

def tasks(dataset: Dataset) -> Dataset

dataset with a training stream that registers each task batch ahead of the step.

Opening the stream again, as a resume does, cancels every rollout the previous stream left in flight.

RolloutScheduler.close

def close() -> None

Cancel every rollout in flight; the source belongs to the caller.