feat(sglang): engine fault tolerance for rollouts

[Kh4L]

Original work: #3187 by @xiuhu17. Re-cut at his request so it can be carried through review on his behalf. His commit is preserved with its original authorship and Signed-off-by.

1 of 3 remaining in the stack. this → #3614 (mxfp8) → #3615 (nvfp4). #3612 has merged, and this branch is rebased directly onto the resulting main. Review the five commits on top of main.

What it adds over #3612

Engine fault tolerance for rollouts: a health monitor that detects a dead SGLang engine, restarts it, and rejoins it to the refit weight-update group, plus a _simulate_crash hook so the path can be exercised deliberately. use_fault_tolerance gates it and defaults off.

The follow-up fixes keep refit armed across recovery, align the configuration with its reader, make stop() safe while the monitor thread is live, and remove the obsolete disable_piecewise_cuda_graph server argument.

Verification

Full CI is green on exact head 0ae89f8b1f8484ee82cebf6613bb10342d08be9a, including unit, functional, H100, and GB200 jobs.

HSG job 6556006 passed the real GPU fault-tolerance path on that exact head: it launched SGLang, deliberately crashed an engine, recovered it, and completed post-recovery generation. Decode CUDA graphs remained enabled; only prefill capture was disabled. Result: PR3613_FAULT_TOLERANCE_REAL_GRAPH_OK job=6556006 commit=0ae89f8b1f8484ee82cebf6613bb10342d08be9a.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Kh4L]

Self-review (/review-pr-team), per the contributor self-review policy

Ran the four reviewer roles the skill specifies — rl-expert, bug-finder, design-reviewer, test-agent — as a real agent team this time (CLAUDE_CODE_EXPERIMENTAL_AGENT_TEAMS=1), so the mandatory devil's-advocate challenge could interrogate the reporting agent across both rounds instead of re-deriving findings alone. Scoped to this PR's increment only, not the refit layer below it.

Three blockers, all reproduced by executing the real source rather than by reading it, all fixed in 5314f53 and be50156. Every one of them is invisible today, because all five shipped sglang recipes set use_fault_tolerance: false and nothing in CI turns it on. The feature as submitted could not be enabled, and would have been broken once enabled.

1. _recover disarmed the weight refit on every refit — silently, on the colocated path

_recover ran unconditionally, including the overwhelmingly common case where nothing died. _start_engines rewrites num_new_engines regardless of whether the previous value has been consumed, so a no-op recovery reset the count to 0 before _refit read it. That count is the only gate on _connect, and _connect is the only place the trainer-side transport is ever built.

The consequence differs by transport, and the difference matters:

• Colocated (weight_transfer_mode: ipc, what all five shipped recipes use) — silent. send_hf_buckets_via_ipc_actor_impl reads _ipc_gather_group / _ipc_gather_src / _ipc_engine_index out of a worker_state that connect_colocate_topology never populated, finds all three None, and takes the placeholder-rank early return on every rank. Training then continues against stale rollout weights, forever, with no error anywhere and no metric that would show it.
• Disaggregated (broadcast) — loud. Rank 0 raises RuntimeError("connect_sglang_rollout_engines_distributed must be called before ...").

Fixed by returning early from _recover when dead_indices is empty.

Worth recording that the first proposed fix — num_new_engines += len(...) at the _start_engines write site — was wrong, and the devil's advocate caught it: __init__ leaves the count at N, the first _recover adds 0, and the assert self.num_new_engines == len(dead_indices) immediately below then fires. The agent that proposed it withdrew it.

2. The boot grace period made the monitor never probe anything

_need_first_wait was re-armed on every resume(). resume() runs once per training step, so any generation phase shorter than rollout_health_check_first_wait (default 60s) left the monitor permanently inside its grace period. Enabled, and dead — it would never health-check an engine, which is the entire feature.

