feat(configs): launcher-owned env-server address overrides

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:<base + i>``. 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 ``<split>/<resolved_name>`` (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:<env_server_base_port + index>`` 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:<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): 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 '<split>/<resolved_name>'"
+            )
+        return self

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(