[faresobeid]

Summary

Fold GenerationSource into TrainEnv, its only owner. Move connect_frozen_client from algo/base.py to clients.py, so algo no longer imports clients lazily. Give TrainEnv setup()/aclose() so orchestrator setup and teardown no longer reach into its parts.

What changed

• orchestrator/generation_source.py is deleted. TrainEnv (envs.py) now holds clients, renderer_config, connected (the frozen source clients it connected) and uses_live_policy. Dropping logprobs for frozen sources happens in TrainEnv.__init__.
• TrainEnv.setup() connects the frozen generation source and runs algorithm.setup() concurrently. TrainEnv.aclose() closes connected and algorithm.connected, never the shared policy clients.
• connect_frozen_client is a plain function in clients.py, next to check_inference_ready. algo/base.py imports it at module level. It is no longer re-exported from prime_rl.orchestrator.algo.
• dispatcher.py reads the TrainEnv directly at the four former generation_source call sites. orchestrator.py calls env.setup() and env.aclose().
• docs/algorithms.md and the algo/__init__.py docstrings refer to TrainEnv.

Behavior change

None.

Checks run

• pytest tests/unit/orchestrator tests/unit/test_configs* -m "not gpu": 297 passed.
• ruff check, ruff format --check: clean.
• Equivalence script, old GenerationSource (origin/main) vs new TrainEnv, on 4 configs: policy, policy with top_p=0.9, sft with a frozen sampling.source, opd with a frozen teacher. Identical sampling_args, requires_sampling_masks and uses_live_policy. After setup(), the same generation clients (policy object, or the frozen endpoint with the same train client config) and the same algorithm clients. Teardown closes the same clients.
• GPU smoke on one node (vLLM 0.31): reverse-text RL (15 steps), sft with a frozen sampling.source (5 steps), opd with a frozen teacher (5 steps). All three ran to "Training finished" and exited 0.

---

Note

Low Risk
Internal orchestrator refactor with stated behavioral equivalence and no changes to training algorithms or dispatch semantics beyond indirection.

Overview
Refactors train rollout sourcing by deleting GenerationSource and moving its responsibilities onto TrainEnv: inference clients, live-vs-frozen policy detection (uses_live_policy), sampling-arg tweaks (e.g. dropping logprobs for frozen endpoints), and frozen-pool connection in setup().

TrainEnv.setup() now connects frozen sampling.source endpoints and runs algorithm.setup() together; TrainEnv.aclose() closes only env-owned clients (connected + algorithm pools), not the shared policy pool. The orchestrator and dispatcher call these surfaces instead of reaching into a nested generation_source.

connect_frozen_client moves from algo/base.py to clients.py (next to check_inference_ready); it is no longer re-exported from prime_rl.orchestrator.algo. Docs and package docstrings now describe TrainEnv as owning the sampling half of the algorithm config.

^{Reviewed by Cursor Bugbot for commit e89ca1a. Bugbot is set up for automated code reviews on this repo. Configure here.}