Replaced with an absolute monotonic deadline armed at start() and after a recovery, never reset by resume(). Time spent paused now counts toward it, and a pause part-way through no longer restarts the clock. Restarted engines re-arm it explicitly, since a freshly booted SGLang server is still loading weights and would be killed by a probe.

Over a representative generate/pause cycle, probes went 0 → 21. The unit-test version of that scenario sees 0 → 2.

3. Flipping use_fault_tolerance: true in any shipped recipe raised KeyError

Config and code disagreed in both directions:

• use_fault_tolerance was declared required, but both readers use .get() and three of the five shipped sglang recipes never set it. The declaration was unenforced — PolicyConfig.generation is typed as the base GenerationConfig, which has no sglang_cfg field, so test_config_validation.py never reaches SglangSpecificArgs at all — but it documented a contract the recipes do not honor. Now NotRequired[bool], matching use_external_router three declarations above it.
• The three rollout_health_check_* knobs were declared NotRequired and read with bare subscripts. Every sglang recipe inherits grpo_math_1B.yaml, which carries no sglang keys, so none of them pick these up from grpo_math_1B_sglang.yaml. Turning the feature on therefore raised KeyError: 'rollout_health_check_interval' from a daemon-thread constructor. They stay NotRequired — they genuinely only matter when the feature is on — but RolloutHealthMonitor now asserts up front, naming every missing key and where the documented values live. This mirrors the existing sglang_router_ip / sglang_router_port handling, which is conditionally required the same way.

Tests

Three regressions added, all of which fail against the pre-fix source and pass after: the missing-key message, _recover leaving num_new_engines alone when nothing died, and _recover re-arming the grace period for restarted engines. test_first_wait_delays_checks_after_resume is renamed to test_first_wait_delays_the_initial_checks — it asserted semantics the fix deliberately removed, and it passed both before and after, so it was never load-bearing.

What the adversarial pass rejected

• A claim that SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION must be set for /health_generate to exist — disputed at confidence 99 and dropped. It would have sent the author chasing a non-issue.
• The unguarded self._health_monitor.pause() in recover_updatable_engines was reported as a null-deref and downgraded to a nit: its only caller is already behind the same use_fault_tolerance gate that creates the monitor, so it is unreachable. Left as the one inline comment below rather than changed unilaterally.

Environment caveats

No GPU on the review machine and no sglang extra, so nothing runtime-dependent was executed against a real engine. The three blockers were instead reproduced by loading the real fault_tolerance.py and the real SGLangGeneration._recover bytes with ray stubbed — that is stronger than reading, weaker than a GPU run. pytest could not be run locally (no ray/torch/omegaconf, no network to install them); the new tests were validated by executing their bodies against both the pre-fix and post-fix source. A GPU re-run with use_fault_tolerance: true would be the real confirmation, and is worth doing before merge.

Generated by Claude Code

[Kh4L]

Addendum to the self-review above — a fourth defect, found after posting.

stop() clears _thread, _stop_event and _pause_event even when the join times out. The thread it failed to reap is still inside _health_monitor_loop, so its next iteration dereferences the now-None event and dies with AttributeError: 'NoneType' object has no attribute 'wait'.

I originally judged this path unreachable and left it out. That was wrong. The join budget is timeout + interval + 5, but one probe is bounded at 2 * timeout by the outer ray.get, and a probe that times out then calls _kill_engine, which spends up to another timeout on the graceful shutdown before ray.kill. With the shipped 60s timeout that is up to 180s of work against a 125s budget — and it happens precisely when an engine is hung, which is the case this feature exists for.

Reproduced by driving stop() against an in-flight health check that outlasts the join, executing the real fault_tolerance.py: pre-fix the thread dies with the AttributeError; post-fix it observes the set stop event and exits cleanly. That scenario is now a regression test, which fails against the pre-fix source.

fixed in 00b8af6

Severity is moderate, not blocker-grade: a stray traceback from a daemon thread during teardown, not wrong training results. Flagging it rather than folding it in silently, since the review above said three findings.

[Kh4L]

