`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`
`log`

API#

nemo_rl.algorithms.single_controller.Generation#: None

nemo_rl.algorithms.single_controller.log#: ‘getLogger(…)’

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
_stall_watchdog_pump: publishes rollout counters and reports stalls or unhealthy environments, which are the failures that otherwise produce no signal at all

Plus _gen_fleet_probe_pump when fleet health is enabled, which probes generation shard liveness on its own, much shorter clock.

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 _maybe_restore_replacement_reserve() → None#

Restore spare prompts diverted before the previous run’s checkpoint.

These were pulled from the dataloader and never dispatched, so the restored dataloader resumes past them. Nothing else in the checkpoint holds them, and without this they are simply gone: one batch of the dataset per divert, plus the training step _clamp_max_num_steps had budgeted for it.

No sampler-name guard, unlike the buffer restore. Spares carry no stamp – they are prompts that never reached admit – so nothing about them depends on which sampler wrote the checkpoint. They are restored even into a run that has since switched to on_dropped_prompt="shrink", where the pool is never drawn on but is still drained back into training at the end of the dataloader.

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. Under on_dropped_prompt=”replace”, divert the batch into the spare pool if the pool is below its low-water mark, and skip admission entirely. Otherwise await sampler.admit(…) to wait until the batch may dispatch and obtain its target_step stamp.

Per prompt:

Acquire _buffer_capacity slot (backpressure)
Acquire sem (cap concurrent in-flight rollouts)
Wait for _rollout_permitted (paused during weight sync)
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)
If the prompt was dropped, substitute a spare and repeat step 4 – for this step, or for whichever step lends this one a finished group in its place (see _take_replacement, _promote_into_step) – or credit the step short so the train pump can close it
Decrement _inflight_rollouts

Once every epoch is done, whatever is left in the spare pool is dispatched as ordinary steps rather than discarded (see _drain_reserve_into_steps).

