Vivek Kashyap
← Back to writing

RL Infra Orchestrator - Part 2: Dispatcher and Train Sink

Note: my notes, cleaned up with AI.

In part 1 we looked at the orchestrator's main loop, the thing that pulls finished rollouts off a queue, hands them to the train sink, and ships full batches to the trainer. But two lines in that loop were basically black boxes: dispatcher.out_q.get(), where a finished rollout showed up, and train_sink.add(episode), which turned rollouts into a batch. Today we open both of them up. Part 1 was more of a map than a deep dive, so this one goes the other way, close to line by line.

Here are the two pieces:

  1. Dispatcher. It decides which prompts to run, sends them off to inference, keeps the engine busy, and drops each finished rollout onto the queue.

  2. Train sink. It takes those finished rollouts and turns them into the batches the trainer learns from, tokenizing them, computing advantages, and filtering out the junk.

Together they're the whole path a rollout takes, from a prompt we want to run to a batch we hand the trainer. We'll start with the train sink, then come back to the dispatcher. Let's dive in.

The train sink

The train sink's job is to take the finished rollouts and turn them into training data. A rollout on its own is just a conversation with a reward attached, and the trainer can't learn from that directly. It needs flat token sequences with advantages attached, packed into a batch. Turning one into the other is what the sink does.

From part 1, the main loop hands each finished episode to the sink:

train_batch = await self.train_sink.add(episode)
if train_batch is not None and not self.draining and not self.stopped.is_set():
    await self.finalize_train_batch(train_batch)

(orchestrator.py:555)

So add is the entry point. It gets called once per finished episode. Most of the time it returns None because it's still collecting, but every so often it returns a full TrainBatch that's ready to ship.

Here's add:

async def add(self, episode):
    group_id = episode[0].group_id
    env_name = episode[0].env_name
    for rollout in episode:
        await self.process_rollout(rollout)
    self.pending_groups[group_id].extend(episode)
    self.pending_group_episodes[group_id] += 1
    if self.pending_group_episodes[group_id] < self.group_size_for(env_name):
        return None                                  # group not complete yet
    await self.process_group(group_id)
    if ready:                                         # batch full?
        return self.process_batch()
    return None

(train_sink.py:127)

process_rollout()

This is the tokenization step. It takes one rollout and turns it into the flat token form the trainer needs.

async def process_rollout(self, rollout):
    if rollout.has_error or not rollout.trainable:
        return
    samples = await asyncio.to_thread(
        trace_to_samples, rollout, env_name=rollout.env_name,
        mm_token_type_ids_mapping=self.mm_token_type_ids_mapping,
    )
    rollout.samples = samples or []

(train_sink.py:153)

trace_to_samples turns the rollout into TrainingSamples: the token ids, plus a mask that flags which tokens the model generated. Those are the ones we train on; the rest is context. This runs on a worker thread so the tokenization overlaps with the dispatcher producing more rollouts, instead of blocking everything.

Errored and untrainable rollouts skip all of this. There's no point tokenizing something we're going to throw away.

process_group()

This fires once a whole group has arrived, all group_size attempts at the same prompt, and it's where they get scored into advantages.

First it pulls the group out and records every rollout in it:

group = self.pending_groups.pop(group_id, [])
for r in group:
    self.pending_rollouts.append(r)

(train_sink.py:176)

That pending_rollouts list is the observation window, everything we've seen since the last batch, kept for metrics. Even rollouts we're about to drop get added, because we still want them in the stats.

Then the errored rollouts are dropped and the survivors go to the algorithm:

await env.algorithm.finalize_group(survivors)

(train_sink.py:212)

This is the line that scores the group into advantages, and it's the piece prime-rl made swappable. env.algorithm might be GRPO or something else, and whichever it is computes the advantages and stamps them onto the samples. The sink just calls finalize_group and lets the algorithm handle it. We'll open up the algorithms in a later post.

After scoring, it copies the env's sampling temperature onto every token (nothing interesting, just a fan-out), then runs the filters:

apply_filters(self.pre_filters, survivors)
for r in survivors:
    if r.is_filtered:
        continue
    self.pending_batch.append(r)

