chore(orchestrator): fold GenerationSource into TrainEnv

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

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",
 ]

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]):

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",

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)

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

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()