fix(orchestrator): make admin-plane timeouts env-overridable

src/prime_rl/entrypoints/orchestrator.py

@@ -10,6 +10,7 @@
 """
 
 import asyncio
+import os
 
 from prime_rl.configs.orchestrator import OrchestratorConfig
 from prime_rl.utils.config import cli
@@ -19,6 +20,11 @@
 def main():
     set_proc_title("Orchestrator")
     config = cli(OrchestratorConfig)
+    # Apply the configured env vars in-process, like `entrypoints/inference.py` does:
+    # the rl/SLURM launchers inject them into the orchestrator subprocess, but a
+    # standalone `orchestrator @ config.toml` (e.g. a platform-hosted run) has no
+    # launcher to do it.
+    os.environ.update(config.env_vars)
     from prime_rl.orchestrator.orchestrator import run_orchestrator
 
     asyncio.run(run_orchestrator(config))

src/prime_rl/eval/online.py

@@ -20,6 +20,7 @@
 from prime_rl.configs.eval import SFTOnlineEvalConfig
 from prime_rl.configs.trainer import FileSystemWeightBroadcastConfig
 from prime_rl.eval.runner import POLL_INTERVAL_S, EvalRunner
+from prime_rl.orchestrator.clients import env_timeout
 from prime_rl.transports.weights import WeightReceiver, setup_weight_receiver
 from prime_rl.utils.config import cli, dump_resolved_config
 from prime_rl.utils.logger import get_logger, setup_logger
@@ -28,7 +29,9 @@
 from prime_rl.utils.utils import clean_exit
 
 # Budget for the trainer's startup broadcast: it is always coming, but only
-# after the trainer has finished loading the model.
+# after the trainer has finished loading the model. Call-time override:
+# `$PRL_STARTUP_BROADCAST_TIMEOUT_S` (a large model's first broadcast can take
+# much longer than the default).
 STARTUP_BROADCAST_TIMEOUT_S = 1200
 
 
@@ -84,7 +87,9 @@ async def watch(self) -> None:
         # its startup broadcast until this receive, and for filesystem it guarantees the
         # served weights match the trainer's incoming policy.
         startup_step = config.resume_step or 0
-        await self.receiver.sync_startup(startup_step, timeout=STARTUP_BROADCAST_TIMEOUT_S)
+        await self.receiver.sync_startup(
+            startup_step, timeout=env_timeout("PRL_STARTUP_BROADCAST_TIMEOUT_S", STARTUP_BROADCAST_TIMEOUT_S)
+        )
         self.runner.policy.version = startup_step
 
         if config.resume_step is None:

src/prime_rl/inference/dynamo.py

@@ -17,6 +17,7 @@
     AdminPlane,
     _admin_post,
     check_health,
+    env_timeout,
     maybe_check_has_model,
     setup_admin_clients,
 )
@@ -352,7 +353,12 @@ async def update_weights(
             self._nccl_initialization_state = "terminal"
 
             try:
-                await _admin_post(self.clients[0], "/pause", params={"mode": "keep", "clear_cache": "false"})
+                await _admin_post(
+                    self.clients[0],
+                    "/pause",
+                    params={"mode": "keep", "clear_cache": "false"},
+                    timeout_s=env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S),
+                )
             except BaseException as failure:
                 raise RuntimeError("Dynamo pause failed; worker state is unknown and restart is required") from failure
 
@@ -367,7 +373,7 @@ async def update_weights(
                 await self._collective_rpc(
                     self.clients[0],
                     method="update_weights_from_path",
-                    timeout=UPDATE_WEIGHTS_TIMEOUT_S,
+                    timeout=env_timeout("PRL_UPDATE_WEIGHTS_TIMEOUT_S", UPDATE_WEIGHTS_TIMEOUT_S),
                     args=[weight_dir.as_posix()],
                 )
             except BaseException as failure:
@@ -377,7 +383,9 @@ async def update_weights(
                 ) from failure
 
             try:
-                await _admin_post(self.clients[0], "/resume", timeout_s=ADMIN_TIMEOUT_S)
+                await _admin_post(
+                    self.clients[0], "/resume", timeout_s=env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S)
+                )
             except BaseException as failure:
                 self._terminalize_nccl()
                 raise RuntimeError("Dynamo resume failed; worker state is unknown and restart is required") from failure

src/prime_rl/orchestrator/clients.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import math
 import os
 from collections.abc import Callable
 from pathlib import Path
@@ -223,7 +224,7 @@ async def update_weights(
                         admin_client,
                         "/update_weights",
                         json={"weight_dir": weight_dir_posix},
-                        timeout_s=UPDATE_WEIGHTS_TIMEOUT_S,
+                        timeout_s=env_timeout("PRL_UPDATE_WEIGHTS_TIMEOUT_S", UPDATE_WEIGHTS_TIMEOUT_S),
                     )
                     for admin_client in self.clients
                 ]
@@ -381,11 +382,35 @@ def _is_retryable_admin_error(exception: BaseException) -> bool:
     return False
 
 
+def env_timeout(name: str, default: float) -> float:
+    """Resolve a timeout in seconds from the environment, falling back to ``default``.
+
+    Read at call time, not import time: entrypoints apply config ``env_vars`` to
+    ``os.environ`` after this module is imported. Invalid values raise instead of
+    silently falling back to the default.
+    """
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    try:
+        value = float(raw)
+    except ValueError:
+        raise ValueError(f"{name} must be a number of seconds, got {raw!r}") from None
+    if not math.isfinite(value):
+        raise ValueError(f"{name} must be a finite number of seconds, got {raw!r}")
+    if value < 0:
+        raise ValueError(f"{name} must be non-negative, got {raw!r}")
+    return value
+
+
 # Per-attempt read timeout for admin ops, overridable per call. The admin
 # AsyncClient uses `timeout=None`, so without this a stuck server would hang the
 # weight update forever: the read timeout converts a hang into a TimeoutException
 # that tenacity retries. Sized for `/pause`, which drains in-flight requests
 # (mode="keep") and so can legitimately take a while.
+# Call-time overrides: `$PRL_ADMIN_TIMEOUT_S`, `$PRL_UPDATE_WEIGHTS_TIMEOUT_S`
+# (via ``env_timeout``) — e.g. a large model whose `/update_weights` collectives
+# run far longer than the default.
 ADMIN_TIMEOUT_S = 300.0
 # `/update_weights` runs a collective NCCL receive across all DP workers, which
 # can take longer than the other admin ops.
@@ -416,8 +441,12 @@ async def _pause_engines(admin_clients: list[AsyncClient], *, step: int) -> None
     """Pause all inference engines, waiting for in-flight requests to drain."""
     logger = get_logger()
     logger.debug(f"Pausing inference engines to update weights to policy v{step}")
+    admin_timeout_s = env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S)
     await asyncio.gather(
-        *[_admin_post(client, "/pause", params={"mode": "keep", "clear_cache": "false"}) for client in admin_clients]
+        *[
+            _admin_post(client, "/pause", params={"mode": "keep", "clear_cache": "false"}, timeout_s=admin_timeout_s)
+            for client in admin_clients
+        ]
     )
     logger.debug("All inference engines paused")
 
@@ -429,7 +458,8 @@ async def _resume_engines(admin_clients: list[AsyncClient]) -> None:
     failures is safe; a dropped /resume would leave engines paused indefinitely.
     """
     logger = get_logger()
