diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index 103b2e6319..7f91299917 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/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. 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") diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 9040cacc1d..a2dc7d12dd 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/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,17 @@ 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 + 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 diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 7e965741eb..4b65b9d347 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -98,10 +98,21 @@ 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``).""" 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 +127,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 +165,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 +208,17 @@ 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``""" 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 +240,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)] diff --git a/src/prime_rl/trainer/parallel_dims.py b/src/prime_rl/trainer/parallel_dims.py index ca30983069..741542f0f8 100644 --- a/src/prime_rl/trainer/parallel_dims.py +++ b/src/prime_rl/trainer/parallel_dims.py @@ -278,12 +278,16 @@ def _is_moe_model(config: ModelConfig) -> bool: return hasattr(model_config, "num_experts") or hasattr(model_config, "n_routed_experts") -def resolve_ep(config: ModelConfig) -> None: +def resolve_ep(config: ModelConfig, world_size: int | None = None) -> None: """Resolve ``ep="auto"`` in-place to a concrete integer. For MoE models, resolves to ``min(fsdp_island_size, 8)`` where ``fsdp_island_size = world_size // dp_replicate``. For non-MoE models, resolves to 1 (no-op). + + ``world_size`` defaults to the initialized process group's size; pass it + explicitly to resolve for a hypothetical allocation without ``dist`` being + initialized (used by the ``rl --check`` preflight). """ if config.ep != "auto": return @@ -294,7 +298,8 @@ def resolve_ep(config: ModelConfig) -> None: get_logger().info(f"EP auto: impl='{config.impl}' does not support EP, resolving ep=1") return - world_size = dist.get_world_size() + if world_size is None: + world_size = dist.get_world_size() if not _is_moe_model(config): config.ep = 1 @@ -316,7 +321,8 @@ 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""" 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 diff --git a/src/prime_rl/utils/doctor.py b/src/prime_rl/utils/doctor.py new file mode 100644 index 0000000000..9afa0a926f --- /dev/null +++ b/src/prime_rl/utils/doctor.py @@ -0,0 +1,522 @@ +"""Preflight checks for the ``rl`` entrypoint (``rl @ config.toml --check``). + +Checks are pure functions ``(RLConfig, HostProbe) -> list[CheckResult]`` so they can +be unit-tested on CPU-only CI with a fake probe. Host/config checks run in +well under a second; the env check spawns each configured env server (bounded +concurrency, per-env timeout) and the endpoint check probes frozen/external +inference servers over the network, so a full run takes seconds to a couple +of minutes. +""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import socket +import tempfile +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING +from urllib.parse import urlparse + +from prime_rl.configs.algorithm import FrozenModelConfig +from prime_rl.utils.logger import get_logger +from prime_rl.utils.pathing import get_all_ckpt_steps, get_ckpt_dir + +if TYPE_CHECKING: + from prime_rl.configs.rl import RLConfig + +# Disk thresholds. FAIL below the floor (can't even hold logs + resolved +# configs safely); WARN below the soft limit since full checkpoints with +# optimizer state commonly run to hundreds of GB. +DISK_FAIL_GB = 5 +DISK_WARN_GB = 100 + +# Per-env timeout for the env spawn check. Deliberately much shorter than the +# orchestrator's ENV_SERVER_SPAWN_TIMEOUT (600s): a slow dataset download is a +# SKIP with a hint, not a FAIL, so --check stays usable in CI. +ENV_CHECK_TIMEOUT = 120.0 + +# Max env servers spawned concurrently. Each is a real child process that +# loads a taskset (RAM + network), so many-env configs shouldn't fork-bomb the +# launcher host; 8 keeps the worst case at a few timeout windows without +# contending badly on downloads. +ENV_CHECK_CONCURRENCY = 8 + +# Timeout for endpoint reachability probes. +ENDPOINT_PROBE_TIMEOUT = 5.0 + + +class CheckStatus(str, Enum): + PASS = "pass" + WARN = "warn" + FAIL = "fail" + SKIP = "skip" + + +@dataclass +class CheckResult: + name: str + status: CheckStatus + detail: str + hint: str | None = None + + +class HostProbe: + """Thin, mockable layer over host state (sockets, disk, env vars).""" + + def port_is_free(self, port: int, host: str = "") -> bool: + # Bind-probe only: listen() is never called, so nothing is exposed and + # no connection can be accepted; the socket closes as soon as the + # answer is known. ``host`` should be whatever the real server will + # bind — "" (all interfaces) mirrors a server binding 0.0.0.0. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + try: + sock.bind((host, port)) + return True + except OSError: + return False + + def disk_free_bytes(self, path: Path) -> int: + # output_dir may not exist yet (--check must not create it); walk up + # to the nearest existing ancestor. + current = path.resolve() + while not current.exists() and current != current.parent: + current = current.parent + return shutil.disk_usage(current).free + + def getenv(self, name: str) -> str | None: + return os.environ.get(name) + + def netrc_has_host(self, host: str) -> bool: + netrc_path = Path(self.getenv("NETRC") or "~/.netrc").expanduser() + try: + return host in netrc_path.read_text() + except OSError: + return False + + def hf_token_present(self) -> bool: + if self.getenv("HF_TOKEN"): + return True + hf_home = Path(self.getenv("HF_HOME") or "~/.cache/huggingface").expanduser() + return (hf_home / "token").is_file() + + +def _is_local_single_node(config: "RLConfig") -> bool: + return config.slurm is None and config.deployment.type == "single_node" + + +### Static tier + + +def check_ports(config: "RLConfig", probe: HostProbe) -> list[CheckResult]: + results: list[CheckResult] = [] + if config.inference is None: + return [ + CheckResult( + "ports", + CheckStatus.SKIP, + "no [inference] block — the external server is probed by the endpoints check", + ) + ] + + # Client base_url port must match the managed server's port (mirrors the + # launch-time validation in entrypoints/rl.py, but before any mutation). + client = config.orchestrator.model.client + if not client.is_elastic: + base_url = client.base_url[0] + client_port = urlparse(base_url).port + server_port = config.inference.server.port + if client_port != server_port: + results.append( + CheckResult( + "client/server port match", + CheckStatus.FAIL, + f"orchestrator client points at port {client_port}, inference server serves on {server_port}", + hint=f"set orchestrator.model.client.base_url to use port {server_port}, " + "or change inference.server.port", + ) + ) + else: + results.append( + CheckResult( + "client/server port match", + CheckStatus.PASS, + f"client base_url and inference server agree on port {server_port}", + ) + ) + else: + results.append(CheckResult("client/server port match", CheckStatus.SKIP, "elastic client — no static URL")) + + # Bind probe is only meaningful when this host is the compute host. + if _is_local_single_node(config): + port = config.inference.server.port + # Probe the exact address the server will bind (vLLM binds all + # interfaces when server.host is unset). + bind_host = config.inference.server.host or "" + if probe.port_is_free(port, bind_host): + results.append(CheckResult("inference port free", CheckStatus.PASS, f"port {port} is free")) + else: + results.append( + CheckResult( + "inference port free", + CheckStatus.FAIL, + f"port {port} is already in use", + hint="a vLLM server from a previous run may still be alive " + "(check `nvidia-smi` / `pgrep -f vllm`), or change inference.server.port", + ) + ) + else: + results.append( + CheckResult( + "inference port free", + CheckStatus.SKIP, + "launcher host is not the compute host (SLURM/multi-node)", + ) + ) + return results + + +def check_parallelism(config: "RLConfig", probe: HostProbe) -> list[CheckResult]: + # tp×dp vs num_infer_gpus and multi-node divisibility are enforced by + # config validators at parse time. Here, mirror the trainer's own startup + # validation (ParallelDims + the CP seq-len divisor) so the crash it would + # produce — after the launcher has cleaned the output dir and spawned all + # processes, with the error buried in logs/trainer.log — happens here + # instead, in milliseconds and with no side effects. + if config.deployment.type != "single_node": + return [CheckResult("parallelism", CheckStatus.SKIP, "multi-node — validated at config parse time")] + + from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep # deferred: imports torch + + model = config.trainer.model + world_size = config.deployment.num_train_gpus + resolved_model = model.model_copy() + ep_note = "" + if not isinstance(model.ep, int): + # Run the trainer's real auto-EP resolution (MoE detection needs the HF + # model config — cached locally or a tiny download) so EP divisibility + # is validated with the value the trainer will actually use. + try: + resolve_ep(resolved_model, world_size=world_size) + ep_note = f" × ep={resolved_model.ep} (auto)" + except ValueError as e: + # Genuine startup failure (e.g. ep_comm_backend requires ep > 1) + return [ + CheckResult( + "parallelism", + CheckStatus.FAIL, + str(e).strip().split("\n")[0], + hint="the trainer would fail at startup with this error", + ) + ] + except Exception: # noqa: BLE001 — HF model config unavailable (offline/gated) + resolved_model.ep = 1 + ep_note = " (ep='auto' unresolvable here — EP divisibility is validated at trainer startup)" + try: + dims = get_parallel_dims(resolved_model, seq_len=model.seq_len, world_size=world_size) + except (AssertionError, ValueError) as e: + return [ + CheckResult( + "parallelism", + CheckStatus.FAIL, + str(e).strip().split("\n")[0], + hint="the trainer would fail at startup with this error — adjust deployment.num_train_gpus, " + "trainer.model.dp_replicate, trainer.model.cp, or seq_len", + ) + ] + detail = ( + f"world={world_size}: dp_replicate={model.dp_replicate} × dp_shard={dims.dp_shard} × cp={model.cp}{ep_note}" + ) + return [CheckResult("parallelism", CheckStatus.PASS, detail)] + + +def check_ckpt(config: "RLConfig", probe: HostProbe) -> list[CheckResult]: + if config.ckpt is None or config.ckpt.resume_step is None: + return [CheckResult("ckpt resume", CheckStatus.SKIP, "not resuming")] + + ckpt_base = config.ckpt.output_dir if config.ckpt.output_dir is not None else config.output_dir + ckpt_dir = get_ckpt_dir(ckpt_base) + steps = get_all_ckpt_steps(ckpt_dir) if ckpt_dir.is_dir() else [] + resume_step = config.ckpt.resume_step + + if resume_step == -1: + if steps: + return [ + CheckResult( + "ckpt resume", CheckStatus.PASS, f"resume_step=-1 resolves to step {steps[-1]} in {ckpt_dir}" + ) + ] + return [ + CheckResult( + "ckpt resume", + CheckStatus.FAIL, + f"resume_step=-1 but no checkpoints found in {ckpt_dir}", + hint="remove ckpt.resume_step to train from scratch, or point ckpt.output_dir/output_dir " + "at the directory holding the checkpoints", + ) + ] + + if resume_step in steps: + return [CheckResult("ckpt resume", CheckStatus.PASS, f"step {resume_step} found in {ckpt_dir}")] + available = ", ".join(map(str, steps[-5:])) if steps else "none" + return [ + CheckResult( + "ckpt resume", + CheckStatus.FAIL, + f"step {resume_step} not found in {ckpt_dir} (latest steps: {available})", + hint="pick an existing step, or use resume_step = -1 for the latest", + ) + ] + + +def check_disk(config: "RLConfig", probe: HostProbe) -> list[CheckResult]: + free_gb = probe.disk_free_bytes(config.output_dir) / 1e9 + if free_gb < DISK_FAIL_GB: + return [ + CheckResult( + "disk", + CheckStatus.FAIL, + f"{free_gb:.1f} GB free at {config.output_dir} (< {DISK_FAIL_GB} GB floor)", + hint="free up space or point output_dir at a larger volume", + ) + ] + if free_gb < DISK_WARN_GB: + return [ + CheckResult( + "disk", + CheckStatus.WARN, + f"{free_gb:.1f} GB free at {config.output_dir}", + hint=f"full checkpoints with optimizer state commonly exceed this — consider ckpt.keep_last, " + f"or a volume with > {DISK_WARN_GB} GB free", + ) + ] + return [CheckResult("disk", CheckStatus.PASS, f"{free_gb:.0f} GB free at {config.output_dir}")] + + +def check_tokens(config: "RLConfig", probe: HostProbe) -> list[CheckResult]: + results: list[CheckResult] = [] + + if config.wandb is not None: + wandb_mode = (probe.getenv("WANDB_MODE") or "").lower() + has_creds = ( + probe.getenv("WANDB_API_KEY") is not None + or probe.netrc_has_host("api.wandb.ai") + or wandb_mode in ("disabled", "offline") + ) + if has_creds: + results.append(CheckResult("wandb auth", CheckStatus.PASS, "credentials found")) + else: + results.append( + CheckResult( + "wandb auth", + CheckStatus.WARN, + "wandb is configured but no credentials were found", + hint="headless launches hang at the interactive wandb login prompt — " + "set WANDB_API_KEY, run `uv run wandb login`, or set WANDB_MODE=disabled", + ) + ) + else: + results.append(CheckResult("wandb auth", CheckStatus.SKIP, "wandb not configured")) + + if probe.hf_token_present(): + results.append(CheckResult("hf auth", CheckStatus.PASS, "HF token found")) + else: + results.append( + CheckResult("hf auth", CheckStatus.PASS, "no HF token found — required only for gated/private models") + ) + return results + + +### Full tier (network + subprocesses) + + +def _probe_endpoint(base_url: str) -> str | None: + """Return None if reachable, else a short error description. + + ``base_url`` is an OpenAI-compatible base like ``http://host:8000/v1``; + any HTTP response (including 401/404) proves a server is listening. + """ + import httpx + + url = base_url.rstrip("/") + "/models" + try: + httpx.get(url, timeout=ENDPOINT_PROBE_TIMEOUT) + return None + except httpx.HTTPError as e: + return f"{type(e).__name__}: {e}" + + +def check_endpoints(config: "RLConfig", probe: HostProbe) -> list[CheckResult]: + """Probe endpoints the rl entrypoint does not manage: frozen/teacher models + always, and the policy server when no [inference] block is configured. + + Upgrades the launch-time "make sure these are serving, otherwise rollouts + will hang" log lines into actual probes. + """ + results: list[CheckResult] = [] + + # Frozen model references (teachers, OPD sources, ...) — never launched by prime-rl. + frozen: dict[str, FrozenModelConfig] = {} + for env in config.orchestrator.train.env: + algo = env.algo + if algo is None: + continue + for ref in (algo.sampling.source, getattr(algo, "teacher", None)): + if isinstance(ref, FrozenModelConfig) and not ref.is_elastic: + frozen.setdefault(ref.name, ref) + for name, ref in frozen.items(): + error = _probe_endpoint(ref.base_url[0]) + if error is None: + results.append(CheckResult(f"frozen model '{name}'", CheckStatus.PASS, f"reachable at {ref.base_url[0]}")) + else: + results.append( + CheckResult( + f"frozen model '{name}'", + CheckStatus.FAIL, + f"unreachable at {ref.base_url[0]} ({error})", + hint="the rl entrypoint does not start frozen models — start this server before " + "launching, or rollouts will hang", + ) + ) + + # Policy inference server, when not managed by this entrypoint. + if config.inference is None: + client = config.orchestrator.model.client + if client.is_elastic: + results.append(CheckResult("external inference", CheckStatus.SKIP, "elastic pool — discovered via DNS")) + else: + error = _probe_endpoint(client.base_url[0]) + if error is None: + results.append( + CheckResult("external inference", CheckStatus.PASS, f"reachable at {client.base_url[0]}") + ) + else: + results.append( + CheckResult( + "external inference", + CheckStatus.FAIL, + f"unreachable at {client.base_url[0]} ({error})", + hint="no [inference] block is configured, so the orchestrator expects a running " + "server at this URL and will hang waiting for it", + ) + ) + return results + + +async def _check_one_env(env_config, log_dir: Path) -> CheckResult: + from prime_rl.orchestrator.envs import Env + + env = Env(env_config) + name = env.name + timed_out = CheckResult( + f"env '{name}'", + CheckStatus.SKIP, + f"did not come up within {ENV_CHECK_TIMEOUT:.0f}s", + hint=f"often a slow first-time dataset download, not a bug — see {log_dir / f'{name}.log'}", + ) + try: + # spawn_timeout does the real bounding: the spawn waits in a worker + # thread (queue.get) that asyncio cancellation cannot interrupt, so an + # outer wait_for alone would block until the thread's own (600s) + # timeout. The wait_for here only guards the pure-async external + # address path (health poll + info), where cancellation works. + await asyncio.wait_for(env.start(log_dir, spawn_timeout=ENV_CHECK_TIMEOUT), timeout=ENV_CHECK_TIMEOUT + 30) + return CheckResult( + f"env '{name}'", + CheckStatus.PASS, + f"loaded: num_tasks={env.num_tasks}, group_scoring={env.requires_group_scoring}", + ) + except (asyncio.TimeoutError, TimeoutError): + return timed_out + except Exception as e: # noqa: BLE001 — report, don't crash the check run + if isinstance(e, RuntimeError) and "did not report its address" in str(e): + return timed_out + detail = str(e).strip().split("\n")[0][:200] or type(e).__name__ + return CheckResult( + f"env '{name}'", + CheckStatus.FAIL, + detail, + hint=f"full server output: {log_dir / f'{name}.log'}", + ) + finally: + # shutdown blocks (join with timeout, kill if stuck) — keep it off the + # event loop so one wedged env doesn't stall the other parallel checks. + await asyncio.to_thread(env.shutdown) + + +def check_envs(config: "RLConfig", probe: HostProbe) -> list[CheckResult]: + """Spawn each configured env server briefly and read its ``info``. + + Catches missing env packages, dataset auth/download failures, and broken + env code before a run ever starts — today these surface minutes in, buried + in the orchestrator's env logs. + """ + env_configs = list(config.orchestrator.train.env) + if config.orchestrator.eval is not None: + env_configs += list(config.orchestrator.eval.env) + if not env_configs: + return [CheckResult("envs", CheckStatus.SKIP, "no environments configured")] + + # Never write into output_dir from --check; env-server logs go to a temp dir. + log_dir = Path(tempfile.mkdtemp(prefix="prime-rl-check-envs-")) + + async def run_all() -> list[CheckResult]: + semaphore = asyncio.Semaphore(ENV_CHECK_CONCURRENCY) + + async def run_one(env_config) -> CheckResult: + async with semaphore: + return await _check_one_env(env_config, log_dir) + + return list(await asyncio.gather(*(run_one(env_config) for env_config in env_configs))) + + return asyncio.run(run_all()) + + +CHECKS = [check_ports, check_parallelism, check_ckpt, check_disk, check_tokens, check_envs, check_endpoints] + +_STATUS_SYMBOL = { + CheckStatus.PASS: "✓", + CheckStatus.WARN: "⚠", + CheckStatus.FAIL: "✗", + CheckStatus.SKIP: "○", +} + + +def run_checks(config: "RLConfig", probe: HostProbe | None = None) -> int: + """Run preflight checks and render a report. Returns the process exit code: + 0 when nothing failed (warnings allowed, so CI can gate on --check without + flaking), 1 when any check failed.""" + logger = get_logger() + probe = probe or HostProbe() + + logger.info("Running preflight checks") + + results: list[CheckResult] = [] + for check in CHECKS: + results.extend(check(config, probe)) + + for result in results: + line = f"{_STATUS_SYMBOL[result.status]} {result.name:<28} {result.detail}" + if result.status == CheckStatus.FAIL: + logger.error(line) + elif result.status == CheckStatus.WARN: + logger.warning(line) + else: + logger.info(line) + if result.hint and result.status in (CheckStatus.FAIL, CheckStatus.WARN): + logger.info(f" ↳ {result.hint}") + + num_failed = sum(r.status == CheckStatus.FAIL for r in results) + num_warned = sum(r.status == CheckStatus.WARN for r in results) + num_passed = sum(r.status == CheckStatus.PASS for r in results) + num_skipped = sum(r.status == CheckStatus.SKIP for r in results) + summary = f"{num_passed} passed, {num_warned} warning(s), {num_failed} failure(s), {num_skipped} skipped" + + if num_failed: + logger.error(f"Preflight failed: {summary}. Fix the failures above before launching.") + return 1 + logger.success(f"Preflight passed: {summary}") + return 0