From 1e3ecc1a8364cfa93f9da25e911c276d8dcd07b5 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sat, 8 Aug 2026 00:13:38 +0000 Subject: [PATCH 1/4] feat(configs): per-source env-server address via serve.address Setting serve.address on a train/eval source marks its env server externally managed: the launcher neither writes its TOML nor spawns a server for it, and the orchestrator connects to the given address. Unset sources keep the derived loopback address, with indices still positional across all sources. Co-Authored-By: Claude Fable 5 --- docs/training.md | 2 +- .../src/prime_rl/configs/orchestrator.py | 23 ++++++---- src/prime_rl/entrypoints/rl.py | 42 +++++++++++-------- 3 files changed, 41 insertions(+), 26 deletions(-) diff --git a/docs/training.md b/docs/training.md index e5da0f2219..9e66a723ac 100644 --- a/docs/training.md +++ b/docs/training.md @@ -37,7 +37,7 @@ This page covers everything you need to launch, observe, checkpoint, and recover | `uv run inference` | vLLM server. | Always use this entrypoint over `vllm serve` — it adds `/update_weights`, `/load_lora_adapter`, and `/init_broadcaster`. | | `uv run trainer` | Standalone trainer process group. | Use only when launching the trainer separately from the orchestrator (e.g. multi-node RL without the `rl` wrapper). | | `uv run orchestrator` | Standalone orchestrator process. | Pair with a separately-launched trainer, inference, and one `env-server` per source. | -| `uv run env-server` | Standalone env server for one environment. | The `rl` launcher starts these automatically (one per train/eval source, at the source's derived `serve.address`); only needed when running the orchestrator standalone. | +| `uv run env-server` | Standalone env server for one environment. | The `rl` launcher starts these automatically (one per train/eval source, at a derived loopback address); only needed when running the orchestrator standalone, or for sources with an explicit `serve.address` — those are externally managed (e.g. their own k8s pod) and the launcher expects the server to already run there. | ## RL Trainer diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 8f5f0e859f..eb604bb672 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -124,15 +124,22 @@ def to_sampling_args(self) -> dict[str, Any]: class ServeConfig(BaseConfig): - """The subset of verifiers' ``ServeConfig`` a source configures — the worker pool - and the per-worker bound. The launcher materializes it into the env server's full - ``[serve]`` block, filling in the source's derived address - (``OrchestratorConfig.env_addresses``).""" + """Verifiers' ``ServeConfig`` as a source configures it — the worker pool, where the + server lives, and the per-worker bound. The launcher materializes it into the env + server's full ``[serve]`` block, filling in the source's derived address + (``OrchestratorConfig.env_addresses``) when ``address`` is unset.""" pool: PoolConfig = Field(default_factory=vf.ElasticPoolConfig) """Worker-pool sizing. ``elastic`` (default) starts at one worker and scales up on demand; ``static`` pre-spawns a fixed ``num_workers``.""" + address: str | None = None + """Where this source's env server is reachable. Unset (default) means the launcher + serves this source at the derived ``tcp://127.0.0.1:`` + address. Setting it marks the server externally managed: the launchers neither write + its env-server TOML nor spawn a server for it, and the orchestrator connects to the + given address — e.g. a k8s deployment running env servers in their own pods.""" + max_concurrent: int | None = Field(None, ge=1) """Episodes in flight per worker (None = unbounded; the dispatcher's ``max_inflight_episodes`` is the run's bound).""" @@ -147,7 +154,7 @@ class EnvConfig(BaseConfig): """The verifiers environment — which env, its seed taskset, each agent, its knobs. Narrowed to the selected env's config class by the env id, else the taskset id.""" serve: ServeConfig = ServeConfig() - """How this source's env server is sized. Consumed by the launcher (which writes each source's env-server config), not by the orchestrator — the orchestrator only connects.""" + """How this source's env server is hosted. The sizing knobs are consumed by the launcher (which writes each source's env-server config); ``address`` is read by the orchestrator, which only connects.""" legacy: vf.LegacyEnvConfig = vf.LegacyEnvConfig() """A classic (v0) environment to run through the bridge instead of ``env``.""" @@ -535,7 +542,7 @@ class OrchestratorConfig(BaseConfig): """Rate limit per environment worker, in tasks per minute. Recommended for sandbox-backed environments to prevent sandbox-not-ready errors during autoscaling. With multiple workers, the effective total rate is ``workers × this value``. None disables rate limiting.""" env_server_base_port: int = Field(5000, ge=1, le=65535) - """First port of the env-server port range: the source at position ``i`` (train, then eval) is served at ``tcp://127.0.0.1:``. Give concurrent runs on one host distinct bases (e.g. one per multi-run orchestrator).""" + """First port of the env-server port range: the source at position ``i`` (train, then eval) is served at ``tcp://127.0.0.1:``. Sources with an explicit ``serve.address`` keep it instead, without shifting the other sources' ports (indices stay positional). Give concurrent runs on one host distinct bases (e.g. one per multi-run orchestrator).""" batch_size: int | None = Field(None, ge=1) """Samples to train on per step (rollout-based batching). Set this OR ``token_batch_size``.""" @@ -746,11 +753,13 @@ def env_sources(self) -> list[tuple[str, EnvConfig]]: @property def env_addresses(self) -> dict[tuple[str, str], str]: """Where each source's env server lives, keyed by ``(split, resolved_name)``: + the source's own ``serve.address`` when set (an externally managed server), else ``tcp://127.0.0.1:`` with ports from ``env_server_base_port`` in ``env_sources`` order. The launcher binds env servers at exactly these addresses and the orchestrator connects to them, so both sides agree from the config alone.""" return { - (split, source.resolved_name): f"tcp://127.0.0.1:{self.env_server_base_port + index}" + (split, source.resolved_name): source.serve.address + or f"tcp://127.0.0.1:{self.env_server_base_port + index}" for index, (split, source) in enumerate(self.env_sources) } diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 279f91e00a..8104629c22 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -49,14 +49,24 @@ def env_servers(config: RLConfig) -> list[tuple[str, EnvConfig, str]]: - """``(split, source, address)`` for every train/eval source. The launcher runs one - env server per source at its deterministic address; the orchestrator connects there.""" + """``(split, source, address)`` for every launcher-managed train/eval source. The + launcher runs one env server per source at its deterministic address; the + orchestrator connects there. A source with ``serve.address`` set is externally + managed — its server runs elsewhere and only the orchestrator connects to it — so + the launcher neither writes its TOML nor spawns a server for it.""" addresses = config.orchestrator.env_addresses return [ - (split, source, addresses[(split, source.resolved_name)]) for split, source in config.orchestrator.env_sources + (split, source, addresses[(split, source.resolved_name)]) + for split, source in config.orchestrator.env_sources + if source.serve.address is None ] +def env_server_names(config: RLConfig, split: str) -> list[str]: + """Names of the launcher-managed env servers for one split.""" + return [source.resolved_name for source_split, source, _ in env_servers(config) if source_split == split] + + def get_physical_gpu_ids() -> list[int]: """Return physical GPU IDs visible to the launcher.""" raw_visible = os.environ.get("CUDA_VISIBLE_DEVICES") @@ -93,9 +103,10 @@ def write_subconfigs(config: RLConfig, output_dir: Path) -> None: with open(output_dir / INFERENCE_TOML, "wb") as f: tomli_w.dump(inference_dict, f) - # One EnvServerConfig TOML per source: `env-server @ ` binds at the source's - # deterministic address, where the orchestrator connects. The source's env/serve/legacy - # blocks carry over; its other knobs (sampling, algo, name, ...) are orchestrator-side. + # One EnvServerConfig TOML per launcher-managed source: `env-server @ ` binds + # at the source's deterministic address, where the orchestrator connects. The source's + # env/serve/legacy blocks carry over; its other knobs (sampling, algo, name, ...) are + # orchestrator-side. for split, source, address in env_servers(config): env_dir = output_dir / ENVS_DIR / split env_dir.mkdir(parents=True, exist_ok=True) @@ -453,10 +464,9 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> else {} ) - # Env servers launch next to the orchestrator, one per train/eval source. - sources = config.orchestrator.env_sources - train_env_names = [source.resolved_name for split, source in sources if split == "train"] - eval_env_names = [source.resolved_name for split, source in sources if split == "eval"] + # Env servers launch next to the orchestrator, one per launcher-managed train/eval source. + train_env_names = env_server_names(config, "train") + eval_env_names = env_server_names(config, "eval") if config.deployment.type == "single_node": script = template.render( @@ -558,10 +568,8 @@ def rl_slurm(config: RLConfig): write_config(config, config_dir, exclude={"slurm", "dry_run", "clean_output_dir"}) logger.info(f"Wrote config to {config_dir / RL_TOML}") - train_env_names = [env.resolved_name for env in config.orchestrator.train.source] - eval_env_names = ( - [source.resolved_name for source in config.orchestrator.eval.source] if config.orchestrator.eval else [] - ) + train_env_names = env_server_names(config, "train") + eval_env_names = env_server_names(config, "eval") log_message = format_log_message( log_dir=log_dir, @@ -575,10 +583,8 @@ def rl_slurm(config: RLConfig): write_subconfigs(config, config_dir) logger.info(f"Wrote subconfigs to {config_dir}") - train_env_names = [env.resolved_name for env in config.orchestrator.train.source] - eval_env_names = ( - [source.resolved_name for source in config.orchestrator.eval.source] if config.orchestrator.eval else [] - ) + train_env_names = env_server_names(config, "train") + eval_env_names = env_server_names(config, "eval") has_infer = config.deployment.infer_nodes_per_replica > 0 log_message = format_log_message( From bfc2e3101628572778d104cc6f45173ce0112c12 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sat, 8 Aug 2026 17:48:19 +0000 Subject: [PATCH 2/4] refactor: subclass vf.ServeConfig instead of mirroring its fields Co-Authored-By: Claude Fable 5 --- .../src/prime_rl/configs/orchestrator.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index eb604bb672..17c74a7ac9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -4,7 +4,6 @@ import verifiers.v1 as vf from pydantic import Field, SerializeAsAny, model_validator from renderers import AutoRendererConfig, RendererConfig -from verifiers.v1.configs.serve import PoolConfig from prime_rl.configs.algorithm import ( AlgoConfig, @@ -123,15 +122,10 @@ def to_sampling_args(self) -> dict[str, Any]: return args -class ServeConfig(BaseConfig): - """Verifiers' ``ServeConfig`` as a source configures it — the worker pool, where the - server lives, and the per-worker bound. The launcher materializes it into the env - server's full ``[serve]`` block, filling in the source's derived address - (``OrchestratorConfig.env_addresses``) when ``address`` is unset.""" - - pool: PoolConfig = Field(default_factory=vf.ElasticPoolConfig) - """Worker-pool sizing. ``elastic`` (default) starts at one worker and scales up on - demand; ``static`` pre-spawns a fixed ``num_workers``.""" +class ServeConfig(vf.ServeConfig): + """Verifiers' ``ServeConfig``, with the address optional: a source that leaves it + unset is served by the launcher, which materializes the env server's full ``[serve]`` + block with the derived address (``OrchestratorConfig.env_addresses``) filled in.""" address: str | None = None """Where this source's env server is reachable. Unset (default) means the launcher @@ -140,10 +134,6 @@ class ServeConfig(BaseConfig): its env-server TOML nor spawn a server for it, and the orchestrator connects to the given address — e.g. a k8s deployment running env servers in their own pods.""" - max_concurrent: int | None = Field(None, ge=1) - """Episodes in flight per worker (None = unbounded; the dispatcher's - ``max_inflight_episodes`` is the run's bound).""" - class EnvConfig(BaseConfig): """One environment a run pulls from: the verifiers blocks it composes (``env`` — what From a069da2ac681bddf22a4d7d4eac47eee822f4d07 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sat, 8 Aug 2026 17:55:08 +0000 Subject: [PATCH 3/4] refactor: consume vf.ServeConfig as-is Companion verifiers bump makes ServeConfig.address optional with no default, so the per-source serve block needs no prime-rl subclass. Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- .../src/prime_rl/configs/orchestrator.py | 17 ++--------------- 2 files changed, 3 insertions(+), 16 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 0a4d872f02..4aba2da7c2 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 0a4d872f021022310a08ec213a25f4efb4a0244a +Subproject commit 4aba2da7c27907ba46ea2181ceec9ea2a2d2694f diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 17c74a7ac9..478fd2f089 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -122,19 +122,6 @@ def to_sampling_args(self) -> dict[str, Any]: return args -class ServeConfig(vf.ServeConfig): - """Verifiers' ``ServeConfig``, with the address optional: a source that leaves it - unset is served by the launcher, which materializes the env server's full ``[serve]`` - block with the derived address (``OrchestratorConfig.env_addresses``) filled in.""" - - address: str | None = None - """Where this source's env server is reachable. Unset (default) means the launcher - serves this source at the derived ``tcp://127.0.0.1:`` - address. Setting it marks the server externally managed: the launchers neither write - its env-server TOML nor spawn a server for it, and the orchestrator connects to the - given address — e.g. a k8s deployment running env servers in their own pods.""" - - class EnvConfig(BaseConfig): """One environment a run pulls from: the verifiers blocks it composes (``env`` — what runs, ``serve`` — how it's hosted, ``legacy`` — a classic v0 env instead) plus this @@ -143,8 +130,8 @@ class EnvConfig(BaseConfig): env: SerializeAsAny[vf.EnvConfig] = vf.SingleAgentEnvConfig() """The verifiers environment — which env, its seed taskset, each agent, its knobs. Narrowed to the selected env's config class by the env id, else the taskset id.""" - serve: ServeConfig = ServeConfig() - """How this source's env server is hosted. The sizing knobs are consumed by the launcher (which writes each source's env-server config); ``address`` is read by the orchestrator, which only connects.""" + serve: vf.ServeConfig = vf.ServeConfig() + """How this source's env server is hosted. The sizing knobs are consumed by the launcher, which writes each source's env-server config with an unset ``address`` filled in as the derived ``tcp://127.0.0.1:``. Setting ``address`` marks the server externally managed: the launchers neither write its env-server TOML nor spawn a server for it, and the orchestrator connects to the given address — e.g. a k8s deployment running env servers in their own pods.""" legacy: vf.LegacyEnvConfig = vf.LegacyEnvConfig() """A classic (v0) environment to run through the bridge instead of ``env``.""" From 72a28e363b8da6f1e5cfaa155fa01a2c56d51593 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sat, 8 Aug 2026 18:14:32 +0000 Subject: [PATCH 4/4] chore: bump verifiers to merged optional-serve-address rev Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 4aba2da7c2..29e3a0f7a3 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 4aba2da7c27907ba46ea2181ceec9ea2a2d2694f +Subproject commit 29e3a0f7a3027d251a8cf8230db6d347357b2f85