feat(sc): track generation fleet health and route NeMo-Gym through NeMo-RL - #3471
Open
asolergi-nv wants to merge 22 commits into
Open
feat(sc): track generation fleet health and route NeMo-Gym through NeMo-RL#3471asolergi-nv wants to merge 22 commits into
asolergi-nv wants to merge 22 commits into
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 NVIDIA-NeMo#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>
asolergi-nv
changed the title
4 tasks
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>
asolergi-nv
force-pushed
the
feat/sc-resiliency-02-fleet-health-router
branch
from
August 4, 2026 15:02
1f1676e to
3fe0d11
Compare
Merging main brought in NVIDIA-NeMo#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 NVIDIA-NeMo#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 NVIDIA-NeMo#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>
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>
asolergi-nv
force-pushed
the
feat/sc-resiliency-02-fleet-health-router
branch
from
August 4, 2026 18:29
3fe0d11 to
7a3715d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part 2/4 of #3454
What this fixes
Containment (Part 1/4, #3470) 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 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.GRPO shard selection now picks among healthy shards (least-outstanding), and generation failures are reported back to the monitor — the routing adapters see failures a liveness probe cannot, such as a shard that answers
/healthand errors on every generation.A NeMo-RL-owned router for NeMo-Gym. Rather than change Gym, hand it a single URL that we own. Gym's
base_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 server thread. 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 "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.
Tests
Fleet-health state machine tests, including that
DEAD → HEALTHYis unreachable and that DEAD/RESTARTING/STALE ignore successful probes.19 router tests run against real aiohttp servers, not fakes — a proxy is precisely the component fakes flatter, since header handling, streaming and status propagation only misbehave over a real socket. Covers all four endpoints Gym calls, including
/tokenize, which is not under/v1because Gym'screate_tokenizestrips the suffix; plus a 512 KB streamed response and status propagation.8 config-validation tests, including that every status in Gym's retry set is rejected.
End-to-end:
grpo_async_gym_single_controller.shwith the router enabled. Gym's only endpoint was the router (:6081) while vLLM served on:3038, and the run completed. The strongest signal isgen_kl_error, which compares vLLM's logprobs against the trainer's recomputation — a proxy corrupting or truncating a response blows it up. It did not move. Now registered inL1_Functional_Tests_SingleController.sh(full mode): the router previously had no functional coverage at all, so a regression in the proxy fronting every Gym rollout would have shipped silently.The chaos test shows the health layer earning its keep. Same scenario and box as PR 1: kill the only generation shard mid-run. PR 1 fails in 222s; this PR fails in 62s, and the log says why:
feat(sc): contain rollout failures in the SingleController path #3470 has no monitor, 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. The monitor here probes the shard directly and trips themin_healthy_shardsfloor, reportingGenerationFleetExhausted, which names 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 to
idle, so this is reproducible rather than an artifact of which state the kill happened to catch. Ray renames the process title during a call, and a worker killed mid-generate_asyncfails much faster by a different route; feat(sc): contain rollout failures in the SingleController path #3470 makes that a separate registered variant instead of a coin flip.Full suite on 2xA6000: 644 unit passed / 3 skipped, plus
grpo_dp_single_controller.sh, both Gym variants, and the chaos test. Gym with defaults confirms inertness —PolicyRouterConfig(enabled=False)and zero router startups.Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information