Skip to content

Feat/rl check preflight#2970

Open
anravich13-cloud wants to merge 11 commits into
mainfrom
feat/rl-check-preflight
Open

Feat/rl check preflight#2970
anravich13-cloud wants to merge 11 commits into
mainfrom
feat/rl-check-preflight

Conversation

@anravich13-cloud

@anravich13-cloud anravich13-cloud commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

feat: rl --check preflight doctor

Adds a preflight mode to run some basic checks before starting runs:

uv run rl @ config.toml --check

validates the run and exits — without launching anything and without writing
to the output directory (unlike --dry-run, which writes resolved configs).
Exit code is 0 when nothing failed (warnings allowed, so CI can gate on it)
and 1 otherwise.

What it checks

Check What / why
client/server port match Client base_url port must equal inference.server.port — a mismatch means the orchestrator hangs against a dead port.
inference port free Bind-probes the server's port to catch zombie vLLMs before minutes of model loading.
parallelism Runs the trainer's real ParallelDims validation (incl. auto-EP resolution and seq-len divisor) so it fails here instead of after launch prep.
ckpt resume Verifies resume_step exists on disk; lists available steps if not.
disk Free space at output_dir — FAIL < 5 GB, WARN < 100 GB (checkpoints die at first write, hours in).
wandb / hf auth WARN if wandb is configured without credentials (headless launches hang at the login prompt); HF token is informational.
envs Spawns each configured env server and reads its info — catches broken/missing envs.
endpoints HTTP-probes frozen/teacher models and external inference servers .

Design notes

  • Checks are pure functions (RLConfig, HostProbe) -> list[CheckResult]; the
    HostProbe seam (sockets, disk, env vars) exists so checks are unit-testable
    on CPU without real host state.
  • Runs at the top of rl(), before validate_output_dir/clean_future_steps
    — preflight must never mutate the output dir (today the first GPU-count
    validation fires after the dir has been cleaned).
  • Checks that don't apply (elastic client, multi-node bind probe, not
    resuming, no envs) report SKIP with a reason rather than disappearing.

Test plan

  • Verified on a GPU box against examples/reverse_text/rl.toml: all-green
    report, wandb WARN correctly fired on a box without credentials, exit codes
    0/1 as expected.
  • Failure modes exercised: occupied port 8000 → FAIL + hint, exit 1;
    --ckpt.resume-step 999 → FAIL listing available steps.
  • --dry-run and normal launch paths untouched (flag short-circuits before
    them; check excluded from written configs).

Follow-ups (out of scope)

  • Refactor rl_local to reuse these check functions at launch time so the
    launch path and --check can't drift.
  • Same treatment for the sft entrypoint.
  • Config migration hints for renamed/removed keys (pydantic-config hook).
  • CPU unit tests for the check functions (removed from this PR to keep it
    minimal; the HostProbe seam makes them cheap to add back).

Note

Medium Risk
Touches the main rl() entry path and spawns real env subprocesses with network probes, but normal launches are unchanged when --check is off; optional world_size on parallel helpers is backward compatible.

Overview
Adds rl @ config.toml --check, a preflight mode that validates the run and exits before output_dir is created or cleaned. Exit 0 if no failures (warnings OK); 1 on any failure.

New RLConfig.check and prime_rl.utils.doctor run host/config checks (client/inference port alignment, local inference bind probe, trainer ParallelDims / seq-len / auto-EP via optional world_size, checkpoint resume paths, disk space, wandb/HF credentials) plus live probes (HTTP for external/frozen inference URLs; bounded parallel env-server spawns with 120s timeout, SKIP on slow startup).

rl() short-circuits to run_checks at the top; check is excluded from dumped TOML. Env start/shutdown gain configurable spawn_timeout and stronger server teardown for the doctor path.

Reviewed by Cursor Bugbot for commit f51dead. Bugbot is set up for automated code reviews on this repo. Configure here.