/ok to test 487027b

[Kh4L]

/ok to test 25f3a01

[Kh4L]

/ok to test 0ae89f8

[Kh4L]

HSG real-GPU validation passed at 0ae89f8b1f8484ee82cebf6613bb10342d08be9a.

• Qwen3-0.6B with 2 live SGLang TP=1 engines
• Injected engine crash was detected; the worker was killed and recreated
• Replacement health check and post-recovery generation both returned 200
• SGLang 0.5.17.dev30+g3003d70f6; decode CUDA graphs enabled, only prefill capture disabled
• Test passed and the Slurm job completed with exit 0

[yuki-97]

Re-cut of #3187, so I re-checked all 14 of my review threads there against this tree. This pass covers the #3187 carry-over only; findings on the new code are held for a follow-up pass.

#3187 carry-over

11 of the 14 are addressed here: the module logger, the ray.get-level probe timeout, the config placement, both test asks, and the four "where is this called" questions.

One is now obsolete — please treat it as closed. r3655023291 asked for rollout_engine_lock to be built only under use_fault_tolerance, the way _health_monitor is. Since #3612 merged, the lock is unconditional on main and every refit acquires it (policy/utils.py:1071), fault tolerance or not — gating it would break the non-FT refit path.

Two still apply, one comment each:

• r3654957493 — use_fault_tolerance was applied as required, then reverted to NotRequired by aeb93c081. The revert rests on the shipped recipes not picking the knobs up from the sglang exemplar, and they do: all six chain to grpo_math_1B_sglang.yaml:42-45. Comment on config.py, listing the three places whose wording rests on that claim.
• r3654976078, second half — recovery is reachable only from _refit, so an engine that dies during generation is not restarted until the next one. Comment on sglang_weight_synchronizer.py.

Nothing here blocks the GPU evidence: HSG 6556006 and the green CI on this exact head both stand.

[yuki-97]

Carried over from the #3187 review (#3187 (comment)), second half — still open.

This is the only production caller of recover_updatable_engines, so recovery is reachable only from _refit: an engine that dies during generation stays dead until the next one.

The in-between state is not benign. _kill_engine sets the slot to None (fault_tolerance.py:268), and the lifecycle calls filter e is not None (sglang_generation.py:406-427), so those degrade to the survivors. Generation does not go through them — it POSTs to the router, which is never told the engine is gone. A request that does fail is retried 3x and then re-raised (http_utils.py:208), and the asyncio.gather at sglang_generation.py:657 has no return_exceptions, so it propagates out of generate() and takes the training step down.

[yuki-97]

Working through where recovery can actually run, the reachable cases come out like this. Only _kill_engine ever sets all_engines[i] = None (fault_tolerance.py:268), and the monitor only probes between prepare_for_generation and finish_generation, so detection and survival need different things from the same window:

engine dies	monitor probing?	outcome
generation window, nothing in flight	yes	detected, restarted at the next refit — the case test_fault_tolerance_real.py constructs
generation window, requests in flight	yes	depends on the router, which I have not verified — could you confirm what the installed sglang-router does with requests already dispatched to the dead engine?
training, or the refit's weights-only stage	no	never detected: the slot stays non-None, _recover finds dead_indices empty and returns, and the refit then broadcasts to a dead handle. No RayActorError handling on that path

Row 1 is the one the design actually handles, and it is also the least likely to occur: it needs the engine to die inside the generation window and with nothing in flight. During generation the engines are busy, and outside it the monitor is paused.

Is that the intended scope for this cut? The PR description — "a health monitor that detects a dead SGLang engine, restarts it, and rejoins it to the refit weight-update group" — reads as covering all three, and the HSG run backing it crashes the engine with _simulate_crash, which is row 1.

