Feat/rl check preflight

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -234,6 +234,9 @@ class RLConfig(BaseConfig):
     dry_run: bool = False
     """Only validate and dump resolved configs, then exit early."""
 
+    check: bool = False
+    """Run preflight checks and exit without launching or writing to the output directory. Validates ports, parallelism, checkpoint resume, disk, and credentials; spawns each configured env server to verify it loads; and probes frozen-model/external inference endpoints."""
+
     ### Validate configs (e.g. raise for unsupported (combinations of) configs)
 
     @model_validator(mode="after")

src/prime_rl/entrypoints/rl.py

@@ -477,7 +477,7 @@ def rl_slurm(config: RLConfig):
     log_dir = get_log_dir(config.output_dir)
 
     if config.deployment.type == "single_node":
-        write_config(config, config_dir, exclude={"slurm", "dry_run", "clean_output_dir"})
+        write_config(config, config_dir, exclude={"slurm", "dry_run", "clean_output_dir", "check"})
         logger.info(f"Wrote config to {config_dir / RL_TOML}")
 
         train_env_names = [env.resolved_name for env in config.orchestrator.train.env]
@@ -528,6 +528,18 @@ def rl_slurm(config: RLConfig):
 
 
 def rl(config: RLConfig):
+    if config.check:
+        # Preflight mode: validate the run can start, then exit. Runs before
+        # validate_output_dir/clean_future_steps on purpose — --check must
+        # never mutate the output directory.
+        from prime_rl.utils.doctor import run_checks
+
+        setup_logger(
+            config.log.level or os.environ.get("PRIME_LOG_LEVEL", "info"),
+            json_logging=config.log.json_logging,
+        )
+        sys.exit(run_checks(config))
+
     resuming = config.ckpt is not None and config.ckpt.resume_step is not None
     clean = config.clean_output_dir and not os.environ.get("NEVER_CLEAN_OUTPUT_DIR")
     ckpt_output_dir = config.ckpt.output_dir if config.ckpt else None

src/prime_rl/orchestrator/envs.py

@@ -98,10 +98,24 @@ def env_client(self) -> EnvClient:
             raise RuntimeError(f"Env {self.name} not started — call start() first.")
         return self._env_client
 
-    async def start(self, log_dir: Path, log_level: str | None = None, json_logging: bool = False) -> None:
-        """Spawn the env server (if needed), connect, and cache its ``info``."""
+    async def start(
+        self,
+        log_dir: Path,
+        log_level: str | None = None,
+        json_logging: bool = False,
+        spawn_timeout: float | None = None,
+    ) -> None:
+        """Spawn the env server (if needed), connect, and cache its ``info``.
+
+        ``spawn_timeout`` caps the wait for a spawned server to report its
+        address (default ``ENV_SERVER_SPAWN_TIMEOUT``). The wait blocks a
+        worker thread that asyncio cancellation cannot interrupt, so callers
+        needing a shorter budget (e.g. preflight) must bound it here rather
+        than with ``asyncio.wait_for``."""
         external = self.config.address is not None
-        address = self.config.address or await self._spawn(log_dir, log_level or "INFO", json_logging)
+        address = self.config.address or await self._spawn(
+            log_dir, log_level or "INFO", json_logging, spawn_timeout=spawn_timeout
+        )
         get_logger().debug(f"Connecting {self.name} to env server {address}")
         self._env_client = EnvClient(address=address)
         # A spawned server already reported its address *after* binding + loading,
@@ -116,7 +130,9 @@ async def start(self, log_dir: Path, log_level: str | None = None, json_logging:
             f"Env {self.name} ready: num_tasks={self.num_tasks} group_scoring={self.requires_group_scoring}"
         )
 
