feat(sc): track generation fleet health and route NeMo-Gym through NeMo-RL - #3590
Conversation
Splits rollout failures into infrastructure (the prompt is fine, the fleet is not) and data (deterministic and prompt-specific). That split drives the retry policy landing in follow-up commits: infra failures re-dispatch the prompt onto another generation shard, data failures get a small budget because another shard would fail identically. HTTP errors are classified by status rather than by type. vLLM answers an over-long prompt with 400 "This model's maximum context length is ...", so treating every aiohttp ClientError as infrastructure would spend the retry budget on a prompt no shard can serve. 5xx plus 408/429 are infra, which matches the set NeMo-Gym itself retries in nemo_gym/openai_utils.py. Anything unrecognized is classified as data so that unexpected exceptions fail loudly rather than being retried into silence. Adds the async_rl.rollout_failure and async_rl.watchdog config blocks plus the three timeout fields. Every default is inert: the timeouts default to null, so a config that does not mention these fields behaves exactly as it did before. The validators reject combinations that would silently do nothing -- notably on_data_exhausted=skip with a zero skip budget, which would otherwise behave exactly like fail_fast while reading as though bad prompts were tolerated. No behaviour change yet: nothing consumes the taxonomy or the config until the following commits. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
_run_single_rollout caught every exception from _generate_response, printed it,
and broke out of the turn loop. Execution then fell through and built a
Completion holding only the prompt with reward=0.0, which generate_and_push
committed as a training row.
It did not even crash downstream. add_grpo_token_loss_masks_and_generation_logprobs
zero-fills a missing generation_logprobs and sets token_loss_mask=0, so the row is
well formed and contributes no gradient -- but its zero reward does enter the
per-prompt GRPO baseline. With use_leave_one_out_baseline a single dead vLLM
worker silently shifts the advantage of every sibling in the group, for the rest
of the run, with nothing in the logs but a print.
Generation failures are now classified and raised: infrastructure errors become
GenerationUnavailable, everything else RolloutDataFailure, both carrying the
prompt and trajectory coordinates a raw traceback lacks. The try body is narrowed
to the generation call so the surrounding bookkeeping can no longer be skipped
half-done, and CancelledError still passes through untouched because it is not an
Exception.
Two adjacent fixes the propagation exposes:
- asyncio.gather propagates the first exception but leaves the remaining
awaitables running detached, so a failed group left N-1 generations queued
against the fleet for a result already being discarded.
_gather_cancelling_siblings cancels and drains them first. This was latent
before -- calculate_rewards could already raise -- and matters more once
retries land.
- The gen_leader_worker_idx catch is narrowed from Exception to
(IndexError, TypeError, ValueError). It is a load-accounting metric and must
not fail a rollout, but the catch should not hide unrelated errors.
Behaviour change: a dead generation shard now fails the run loudly instead of
quietly degrading the batch. The re-dispatch policy that keeps the run alive
lands in a follow-up commit; correctness first.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Nothing in the SingleController rollout path had a deadline. NeMo-Gym compounds this: it passes aiohttp ClientTimeout() -- every field None, so all timeouts disabled -- and retries ClientOSError and ServerDisconnectedError in uncapped 0.5s loops. A dead vLLM endpoint therefore parks a rollout forever, and each parked rollout permanently holds one _buffer_capacity and one max_inflight_prompts permit. Enough of them and the rollout pump blocks on sem.acquire() while the train pump spins, with no exception raised anywhere. Adds deadlines at the three waits, resolved from async_rl into a RolloutTimeouts dataclass so async_rl remains the single place a default lives: - the whole NeMo-Gym prompt-group stream (rollout_timeout_s) - one generate_async turn on the native path (generation_timeout_s) - one calculate_rewards environment step (env_timeout_s) The gym deadline deliberately spans the entire stream rather than each await. Gym yields rows as they finish, so a per-await budget would reset every time a fast row landed and never fire for the slow row actually holding the group up. The env deadline frees the rollout, not the thread: Python cannot kill a running thread, so a hung env call keeps its thread-pool slot until its own ray.get returns. Unblocking the rollout is still the point. _Deadline wraps asyncio.timeout and consults expired() before relabelling, so a TimeoutError raised by the wrapped code is not reported with our deadline's duration -- that would send anyone debugging it to the wrong knob. Expiry surfaces as RolloutTimeout, an infra failure, so the retry policy will treat it as retriable. Outer cancellation still propagates as CancelledError. Also reclassifies the truncated-gym-stream error from bare RuntimeError to GymTransportError and names the missing rows. Rows going missing is a transport problem and must be retriable rather than reading as a bad prompt. All three default to null, so a config that does not set them behaves exactly as before. Two pre-existing object.__new__ test fakes are updated for the new attribute. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
P0.2 made a dead generation shard fail the run loudly rather than quietly
corrupting the batch. This restores availability without giving the corruption
back: generate_and_push now retries, and no prompt is discarded for
infrastructure reasons.
An infra failure says the fleet is unwell, not the prompt, so the attempt is
retried. The retry re-enters generation-shard selection, which is what makes it
land somewhere else -- generate_and_push itself knows nothing about shard health,
and does not need to. Exhausting the infra budget therefore means the same
failure followed the prompt across repeated selections, which is reported as
RolloutRedispatchExhausted rather than absorbed.
Deterministic failures get their own, much smaller budget. Another shard rejects
the prompt identically, so retrying mostly burns time -- but one retry is still
worth taking, because a shard under memory pressure can return an empty
generation that looks deterministic and is not. On exhaustion the default is to
fail the run, since a genuinely deterministic failure is almost always a config
bug (usually max_total_sequence_length versus the engine's max_model_len).
on_data_exhausted=skip exists for long soaks and is bounded by a run-wide
max_skipped_prompts.
Details worth knowing when reading this:
- The slot is reserved inside the loop, so each attempt owns a fresh group_id
and a failed attempt's rows cannot collide with the retry's.
- The loop condition is the infra budget, so exhaustion exits through a normal
terminal rather than raising from inside the handler. RolloutRetryPolicy
rejects a zero budget, which is what makes that terminal's invariant hold.
- Data failures do not back off. Waiting cannot help a deterministic failure.
- Cancellation is caught by a separate `except BaseException` and never
retried: tearing down the controller must not look like a transient fault.
- A SKIPPED prompt never reaches the buffer, so the train pump will never
release its backpressure permit. _dispatch_one_prompt releases it directly;
getting this wrong leaks one slot per skipped prompt until the pump wedges.
There is a parametrized test pinning both branches.
RolloutStats counts commits, skips, re-dispatches and data failures by reason.
A rising re-dispatch count is the only externally visible sign that the fleet is
degrading, so it is not optional bookkeeping. Wiring it into the SC logger comes
with the watchdog.
Known gap, deliberately deferred: a gym retry currently redoes the whole prompt
group rather than only the rows that never arrived. _run_rollouts already tracks
received_row_indices and trajectory_collector.py has the prior art, so partial
re-dispatch is a follow-up rather than a redesign. Correct and wasteful before
efficient.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Every other guard in this phase reacts to something raising. The wedge this work exists to prevent raises nothing at all: rollouts sit in NeMo-Gym's uncapped retry loop, the train pump spins on sleep(0.005), and Ray reports everything healthy. The only way to see it is to notice that committed groups stopped moving while rollouts are still in flight. _watchdog_pump runs as a third asyncio task and does two things. It publishes the RolloutStats counters plus in-flight and idle-time gauges, so a degrading fleet is visible before it wedges -- rollout/idle_s is the leading indicator. And it reports a stall, defined as in-flight rollouts with no commits for stall_timeout_s. Progress is measured by the committed counter rather than a timestamp because "no group has landed" is the property that matters, whatever the cause. An idle controller with nothing in flight is explicitly not a stall; between epochs there is legitimately no work. stall_action defaults to warn so the threshold can be tuned against a real workload before it is allowed to end a run. NemoGym gains health_check(), a thin wrapper over NeMo-Gym's own RunHelper.poll(). Gym already implements that check and calls it every 60s from run_forever(); NeMo-RL only ever called rh.start(), so it never ran. Without it a dead tool server surfaces as unexplained rollout timeouts rather than a named process. The watchdog polls every environment handle that exposes the method and skips those that do not -- only NeMo-Gym has subprocess servers to lose. While there: NemoGym now declares rh/rch/head_server_config/node_ip/ head_server_port in __init__ and guards the methods that need them. Ray recreates a restarted actor through __init__ alone, which does not start the Gym servers, so a restarted NemoGym previously surfaced that state as an AttributeError from deep inside a rollout. It now says what actually happened. shutdown() became a no-op in that state too, since it runs in a finally block and must not mask a real training error. run() awaits the watchdog first when several tasks finish together: it only completes by raising, and its diagnosis is more specific than the pumps', whose own symptom would just be "waiting". Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
Fault injection found this: killing a vLLM generation worker wedged the loop for
six minutes while the watchdog watched and said nothing.
The stall condition required rollouts to be in flight, on the reasoning that an
idle controller has legitimately no work. The wedge has none. _sync_weights
clears _rollout_permitted on entry and only sets it on exit, so a hung weight
sync leaves the gate shut and the rollout pump parked at
`await self._rollout_permitted.wait()` -- before dispatch, before anything can
fail. Nothing is in flight to count, nothing fails, and the run sits there.
The watchdog's own metrics from that run diagnose it exactly:
rollout/inflight: 0.0 at every sample
rollout/idle_s: 371.8 <- it saw the idleness
rollout/redispatch_total: 0.0 <- no rollout ever failed
rollout/committed_total: 10, frozen
and train/loss stopping one step short of the watchdog's own step counter places
the hang inside _sync_weights: _train_pump increments _train_steps, then syncs,
then logs. The sync is a NCCL broadcast to an inference rank that no longer
exists, and ray.get on the trainer future comes before the one that would have
raised ActorDiedError.
Progress is now the pair (committed groups, completed train steps), and what
separates a stall from an idle gap is whether work remains, not whether anything
is in flight. rollout/train_steps is published alongside so the two can be told
apart from outside.
Also adds tests/functional/grpo_dp_single_controller_chaos.sh, the harness that
found this. Two things it has to do that are not obvious:
- Teardown reaps VLLM::EngineCore and the policy workers, not just the driver.
vLLM runs its engine in a child process that survives its parent actor being
killed -- which is exactly what this test does on purpose -- and the orphan
holds tens of GB of device memory. A leaked run made the next one fail in
placement-group setup, reading as an unrelated flake. It is hard to spot
because nvidia-smi in a container reports host pids, so the offender is
invisible to `ps -p`.
- Startup refuses to run on a dirty GPU, so a leftover allocation cannot be
misreported as a failure of the code under test.
On the run that passes, the re-dispatch path catches the kill in ~10s:
ActorDiedError -> GenerationUnavailable -> three attempts -> exhausted. Ray
reports actor death immediately, so the generation deadline never has to fire.
Not in any CI lane: it needs GPUs and is timing-sensitive.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
A gym retry redid the whole prompt group. NeMo-Gym's stream dies on its first
failing row, so one bad row takes every later row with it -- at
num_generations_per_prompt=16 that means paying 16 generations to recover
however many were actually lost.
_run_rollouts now keeps completed rows across attempts and re-sends only the
pending ones, which is the same shape as the legacy collector's pending-group
retry in trajectory_collector.py. A 5-row group that loses rows 3-4 costs 7
row-generations instead of 10.
Three details that shape the implementation:
- The deadline still belongs to the prompt group, not to each attempt.
Wrapping each attempt instead would silently multiply rollout_timeout_s by
max_gym_row_attempts.
- Only infrastructure failures are re-dispatched. A prompt NeMo-Gym cannot
serve fails the same way every time, and retrying it here would also
multiply against the outer data budget in generate_and_push.
- Row indices are validated against the original group rather than the pending
subset, because a re-dispatched row keeps its original _rowidx so results
stay ordered. That makes "_rowidx equals position" a contract of
_run_rollouts, so it is now checked up front -- the alternative is a KeyError
several frames deeper.
max_gym_row_attempts defaults to 3, matching the legacy
_MAX_NEMO_GYM_STREAM_RETRIES. The RolloutRetryPolicy default stays 1 so a
directly-constructed RolloutManager does not silently gain retries.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
The harness selected its victim with pgrep -f 'ray::.*[Gg]eneration'. In a run against vLLM 0.25.1 that matched a process which was NOT the serving generation worker: the kill landed, generation carried on, training reached step 3/50, and the test was on course to report a wedge that never happened -- a false failure of exactly the containment behaviour it exists to prove. Widen the pattern to match both forms the worker appears as (ray::-titled and isolated-venv path; both were observed), but the assertions matter more than the pattern: * print the victim's cmdline, so a mis-target is visible in the log rather than masquerading as a hang; * refuse to proceed unless the victim looks like a generation worker; * verify it actually died, rather than assuming SIGKILL landed. Also raise DEATH_DEADLINE_S from 300s to 600s. An observed run took 222s to fail -- 5 re-dispatch attempts with capped exponential backoff, which is the designed behaviour -- leaving only 26% headroom on a workstation that is faster than CI. Verified on 2xA6000 with TransformerEngine rebuilt for sm_86: [chaos] killing generation worker pid=69107 [chaos] cmdline: ray::VllmAsyncGenerationWorker [chaos] PASS: bounded, attributable failure 222s after the kill Signed-off-by: asolergibert <asolergibert@nvidia.com>
The harness was excluded as 'inherently timing-sensitive'. That no longer holds: the death deadline is 600s against an observed 222s, and the victim is now asserted rather than assumed. The recovery tests already in this lane kill processes the same way, so excluding this one was also inconsistent. It matters more than the others. A wedge raises no exception and fails no assertion anywhere else, so without this test a regression that restores the silent hang -- the failure this whole series exists to remove -- would be caught by nothing. Full mode only; it costs ~10 minutes and deliberately drives the job to a non-zero exit. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Adding `timeouts` as a required positional argument broke every direct construction of AsyncRolloutImpl. That went unnoticed because production always builds it through the config path, which passes the field -- the only direct construction is in tests. Upstream's new tests/unit/experience/test_rollout_manager_router_replay.py (from #3378) is one, so after syncing onto current main it failed with TypeError: __init__() missing 1 required positional argument: 'timeouts'. The rebase was textually clean; the incompatibility is in the signature, not in any line either side edited, so nothing flagged it. RolloutTimeouts is a frozen dataclass whose fields all default to None, meaning "no deadline, wait indefinitely" -- the historical behaviour. So defaulting the parameter restores the old semantics for callers that do not ask for deadlines, rather than silently imposing one. Same reasoning already applied to RolloutRetryPolicy, whose defaults reproduce single-attempt behaviour so a directly-constructed manager does not silently gain retries. No separate regression test: upstream's file is the regression test, and it runs in our lane. Signed-off-by: asolergibert <asolergibert@nvidia.com>
…ich state
Victim selection was a coin flip. `pgrep -f` on a loose
'[Gg]enerationWorker' substring, then `head -1`, which picks by pid
order. Sampling a live run 378 times shows what that actually matched:
/opt/ray_venvs/...VllmAsyncGenerationWorker/bin/python a CHILD, not the actor
bash -c exec /opt/ray_venvs/...GenerationWorker... the launcher shell
ray::VllmAsyncGenerationWorker the actor, between calls
ray::vllm_policy-0-0:VllmAsyncGenerationWorker.__init__ the actor, constructing
ray::VllmAsyncGenerationWorker.generate_async the actor, serving a rollout
ray::VllmAsyncGenerationWorker.init_collective_async the actor, setting up refit
ray::VllmAsyncGenerationWorker.shutdown the actor, tearing down
Three distinct processes across five states, because Ray retitles a
worker with setproctitle for the exact duration of each call.
It never failed, and that is the point: every one of those scenarios does
end in a bounded attributable failure, which is all the test asserted. The
divergence was only visible as wall-clock across branches -- 7s on one,
222s on another, for what was supposed to be the same test.
The existing guard did not help. `case $VICTIM_CMD in *[Gg]enerationWorker*`
was added earlier to catch a mis-targeted kill, and it passes for all
seven forms above, including the launcher shell and the venv child. It
checked WHAT was killed and never IN WHAT STATE, so it was structurally
blind to this while reading like protection against it.
Now the actor is matched structurally -- anchored ray:: prefix, optional
<name>: infix -- so the child and the shell cannot match, and the state is
pinned:
idle (default) no method suffix. Nothing is in flight, so the loss
must be DETECTED, by health probe or by the stall detector.
serving .generate_async. An in-flight rollout RPC dies with the worker
and surfaces immediately, detection doing no work.
Not "any method suffix": __init__, init_collective_async and shutdown are
three further distinct scenarios, and folding them in would reintroduce
the ambiguity being removed. The title is re-read immediately before the
kill and the run fails if it changed, so the sub-millisecond window
between scan and kill cannot quietly restore the coin flip.
Both modes registered in the lane. Pinning to idle alone would silently
drop a scenario the old selection used to hit by chance, and the serving
path costs seconds.
Verified on 2xA6000: idle twice -> ray::VllmAsyncGenerationWorker, 222s
both times; serving -> ray::VllmAsyncGenerationWorker.generate_async, 12s,
GenerationUnavailable + RolloutRedispatchExhausted. Regexes checked
against all seven observed titles: each mode matches exactly one.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
The pre-flight check took the max memory.used across every GPU and aborted above 1GiB. That is an OR over the whole host, so it gets likelier to trip the bigger the machine: this test pins itself to 2 GPUs, so on an 8-GPU CI runner a single unrelated process on one GPU aborts it with six sitting idle. The message says 'clean up before running', which reads like a real problem with the machine rather than an over-strict check. Count free GPUs and require $GPUS of them instead. It still catches the case the check was written for -- a previous test in the lane leaking a VLLM::EngineCore -- because that drops the free count below the requirement. Verified across simulated host sizes: 2-GPU box, both idle free=2 proceed 2-GPU box, leftover EngineCore free=1 ABORT (still caught) 8-GPU node, all idle free=8 proceed 8-GPU node, one GPU busy free=7 proceed (was: ABORT) 8-GPU node, six busy free=2 proceed 8-GPU node, seven busy free=1 ABORT GPUS is now defined once and feeds both cluster.gpus_per_node and the check, so the two cannot drift apart. Re-ran the chaos test after the change: killed ray::VllmAsyncGenerationWorker in state idle, bounded failure at 222s -- identical to the two runs before it. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Two files imported the same module twice instead of merging the names,
which ruff's isort rules (I001) reject:
from nemo_rl.experience.rollout_manager import RolloutManager, RolloutTimeouts
from nemo_rl.experience.rollout_manager import (RolloutRetryPolicy,)
`ruff check` did not catch this locally, and neither did the SLURM run --
both reported "All checks passed!". The repo's .pre-commit-config.yaml
registers the ruff hook twice:
- id: ruff
args: ["--fix"]
- id: ruff
args: ["check", "--select", "I", "--fix"]
Only the second selects `I`, and `I` is not in the default selection that
a plain `ruff check` uses. So the rule that failed in CI was never being
run locally. Verified by reproducing it: `ruff check --select I` reports
exactly the two errors GitHub reported, and the fix it applies is
byte-identical to the diff in the CI log.
Both files are ones this branch introduced changes to; no unrelated files
were touched.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Merging main brought in #2518, which moves GRPOConfig from TypedDict to a pydantic BaseModel, so `master_config.grpo["key"]` becomes `master_config.grpo.key`. The merge was textually clean and left three sites broken. Two were visible: test_rollout_pump.py fed `grpo={"max_num_epochs": 1}` into _rollout_pump, which #2518 migrated to attribute access, giving "AttributeError: 'dict' object has no attribute 'max_num_epochs'" on both parametrisations. The third was not visible, and is the one that mattered. single_controller.py read `self._master_config.grpo["max_num_steps"]` while test_watchdog_pump.py passed a dict -- self-consistent, so the unit tests stayed green while production, which now receives a real pydantic config, would have raised at runtime in the watchdog. A functional test would have caught it; no unit test could, because the fake and the code agreed with each other and both disagreed with reality. Fixtures now build GRPOConfig.model_construct(...), matching how #2518 migrated upstream's own fixtures in the same file. 315 passed across tests/unit/single_controller and tests/unit/experience. Signed-off-by: asolergibert <asolergibert@nvidia.com>
GRPOConfig.model_construct(...) made two assignments exceed the line limit, so ruff-format wanted to wrap them. CI's ruff-format hook would have rejected this. It was reported as clean locally because the check itself was wrong: `ruff format --check ... | grep -o 'already formatted'` matches the tail of "2 files would be reformatted, 542 files already formatted". Checking the exit code instead of grepping the message. Signed-off-by: asolergibert <asolergibert@nvidia.com>
The chaos harness picked its victim by matching `ray::` process titles. That worked on a workstation and found ZERO generation actors on a GB200 cluster (job 5861743): the poll never matched, the run simply finished, and both chaos variants reported "job died before the kill" -- a green-looking failure that never killed anything. Process titles are a runtime implementation detail; the GCS actor table is the runtime's own record, so ask that instead. Titles are still used for the idle-vs-serving distinction, because Ray only exposes the running method through the dashboard state API and init_ray disables the dashboard. So: discover authoritatively, refine by title where titles exist, and say so plainly where they do not -- VICTIM_STATE=any then kills a generation actor without pinning the state, which still asserts a bounded attributable failure but stops separating the detection path from the in-flight-RPC path. Also unbuffers the driver. The harness detects progress by grepping the run log, and the driver's stdout is a redirected file, so Python block-buffered it: job 5892910 wrote "train step 3/24" at 10:40:38 and the harness did not see it until 10:48:21, by which time there was nothing left to kill. The actor prints with flush=True, but that only reaches the driver. The remaining changes are all about being able to diagnose the next failure: ACTOR_WAIT_S separates "no actors yet" from "actors never reached the state", which were previously indistinguishable in the log; ACTOR_QUERY_TIMEOUT_S bounds one discovery attempt; and the helper's stderr is printed inline rather than to a side file, because two attempts at a side file never produced one that survived to be read. Developed while getting the recovery tests to run on the cluster, and belongs here: none of it depends on anything the later branches add, and without it this branch's chaos test cannot pass on that hardware at all. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
…01-containment # Conflicts: # nemo_rl/environments/nemo_gym.py
The main sync brought upstream's test_logs_setup_timing_metrics, which builds
actor_args as a SimpleNamespace listing only the fields upstream's own __init__
reads. env_handles is not among them -- upstream declares it on
SingleControllerActorArgs but never touches it in __init__ -- so our
environment-health-check read turned a field nobody else needs into a
construction requirement, and the test died with AttributeError.
getattr with a {} default. Every production caller goes through the dataclass and
still supplies it, so behaviour is unchanged; only test fixtures that do not
exercise environment health are let off.
Identical in shape to the `timeouts` regression from the previous main sync: a
required parameter we added, discovered by an upstream test constructing the
object directly. Worth remembering as the recurring cost of reading new fields
in __init__ rather than where they are used.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Both SC rollout paths pick a generation shard by static round-robin with no idea
whether the shard is alive -- the native path in
VllmGeneration._async_generate_base, and the NeMo-Gym path inside Gym's own
_resolve_client. This adds the missing half: which shards are eligible to serve,
and why. Nothing consumes it yet.
GenerationFleetMonitor is deliberately a pure state machine. Probing, restarting
and pushing membership are I/O and belong to the caller, which keeps every
transition testable without Ray, a network or a GPU, and keeps one description of
"eligible" that both routing adapters will read.
The transition that carries the most weight is the one that does not exist: a
shard cannot go from DEAD back to HEALTHY on its own. A restarted engine holds
whatever weights it loaded at init, so re-admitting it because it started
answering probes again would feed training rollouts generated from a checkpoint
hundreds of steps stale -- invisible, and worse than the outage that caused it.
Recovery must pass through STALE and a completed refit, and DEAD/RESTARTING/STALE
all ignore successful probes to make that hard to get wrong by accident.
Other decisions worth knowing:
- SUSPECT still serves traffic. Draining on a single failed probe would make a
transient blip cost a shard's worth of throughput; only unhealthy_threshold
consecutive failures condemn a shard.
- The membership epoch advances only when the *serving set* changes, so
HEALTHY -> SUSPECT does not disturb it. Downstream reconciliation can then be
an integer comparison in the common case.
- Retirement is terminal and bounded by max_restart_attempts_per_shard, with
min_healthy_shards as the floor below which the run stops being worth
continuing.
- HealthyShardSelector uses least-outstanding rather than round-robin: it
steers away from a shard that is merely slow or wedged without that having to
be diagnosed first.
async_rl.fleet_health declares only the knobs P1 consumes. on_dead_shard is a
Literal accepting just "fail_fast" so the recovery modes that need the
communicator rebuild are rejected rather than silently doing nothing.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
_async_generate_base advanced a counter modulo dp_size and sent the request wherever it landed, with no idea whether that shard was alive. A dead shard therefore kept receiving its full 1/N share of traffic for the rest of the run, and each request rediscovered the same corpse. Selection now goes through HealthyShardSelector when fleet health is enabled: a quarantined shard is skipped, and least-outstanding steers away from one that is merely slow without that having to be diagnosed first. With no selector attached the historical round-robin is reproduced exactly, so an unconfigured run is unchanged. The other half is reporting. A dead worker surfaces as ray.exceptions.RayError, which is now fed to the monitor before being re-raised as GenerationUnavailable. Reporting is what lets the *next* request skip the shard; the retype is what tells the rollout retry policy the prompt is fine and worth re-dispatching elsewhere. Together they close the loop that P0.4 could only half-close -- re-dispatch already existed, but nothing steered it away from the failure. The generation body moved into _generate_on_shard so the acquire/release of the in-flight count can sit in a finally around the whole stream. A leaked count would permanently bias selection away from a shard that is actually fine. setup_single_controller builds the monitor when async_rl.fleet_health.enabled and hands it to the SingleController through actor args; a backend without attach_fleet_health raises rather than silently ignoring the request. Nothing drives the probe loop yet, so detection is currently failure-driven only. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
…reeze gap
Wires the P1.1 monitor into the SingleController and adds the external observer
the in-actor watchdog cannot be.
The watchdog tick now probes every serving shard for Ray actor liveness, folds
the result into the monitor, publishes fleet state, and raises
GenerationFleetExhausted once too few shards remain to be worth continuing. Only
serving shards are probed: a quarantined shard answering again says nothing about
whether its weights are current, and the monitor ignores such probes anyway. Ray
liveness is the cheap authoritative signal for "the process is gone" but misses a
vLLM engine core dying under a live worker, which is why the routing adapters
also report what they observe -- both feed the same counters.
Two loop-freeze fixes, which is the residual risk once the monitor lives on the
SC actor:
- invalidate_kv_cache moves to asyncio.to_thread. It was the one call into the
workers made directly on the event loop; run there against a wedged worker it
would freeze the loop itself, taking the watchdog -- an asyncio task on that
same loop -- down with it.
- The driver polls ping() around the run. The in-actor watchdog cannot observe
its own loop being blocked, and the driver is already a separate process
holding the handle, which makes it the cheapest possible external observer.
ping() has existed since the SC landed and had no caller until now.
Environment teardown is also bounded, with a ray.kill fallback, matching what the
legacy GRPO path already does. It runs in a finally block, so a hung shutdown
would otherwise replace a real training error with an indefinite wait.
Validated on 2xA6000 by the chaos harness with fleet health enabled: the monitor
takes shard 0 healthy -> suspect -> dead and the job stops 10s after the kill
with both GPUs released. Two caveats worth stating rather than implying
otherwise. The threshold was crossed largely by adapter-reported failures rather
than by probes, since the rollout path reaches the dead shard before the third
probe lands. And with 1 training + 1 inference GPU dp_size is 1, so there is
nowhere to fail over to -- what P1 actually buys, traffic moving to a surviving
shard, is covered by unit tests only and needs >=3 GPUs to show end to end.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
NeMo-Gym picks a policy endpoint by static round-robin over a list fixed at
process start, never fails over, and retries a refused connection in an uncapped
loop with no HTTP timeout. A dead vLLM endpoint therefore keeps receiving ~1/N of
new rollouts for the rest of the run.
Rather than change Gym, hand it a single URL NeMo-RL owns. Gym's
VLLMModelConfig.base_url accepts one string, so its round-robin becomes a no-op
and the routing decision moves next to the fleet health that already knows which
shards serve. The switch is one expression in setup_single_controller.
Three decisions carry this:
- The URL never changes. The port is reserved once and passed in, so Ray
recreating a restarted actor rebinds the same address. Gym is never
reconfigured and never has to fail over, which matters because failing over
is exactly what it cannot do. If the router picked a fresh free port on
restart -- the way everything else here allocates ports -- Gym would hold a
dead URL forever, since it never re-resolves.
- Every piece of state is built in __init__, including the server thread. A
restarted actor is immediately usable. This is the deliberate inverse of the
NemoGym mistake, where servers were started from a _spinup that Ray never
re-runs.
- The no-healthy-backend status must stay outside Gym's retry set. Gym retries
429/500/502/503/504/520, and for the rate-limit subset it raises its own
retry ceiling per attempt, so answering with one of those would spin forever
-- recreating the hang the router exists to prevent. It defaults to 409 and
the config validator rejects the retried codes outright rather than leaving
it to be discovered in production.
Not a redirect, though aiohttp does follow 307 with the body intact. A redirect
puts Gym's socket back on a vLLM endpoint directly, so a backend dying
mid-request drops it into that same uncapped retry loop.
Membership is pushed from the SingleController watchdog as the full serving set
whenever the fleet's membership epoch moves. Full sets rather than deltas mean a
dropped, reordered or post-restart update converges on the next tick without
sequence numbers or replay. A restarted router comes up believing every backend
serves, which is self-correcting and strictly better than serving nothing.
Bodies stream through in both directions rather than being buffered; a completion
carrying per-token logprobs is large and this sits on every rollout's critical
path.
Tested against real aiohttp servers rather than fakes -- header handling,
streaming and status propagation only misbehave over an actual socket -- plus a
live Ray-actor check confirming the proxy, the 409 drain path and the metrics.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: asolergibert <asolergibert@nvidia.com>
policy_router.enabled was enabled by no test anywhere. The branch's other three features ride on the chaos test, whose branch-2 version turns on fleet_health.enabled, but the router did not -- so a regression in the proxy that fronts every NeMo-Gym rollout would have shipped silently. Registers the existing Gym functional test a second time with the router and fleet health on. gen_kl_error is what earns its keep here: it compares vLLM's logprobs against the trainer's recomputation, so a proxy that corrupts or truncates a response blows it up. A run that merely completes would not prove the payload survived the extra hop. Verified on 2xA6000 before registering: exit 0, reward 0.5, Gym's only endpoint the router on :6081 while vLLM served on :3038, and the resolved config showing PolicyRouterConfig(enabled=True) -- an override for a field that does not exist is accepted silently, and the test would then pass while exercising nothing. Full mode only; it is a second ~20 minute Gym run. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Same rule as the fix on the containment branch: ruff's `I` rules are only selected by the second ruff hook in .pre-commit-config.yaml, not by a plain `ruff check`, so these were invisible locally and failed in CI. Three orderings: third-party `ray` before `torchdata`/`transformers`, `policy_router` after `interfaces`, and a stray blank line splitting an import block. Signed-off-by: asolergibert <asolergibert@nvidia.com>
probe_interval_s was decorative. The probe ran inside _watchdog_pump, so its real cadence was watchdog.interval_s (30s), and nothing anywhere read the configured value -- it existed only to be validated against probe_timeout_s. That put dead-shard detection at unhealthy_threshold * 30s = 60-90s, against a config comment promising ~15s. Worse, it is longer than refit_timeout_s: a refit hung on a dead rank always aborted before the monitor had decided which rank was gone, so the rebuild the abort exists to trigger saw an empty absent set and did nothing, and the retry died on a communicator that had been aborted and never rebuilt. Arithmetic rather than a race -- it could not have worked. Job 5925668. The probe now runs in its own pump at probe_interval_s, created only when fleet health is enabled so a default run gains no timer. Shards are probed concurrently: sequentially a round costs probe_timeout_s per shard, so any fleet larger than probe_interval_s / probe_timeout_s could not finish a round inside its own interval, and the config validator compares only those two -- silently assuming one probe per tick. The router membership push moves with it, so a membership change reaches the router at detection speed rather than on the watchdog's clock. It is epoch-gated, so an unchanged serving set still costs nothing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
chaos-serving failed in GitHub CI without running a line of product code: 19:51:42.05 [chaos] PASS (chaos-idle) 19:51:43.40 chaos-idle's bash exits, cleanup trap done 19:51:43.56 [chaos] FAIL: need 2 free GPUs, found 1 19:51:43.58 0, 50479 MiB <- still held, by a live pid Three things compound. cleanup() signals and returns, but SIGKILL does not free device memory synchronously -- a worker holding ~50GB stays in nvidia-smi for seconds while the driver tears its context down. run_test is just `time "$@"`, so the next test starts with no gap. And the pre-flight was a single sample with no retry, on a 2-GPU runner where one leaked GPU is the whole margin. It passed on the cluster only because submit-ci.sh sleeps 5s between tests. That hygiene belongs in the test that makes the mess, not in one driver: cleanup() now waits for the memory to come back, and the pre-flight waits rather than sampling, so it is robust to anything else on the box too. Also moves the pre-flight above the training launch. It used to run after, so a refusal was not a refusal -- the driver was already up and the EXIT trap shot it down, which is where the stray "line 129: <pid> Killed" in the CI log came from. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: asolergibert <asolergibert@nvidia.com>
…ncy-02-fleet-health-router
…ncy-02-fleet-health-router
…01-containment Signed-off-by: asolergibert <asolergibert@nvidia.com> # Conflicts: # nemo_rl/algorithms/single_controller.py # nemo_rl/environments/nemo_gym.py
…ncy-02-fleet-health-router Signed-off-by: asolergibert <asolergibert@nvidia.com>
…02-fleet-health-router Signed-off-by: asolergibert <asolergibert@nvidia.com> # Conflicts: # nemo_rl/algorithms/single_controller.py # nemo_rl/algorithms/single_controller_utils/config.py # nemo_rl/algorithms/single_controller_utils/setup.py # tests/functional/L1_Functional_Tests_SingleController.sh # tests/functional/grpo_dp_single_controller_chaos.sh # tests/unit/single_controller/test_resiliency_config.py # tests/unit/single_controller/test_single_controller.py # tests/unit/single_controller/test_watchdog_pump.py
test_fleet_health.py and test_policy_router.py -- 50 tests -- were collected by zero shards: Models_1..4 pass --ignore=unit/models/generation/, Other passes --ignore=unit/models/, the Vllm base runs glob only test_vllm*.py plus test_openai_server_utils.py by name, and every --*-only catch-all keeps marked items only. Named in the Vllm base runs rather than marked @pytest.mark.vllm: the import chain is vLLM-free and all 50 pass in a CPU-only environment, so the absent marks are correct. Verified by collection count -- the base run goes 387 -> 437. Signed-off-by: asolergibert <asolergibert@nvidia.com>
…ohttp choose Four interlocking holes on the failure path, all of which had to close together. 1. The router had no exception handling, so aiohttp answered for it -- and a wedged backend produces 504. 504 is in NeMo-Gym's rate-limit retry subset, where _request_with_retry raises its own ceiling on every attempt: an unbounded retry loop at backend_timeout_s per turn. That is exactly the hang _check_status_is_not_retried_by_gym exists to prevent, arriving through the error path that validator never covered. The router now answers a deliberate 500 -- Gym's *bounded* retry set -- so Gym re-sends the one call, the next pick lands on a healthy shard, and a multi-turn rollout keeps its turns. 2. Least-outstanding then preferred the corpse: _handle's finally returns a fast-failing backend to inflight=0, so it was selected for every subsequent request until a probe caught up -- strictly worse than the round-robin the router replaced, in the window it was built for. The failing backend is now dropped locally until the next membership push restores it. The drop is armed only when fleet health is driving membership. Without a monitor nothing ever pushes, so the drop would be permanent and a few transient blips would drain the fleet with no way back. 3. A wedged engine could never be condemned. report_failure delegated to record_probe(ok=False), sharing the streak that record_probe(ok=True) zeroes -- and a wedged vLLM answers is_alive from a healthy worker process. Reported failures now have their own streak that only a successful generation, or a refit, clears. The router counts failures per backend and the probe tick drains them in, so the ledger learns what only a real request can reveal. 4. A generation timeout raised a bare RuntimeError from inside the try, so the RayError handler that reports to the ledger never saw it -- the one failure mode the fleet-health docstrings cite to justify reactive reporting. It is now reported and typed as GenerationUnavailable (INFRA), not DATA. Also drops the membership-push epoch gate, which made a router restart unrecoverable: a recreated actor rebuilds its serving set as every backend while the epoch has not moved, so the gate blocked the corrective push forever and Gym routed to a quarantined shard for the rest of the run. Both docstrings claimed the opposite. The push is a short list of strings on a probe-interval timer. The push and drain are wrapped: run() awaits the probe task and re-raises, so an unguarded RayActorError from the max_restarts=-1 router being recreated ended training over something the next tick would have healed. _push_router_membership had zero test coverage; it now has ten tests, including the restarted-router regression. Signed-off-by: asolergibert <asolergibert@nvidia.com>
… nothing Every item here is a setting that parsed cleanly and then had no effect, or a failure that surfaced far from its cause. - policy_router.enabled=true on a native run was ignored outright: _maybe_start_policy_router is only reached inside the use_nemo_gym branch. Now rejected, matching how fleet_health treats an unsupported backend. - policy_router.enabled with fleet_health disabled is legitimate -- the stable URL, the backend deadline and least-outstanding all still apply -- but nothing ever pushes a serving set, so it cannot fail over. Warns rather than rejects. - A transposed port range surfaced as "empty range for randrange()" from inside port allocation. Validated where the typo is. - connect_timeout_s is now its own knob, defaulting to 5s. It shared backend_timeout_s (600s), so a black-holed SYN parked a rollout for ten minutes on a handshake that is milliseconds-or-never. - fleet_health.selection accepted "round_robin" and silently gave the caller least_outstanding, because nothing dispatches on it. Narrowed to a Literal of one, like on_dead_shard directly below it. - fleet_health with async_engine=false failed as "base_urls has 1 entries for N shards", which reads like an internal bug. An all-None URL list is now dropped: health tracking needs no URLs, only the router push does. - attach_fleet_health is declared on GenerationInterface rather than discovered with hasattr, so an unsupported backend names itself. Same shape as the refit hooks alongside it. - Neither block appeared in the exemplar YAML, leaving ten knobs discoverable only from Python source. Added under the Resiliency banner P1 established, with values verified equal to the model defaults. Signed-off-by: asolergibert <asolergibert@nvidia.com>
The system now has three observers and, until this commit, only one of them said what it observed. Renames only -- no behaviour changes -- and they land now because a config key is free to rename before release and breaking after. - PolicyRouter -> GenerationRouter (module, config key async_rl.generation_router). "Policy" in NeMo-RL means the trainer, under models/policy/; this lives in models/generation/ and routes to the generation fleet. The name came from Gym's vocabulary, where it fronts policy_base_url, but inside this repo it pointed at the wrong subsystem. - async_rl.fleet_health -> async_rl.generation_fleet_health, self._fleet_monitor -> self._gen_fleet, _fleet_probe_pump -> _gen_fleet_probe_pump, and the fleet/* metric prefix -> gen_fleet/*. "Fleet" means only the generation DP shards, but async_rl also governs trainer, env and rollout knobs, and the bare metric prefix lands in dashboards with no type context. - GenerationFleetMonitor -> GenerationFleetHealth. It holds no Ray handle, no network, not even its own clock -- it is structurally incapable of monitoring, by design. It is the fleet's health model: fed observations, keeps the verdicts. The name now matches its module and its config key. - async_rl.watchdog -> async_rl.stall_watchdog, _watchdog_pump -> _stall_watchdog_pump. It watches four things now; the rename names its namesake duty and leaves the docstring to cover what rides along. - _run_with_liveness_watch -> _run_with_controller_liveness_watch. It watches the SingleController's own event loop -- the one thing the in-actor loops structurally cannot observe. Deliberately left alone: VllmGeneration.fleet_monitor/fleet_selector, GenerationFleetExhausted, FleetHealthConfig/FleetHealthPolicy and fleet_health.py itself. Context already scopes all of them. async_rl.watchdog shipped in the containment PR, so AsyncRLConfig now rejects all three old block names rather than letting extra="allow" swallow them -- a config carrying watchdog: would otherwise lose stall detection silently. Signed-off-by: asolergibert <asolergibert@nvidia.com>
grpo_async_gym_single_controller.sh runs gpus_per_node=1 with tensor_parallel_size=1, so dp_size=1. With one backend _pick_backend has one choice, set_serving_backends never sees a shrinking set, and the no-healthy-backend path never fires -- the lane proves pass-through and nothing else. "NeMo-Gym never has to fail over" is the entire reason the router exists, and it had no functional coverage at all, which is how the error-path defects in _handle could have shipped. Adds the inverse of the chaos harness. That one runs a single-shard fleet and asserts a bounded FAILURE, because there is nowhere to fail over to. This runs two generation shards, kills one mid-run, and asserts SURVIVAL: the run finishes on the survivor, a shard is quarantined in the log, and gen_kl_error stays flat so the router demonstrably did not corrupt a payload while re-routing. Needs >= 3 GPUs (2 generation + 1 trainer) and self-skips below that rather than passing vacuously -- a green tick on a 2-GPU runner would be worse than no test. Verified the skip path on a 2-GPU box; the failover path needs the cluster. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 6251221 reported "the run died after losing one of two shards" -- the exact message for the property under test failing -- and it was misleading. The run had executed zero training steps: it died inside spinup_nemo_gym_actor with "Process `policy_model` finished unexpectedly!", and no shard was ever quarantined. The cause was the harness. Generation actors exist as soon as the inference cluster is built, but NeMo-Gym's _spinup runs after that and takes minutes. A 180s timer started at actor discovery therefore fired while the run was still in SETUP; the kill took down the vLLM server Gym's policy_model process depends on, Gym's own poll() failed the run, and the harness attributed that to failover. Waits for `train step N/` in the run log instead, which is only printed once the train pump has completed a step -- i.e. rollouts really are flowing through the router. Reaching steady state is a precondition of this test, not what it measures, so failing to reach it now reports itself as such rather than as a failover failure. Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 6251951 got further and still proved nothing: it reached the step gate with the run ALREADY at 40/40, so the kill hit a shutting-down job. The "no shard was quarantined" guard caught it rather than reporting a false pass, which is the one part that worked. Two causes. Actor discovery ran before the step gate and cost more wall-clock than the whole run. Its first ray.init resolved a stale GCS address and spent two minutes timing out before retrying against the live one and succeeding instantly. Discovery now runs after the step gate, where the actors certainly exist -- it is not on the critical path to the kill, and nothing about it needs to be. And the run was simply too short: 40 steps at ~4s/step is under three minutes end to end, so any discovery cost at all consumed it. max_num_steps is now 150, which leaves real work after the kill. The runway check is re-read at the kill site rather than at the step gate, since discovery sits between the two, and it reports itself as a precondition failure naming the step counts -- not as a failover failure. Signed-off-by: asolergibert <asolergibert@nvidia.com>
… deliver
Job 6258553 ran the test as intended for the first time and produced a real
result: the kill landed in steady state, the ledger took the shard
healthy -> suspect -> dead on a genuine ActorDiedError, and then the run wedged.
WARNING: rollout stall -- no rollout committed and no train step completed
in 2018s (0 rollouts in flight, 12 groups committed, step 3/150)
That is the P0.6 signature exactly -- nothing in flight, committed frozen, step
frozen -- because _sync_weights parks the rollout pump behind _rollout_permitted
and never returns: the refit broadcast still addresses the rank that just died.
Rebuilding the communicator over the survivors is the elastic-recovery PR's job,
not this one's. The test was asserting a property this code does not implement.
Split by EXPECT. The default, quarantine, asserts what this part actually
delivers: the shard is condemned and router/serving_backends drops to 1, so
NeMo-Gym stops being handed the corpse. It then stops the run deliberately rather
than waiting for a completion that cannot come. EXPECT=survival keeps the
completion assertion for the part of the stack that can satisfy it.
Both assert gen_kl_error, which is what shows the extra hop did not corrupt a
payload, and both now assert on router/serving_backends rather than on a log
line -- that counter is what routing actually reads.
Signed-off-by: asolergibert <asolergibert@nvidia.com>
Job 6262310 reached the step gate and saw 0 training steps for 26 minutes, then 150 at once -- so the kill landed on a finished run and the runway guard failed it, correctly, at "step 150 of 150". The run was fine. The harness could not see it. Ray forwards actor output to the DRIVER, and the driver's stdout here is a redirected file, so Python block-buffers it: nothing reaches RUN_LOG until a buffer fills or the process exits. grpo_dp_single_controller_chaos.sh and grpo_sc_generation_shard_recovery.sh both already set PYTHONUNBUFFERED=1, with a comment naming this exact failure and the job that found it (5892910). This harness greps RUN_LOG the same way and was written without it. Signed-off-by: asolergibert <asolergibert@nvidia.com>
In job 6262310 the step gate opened at 150/150 and the next thing to run was actor discovery, whose own liveness check failed first with 'job died before any actor appeared'. True, and it points at Ray rather than at the clock. The runway check now also runs before discovery, where it can say what actually happened. Signed-off-by: asolergibert <asolergibert@nvidia.com>
…s to backend_timeout_s bounds one HTTP call; rollout_failure.nemo_gym.rollout_timeout_s bounds the whole prompt-group stream, which is many of them. Set the inner one larger and it can never fire -- the rollout deadline always expires first, so the timeout the router exists to add is dead config, and the failure surfaces as a group timeout when the router could have named the backend. Raised in review alongside the connect-timeout split; that half landed, this one did not. Signed-off-by: asolergibert <asolergibert@nvidia.com>
|
Thanks — this was a genuinely useful review. Every item is addressed in the PR now; each thread has the detail. Summarising what moved, since the shape of the change is larger than the individual threads suggest. The router's error path was the one real defect, and it was three problems wearing a coat. The functional gap you named turned into the PR's main new evidence. The shipped Gym lane is On the scope note ( One thing changed under a recommendation since you wrote it. You suggested renaming And the finding that made everything else verifiable: both new test files were collected by zero CI shards. 50 tests, none of them running. The Vllm base run goes 387 → 437 with them named. |
|
/ok to test 5c9b813 |
…03-elastic-recovery Part 2 (#3590) landed as a squash, so every file it touched conflicts with itself here. Verified all 17 are byte-identical between the Part 2 branch tip and main, which makes "ours" (that same content plus this branch's delta) the provably correct side. pyrefly.toml was the one real conflict: three commits each appended an allow-list entry, all three kept. The substantive change comes from the colocated Megatron reshard (#3490): mark_stale() is gone from WeightSynchronizer and every implementation, with per-step staleness now owned by the training loop. Nothing on this branch called it, so the removal merged cleanly; the one leftover was a fake in test_reconcile_communicator that still defined it, now dropped so the fake matches the interface it stands in for.
Note
Mirror of #3471, originally opened by @asolergi-nv from the
asolergi-nv/RLfork.The branch was re-pushed to
NVIDIA-NeMo/RLunchanged (same head SHA) so the four partscan be linked as a proper PR stack —
gh stackcannot operate on fork branches.Please review and comment here. The original #3471 is left open for history.
Part 2/4 of #3454
What this fixes
Containment (Part 1/4, #3589) stops a dead engine from wedging or corrupting a run, but the fleet still has no idea which engines are alive, so both rollout paths keep sending work to a dead one:
_resolve_clientnever re-resolves. A dead endpoint keeps receiving roughly 1/N of new rollouts for the rest of the run.What it does
A generation-fleet health model (
fleet_health.py) — a pure state machine with I/O injected, so it is testable without Ray or GPUs:SUSPECTexists so a single failed probe does not cost a shard's throughput.DEAD → HEALTHYis deliberately unreachable: only a completed refit returns a shard to service, because an engine that answers/healthsays nothing about whether its weights are current.The class is
GenerationFleetHealth— it holds no Ray handle, no network and not even its own clock, so it cannot monitor anything. It is the fleet's health model: two observers write into it (the probe loop, and the generate path reporting failures), three readers consult it (the shard picker, the router membership push, the stall watchdog).GRPO shard selection now picks among healthy shards (least-outstanding), and generation failures are reported back — the routing adapters see failures a liveness probe cannot, such as a shard that answers
is_aliveand errors on every generation.A NeMo-RL-owned router for NeMo-Gym (
GenerationRouter). Rather than change Gym, hand it a single URL that we own. Gym'sbase_urltakes one string, so its round-robin becomes a no-op and the routing decision moves next to the health data. Zero NeMo-Gym changes.Three decisions:
The router's URL never changes. The port is reserved once and passed in, so Ray recreating a restarted actor rebinds the same address. This matters because Gym never re-resolves: if the router allocated a fresh free port on restart — the way everything else in this codebase allocates ports — Gym would hold a dead URL forever. There is a test pinning it.
All router state is built in
__init__, including the listening socket. The deliberate inverse of theNemoGym._spinupshape, where the servers are started from a separate method Ray never re-runs on restart — leaving a live actor that cannot serve. The socket is bound synchronously so a port conflict fails actor construction with the port in the traceback, rather than killing a daemon thread whilebase_url()keeps handing Gym an address nobody listens on.The "no healthy backend" status is load-bearing and validated at config load. Gym retries 429/500/502/503/504/520 and raises its own retry ceiling on the rate-limit subset, so answering with a 503 would spin forever and recreate the exact hang this removes. The default is 409, and the config rejects anything in Gym's retry set rather than leaving it to be discovered in production.
What changed after review
The review surfaced one correctness cluster and a set of smaller fail-loud and naming items. Everything below is in this PR now; each inline thread has the detail.
The error path was the serious one.
_handlehad no exception handling, so aiohttp answered for it — and a wedged backend produces 504, which is in Gym's rate-limit retry subset where its ceiling rises on every attempt. That is an unbounded retry loop atbackend_timeout_sper turn: precisely the hang_check_status_is_not_retried_by_gymexists to prevent, arriving through the error path that validator never covered. It now answers a deliberate 500 (Gym's bounded set, so Gym re-sends one call and a multi-turn rollout keeps its completed turns), drops the failing backend locally until the next membership push, and counts the failure so the controller's tick can report it into the ledger.Three things had to close with it:
_handle'sfinallyreturned a fast-failing backend toinflight=0, so it was selected for every subsequent request — worse than the round-robin the router replaced, in the window it was built for.report_failuredelegated torecord_probe, sharing the streak that a successful probe zeroes — and a wedged vLLM keeps answeringis_alive. Reported failures now have their own streak that only a successful generation, or a refit, clears.asyncio.TimeoutErrorbecame a bareRuntimeErrorraised inside thetry, so the handler that reports never saw it. Now reported and typedGenerationUnavailable(INFRA) rather than classifying as DATA.The membership-push epoch gate is gone. It made a router restart unrecoverable: a recreated actor rebuilds its serving set as every backend while the epoch has not moved, so the gate blocked the one push that would have corrected it.
_push_router_membershipalso had zero coverage; it now has ten tests, including that regression.Naming. The system has three observers and only one said what it watched.
PolicyRouter→GenerationRouter("policy" means the trainer in this repo);async_rl.fleet_health→generation_fleet_healthand thefleet/*metrics →gen_fleet/*;async_rl.watchdog→stall_watchdog;GenerationFleetMonitor→GenerationFleetHealth.async_rl.watchdogshipped in Part 1, soAsyncRLConfignow rejects all three old block names rather than lettingextra="allow"swallow them and silently drop stall detection.Tests
DEAD → HEALTHYis unreachable, that DEAD/RESTARTING/STALE ignore successful probes, and that a shard answering probes while failing every generation is still condemned./tokenize(not under/v1, because Gym'screate_tokenizestrips the suffix), a 512 KB streamed response, status propagation, and the full error path: a wedged backend is answered 500 and never 504, the failing backend leaves the serving set, and a port conflict fails construction.Models_*ignoreunit/models/generation/,Otherignoresunit/models/, and the Vllm shards glob onlytest_vllm*.py— so all of this coverage was dead. Verified by collection count: the Vllm base run goes 387 → 437.Functional
grpo_async_gym_single_controller.shwith the router enabled, registered inL1_Functional_Tests_SingleController.sh. Gym's only endpoint is the router while vLLM serves elsewhere, andgen_kl_error— which compares vLLM's logprobs against the trainer's recomputation — does not move, so the proxy is not corrupting or truncating responses.Chaos, single-shard fleet. Same scenario and box as Part 1: kill the only generation shard mid-run. Part 1 fails in 222s; this PR fails in 62s, and the log says why:
Part 1 has no health model, so its only backstop is the stall detector: it reports
RolloutStallonce the rollout has been quiet long enough — true, but that names the symptom, and the timeout has to be generous enough not to fire on a slow-but-healthy rollout. This probes the shard directly and trips themin_healthy_shardsfloor, naming the cause. Both are bounded and attributable; this one is ~4x faster and more actionable. Both numbers are from runs with the victim state pinned toidle, so this is reproducible rather than an artifact of which state the kill happened to catch.Failover, two-shard fleet — new, and the property this PR exists for. The Gym lane above cannot prove it: it runs
gpus_per_node=1, tensor_parallel_size=1, sodp_size=1. With one backend_pick_backendhas one choice,set_serving_backendsnever sees a shrinking set, and the no-healthy-backend path never fires — it demonstrates pass-through and nothing else.grpo_sc_gym_router_failover.shruns two generation shards, kills one mid-run, and asserts the fleet routes around it. Needs ≥3 GPUs (2 generation + 1 trainer) and self-skips below that rather than passing vacuously.Result on 3×GB200 (oci-hsg, job 6263801, 13m34s, exit 0):
median(train/gen_kl_error) < 1.3min(router/serving_backends) == 1max(gen_fleet/shards/dead) >= 1router/serving_backendsdropping 2 → 1 is the load-bearing one: it is the router's own counter, the one_pick_backendreads, so it shows NeMo-Gym genuinely stopped being handed the dead backend rather than that a log line was printed.gen_kl_errorat 0.00128 sits with the healthy baselines (0.00106–0.00118), so re-routing mid-run did not corrupt a payload.What it deliberately does not assert. The run does not survive the loss on this PR, and the test does not pretend otherwise — it stops the run after verifying quarantine. The next weight refit still broadcasts to the rank that died and hangs inside NCCL, so the rollout pump stays parked and the stall watchdog warns (observed directly: job 6258553 sat there for 33 minutes). Rebuilding the communicator over the survivors is feat(sc): keep training when a generation shard dies #3591's job, and the same harness asserts survival there via
EXPECT=survival. This is the P0/P1 boundary the design doc already states: detection and routing here, recovery next.Issues
List issues that this PR closes (syntax):
Usage
Before your PR is "Ready for review"
Pre checks:
Additional Information