Task Scheduler

1 Role in the Rollout Pipeline

The TaskScheduler (task_scheduler.py:127) drives a rollout step. It is not a Ray actor — it runs as plain async Python inside the SchedulerRolloutManager actor; only the SGLang engine holds a persistent Ray allocation.

Every step makes the same four moves — fetch → distribute → batch → run — traced below, with results streaming back through on_task_done as each task completes.

2 Batching Strategies

Batch composition is the scheduler's one pluggable decision. A BatchingStrategy (task_scheduler.py:57) is any callable with the signature below; everything downstream is uniform.

# task_scheduler.py:57 — the strategy is a plain callable, no class hierarchy
BatchingStrategy = Callable[[list[Task], dict[str, EvaluatorSpec]], list[Batch]]

TaskScheduler.__init__ takes an optional batching_strategy, defaulting to fifo_batching (task_scheduler.py:156) — the only shipped strategy.

fifo_batching default

It groups tasks by task_type in encounter order, sealing a batch at the matching spec's batch_size (task_scheduler.py:60):

def fifo_batching(tasks, specs) -> list[Batch]:
    batches, open_batch = [], {}
    for task in tasks:
        ttype = task.config.task_type
        batch_size = specs[ttype].batch_size if ttype in specs else None
        if ttype not in open_batch:                 # start a batch for this type
            b = Batch(task_type=ttype); open_batch[ttype] = b; batches.append(b)
        open_batch[ttype].tasks.append(task)
        if batch_size is not None and len(open_batch[ttype].tasks) >= batch_size:
            del open_batch[ttype]                        # seal → next task opens a new one
    return batches

Task-type-aware interleaving on the eval path

On the eval path, execute_eval() builds every (prompt × copy) sample and interleaves them round-robin — ABCABC… not AABBCC… (task_scheduler.py:746-774) — so concurrent tasks span datasets and smooth sandbox/inference pressure. All share one run_all(), split apart afterward by _eval_dataset_name.

Spec resolution & the config cascade

Each batch resolves its spec via _resolve_spec(): exact evaluator_id match, else a warning and fallback to the first spec (task_scheduler.py:321). Two generation settings resolve through a three-level cascade:

_resolve_stopping — task_scheduler.py:162 · _build_sampling_params — task_scheduler.py:200. The per-dataset overrides ride in on each sample's metadata["_dataset_cfg"], set by MultiDataSource.

3 The Two-Level Resource Pool

PlacementPool (resource.py:137) is admission control in two composing modes: a global pool from the cluster layout gates whole batches; a short-lived local pool hands one token per running sample.

Resource descriptors

Three dataclasses describe supply and demand (resource.py:18-90):

All-or-nothing batch acquisition

try_acquire_batch(budget, per_sample, max_tokens) (resource.py:220) takes tokens one per_sample at a time — under the pool lock, greedily from the node with the most free GPU, cross-node allowed — until budget is spent or max_tokens (the sample count) is hit. If any token can't be placed, all roll back and it returns None, so the batch waits rather than half-commit resources that block others.

Fractional GPU sharing

Each GPU has capacity 1.0; a sample consumes its Resource.gpu fraction, so gpu=0.33 packs three jobs per GPU (__init__.py:48-51). _NodePool.try_allocate best-fits by remaining capacity (resource.py:104); max_tokens_for counts floor(1.0 / per_sample.gpu) slots per GPU. The SWE recipe sets gpu=0 — sandboxes are CPU/memory-bound and student/teacher GPUs belong to SGLang — so SWE batches gate on CPU and memory alone.

Deadlock prevention

Before launching, run_all() asserts the pool can meet the largest per-spec demand it will ever request (task_scheduler.py:358-387): the sample_resource tokens that fit a spec's resource budget (capped by the biggest batch) must not exceed max_tokens_for(sample_resource):

tokens_needed = min(budget_tokens, max_batch_samples[spec.evaluator_id])
pool_capacity = pool.max_tokens_for(spec.sample_resource)
assert pool_capacity >= tokens_needed, (
    f"Deadlock: EvaluatorSpec '{spec.evaluator_id}' needs {tokens_needed} "
    f"placement tokens but the pool can only provide {pool_capacity}…")

It fails fast at startup instead of hanging on an unadmittable batch.

Cluster layout schema

_build_placement_pool (__init__.py:37) reads --env-cluster-layout as inline JSON or a file path. Each node is a GPU-id list (shorthand) or a full dict; missing cpu/memory_gb fall back to --env-num-cpus / --env-memory-gb. Absent the flag, the pool is None and batches launch ungated. SWE scripts generate it per run:

# rl_engine/examples/swe_agent/config/cluster_layout.json — full-dict form
{
  "127.0.0.1": { "gpus": [0, 1, 2, 3], "cpu": 96, "memory_gb": 384 }
}

# shorthand — a bare GPU-id list; cpu/mem come from --env-num-cpus / --env-memory-gb
{ "127.0.0.1": [0, 1, 2, 3] }

4 Evaluator Lifecycle

