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

@@ -194,9 +194,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