(train_sink.py:221)

The filters flag junk like gibberish or repetition, and anything flagged is skipped. Everything clean lands in pending_batch, the list process_batch pulls batches from.

So a group comes in, loses its errored and filtered rollouts, and puts the clean ones on the batch list. Two lists to keep straight: pending_rollouts has everything (for metrics), pending_batch has only the clean ones (for training).

process_batch()

First thing to notice, this one is not async:

def process_batch(self):
    ...

(train_sink.py:253)

Every other step was async because it did IO (tokenizing on a thread, scoring against a teacher over the network). process_batch is pure CPU. It's just slicing lists in memory, so there's nothing to await.

What it does is small. It takes a batch-sized chunk off pending_batch, runs the post-batch filters, and flattens the survivors into samples:

if self.post_filters:
    apply_filters(self.post_filters, cohort)
samples = [sample for r in cohort if not r.is_filtered for sample in r.samples]

(train_sink.py:275)

"Flatten" here means going from rollouts to samples. Remember each rollout was tokenized into one or more TrainingSamples back in process_rollout. This line walks the chosen rollouts, skips any that got filtered, and collects all their samples into one flat list. That flat list is the payload the trainer trains on.

Finally, it resets the observation window and returns the batch:

rollouts = self.pending_rollouts
if samples:
    self.pending_rollouts = TrainRollouts()
return TrainBatch(rollouts=rollouts, samples=samples)

(train_sink.py:288)

pending_rollouts is the window we've been filling up in process_group. It gets handed off with the batch (for metrics) and reset to empty for the next round. The result is a TrainBatch carrying samples (what gets trained on) and rollouts (what gets measured). That's exactly the TrainBatch the main loop's finalize_train_batch picks up and ships to the trainer.

And that's the whole train sink. One rollout comes in and gets tokenized, a full group gets scored into advantages, and enough scored survivors get packed into a batch.

The dispatcher

One thing before we start: the dispatcher is big, around 700 lines, and it does more than we'll get into here. It also cancels stale rollouts when the weights update, schedules evals, and handles shutdown, and we're skipping all of that. We'll just follow the main path, how a rollout gets scheduled, run, and handed back. That's still most of the file.

Back in part 1, start() launched the dispatcher as a background task with asyncio.create_task, so it runs its own loop forever, next to the main loop.

start()

Here's that loop, trimmed:

async def start(self):
    while not self.stopped.is_set():
        await self.fill_inflight()
        done, _ = await asyncio.wait(self.inflight, return_when=asyncio.FIRST_COMPLETED, timeout=0.5)
        for task in done:
            await self.handle_completed_rollout(task)

(dispatcher.py:238)

Two calls matter. fill_inflight is the schedule step, it starts new rollouts. handle_completed_rollout is the collect step, it deals with rollouts that just finished. Everything in the dispatcher hangs off one of these two. Let's follow the schedule side first.

fill_inflight()

fill_inflight keeps starting rollouts until either there's no room left or there's nothing left to start. That's the whole idea, keep the inference engine as busy as it's allowed to be.

All it really does is call try_schedule in a loop:

async def fill_inflight(self):
    while True:
        if self.available_permits <= 0:
            return
        scheduled = await self.try_schedule("train")
        if not scheduled:
            return

(dispatcher.py:306)

So each pass either starts one rollout (through try_schedule) or stops because we're full or out of work. (There's also an eval-vs-train mode in here, but ignore that for now.) The real logic is in try_schedule.

try_schedule()

try_schedule starts exactly one rollout, and its main decision is about the KV cache.

Here's the idea. Rollouts come in groups, several attempts at the same prompt. Since all attempts of a group share that prompt, sending them to the same inference server means the server already has the prompt cached and doesn't recompute it. So try_schedule prefers to continue a group that's already open before starting a brand new prompt:

for gid, group in self.groups.items():
    if group.rollouts_to_schedule > 0:       # this group still has attempts left
        return await self.schedule_group_rollout(gid, group)

fresh = self.next_fresh_group(...)           # nothing open, grab a new prompt
return await self.schedule_group_rollout(gid, fresh)

(dispatcher.py:343)