-    await asyncio.gather(*[_admin_post(client, "/resume") for client in admin_clients])
+    admin_timeout_s = env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S)
+    await asyncio.gather(*[_admin_post(client, "/resume", timeout_s=admin_timeout_s) for client in admin_clients])
     logger.debug("All inference engines resumed")
 
 
@@ -504,7 +534,7 @@ async def initialize(admin_client: AsyncClient, rank_offset: int) -> None:
         await _admin_post(
             admin_client,
             "/init_broadcaster",
-            timeout_s=max(ADMIN_TIMEOUT_S, timeout),
+            timeout_s=max(env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S), timeout),
             json={
                 "host": host,
                 "port": port,

src/prime_rl/orchestrator/orchestrator.py

@@ -1069,8 +1069,12 @@ def main() -> None:
     set_proc_title("Orchestrator")
     import uvloop
 
+    config = cli(OrchestratorConfig)
+    # Apply the configured env vars in-process: a standalone `python -m
+    # prime_rl.orchestrator.orchestrator @ config.toml` has no launcher to do it.
+    os.environ.update(config.env_vars)
     uvloop.install()
-    asyncio.run(run_orchestrator(cli(OrchestratorConfig)))
+    asyncio.run(run_orchestrator(config))
 
 
 if __name__ == "__main__":

tests/unit/orchestrator/test_clients.py

@@ -3,13 +3,16 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
+import pytest
 from verifiers.v1.configs.client import EvalClientConfig
 
 from prime_rl.configs.shared import ClientConfig
 from prime_rl.orchestrator.clients import (
+    ADMIN_TIMEOUT_S,
     AdminPlane,
     _is_retryable_lora_error,
     check_health,
+    env_timeout,
     load_lora_adapter,
     setup_client,
 )
@@ -110,6 +113,35 @@ def test_setup_client_creates_renderer_client():
     assert client.headers["X-Test"] == "test"
 
 
+def test_env_timeout_reads_env_at_call_time(monkeypatch):
+    monkeypatch.setenv("PRL_ADMIN_TIMEOUT_S", "1234")
+    assert env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S) == 1234.0
+
+
+def test_env_timeout_defaults_when_unset(monkeypatch):
+    monkeypatch.delenv("PRL_ADMIN_TIMEOUT_S", raising=False)
+    assert env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S) == ADMIN_TIMEOUT_S
+
+
+def test_env_timeout_raises_on_non_numeric_value(monkeypatch):
+    monkeypatch.setenv("PRL_ADMIN_TIMEOUT_S", "abc")
+    with pytest.raises(ValueError, match="PRL_ADMIN_TIMEOUT_S"):
+        env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S)
+
+
+def test_env_timeout_raises_on_negative_value(monkeypatch):
+    monkeypatch.setenv("PRL_ADMIN_TIMEOUT_S", "-1")
+    with pytest.raises(ValueError, match="PRL_ADMIN_TIMEOUT_S"):
+        env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S)
+
+
+@pytest.mark.parametrize("raw", ["nan", "inf", "1e309"])
+def test_env_timeout_raises_on_non_finite_value(monkeypatch, raw):
+    monkeypatch.setenv("PRL_ADMIN_TIMEOUT_S", raw)
+    with pytest.raises(ValueError, match="PRL_ADMIN_TIMEOUT_S"):
+        env_timeout("PRL_ADMIN_TIMEOUT_S", ADMIN_TIMEOUT_S)
+
+
 def test_check_health_retries_non_success_status():
     client = AsyncMock()
     unavailable = httpx.Response(503, request=httpx.Request("GET", "http://worker/health"))