_divert_batch_to_reserve( prompt_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → bool#

Consume a whole batch as spare prompts instead of admitting it as a step.

Returns whether the batch was taken, in which case the caller must not admit it. Diverting before admit is what keeps the stamp sequence honest: admitting a batch and then dispatching nothing for it would leave a target step that no group is ever generated for, which is exactly the hang the shortfall accounting exists to prevent.

A whole batch at a time because the dataloader only yields batches. The spares that go unused are not wasted work – nothing has been generated for them – and they stay in the pool for later steps.

Nothing is diverted until the sampler has actually stamped a batch, so a run whose sampler never stamps does not lose a batch of prompts to a pool it can never draw on. The cost is that the first batch is always admitted rather than diverted; in practice the pool is filled while that first batch’s rollouts are still running, so it is available by the time any of them can be given up on.

async _drain_reserve_into_steps( launch: Callable[[nemo_rl.data.interfaces.DatumSpec, Optional[int]], Awaitable[None]], ) → None#

Train on the leftover spares once the dataloader has nothing more to give.

Spares were consumed from the dataset like any other prompt, so leaving them in the pool at the end of the last epoch throws away data the run already paid for.

It also restores the step count. _clamp_max_num_steps derives max_num_steps from len(dataloader), and every diverted batch is one fewer batch the loop can admit – so without this a replace-mode run quietly finishes one step short of the budget it was configured with, per divert.

Whole steps only. A partial pool dispatched as a step is short by construction, and min_step_batch_fraction would then reject it and fail a run that had otherwise completed cleanly. In the ordinary case the pool holds exactly one batch (the dataloader uses batch_size=num_prompts_per_step), so the common outcome is that the whole thing is recovered.

Not gated on on_dropped_prompt: an empty pool makes this a no-op anyway, and only “replace” ever fills one, so the gate would buy nothing while stranding a pool restored from a checkpoint into a run that has since switched to “shrink”.

_take_replacement( target_step: Optional[int], replacements_used: int, ) → Optional[nemo_rl.data.interfaces.DatumSpec]#

A spare prompt to stand in for a dropped group, or None to shrink instead.

None covers the four ways a replacement can be unavailable: it was not asked for, the sampler did not stamp this prompt so no step is waiting on it, the per-slot budget is spent, or the pool is empty because the dataloader is exhausted. Every one of them falls back to on_dropped_prompt="shrink" rather than waiting, because a step whose replacements keep failing still has to close.

_promote_into_step( target_step: Optional[int], ) → Optional[int]#

Fill a dropped step from a later step’s finished work, and name the lender.

Where a replacement goes, rather than whether one happens. The lost step closes on generation that already exists instead of waiting out a rollout with the trainer idle, and the caller redirects its spare prompt to the lender, which is due a training step later and has the slack to absorb the wait. The same prompt is generated either way.

Only ever reached with a spare already in hand, which is what makes the borrow safe to take: an unrepaid loan is the same hole one step later.

Returns None – leaving the caller filling the dropped step directly – when nothing is stamped so no step is stranded, when the trainer has already moved past this step (a second drop can land after the first one closed it short, and a group stamped for a finished step would only be evicted), or when no later step has a finished group to lend. The last is always the case at in_order.max_lookahead_versions=0, where the next batch is not dispatched until this step trains.

Returns:: The step that lent the group, which the caller now owes a rollout, or None.

_credit_shortfall(target_step: Optional[int]) → None#: Record that a stamped step will never receive a group it is waiting for.

_target_groups_for_step(step: int) → int#

How many prompt groups this step should train on, after dropped prompts.

num_prompts_per_step is the target; groups stamped for this step that were given up on are subtracted, because they are never arriving and a sampler that matches batches to steps exactly cannot substitute another step’s groups for them. Without this the pump waits on a group no one is generating.

The step trains on fewer samples than configured, which is the point: a smaller step beats a stalled run. The count is logged as dropped_prompt_groups so the batch size a step actually used is recoverable afterwards.

How much smaller is bounded by min_step_batch_fraction, and that bound has to live here because neither drop budget provides it. Both budgets are run-scoped – the consecutive counter is cleared by any commit, including commits for other steps – so drops landing on one step while other steps succeed can shrink it without ever tripping them.

Raises:: RuntimeError – The step fell below min_step_batch_fraction of num_prompts_per_step. Training a fraction of a batch is a silent change to the gradient estimate, so it is refused rather than absorbed.

async _train_pump() → None#

Per-prompt-group streaming train loop.

Per step:

sampler.evict drops stale groups from the buffer and clears their TQ rows.
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.
_advantage_stage(train_meta).
trainer.train_microbatches_from_meta + finish_train_step.
dp_client.clear_samples on consumed sample_ids; release _buffer_capacity per dropped group, then sync.

async _stall_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 _gen_fleet_probe_pump() → None#

Probe the generation fleet on its own clock.

Separate from the watchdog because the two cadences answer different questions. The watchdog publishes counters and notices a stalled run, which is a minutes-scale concern; liveness detection is the input to every recovery decision and has to be seconds-scale.

Sharing the watchdog’s loop made probe_interval_s decorative – probes ran at watchdog.interval_s and nothing read the configured value. With the shipped defaults that put detection at 30s * unhealthy_threshold, i.e. 60-90s, which is longer than the refit deadline: by the time a hung refit aborted, the monitor still had the dead shard as SUSPECT, so the rebuild that abort exists to trigger saw an empty absent set and did nothing. Arithmetic, not a race – it could never have worked. Job 5925668.

async _probe_generation_fleet() → None#

Ask every serving generation shard whether it is still alive.

Ray actor liveness is the cheap authoritative signal for “the process is gone”, and it is what the probe uses. It does not catch every failure – a vLLM engine core can die while the worker process and its HTTP thread survive – which is why the routing adapters also report the failures they observe. The two signals feed the same counters.

Only serving shards are probed: a quarantined shard answering again says nothing about whether its weights are current, and the monitor ignores such probes anyway.

Shards are probed concurrently. Sequentially, a tick costs up to probe_timeout_s per shard, so a fleet of four would take 8s to complete a round the config promises every 5s – and config validation only checks probe_timeout_s < probe_interval_s, which silently assumes one probe per tick. Concurrent, a round is bounded by probe_timeout_s at any fleet size.

async _push_router_membership() → None#

Tell the NeMo-Gym router which backends are currently serving.

Pushed as the full set rather than a delta, so a dropped or reordered update – or a restarted router, which comes up believing every backend serves – converges on the next tick without sequence numbers or replay.

Pushed unconditionally, not gated on the membership epoch moving. The gate looked free – an unchanged serving set costs nothing to skip – but it made the router’s own restart unrecoverable: a recreated actor rebuilds _serving as every backend, while the epoch it was last pushed at has not moved, so the gate blocked every corrective push and Gym routed to a quarantined shard for the rest of the run. The payload is a short list of strings on a probe-interval timer; the gate bought nothing and cost the guarantee both docstrings advertised.

It is also what makes the router’s reflex drop safe: dropping a failing backend locally is only correct because a later push puts it back.

async _drain_router_failures() → None#

Fold the router’s observed backend failures into the fleet ledger.

The router is the only component that sees a wedged engine: it answers is_alive from a healthy worker process, so no probe can condemn it. The router holds no monitor reference by design – membership flows one way – so it counts failures per backend URL and this drains them here, on the tick that already talks to it.

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:

_rollout_permitted.clear() — no new dispatches
Optionally calibrate FP8 KV-cache scales.
weight_synchronizer.sync_weights(kv_scales=…)
_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, ) → tuple[nemo_rl.data_plane.KVBatchMeta, bool]#

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.

Returns:: The updated batch metadata and whether the batch contains at least one valid training token.

_advantage_input_fields() → list[str]#

nemo_rl.algorithms.single_controller#

Module Contents#

Classes#

Data#

API#

`nemo_rl.algorithms.single_controller`#