So if a group is open (still has attempts to launch), continue it. If not, open a fresh one. Either way it ends by calling schedule_group_rollout, which is where the rollout gets fired off.

schedule_group_rollout()

This is the one that launches a rollout. Three things worth pulling out.

The pinned client. The first attempt of a group picks an inference server, and every later attempt reuses the same one:

if group.pinned_client is None:
    group.pinned_client = await pool.select_train_client(...)
client = group.pinned_client

(dispatcher.py:415)

This is the cache idea from try_schedule made real. Pinning the whole group to one server (pinned_client) is what makes all its attempts hit the same cached prompt.

Running the environment. Then it reserves capacity and starts the rollout:

await self.acquire(permits)
task = asyncio.create_task(env.run(client=client, model_name=model_name, **addressing))

(dispatcher.py:463)

env.run is where the rollout actually runs. The env drives the agent and calls the inference server to generate, and create_task launches it in the background so the dispatcher can move on. There are two flavors: modern (v1) envs run one rollout at a time with env.run, while older (v0) envs run the whole group in one shot with env.run_group. Same idea, one just batches the group together.

Permits. The acquire call reserves permits. A permit is just a slot for one in-flight rollout. The dispatcher has a fixed number of them, so permits are how it caps how many rollouts run at once. Start a rollout, take a permit. Finish one, give it back.

next_fresh_group()

The one piece we skipped: when there's no open group to continue, next_fresh_group grabs a new prompt from the source.

example = source.next_example(self.available_permits)
...
policy_version_at_start=self.policy.version,

(dispatcher.py:369)

source.next_example just hands over the next prompt to run. The new group gets stamped with the current policy version, which is the policy_version tag we used in part 1 for off-policy staleness. Then it goes back to try_schedule, which passes it to schedule_group_rollout like any other group.

That's the whole schedule side: fill_inflight loops, try_schedule picks a group (open or fresh), and schedule_group_rollout fires the rollout at a pinned server.

handle_completed_rollout()

Now the collect side. When a rollout finishes, start() calls this on it. Its job is to emit every finished rollout exactly once onto out_q, the queue the main loop reads from.

First it takes the rollout out of the in-flight ledger and gives back its permit:

meta = self.inflight.pop(task, None)
if meta is None:
    return
self.release(meta.rollout_count)

(dispatcher.py:504)

release returns the permit we took when we started the rollout, freeing a slot for a new one. (The meta is None check just means "someone already handled this," so skip it.)

Then it gets the result, and this is where the different endings are handled:

try:
    result = task.result()
    rollouts = result if isinstance(result, list) else [result]
except asyncio.CancelledError:
    return
except Exception as exc:
    rollouts = [Rollout(...) for _ in range(meta.rollout_count)]
    for r in rollouts:
        r.capture_error(exc)

(dispatcher.py:511)

Three cases. If it succeeded, we get the real rollouts. If it was cancelled, we just return, someone else already cleaned it up. If it crashed, we build placeholder rollouts and stamp the error on them, so the failure still flows through instead of vanishing.

There's also an empty-rollout case just below:

for r in rollouts:
    if not r.has_error and r.num_turns == 0:
        r.errors.append(vf.Error(type="EmptyTrajectory", ...))
        r.ok = False

(dispatcher.py:529)

A rollout that came back with nothing (no error, but zero turns) gets marked as an error too. The point of all this is uniformity: success, crash, cancel, or empty, everything ends up as a rollout with a clear status, so nothing the dispatcher started ever silently disappears. Then it hands the rollouts to emit_episode.

emit_episode()

Last step, and it's short. It stamps the rollouts with their tags and puts them on the queue:

await self.out_q.put(rollouts)

(dispatcher.py:572)

That out_q.put is the handoff. This is the exact queue the main loop pulls from with dispatcher.out_q.get() back in part 1. So the moment this line runs, the rollout leaves the dispatcher's world and becomes the main loop's problem, which feeds it to the train sink we just covered.

And that's the round trip. fill_inflight starts rollouts, they run in the background, handle_completed_rollout collects them when they finish, and emit_episode drops them on the queue for the main loop.

Thanks for reading — if you made it all the way down here, I appreciate it.