An Evaluator (evaluator.py:197) is created for one batch, drains its queue, and is discarded. It holds no domain logic — a BaseAgent handler owns the episode; its EvaluatorSpec binds that handler to a resource model:

# evaluator_config_dagger.py:80 — the SWE-DAgger spec (real values)
EvaluatorSpec(
    evaluator_id="swe",
    handler=DAggerSWEAgent(…),
    resource=Resource(gpu=0, concurrency=env_concurrency, cpu=96, memory_gb=320),  # per-evaluator budget
    run_resource=Resource(gpu=0, concurrency=1, cpu=4, memory_gb=16),   # per-sample container
    eval_resource=Resource(gpu=0, concurrency=1, cpu=4, memory_gb=16),  # per-sample verifier
    batch_size=args.evaluator_batch_size or 1,
    default_stopping=StoppingCriteria(max_turns=100, max_context_len=65536),
)

Since run and eval never overlap for a sample, the spec's sample_resource (evaluator.py:68) is the per-field max(run_resource, eval_resource) — the peak one sample holds, pre-acquired from the global pool. Per-evaluator concurrency is Resource.concurrency: the SWE recipe sets it from STACX_DAGGER_ENV_CONCURRENCY (default 4), the installed-scaffold path from --evaluator-concurrency (evaluator_factory.py:225). Either becomes the asyncio.Semaphore ceiling.

Per-sample execution

_execute_sample (evaluator.py:338) is the core. Under the semaphore it acquires one token and holds it across the whole episode — never releasing between phases — so at most concurrency containers exist at once:

After eval, _populate_token_data copies token IDs, log-probs, and the loss mask off the AgentResult onto the sample for the trainer. Samples in a group run via asyncio.gather; Evaluator.run() fires on_task_done per task so tqdm and the collector update live.

Abort classification

When a sample raises, the except block marks it ABORTED, sets reward={"score": 0.0}, and injects a 2-token dummy with remove_sample=True — zero gradient, but Megatron's batch shape stays intact (evaluator.py:429-452). It preserves the agent's own termination reason before overwriting it with a transport classification:

_agent_reason = sample.metadata.get("_termination_reason")
if _agent_reason is not None:
    sample.metadata["_termination_reason_agent"] = _agent_reason   # kept for diagnostics
sample.metadata["_termination_reason"] = classify_abort(str(e))   # agent_timeout / infra_*

classify_abort (termination_summary.py:143) maps the exception to agent_timeout / infra_conn / infra_sandbox / infra_other; classify_termination folds (reward, reason) into one of eight failure classes below. _annotate_abort attaches diagnostics (phase, window, node IP, turns) for the summary.

5 Training-Loop Integration

The scheduler alone knows the current rollout index, so it publishes per-iteration state as env vars the agent and loss code read — set by execute_rollout() before the fetch loop (task_scheduler.py:481-497):

The t_max_step > 0 guard avoids clobbering a statically-set tail length with T_max=0. execute_eval() pops STACX_DAGGER_BETA, so eval runs student-only at β=0 (re-injectable via STACX_DAGGER_EVAL_BETA). The schedule lives in DAggerSchedule (rollout/utils/dagger_schedule.py) — linear/cosine/exponential/step β decay plus a linear t_max ramp — built at train.py:281. See Algorithm Recipes.

Rejection sampling & the finish-only filter

Two independent survivorship filters compose — rejection first, finish-only on the survivors (rollout_manager.py:446-467).

Reuse & resume

Two resume paths coexist, both keyed off the saved trajectory .pt (rollout_manager.py:201-243):

• Fresh (live DAgger / AggreVaTe / OPD): STACX_FORCE_FRESH_ROLLOUT=1 bypasses reuse; every iter rolls out with the improving student, saving trajectory_sft_samples.pt plus a per-iter …_iter_NNN.pt (gated by STACX_SAVE_PER_ITER_PT, default on).
• Reuse (iter-SFT): STACX_REUSE_ROLLOUT=1 always loads the .pt; STACX_REUSE_PER_ITER_DIR replays a directory of per-iter dumps in sequence. Without either, rollout_id > 0 auto-reuses iter-0's data (multi-epoch SFT).

execute_rollout() also writes a write-only partial checkpoint per group (task_scheduler.py:556-560); it is never loaded on resume — a resumed rollout re-runs from scratch, since the data source's cursor and RNG aren't checkpointed.

Eval outputs

With --save-eval-details, execute_eval() streams three artifacts per iteration: a per-sample JSONL (eval_details/{id}.jsonl), a per-turn placement log (node IP + GPU), and a failure-class summary.json. Rollout steps emit the same at eval_details/train_{id}.summary.json. Per-dataset metrics are eval/{name}/accuracy and, when n_samples_per_eval_prompt > 1, eval/{name}/pass@k.

6 Knob Reference

Flags and variables on the scheduler path. CLI flags are set once at launch; code-set STACX_* variables are written per iteration — never set those by hand. Full catalog in scripts/train/swe/ENV_REFERENCE.md.

7 Class & Function Reference

Related pages: Rollout Architecture · Inference Backend · Sandbox Management · Algorithm Recipes.