On row 3, to be clear that the gap is structural rather than an oversight: the probe is health_generate, a real generation request, so it only means anything against an engine that can serve. finish_generation has released the weights and KV cache, and the refit's weights-only stage is mid-write — a probe in either window fails, and a failed probe is _kill_engine, so an unpaused monitor would kill healthy engines every training phase. The pause is load-bearing, which is exactly why row 3 cannot be closed by dropping it. It needs a liveness signal that does not depend on serving — Ray already reports actor death, and that holds while offloaded and mid-refit alike, with health_generate left to answer the narrower "can it still serve" question.

[yuki-97]

Every sglang recipe resolves this key. All six chain to grpo_math_1B_sglang.yaml:42-45 — three inherit it directly, three through a sibling recipe — and ec7e0a6cc is the commit that adds use_fault_tolerance: false plus the three knobs at 60 there. Nothing in the tree relies on the key being absent, so declare it required and let the exemplar be the single source of the default.

Suggested change

	use_fault_tolerance: NotRequired[bool]
	use_fault_tolerance: bool

The three knobs below can stay NotRequired — they only matter when this is on, and the constructor asserts on them. Both readers can keep their bare .get(), which is what lets a pre-existing out-of-tree config load with fault tolerance off.

Three comments justify the current shape with the opposite claim, and need updating with it:

• lines 101-104, directly above — "Absent is equivalent to False, which is what the recipes that predate this feature rely on".
• fault_tolerance.py:46-49 — "the recipes that ship with it off do not carry them". It also says a missing key would be "a bare KeyError from a thread", but the monitor is constructed at sglang_generation.py:145 on the calling thread and start() runs after it, so that would be an ordinary init-time KeyError with a full traceback.
• test_fault_tolerance.py:488-491 — the test_monitor_names_the_missing_tuning_keys docstring. Every sglang recipe does inherit grpo_math_1B.yaml, transitively, but the chain runs through grpo_math_1B_sglang.yaml, so "flipping use_fault_tolerance: true in one of them" does not reach the assert. The scenario the test covers is a hand-written or out-of-tree config.

The assert itself is worth keeping either way: it names all three missing keys at once where a KeyError names only the first. Only the stated reason for it is wrong.

[yuki-97]

Scope: the code this PR adds on top of #3612 — fault_tolerance.py, the recovery path in sglang_generation.py, and its single wiring point in sglang_weight_synchronizer.py. Separate pass from the #3187 carry-over review I just submitted.

The one thing worth resolving before the rest: #3613 (comment) works through when recovery is actually reachable, and lands on three cases.

• Row 1 — the engine dies inside the generation window with nothing in flight — is the one the design handles, and it is what the GPU run exercises.
• Row 2, a death with requests in flight, is unresolved, and it decides how much this feature is worth: on our side nothing tolerates a failed request, so whether the run survives rests entirely on the router re-routing those requests, which I could not confirm for the installed sglang-router version. If it does not, then the only deaths this feature recovers from are the ones that were not going to take the run down anyway — which is close to no coverage at all. Please check that against a version you can verify.
• Row 3, a death while the monitor is paused, is never detected.

If row 2 does re-route, row 2 is the case a functional test should pin — see #3613 (comment).

Also checked:

• Router registration lifecycle across kill and restart — followed DELETE /workers/{id} to its only call site. Same version gap as above: the router's retry, health-check and circuit-breaker defaults differ between the sglang rev this branch pins and the PyPI package it installs, so that comment states the asymmetry and leaves the severity to you.
• Test coverage of the new wiring, against CI's own coverage report at this head. This includes _recover's offload/onload block, uncovered on both tiers: I walked the sequence against prepare_for_generation's tag gating and it is correct.
• Production consumers of every new symbol — writes vs reads, source vs tests.
• Restart policy against the existing fleet_health.py model.
• _recover's remaining guards: the redundant and dead_indices and the assert that cannot fire by construction are both harmless. The inner health_generate timeout is passed but never asserted; only the outer one is covered.

The remaining comments are lower severity — one [TEST], a test-only method on the production worker, two symbols with no production consumer, an outdated comment, an unbounded restart loop, a return-shape question, and a .get() nit contingent on the use_fault_tolerance: bool change in the carry-over review.

