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..3671383beb 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py
@@ -537,6 +537,9 @@ class OrchestratorConfig(BaseConfig):
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)."""
+ env_server_addresses: dict[str, str] = Field(default_factory=dict)
+ """Launcher-owned overrides of where each source's env server lives, keyed ``/`` (e.g. ``train/wordle``). A listed source is externally managed: the launchers neither write its env-server TOML nor spawn a server for it, and the orchestrator connects to the given address instead of the derived loopback one. Unlisted sources keep the derived ``tcp://127.0.0.1:`` address (indices stay positional across ALL sources, so overriding one source never shifts another's port). Sources themselves stay deployment-agnostic — this block is the launcher recording where it chose to run each server, not user intent: launchers whose orchestrator and env servers cannot share a host (e.g. the k8s chart running env servers in separate pods) inject their addresses here."""
+
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 +749,32 @@ 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)``:
+ an ``env_server_addresses`` override when the launcher set one, 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): self.env_server_addresses.get(f"{split}/{source.resolved_name}")
+ or f"tcp://127.0.0.1:{self.env_server_base_port + index}"
for index, (split, source) in enumerate(self.env_sources)
}
+
+ @model_validator(mode="after")
+ def validate_env_server_addresses(self):
+ """Reject override keys that match no source — a typo would otherwise silently
+ fall back to the derived loopback address and the run would hang polling a
+ server nobody runs. Skipped when no sources are present: external render
+ paths (e.g. rl-k8s's validator) validate with the source sections stripped."""
+ if not self.env_server_addresses:
+ return self
+ known = {f"{split}/{source.resolved_name}" for split, source in self.env_sources}
+ if not known:
+ return self
+ unknown = sorted(set(self.env_server_addresses) - known)
+ if unknown:
+ raise ValueError(
+ f"env_server_addresses keys {unknown} match no train/eval source "
+ f"(known: {sorted(known)}); overrides are keyed '/'"
+ )
+ return self
diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py
index 279f91e00a..d880b4f47f 100644
--- a/src/prime_rl/entrypoints/rl.py
+++ b/src/prime_rl/entrypoints/rl.py
@@ -49,11 +49,17 @@
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. Sources with an ``env_server_addresses`` override are
+ externally managed (another launcher runs their server at the given address) and
+ are skipped — no TOML is written and no server is spawned for them."""
addresses = config.orchestrator.env_addresses
+ overridden = set(config.orchestrator.env_server_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 f"{split}/{source.resolved_name}" not in overridden
]
@@ -453,10 +459,12 @@ 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 (externally-managed sources — env_server_addresses
+ # overrides — are skipped, same as env_servers()).
+ launcher_sources = [(split, source) for split, source, _ in env_servers(config)]
+ train_env_names = [source.resolved_name for split, source in launcher_sources if split == "train"]
+ eval_env_names = [source.resolved_name for split, source in launcher_sources if split == "eval"]
if config.deployment.type == "single_node":
script = template.render(