feat: per-source env-server address via serve.address

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
 

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,21 +122,6 @@ def to_sampling_args(self) -> dict[str, Any]:
         return args
 
 
-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``)."""
-
-    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``."""
-
-    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
     runs, ``serve`` — how it's hosted, ``legacy`` — a classic v0 env instead) plus this
@@ -146,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 sized. Consumed by the launcher (which writes each source's env-server config), not by the orchestrator — the orchestrator 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:<env_server_base_port + index>``. 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``."""
@@ -535,7 +519,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:<base + i>``. 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:<base + i>``. 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 +730,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:<port>`` 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)
         }

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 @ <path>` 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 @ <path>` 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(