[yuki-97]

shutdown issues DELETE /workers/{id} before killing the process (sglang_worker.py:239-253); this path does not. A failed health probe is exactly when shutdown() is wedged, and ray.kill does not deregister, so an engine the monitor kills stays in the router's worker table and eviction is left to the router's own health check. The HSG run does not exercise this either — _simulate_crash is self.shutdown(), so its engine deregisters on the way out.

Suggest fix: send DELETE /workers/{id} from the monitor before ray.kill, so both teardown paths leave the same router state.

What it costs at runtime I could not settle. _start_router passes a bare RouterArgs(), and at the sglang rev this branch pins those defaults carry retry_max_retries=5 with per-attempt worker re-selection plus a circuit breaker, which would absorb most of a stale entry — but the installed package is PyPI sglang-router==0.3.2, which is not that tree. If you know how that version behaves, it decides whether this is cosmetic.

[yuki-97]

[TEST] The use_fault_tolerance: true arm here is executed by no test. CI's own coverage report at this head names line 137 as one of only two uncovered statements in the file: sglang_weight_synchronizer.py 101 2 98% 96, 137.

_mock_sglang_generation (test_weight_synchronizer.py:245) builds sglang_cfg as a plain dict without the key, so all 20 synchronizer tests take the false branch, and the GPU test calls recover_updatable_engines() directly rather than through sync_weights.

The monitor and _recover themselves are well covered; this is the one line that wires them into the refit path, so it can break with every existing test still green. If it regressed, the next refit dies with AttributeError: 'NoneType' object has no attribute 'update_weights_from_tensor' inside policy/utils.py — nowhere near anything named fault tolerance.

Suggest fix: give _mock_sglang_generation a use_fault_tolerance kwarg and assert recover_updatable_engines runs before get_updatable_engines_and_lock. It is already a MagicMock, so no fixture change is needed.

[yuki-97]

This 5-tuple is returned to nobody. The only production caller discards it (sglang_weight_synchronizer.py:137) and re-reads the identical tuple eight lines later via get_updatable_engines_and_lock (:145).

The two return expressions are textually identical, and _recover is synchronous with the monitor paused throughout, so the two reads cannot differ. Only test_recover_updatable_engines_reports_engine_state consumes it, which pins a shape nothing else uses.

Suggest fix: annotate -> None and end the method at self._recover(); the sibling test just below already covers num_new_engines surviving a no-op recovery.

[yuki-97]

Outdated — this is #3613, and recovery never takes this lock. The only acquire/release is policy/utils.py:1071 / :1104, trainer rank 0 only. Suggest:

        # Serializes weight-update broadcasts. Acquired per bucket by trainer
        # rank 0 in policy/utils.py; nothing on the generation side takes it.

Same correction applies to that function's docstring (policy/utils.py:1050-1052): it credits the lock with keeping health-check pings off the broadcast, but that comes from the monitor's pause.

[yuki-97]

An engine that boots fine and then dies again every step is restarted forever — there is no attempt cap and no floor on surviving engines.

A hard failure propagates out of _refit and ends the run, so that case is covered. The unbounded one is a flaky node: every step pays a full actor create, weight load and _connect NCCL rebuild, with no escalation and no abort. fleet_health.py already models this — max_restart_attempts_per_shard / min_healthy_shards (:105-106), the check at :304, and a terminal RETIRED state. Reusing it, or just capping and aborting here, would avoid a silent throughput cliff.

[yuki-97]

Nit, and only if you take the use_fault_tolerance: bool change in #3613 (comment): once the field is required, read it directly. config-conventions allows a bare .get(key) precisely for NotRequired fields, and asks for direct access once a key is required.

Suggested change

	if sglang_cfg["sglang_cfg"].get("use_fault_tolerance"):
	if sglang_cfg["sglang_cfg"]["use_fault_tolerance"]:

