nemo_rl.algorithms.single_controller#

SingleController: asyncio orchestrator for the RL training loop.

CPU-only Ray actor that runs two concurrent pumps plus a watchdog, and coordinates the other actors via lightweight RPCs. SC sends control signals and reads metadata only — model tensors still move through DataPlane or NCCL.

Data flow: _rollout_pump → gen.generate_and_push(prompt, dp_client) ← RPC to GenWorker GenWorker → dp_client.put_samples(…) _train_pump → sampler.evict/select against TQReplayBuffer → _advantage_stage(meta) → dp_client.get_samples(…) → adv_estimator.compute_advantage(…) → dp_client.put_samples(…) → trainer.begin/train_microbatches/finish_train_step (split API, driver-side TQPolicy via asyncio.to_thread) Trainer → dp_client.get_samples(…) (via its own client) → dp_client.clear_samples(…) ← SC clears after train _sync_weights → WeightSynchronizer.sync_weights()

Module Contents#

Classes#

SingleControllerActor

CPU-only Ray actor that orchestrates the RL training loop.

Data#

Generation

API#

nemo_rl.algorithms.single_controller.Generation#

None

class nemo_rl.algorithms.single_controller.SingleControllerActor(
master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig,
actor_args: nemo_rl.algorithms.single_controller_utils.setup.SingleControllerActorArgs,
setup_timing_metrics: nemo_rl.algorithms.metric_utils.SetupTimingMetrics,
)#

CPU-only Ray actor that orchestrates the RL training loop.

Owns three concurrent asyncio tasks:

  • _rollout_pump: dispatches prompts to GenerationWorkerActor

  • _train_pump: claims DataPlane meta, trains, clears consumed rows, then runs _sync_weights (drain gate + weight synchronization) inline after each optimizer step

  • _watchdog_pump: publishes rollout counters and reports stalls or unhealthy environments, which are the failures that otherwise produce no signal at all

All other actors are passive — they expose methods and wait to be called.

Initialization

Initialize the SingleController actor.

Parameters:

  • master_config – SC MasterConfig.

  • actor_args – Pre-built actor args from setup_single_controller.

  • setup_timing_metrics – Driver-side setup timings; logged here (Logger isn’t cloudpickleable).

async run() → dict[str, Any]#

Main entry point. Runs until max_train_steps is reached.

async ping() → dict[str, Any]#

Liveness check — returns immediately if event loop is running.

async _maybe_restore_replay_buffer() → None#

Restore replay-buffer groups from the previous run’s checkpoint.

Skipped with a warning when the checkpoint was written under a different sampler: restored groups carry the saving sampler’s weight/target-step stamps, which another policy may never select.

async _ray_get(obj_ref: Any) → Any#

Await a Ray ObjectRef without blocking the asyncio event loop.

async _call_dp(method_name: str, **kwargs) → Any#

Call a DataPlaneClient method or a Ray actor exposing that method.

async _rollout_pump() → None#

Continuously dispatch rollout tasks until cancellation.

Per batch: 0. await sampler.admit(…) to wait until the batch may dispatch and obtain its target_step stamp.

Per prompt:

  1. Acquire _buffer_capacity slot (backpressure)

  2. Acquire sem (cap concurrent in-flight rollouts)

  3. Wait for _rollout_permitted (paused during weight sync)

  4. Call rollout_manager.generate_and_push(prompt) — local async RolloutManager reserves a slot, runs the rollout, then commits the group via TQReplayBuffer (→ dp_client.put_samples + mark ready)

  5. Decrement _inflight_rollouts

async _train_pump() → None#

Per-prompt-group streaming train loop.

Per step:

  1. sampler.evict drops stale groups from the buffer and clears their TQ rows.

  2. sampler.select returns K prompt groups (or None) and drops them from the buffer; DP rows survive so the trainer can read them. Already trainable — buffer wrote training-shaped rows at rollout time.

  3. _advantage_stage(train_meta).

  4. trainer.train_microbatches_from_meta + finish_train_step.

  5. dp_client.clear_samples on consumed sample_ids; release _buffer_capacity per dropped group, then sync.

async _watchdog_pump() → None#

Report rollout health, and detect stalls nothing else catches.

Progress is the pair (committed groups, completed train steps) rather than a timestamp: both counters already exist, and “neither has moved” is the property that actually matters.

Deliberately not conditioned on rollouts being in flight. An earlier version required that, on the reasoning that an idle controller has legitimately no work – and a fault-injection run walked straight through the gap. Killing a generation worker wedged the loop with zero rollouts in flight and zero failures recorded: the rollout pump was blocked on backpressure behind a train pump that could no longer finish a step, so nothing was in flight to count. The watchdog observed six minutes of idleness and said nothing.

What separates a real stall from an idle gap is whether work remains, so that is what is checked instead.

async _check_env_health(timeout_s: float) → list[str]#

Ask each environment actor that exposes a health check whether it is whole.

Returns the problems found, empty when everything is well. It reports rather than raises so the caller can route the verdict through stall_action, the same way the stall path does. Raising here bypassed stall_action entirely: under the documented default ("warn", which promises to “only report”), and with gym_subprocess_check defaulting to true, an unhealthy environment killed the run – a run-ending path switched on by default, in a feature whose whole posture is inert-by-default.

Each probe is bounded. NemoGym is an asyncio actor, so a wedged environment – precisely the case this check exists to catch – left the await hanging forever, the pump never ticked again, and stall detection was dead exactly when it was needed. A probe that does not answer within one tick IS the unhealthy signal; it is not a reason to stop watching.

Environments without the method are skipped rather than treated as unhealthy; only NeMo-Gym has subprocess servers to lose.

async _abort_stale_inflight() → int#

Abort in-flight rollouts that the sampler can no longer select.

async _save_checkpoint(step_metrics: dict[str, Any]) → None#

Write a full checkpoint for the just-finished train step.

Everything except the (possibly async) policy weight write must be on disk before begin_finalization; rollouts keep running throughout.

async _sync_weights(
*,
calibration_data: Optional[nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any]] = None,
) → int#

Pause new rollout dispatches, synchronize weights, resume.

SC owns the pause gate; in-flight generations continue through the refit — vLLM V1 async engine supports weight updates during pending requests.

Flow:

  1. _rollout_permitted.clear() — no new dispatches

  2. Optionally calibrate FP8 KV-cache scales.

  3. weight_synchronizer.sync_weights(kv_scales=…)

  4. _rollout_permitted.set() — resume

Parameters:

calibration_data – Optional data used to calibrate FP8 KV-cache scales before synchronizing weights.

Returns:

The number of stale in-flight rollout groups aborted before the weight synchronization.

async _advantage_stage(
meta: nemo_rl.data_plane.KVBatchMeta,
) → nemo_rl.data_plane.KVBatchMeta#

Fetch advantage inputs, compute advantages, and write them back.

SC owns the prompt-group-scoped advantage stage because the selected KVBatchMeta still contains complete prompt groups before trainer DP sharding. Tensor payloads still move through DataPlane: SC fetches only the configured advantage input columns and writes the computed advantages column back under the same sample_ids.

_advantage_input_fields() → list[str]#