-    async def _spawn(self, log_dir: Path, log_level: str, json_logging: bool) -> str:
+    async def _spawn(
+        self, log_dir: Path, log_level: str, json_logging: bool, spawn_timeout: float | None = None
+    ) -> str:
         """Spawn a v1 EnvServer child process (it loads the env; we never do).
         The server binds an OS-assigned port (``:0``) and reports the concrete
         address back over a queue — no free-port guess, no TOCTOU race. Its output
@@ -152,10 +168,11 @@ async def _spawn(self, log_dir: Path, log_level: str, json_logging: bool) -> str
         )
         process.start()
         self._env_server_process = process
+        timeout = spawn_timeout if spawn_timeout is not None else ENV_SERVER_SPAWN_TIMEOUT
         try:
-            address = await asyncio.to_thread(address_queue.get, timeout=ENV_SERVER_SPAWN_TIMEOUT)
+            address = await asyncio.to_thread(address_queue.get, timeout=timeout)
         except queue.Empty:
-            raise RuntimeError(f"Env server {self.name} did not report its address within {ENV_SERVER_SPAWN_TIMEOUT}s")
+            raise RuntimeError(f"Env server {self.name} did not report its address within {timeout}s")
         finally:
             address_queue.close()
             address_queue.join_thread()
@@ -194,9 +211,18 @@ async def run_group(
         return [ROLLOUT_TYPE.model_construct(**dict(wire)) for wire in wires]
 
     def shutdown(self) -> None:
+        """Terminate the spawned env server and reap it (terminate → join →
+        kill), mirroring ``Envs.shutdown``. Blocks up to ~30s for a stuck
+        process; call via ``asyncio.to_thread`` from async contexts."""
         if self._env_server_process is None:
             return
-        self._env_server_process.terminate()
+        process = self._env_server_process
+        process.terminate()
+        process.join(timeout=25)
+        if process.is_alive():
+            get_logger().warning(f"Env server {process.pid} did not exit after 25s, force killing")
+            process.kill()
+            process.join(timeout=5)
         self._env_server_process = None
 
 
@@ -218,8 +244,16 @@ def __init__(self, config: EvalEnvConfig):
         self.sampling_args = config.sampling.to_sampling_args()
         self.examples: list[dict] = []
 
-    async def start(self, log_dir: Path, log_level: str | None = None, json_logging: bool = False) -> None:
-        await super().start(log_dir=log_dir, log_level=log_level, json_logging=json_logging)
+    async def start(
+        self,
+        log_dir: Path,
+        log_level: str | None = None,
+        json_logging: bool = False,
+        spawn_timeout: float | None = None,
+    ) -> None:
+        await super().start(
+            log_dir=log_dir, log_level=log_level, json_logging=json_logging, spawn_timeout=spawn_timeout
+        )
         n = self.num_tasks if self.config.num_examples < 0 else min(self.config.num_examples, self.num_tasks)
         self.examples = [{"task_idx": i} for i in range(n)]
 

src/prime_rl/trainer/parallel_dims.py

@@ -316,7 +316,13 @@ def resolve_ep(config: ModelConfig) -> None:
         )
 
 
-def get_parallel_dims(config: ModelConfig, seq_len: int | None = None) -> ParallelDims:
+def get_parallel_dims(config: ModelConfig, seq_len: int | None = None, world_size: int | None = None) -> ParallelDims:
+    """Build (and validate) ParallelDims from a model config.
+
+    ``world_size`` defaults to the initialized process group's size; pass it
+    explicitly to validate a hypothetical allocation without ``dist`` being
+    initialized (used by the ``rl --check`` preflight).
+    """
     assert isinstance(config.ep, int), (
         f"config.ep must be resolved to an int before get_parallel_dims; got {config.ep!r}. "
         "Call resolve_ep(config) first."
@@ -329,7 +335,7 @@ def get_parallel_dims(config: ModelConfig, seq_len: int | None = None) -> Parall
         cp=config.cp,
         pp=1,
         ep=config.ep,
-        world_size=dist.get_world_size(),
+        world_size=world_size if world_size is not None else dist.get_world_size(),
     )
 
     # Validate sequence length against parallel dimensions requirements