Same read at sglang_weight_synchronizer.py:136.

This supersedes the "both readers can keep their bare .get()" line in that thread — I hedged it for out-of-tree configs, but the exemplar is what supplies the default, and the convention is explicit that the call site should not carry one.

[yuki-97]

Dead state, both halves. The accessor is read only by the two new test files, and once it goes the field itself has four writes and no reader at all.

It is also derivable: checking is enabled exactly when _pause_event is cleared and _stop_event is not set. That holds at every write site — start leaves the pause event set, resume clears it, pause sets it, and stop clears the pause event but sets the stop event first. test_stop_leaves_events_intact_when_the_join_times_out shows this directly: the line above its is_checking_enabled() assertion already asserts _stop_event.is_set(), which is what makes the flag False.

Suggest dropping the accessor here, plus self._is_checking_enabled and its four writes (lines 73, 131, 157, 170). The tests already assert on _thread, _stop_event and _pause_event, so they can assert the two events instead.

Suggested change

	def is_checking_enabled(self) -> bool:
	"""Return whether health checking is currently enabled (not paused)."""
	return self._is_checking_enabled

[yuki-97]

Working through where recovery can actually run, the reachable cases come out like this. Only _kill_engine ever sets all_engines[i] = None (fault_tolerance.py:268), and the monitor only probes between prepare_for_generation and finish_generation, so detection and survival need different things from the same window:

engine dies	monitor probing?	outcome
generation window, nothing in flight	yes	detected, restarted at the next refit — the case test_fault_tolerance_real.py constructs
generation window, requests in flight	yes	depends on the router, which I have not verified — could you confirm what the installed sglang-router does with requests already dispatched to the dead engine?
training, or the refit's weights-only stage	no	never detected: the slot stays non-None, _recover finds dead_indices empty and returns, and the refit then broadcasts to a dead handle. No RayActorError handling on that path

Row 1 is the one the design actually handles, and it is also the least likely to occur: it needs the engine to die inside the generation window and with nothing in flight. During generation the engines are busy, and outside it the monitor is paused.

Is that the intended scope for this cut? The PR description — "a health monitor that detects a dead SGLang engine, restarts it, and rejoins it to the refit weight-update group" — reads as covering all three, and the HSG run backing it crashes the engine with _simulate_crash, which is row 1.

On row 3, to be clear that the gap is structural rather than an oversight: the probe is health_generate, a real generation request, so it only means anything against an engine that can serve. finish_generation has released the weights and KV cache, and the refit's weights-only stage is mid-write — a probe in either window fails, and a failed probe is _kill_engine, so an unpaused monitor would kill healthy engines every training phase. The pause is load-bearing, which is exactly why row 3 cannot be closed by dropping it. It needs a liveness signal that does not depend on serving — Ray already reports actor death, and that holds while offloaded and mid-refit alike, with health_generate left to answer the narrower "can it still serve" question.

[yuki-97]

Which of these two is intended? The lock comes back in the tuple and is discarded here, and then _send_buckets reads it off the generation object instead (:288):

rollout_engine_lock=self._generation.rollout_engine_lock,

• If the attribute read is the intended source, the lock should come out of the return shape — a method named get_updatable_engines_and_lock whose only caller throws the lock away is misleading, and the same slot is dead in recover_updatable_engines too.
• If the returned one is meant to be used, drop the leading underscore and pass it into _send_buckets rather than re-reading the attribute.

They are the same object today, so this is a readability question rather than a bug. Related to the return-shape comment on sglang_generation.py.

[yuki-97]

would be good to add a functional test to show this feature works well, like tests/functional/grpo_dp_single_controller_chaos.sh.

it depends on #3613 (comment) though — of the three cases there, row 1 is what the GPU test already covers, row 3 isn't implemented, and row 2's expected outcome isn't settled. once we know whether a mid-generation death is supposed to survive, the test can assert that, the same way the SC lane splits EXPECT=survival from the bounded-failure variant.