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(