`nemo_rl.algorithms.single_controller_utils.setup`#

Driver-side factory for the SingleController (async-RL) training path.

setup builds the full SingleControllerActorArgs on the driver and the caller passes it to SingleControllerActor.remote. Everything lives on the driver because driver-side TQPolicy owns the worker group directly — running this inside another Ray actor nests runtime_envs and breaks Ray’s resource resolution (see the PR #2692 follow-up).

Module Contents#

Classes#

SingleControllerActorArgs

All inputs SingleControllerActor needs, built driver-side by setup_single_controller().

Functions#

`_build_clusters`	Allocate train + inference clusters; one shared cluster when colocated.
`_build_generation`	Spin up the generation backend (vLLM or SGLang).
`_finish_deferred_generation`	Finish loading and starting the deferred generation.
`_build_trainer`	Build the TQ-mediated trainer (driver-side TQPolicy).
`_spinup_gym`	Spin up the NeMo-Gym actor against the reserved vLLM URLs.
`_generation_max_seq_len`	Return the per-backend max sequence length.
`_clamp_max_num_steps`	Clamp grpo.max_num_steps to max_num_epochs * len(dataloader).
`_maybe_inject_megatron_train_iters`	Set train_iters from max_num_steps after its dataloader clamp.
`_maybe_attach_fleet_health`	Route generation through fleet health, when it is enabled and supported.
`_shard_base_urls`	Per-shard OpenAI base URLs, or None when the backend exposes no servers.
`_maybe_start_generation_router`	Start the NeMo-Gym-facing router, if enabled.
`_build_retry_policy`	Translate `async_rl.rollout_failure` into the rollout layer’s policy object.
`setup_single_controller`	Build the full SC actor args driver-side.

API#

class nemo_rl.algorithms.single_controller_utils.setup.SingleControllerActorArgs#

All inputs SingleControllerActor needs, built driver-side by setup_single_controller().

Passed as a single arg to SingleControllerActor.remote so the actor’s init does no construction work — every heavy object is cloudpickled in.

gen_handle: Any#: None

trainer_handle: Any#: None

env_handles: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface]#: None

train_cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster#: None

inference_cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster#: None

dp_client: nemo_rl.data_plane.DataPlaneClient#: None

dataloader: torchdata.stateful_dataloader.StatefulDataLoader#: None

weight_synchronizer: nemo_rl.weight_sync.WeightSynchronizer#: None

advantage_estimator: Any#: None

loss_fn: nemo_rl.algorithms.loss.interfaces.LossFunction#: None

rollout_manager: nemo_rl.experience.rollout_manager.RolloutManager#: None

tq_buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer#: None

partition_id: str#: None

save_state: nemo_rl.algorithms.grpo.GRPOSaveState#: None

last_checkpoint_path: Optional[str]#: None

fleet_monitor: Optional[nemo_rl.models.generation.fleet_health.GenerationFleetHealth]#: None

generation_router: Optional[ray.actor.ActorHandle[nemo_rl.models.generation.generation_router.GenerationRouterImpl]]#: None

nemo_rl.algorithms.single_controller_utils.setup._build_clusters( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, ) → tuple[nemo_rl.distributed.virtual_cluster.RayVirtualCluster, nemo_rl.distributed.virtual_cluster.RayVirtualCluster]#: Allocate train + inference clusters; one shared cluster when colocated.

nemo_rl.algorithms.single_controller_utils.setup._build_generation( inference_cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster, master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, *, defer_model_load: bool = False, ) → tuple[Any, float]#

Spin up the generation backend (vLLM or SGLang).

Parameters:

inference_cluster – Ray virtual cluster the generation workers run on.
master_config – SC MasterConfig.
defer_model_load – If True (for the NeMo-Gym flow), reserve OpenAI server URLs without loading weights; caller runs gen.load_and_start() later.

Returns:

A tuple of (generation object, wall time spent in this call). The generation object is a VllmGeneration or SGLangGeneration.

nemo_rl.algorithms.single_controller_utils.setup._finish_deferred_generation( generation: Any, ) → tuple[Any, float]#

Finish loading and starting the deferred generation.

Parameters:: generation – The deferred generation object.
Returns:: A tuple of (finished generation object, wall time spent in this call).

nemo_rl.algorithms.single_controller_utils.setup._build_trainer( train_cluster: nemo_rl.distributed.virtual_cluster.RayVirtualCluster, master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, tokenizer, processor, *, weights_path: Optional[pathlib.Path], optimizer_path: Optional[pathlib.Path], ) → tuple[Any, float]#

Build the TQ-mediated trainer (driver-side TQPolicy).

Parameters:

train_cluster – Ray virtual cluster the trainer workers run on.
master_config – SC MasterConfig.
tokenizer – Tokenizer used by the policy.
processor – Optional AutoProcessor for VLM paths.
weights_path – Checkpointed policy weights to resume from, or None.
optimizer_path – Checkpointed optimizer state to resume from, or None.

Returns:

A tuple of (TQPolicy trainer, wall time spent in this call).

nemo_rl.algorithms.single_controller_utils.setup._spinup_gym( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, base_urls: list[str], tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase, ) → tuple[Any, float]#

Spin up the NeMo-Gym actor against the reserved vLLM URLs.

Parameters:

master_config – SC MasterConfig.
base_urls – Reserved vLLM OpenAI server URLs.
tokenizer – Installed on the actor at spinup rather than passed per rollout call. See NemoGym.set_tokenizer.

Returns:

A tuple of (NeMo-Gym actor, wall time spent in this call).

nemo_rl.algorithms.single_controller_utils.setup._generation_max_seq_len(generation_config) → int#

Return the per-backend max sequence length.

vllm uses vllm_cfg.max_model_len; sglang uses sglang_cfg.context_length; megatron generation has no dedicated field and routes max_new_tokens through as max_sequence_length on the inference worker.

nemo_rl.algorithms.single_controller_utils.setup._clamp_max_num_steps( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, dataloader: torchdata.stateful_dataloader.StatefulDataLoader, ) → None#: Clamp grpo.max_num_steps to max_num_epochs * len(dataloader).

nemo_rl.algorithms.single_controller_utils.setup._maybe_inject_megatron_train_iters( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, ) → None#: Set train_iters from max_num_steps after its dataloader clamp.

nemo_rl.algorithms.single_controller_utils.setup._maybe_attach_fleet_health( generation: Any, master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, ) → Optional[nemo_rl.models.generation.fleet_health.GenerationFleetHealth]#

Route generation through fleet health, when it is enabled and supported.

Returns:: The monitor the SingleController should drive, or None when fleet health is disabled or the backend does not support it.

nemo_rl.algorithms.single_controller_utils.setup._shard_base_urls( generation: Any, ) → Optional[list[Optional[str]]]#: Per-shard OpenAI base URLs, or None when the backend exposes no servers.

nemo_rl.algorithms.single_controller_utils.setup._maybe_start_generation_router( generation: Any, master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, ) → Any#

Start the NeMo-Gym-facing router, if enabled.

Returns:: The router actor handle, or None when the router is disabled.

nemo_rl.algorithms.single_controller_utils.setup._build_retry_policy( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, ) → nemo_rl.experience.rollout_manager.RolloutRetryPolicy#: Translate async_rl.rollout_failure into the rollout layer’s policy object.

nemo_rl.algorithms.single_controller_utils.setup.setup_single_controller( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase, *, processor: Optional[transformers.AutoProcessor] = None, partition_id: str = 'rollout_data', ) → tuple[nemo_rl.algorithms.single_controller_utils.setup.SingleControllerActorArgs, nemo_rl.algorithms.metric_utils.SetupTimingMetrics]#

Build the full SC actor args driver-side.

Parameters:

master_config – SC MasterConfig.
tokenizer – Tokenizer used by the policy.
processor – Optional AutoProcessor for VLM paths.
partition_id – TQ partition the rollout writer + sampler share.

Returns:

A tuple of (pre-built SC actor args, driver-side per-phase timings logged by the SC actor).

nemo_rl.algorithms.single_controller_utils.setup#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.algorithms.single_controller_utils.setup`#