Which process owns which piece of the SingleController checkpoint

Architecture companion for NVIDIA-NeMo/RL #3480 · green = added by this PR · all permalinks pin the PR head 608e9bd · see also what a restart loses with in_order.

Two things about this layout surprise people. First, almost nothing is built inside the controller actor: setup_single_controller() builds every object on the driver and hands them over in one argument, so the actor's __init__ does no construction (why). Second, there are two data-plane clients, and this PR gives them different jobs: the trainer's client restores TransferQueue during setup, and the controller's client writes the snapshot later. TransferQueue is the only place rollout tensors are ever stored.

DRIVER PROCESS — run_grpo_single_controller.py setup_single_controller() TQPolicy — owns trainer workers dp_client(bootstrap=True) load_data_plane_checkpoint() NEW restore happens here, before the actor exists SingleControllerActor — one Ray actor, one asyncio loop rollout pump train pump sampler TQReplayBuffer DataPlaneCheckpointBarrier NEW dp_client(bootstrap=False) _validate_replay_inventory() + _save_data_plane_checkpoint() NEW OTHER RAY ACTORS trainer workers generation (vLLM) TransferQueue controller + storage the only place rollout tensors are stored 1 2 3 4 5 6 CHECKPOINT BUNDLE ON DISK — step_N/ policy/ · train_dataloader.pt data_plane/ NEW replay_buffer_metadata.pt NEW data_plane/ holds the tensors and is written by TransferQueue itself. replay_buffer_metadata.pt holds only which rows belong to which group — no tensors.

The six flows, in the order they happen

These are every data-plane call the controller and the trainer make. Numbered by time, so 1 really does come before 2 — the restore finishes before the actor is created. Not drawn: weight sync from trainer to generation, and generation traffic, neither of which touches TransferQueue.

What movesDetail
1 New. Trainer's client → TransferQueue: restore TQPolicy.load_data_plane_checkpoint() → dp_client.load_checkpoint() → tq.load_checkpoint(). This runs on the driver at setup.py:697, before the connect-only client is built and before the actor exists. It has to: TransferQueue can only be restored while it is still untouched, and the adapter refuses a restore after any read or write through that client (_require_clean_for_load).
2 Driver → controller actor Everything the actor needs, cloudpickled as one SingleControllerActorArgs. This PR adds one field to it — data_plane_checkpoint_metadata, the metadata the driver read back during restore, which the actor later checks its own sidecar against.
3 Replay buffer → TransferQueue: finished rollouts Unchanged in shape — put_samples on commit, clear_samples when a group is dropped. What is new is that both now run inside the barrier, so they cannot overlap a snapshot.
4 Controller → TransferQueue: the advantage stage Unchanged, and easy to miss. Before a batch trains, the controller reads its rows back (get_samples) and writes the computed advantages onto the same rows (put_samples). It only overwrites fields on keys that already exist, so it never changes which rows are there — which is why the checkpoint's inventory check is unaffected by it. It takes no barrier slot; it is safe because it runs in the same task as the save.
5 New. Controller → TransferQueue: snapshot Three new client methods: list_sample_ids() to check the stored rows match the index, save_checkpoint(), and load_checkpoint(). All three are on the abstract client, so every adapter implements them.
6 New. Controller → disk: the index file The controller writes replay_buffer_metadata.pt itself from metadata_state_dict(). TransferQueue writes data_plane/ separately. A digest ties the two together so a mismatched pair is caught on resume.

Save and restore, in order

This change touches both halves, so here are both. Ordering is the part that is easy to get wrong in each: on save, the index is captured and checked before the snapshot is written, all inside one exclusive window. On restore, TransferQueue is loaded before the controller actor exists at all.

SAVE — runs every save_period steps, inside the controller actor train pump replay buffer TransferQueue step_N/ on disk note the loader place make tmp_step_N/ policy weights + optimizer train_dataloader.pt — the loader place barrier held: finished rollouts queue at commit until this window closes metadata_state_dict() — which rows form which group list_sample_ids() — do the stored rows match that index? save_checkpoint() — TransferQueue writes data_plane/ itself replay_buffer_metadata.pt — the index, no tensors rename tmp_step_N/ to step_N/ — nothing is resumable until this happens RESTORE — runs once, at startup setup (driver) TQPolicy TransferQueue SC actor replay buffer read step_N/ trainer state, loader place load_data_plane_checkpoint() bootstrap client: load_checkpoint() rows and consumer places come back digest, group count, mode=authoritative TransferQueue is restored above this line. Nothing has touched it yet. build connect-only client, then the buffer SingleControllerActor.remote(actor_args + digest) run() load_state_dict(sidecar) index only - no writes to TQ list_sample_ids(): do the rows match the index? start the rollout and train pumps

Three orderings are doing real work. On save, the replay index is captured and checked against the stored rows before the snapshot is written — so a mismatch stops the checkpoint instead of producing a bundle that looks fine and cannot be resumed. On restore, TransferQueue is loaded before the connect-only client exists, because a restore is only safe onto an untouched client; the adapter refuses one after any read or write. And the index is rebuilt before the pumps start, so no live writer can race the check.

The shaded band on the save side is the one thing with a run-time cost: while it is held, finished rollouts wait at commit. Generation keeps going — only the hand-off pauses.

What the barrier is for

Generation keeps running while a checkpoint is being written, so without something holding them apart the snapshot and the index would describe different sets of rows. DataPlaneCheckpointBarrier is the new piece that stops that. It allows many writes at once but only one checkpoint, and a checkpoint waits for the writes already in progress to finish before it starts. Finished rollouts still arrive during a save — they just wait at the commit step.

It lives in the controller actor and is handed to the replay buffer once at startup (single_controller.py:207), which is why both the buffer's writes and the controller's snapshot go through the same object. It is plain asyncio with no Ray or torch in it, so it can be tested on its own in milliseconds — worth keeping that way.

The part worth remembering. Restore and save happen in different processes through different clients: the trainer's client restores TransferQueue on the driver during setup, and the controller's client writes the snapshot later from inside the actor. That ordering is held together by the order of statements in one setup function, with nothing asserting it. If the connect-only client or the replay buffer were ever built before setup.py:697, the restore would land on a TransferQueue that had already been touched, and the guard inside the adapter could not see it — it only knows about its own client.