Adds a preflight mode to the rl entrypoint that validates a run can start
without launching it or writing to the output directory.

Static tier (default, sub-second): client/server port match + bind probe,
trainer GPU/cp divisibility (previously floor-divided silently), checkpoint
resume-step existence, disk space thresholds, wandb/HF credential presence.

Full tier (--check-level full): spawns each configured env server briefly to
validate it loads (num_tasks, group scoring), and probes frozen-model and
external inference endpoints instead of just logging 'make sure these are
serving'.

Checks are pure functions (RLConfig, HostProbe) -> [CheckResult] with a
mockable HostProbe, unit-tested on CPU. Exit code: 0 if no failures (warnings
allowed, CI-friendly), 1 otherwise.
…wns (cap 8)

One flag, one behavior: the env and endpoint checks are the most valuable
part of preflight, so they shouldn't hide behind an opt-in tier. Env servers
now spawn with bounded concurrency (ENV_CHECK_CONCURRENCY = 8) so many-env
configs finish in a few timeout windows without fork-bombing the launcher
host; per-env timeout -> SKIP semantics unchanged.
Comment thread src/prime_rl/utils/doctor.py Fixed
…ert already catches this; the win is failing before side effects, not preventing stranded GPUs
Comment thread src/prime_rl/utils/doctor.py Outdated
Comment thread src/prime_rl/utils/doctor.py
Env.shutdown only sent SIGTERM and moved on, so a stuck env server could
outlive --check (or hang preflight at exit, since the spawn is non-daemon
and multiprocessing joins children at interpreter exit). Mirror
Envs.shutdown's terminate -> join(25) -> kill -> join(5) sequence, and call
it via asyncio.to_thread in the doctor so one wedged env doesn't stall the
other parallel env checks.
…ll interfaces

port_is_free now takes the host and check_ports passes
inference.server.host (falling back to all interfaces, matching vLLM's
default bind). Probing the exact address the server will bind keeps the
check faithful; the probe never listens, so nothing is exposed.
Comment thread src/prime_rl/utils/doctor.py Dismissed
Comment thread src/prime_rl/utils/doctor.py Outdated
Comment thread src/prime_rl/utils/doctor.py Outdated
…check to full ParallelDims validation

Env timeout: asyncio.wait_for cannot interrupt the spawn's queue.get worker
thread, so the 120s budget silently stretched to the thread's own 600s
timeout (and asyncio.run joins executor threads at shutdown, delaying exit
too). Env.start/_spawn now take spawn_timeout (default unchanged:
ENV_SERVER_SPAWN_TIMEOUT) so preflight bounds the wait at its source; the
outer wait_for remains only for the pure-async external-address path.

Parallelism: the check only tested num_train_gpus % cp, but trainer startup
asserts dp_replicate * dp_shard * cp * pp == world_size and the CP seq-len
divisor. Mirror it by constructing ParallelDims (deferred torch import) with
the same mapping as get_parallel_dims, plus the seq_len % (2*cp) rule.
Explicit integer ep is validated; ep='auto' resolves at trainer startup and
is reported as deferred.
…pping

get_parallel_dims gains an optional world_size (defaults to
dist.get_world_size(), so trainer callers are unchanged), letting preflight
validate a hypothetical allocation through the exact function the trainer
runs — no duplicated dp_shard/pp mapping or seq-len rule to drift.

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit 742fdb2. Configure here.

Comment thread src/prime_rl/utils/doctor.py
resolve_ep gains an optional world_size (defaults to dist.get_world_size();
trainer callers unchanged) so the parallelism check resolves ep='auto' with
the trainer's own logic — MoE detection via the HF model config included —
and validates EP divisibility with the value the trainer will actually use.
When the HF config is unreachable (offline/gated), the check degrades to
ep=1 and says so explicitly in the report instead of implying full
validation. Closes the false-pass where e.g. a 12-GPU island with cp=3
auto-resolves ep=8 and crashes ParallelDims after launch prep.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants