diff --git a/docs/algorithms.md b/docs/algorithms.md index 7263f5ed866..594c1c4b395 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -139,7 +139,7 @@ algo.type = "echo" ### The Algorithm Classes -At runtime, each env's resolved config builds two objects: a `GenerationSource` (`prime_rl.orchestrator.generation_source`) that resolves the `sampling.source` model into the inference pool used for train episodes, and one of the named algorithm classes in `prime_rl.orchestrator.train.algo` (one module per algorithm: `algo/grpo.py`, `algo/opd.py`, …) from the algorithm config. Algorithm dispatch is keyed on `algo.type` — it names the algorithm, and each config class's defaults are its vetted parameterization: +At runtime, each env's `TrainEnv` (`prime_rl.orchestrator.envs`) resolves the `sampling.source` model into the inference pool used for train episodes, and its resolved config builds one of the named algorithm classes in `prime_rl.orchestrator.train.algo` (one module per algorithm: `algo/grpo.py`, `algo/opd.py`, …) from the algorithm config. Algorithm dispatch is keyed on `algo.type` — it names the algorithm, and each config class's defaults are its vetted parameterization: | `algo.type` | Class | hook(s) — stage | |---|---|---| diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index f1eb4050c81..a2f1dd6ff7e 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/src/prime_rl/orchestrator/algo/__init__.py @@ -2,8 +2,8 @@ The config side (``prime_rl.configs.algorithm``) defines *what* an algorithm is — a bundle of sampling and the per-token training signal. This package -turns the signal half into runtime objects (the sampling half is the env's -:class:`~prime_rl.orchestrator.generation_source.GenerationSource`): +turns the signal half into runtime objects (the sampling half is interpreted by +:class:`~prime_rl.orchestrator.envs.TrainEnv`): - one module per algorithm (``grpo``, ``echo``, ``max_rl``, ``rae``, ``hierarchical_grpo``, ``opd``, ``opsd``, ``sft``, ``debug``) — each named @@ -22,7 +22,7 @@ class owns its scoring hooks from typing import TYPE_CHECKING -from prime_rl.orchestrator.algo.base import Algorithm, connect_frozen_client +from prime_rl.orchestrator.algo.base import Algorithm from prime_rl.orchestrator.algo.debug import DebugAlgorithm from prime_rl.orchestrator.algo.echo import EchoAlgorithm from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm @@ -56,7 +56,7 @@ class owns its scoring hooks def build_algorithm(config: AlgoConfig, clients: InferenceClient) -> Algorithm: cls = ALGORITHM_CLASSES[config.type] # The Algorithm is the runtime of the algorithm config's training signal - # (its sibling GenerationSource interprets the sampling half). Every algorithm is + # (its TrainEnv interprets the sampling half). Every algorithm is # handed the live policy pool — opsd self-distills against it, others may # judge against it or ignore it. Other models (a frozen teacher, a hint # renderer) are built from the algorithm's own config in setup(). @@ -75,7 +75,6 @@ def build_algorithm(config: AlgoConfig, clients: InferenceClient) -> Algorithm: "RAEAlgorithm", "SFTDistillAlgorithm", "build_algorithm", - "connect_frozen_client", "assign_advantages", "stamp_loss_routing", ] diff --git a/src/prime_rl/orchestrator/algo/base.py b/src/prime_rl/orchestrator/algo/base.py index e2143e289d9..88ea92ab2f4 100644 --- a/src/prime_rl/orchestrator/algo/base.py +++ b/src/prime_rl/orchestrator/algo/base.py @@ -2,34 +2,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import verifiers.v1 as vf from prime_rl.configs.algorithm import AlgoConfig, FrozenModelConfig -from prime_rl.utils.logger import get_logger - -if TYPE_CHECKING: - from renderers import RendererConfig - - from prime_rl.orchestrator.clients import InferenceClient - - -async def connect_frozen_client( - config: FrozenModelConfig, *, renderer_config: RendererConfig | None = None -) -> InferenceClient: - """Connect to an externally hosted frozen model and wait for it.""" - from prime_rl.orchestrator.clients import InferenceClient, check_inference_ready - - get_logger().info(f"Initializing frozen model pool (model={config.name}, base_url={config.base_url})") - if renderer_config is not None: - clients = InferenceClient( - config, model_name=config.name, train_client_type="renderer", renderer_config=renderer_config - ) - else: - clients = InferenceClient(config, model_name=config.name) - await check_inference_ready(config, config.name) - return clients +from prime_rl.orchestrator.clients import InferenceClient, connect_frozen_client def iter_trainable_traces(episodes: list[vf.Episode]): diff --git a/src/prime_rl/orchestrator/clients.py b/src/prime_rl/orchestrator/clients.py index 403e8c29b56..d9e1cca05f2 100644 --- a/src/prime_rl/orchestrator/clients.py +++ b/src/prime_rl/orchestrator/clients.py @@ -20,6 +20,7 @@ resolve_headers, ) +from prime_rl.configs.algorithm import FrozenModelConfig from prime_rl.configs.shared import ClientConfig from prime_rl.utils.logger import get_logger @@ -228,6 +229,21 @@ async def check_inference_ready(client_config: ClientConfig, model_name: str) -> await admin.aclose() +async def connect_frozen_client( + config: FrozenModelConfig, *, renderer_config: RendererConfig | None = None +) -> InferenceClient: + """Connect to an externally hosted frozen model and wait for it.""" + get_logger().info(f"Initializing frozen model pool (model={config.name}, base_url={config.base_url})") + if renderer_config is not None: + clients = InferenceClient( + config, model_name=config.name, train_client_type="renderer", renderer_config=renderer_config + ) + else: + clients = InferenceClient(config, model_name=config.name) + await check_inference_ready(config, config.name) + return clients + + def setup_client( client_config: ClientConfig, client_type: str = "openai_chat_completions", diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index b9410915cea..e53fec4cd21 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -214,10 +214,10 @@ def _train_generation_for(self, env_name: str) -> tuple[InferenceClient, str, bo """``(clients, model_name, is_live)`` for *train* rollouts of this env — eval always uses the policy.""" assert self.train_envs is not None # train groups only exist when train is configured - source = self.train_envs.get(env_name).generation_source - if source.uses_live_policy: - return source.clients, self.policy.model_name, True - return source.clients, source.clients.model_name, False + env = self.train_envs.get(env_name) + if env.uses_live_policy: + return env.clients, self.policy.model_name, True + return env.clients, env.clients.model_name, False @property def inflight_train_count(self) -> int: @@ -316,7 +316,7 @@ def inflight_staleness(self) -> list[int]: return [ (self.progress.step - 1) - meta.policy_version for meta in self.inflight.values() - if meta.kind == "train" and self.train_envs.get(meta.env_name).generation_source.uses_live_policy + if meta.kind == "train" and self.train_envs.get(meta.env_name).uses_live_policy ] # ── lifecycle ────────────────────────────────────────────────────────── @@ -423,7 +423,7 @@ async def on_version_pending(self, step: int) -> None: gid for gid, group in self.groups.items() if group.kind == "train" - and self.train_envs.get(group.env_name).generation_source.uses_live_policy + and self.train_envs.get(group.env_name).uses_live_policy and group.policy_version_at_start < min_version ] cancelled = 0 @@ -722,7 +722,7 @@ async def emit_episode( live_policy = meta.kind == "eval" if meta.kind == "train": assert self.train_envs is not None - live_policy = self.train_envs.get(meta.env_name).generation_source.uses_live_policy + live_policy = self.train_envs.get(meta.env_name).uses_live_policy policy = vf.PolicySpan(start=policy_version, end=self.policy.version) if live_policy else None work: vf.WorkInfo = ( vf.EvalWorkInfo(step=meta.step, policy=policy) diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 0374aaf9548..b0053a0411a 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -26,11 +26,13 @@ from typing import Generic, TypeVar import verifiers.v1 as vf +from renderers import RendererConfig from verifiers.v1.serve import EnvClient +from prime_rl.configs.algorithm import FrozenModelConfig from prime_rl.configs.orchestrator import EnvConfig, EvalSourceConfig, TrainSourceConfig from prime_rl.orchestrator.algo import Algorithm, build_algorithm -from prime_rl.orchestrator.generation_source import GenerationSource +from prime_rl.orchestrator.clients import InferenceClient, connect_frozen_client from prime_rl.utils.logger import format_time, get_logger from prime_rl.utils.pathing import env_address_file @@ -150,17 +152,38 @@ def __init__( config: TrainSourceConfig, address: str | None, address_file: Path, - generation_source: GenerationSource, + clients: InferenceClient, + renderer_config: RendererConfig | None, algorithm: Algorithm, ): super().__init__(config, address, address_file) - self.generation_source = generation_source + # Train rollouts are generated from `clients`: the policy, or the frozen + # `sampling.source` connected in setup() with the renderer (token-in/out) client. + self.clients = clients + self.renderer_config = renderer_config + self.connected: InferenceClient | None = None self.algorithm = algorithm - self.sampling_args = generation_source.sampling_args(config.sampling.to_sampling_args()) + self.uses_live_policy = config.algo.sampling.source == "policy" + self.sampling_args = config.sampling.to_sampling_args() + if not self.uses_live_policy: + # Logprobs only feed importance ratios on policy-sampled tokens; frozen endpoints may reject the knob. + self.sampling_args.pop("logprobs", None) # Truncated policy sampling must ship the sampling masks the trainer replays. - self.requires_sampling_masks = ( - config.sampling.truncates_distribution() and config.algo.sampling.source == "policy" - ) + self.requires_sampling_masks = config.sampling.truncates_distribution() and self.uses_live_policy + + async def setup(self) -> None: + async def connect_source() -> None: + source = self.config.algo.sampling.source + if isinstance(source, FrozenModelConfig): + self.connected = await connect_frozen_client(source, renderer_config=self.renderer_config) + self.clients = self.connected + + await asyncio.gather(connect_source(), self.algorithm.setup()) + + async def aclose(self) -> None: + for clients in (self.connected, self.algorithm.connected): + if clients is not None: + await clients.aclose() class EvalEnv(Env): @@ -218,9 +241,8 @@ async def start(self) -> None: class TrainEnvs(Envs[TrainEnv]): - """Collection of training environments, each paired with its - :class:`GenerationSource` and runtime :class:`Algorithm`, built from the env's - resolved algorithm config.""" + """Collection of training environments, each with its runtime + :class:`Algorithm`, built from the env's resolved algorithm config.""" def __init__( self, @@ -238,7 +260,8 @@ def __init__( config, addresses[("train", config.resolved_name)], env_address_file(config_dir, "train", config.resolved_name), - GenerationSource(config.algo.sampling, clients, renderer_config), + clients, + renderer_config, build_algorithm(config.algo, clients), ) self._envs[env.name] = env diff --git a/src/prime_rl/orchestrator/generation_source.py b/src/prime_rl/orchestrator/generation_source.py deleted file mode 100644 index 344acbfd6c1..00000000000 --- a/src/prime_rl/orchestrator/generation_source.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Resolve the model source used to generate an env's train episodes.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -from prime_rl.configs.algorithm import FrozenModelConfig, SamplingConfig -from prime_rl.orchestrator.algo import connect_frozen_client - -if TYPE_CHECKING: - from renderers import RendererConfig - - from prime_rl.orchestrator.clients import InferenceClient - - -class GenerationSource: - """One env's generation model and inference clients. - - ``clients`` is what train rollouts are generated from: the policy clients, - swapped for a connected frozen endpoint in :meth:`setup` when the source is - an inline frozen model. A frozen source *generates* rollouts, so it uses - the renderer (token-in/out) client (built from ``renderer_config``) — the - rollout must carry tokens for training.""" - - def __init__(self, config: SamplingConfig, clients: InferenceClient, renderer_config: RendererConfig | None = None): - assert config.source is not None, "sampling.source must be resolved by config validation" - self.config = config - self.clients: InferenceClient = clients - self.renderer_config = renderer_config - self.connected: InferenceClient | None = None # frozen clients connected in setup(); closed at shutdown - - async def setup(self) -> None: - """Connect clients to a frozen generation source and wait for - readiness. Must run before dispatching.""" - if isinstance(self.config.source, FrozenModelConfig): - self.clients = await connect_frozen_client(self.config.source, renderer_config=self.renderer_config) - self.connected = self.clients - - @property - def uses_live_policy(self) -> bool: - return self.config.source == "policy" - - def sampling_args(self, args: dict) -> dict: - """Source-specific sampling-arg overrides. Sampling logprobs are only - needed for importance ratios on policy-sampled tokens — frozen - endpoints may reject the knob.""" - if not self.uses_live_policy: - args.pop("logprobs", None) - return args diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 89e8df70439..9b0bf8da8f1 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -282,10 +282,7 @@ async def setup(self) -> None: get_logger().success(f"Policy inference pool ready after {format_time(time.perf_counter() - t0)}") # Build + ready pools for each env's frozen generation source and the # algorithm's frozen reference model - await asyncio.gather( - *(env.generation_source.setup() for env in self.train_envs), - *(env.algorithm.setup() for env in self.train_envs), - ) + await asyncio.gather(*(env.setup() for env in self.train_envs)) get_logger().info(f"Initializing weight broadcast ({config.weight_broadcast})") t0 = time.perf_counter() @@ -1024,11 +1021,9 @@ async def teardown() -> None: if self.admin_plane is not None: await self.admin_plane.aclose() if self.train_envs is not None: - get_logger().debug("Stopping generation source and algorithm clients") + get_logger().debug("Stopping train env clients") for env in self.train_envs: - for clients in (env.generation_source.connected, env.algorithm.connected): - if clients is not None: - await clients.aclose() + await env.aclose() get_logger().info("Stopping orchestrator components") t0 = time.perf_counter()