fix(sglang): rebind checkpoint-engine refit peers across engine restarts - #3944
Draft
Kh4L wants to merge 27 commits into
Draft
fix(sglang): rebind checkpoint-engine refit peers across engine restarts#3944Kh4L wants to merge 27 commits into
Kh4L wants to merge 27 commits into
Conversation
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Three gaps found reviewing this branch: - tests/functional/grpo_sglang_nixl_non_colocated.sh was referenced by nothing, so it had never run. The CI guard only checks that L1_Functional*.sh shards appear in the workflow matrix, not that leaf scripts are called. Register it with L1_Functional_Tests_SGLang.sh, which is already in all three matrices. - The only test of _aligned_checkpoint_engine_batches asserted weight names only. The aligner itself enforces name equality across ranks, so that assertion is invariant under any rank permutation. Mutating aligned[rank] -> aligned[len(pending)-1-rank] delivered every shard to the wrong TP rank and the suite stayed green. Assert the tensors too; the mutant now fails. - docs/design-docs/checkpoint-engines.md and docs/guides/checkpoint-engine-refit.md both still said SGLang has no checkpoint-engine refit. Update both, record the actual limits (one node per logical engine, no shard_expert_weights), and generalize the 'Adding Another Backend' timing-line step, which named the vLLM line only. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Five properties that had no coverage; each was verified by mutation, i.e. the test fails when the corresponding behaviour is broken and passes when it is not. - Recycled receive buffers. nixl.py hands out views into a rotating buffer pool, so a fake engine that allocates fresh tensors every batch makes that entire bug class invisible. _RecyclingEngine poisons a buffer once it is recycled. The aligner is correct today because it only advances a rank whose deque is empty; making it prefetch turns the yielded tensors into NaN. - Multi-dtype batches. NIXL packs buckets by bytes, not dtype, so a mixed bf16/fp32 batch is the normal production shape, but nothing drove more than one dtype group through the update path. - weight_version across refits. Its semantics were pinned only for the first refit, so bumping per POST instead of per refit went undetected. - base_gpu_id remapping. _to_local_gpu_id was stubbed to identity with base_gpu_id=0, which is exactly the case where remapping and doing nothing are indistinguishable. - Payload index to SGLang rank. SGLang indexes serialized_named_tensors by its own TP rank, so a transposed list loads every shard onto the wrong GPU and still reports success. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
refit.md still stated flatly that non-colocated SGLang generation is not supported, which now contradicts checkpoint-engine-refit.md on the same click (refit.md links to it). The statement is still true for every transport other than checkpoint-engine refit, which factory.py:107-110 rejects, so narrow it rather than delete it, and add SGLang to the NIXL full-weights row. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Two follow-ups on this PR's own changes.
The dp_size=1 guard (weight_sync/factory.py:90) was documented only in
checkpoint-engine-refit.md. design-docs/checkpoint-engines.md still said
'Both are rejected' for what is now three constraints, and refit.md's
constraint table did not mention it at all.
MetricSetupTiming.vllm_checkpoint_engine_init_time_s is no longer written
by anything: grpo.py:1504 moved to extras[f'{backend}_checkpoint_engine_
init_time_s']. For vLLM that formats to the same string, and to_dict()
merges extras over the typed fields, so the emitted metric name is
unchanged -- the field is just dead state now.
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
The factory guarded dp_size and shard_expert_weights but not pp_size. The mixin creates one receiver per engine GPU, while SGLang indexes serialized_named_tensors by TP rank -- and sglang_worker.py:386 asserts tp_size == num_gpus_per_engine // pp_size, so with pp_size>1 the payload list is pp_size times longer than the engine expects. Fail loudly at setup instead. Recorded in all three docs alongside the other two limits. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Mirrors the dp_size case immediately above it. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
A hung or dead SGLang engine currently wedges the whole run: the refit's `ray.get` on its actor never returns, and the router keeps sending `/generate` to it. Add `RolloutHealthMonitor`, a daemon thread that polls `/health_generate` on every node-0 engine and, on failure, kills and restarts the actor. The refit picks the survivors up through the engine registry this PR's base already exposes (`get_updatable_engines_and_lock`), so a recovered engine is reconnected on the next weight update. The monitor must not run while the engines are offloaded: `/health_generate` always executes a real one-token generation (the `SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=false` bypass only covers `/health`), so probing a released engine reports a false failure and kills a live actor. It is therefore wired to the generation lifecycle — resumed by `prepare_for_generation` once the KV cache is back, paused by `finish_generation` and for the duration of a refit. `pause()` blocks on an in-flight probe rather than only setting a flag, so a probe cannot overlap `release_memory_occupation`. `_kill_engine` bounds its graceful-shutdown `ray.get` and kills the actor even when that shutdown fails, which is the exact case the monitor exists to handle. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu> (cherry picked from commit f563fa9)
…grace Two defects made the fault-tolerance path unusable once enabled. Both are invisible today because every shipped recipe sets `use_fault_tolerance: false`, so `_recover` is never reached. `_recover` ran unconditionally on every refit, including the common case where nothing died. `_start_engines` rewrites `num_new_engines` without regard for whether the previous value has been consumed, so the no-op recovery reset the count to 0 before the refit read it. That count is the only gate on `_connect` in the weight synchronizer, and `_connect` is the only place the trainer-side transport is ever built. What that costs depends on the transport. Colocated (`weight_transfer_mode: ipc`, which is what all five shipped sglang recipes use) fails silently: `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 with no error anywhere. Disaggregated (`broadcast`) at least fails loudly, since rank 0 raises RuntimeError from `update_weights_to_sglang_distributed`. Return early when there are no dead engines. The boot grace period was a `_need_first_wait` boolean re-armed on every `resume()`. `resume()` runs once per training step, so any generation phase shorter than `rollout_health_check_first_wait` left the monitor permanently inside its grace period and no probe ever ran -- the feature was enabled but dead. Replace the flag with an absolute monotonic deadline that is armed at `start()` and after a recovery, and is not reset by `resume()`. Time spent paused now counts toward the deadline, and a pause part-way through no longer restarts the clock. Restarted engines re-arm it explicitly, since a freshly booted server is still loading weights and a probe would kill it. Verified by executing the real source: `num_new_engines` survives at 4 across a no-op recovery, and health probes over a representative generate/pause cycle went from 0 to 21. Signed-off-by: Serge Panev <spanev@nvidia.com>
`use_fault_tolerance` was declared required while both readers used `.get()`, and three of the five shipped sglang recipes never set it at all. The declaration was unenforced -- `PolicyConfig.generation` is typed as the base `GenerationConfig`, which has no `sglang_cfg` field, so `test_config_validation.py` never validates `SglangSpecificArgs` -- but it documented a contract the recipes do not honor. Declare it `NotRequired[bool]`, matching `use_external_router` in the same file and the `.get()` truthiness reads that already exist. The three `rollout_health_check_*` knobs had the inverse problem: declared `NotRequired`, read with bare subscripts. Every sglang recipe inherits `grpo_math_1B.yaml`, which carries no sglang keys, so none of them pick the knobs up from the sglang exemplar -- flipping `use_fault_tolerance: true` in any shipped recipe raised `KeyError: 'rollout_health_check_interval'` from a daemon-thread constructor. They stay `NotRequired`, since they only matter when the feature is on, but `RolloutHealthMonitor` now asserts on them up front and names every missing key plus where the documented values live. This follows the existing precedent for `sglang_router_ip`/`sglang_router_port`, which are conditionally required in the same way. Tests cover all three blockers fixed on this branch: the missing-key message, `_recover` leaving `num_new_engines` alone when nothing died, and `_recover` re-arming the boot grace period for restarted engines. Against the pre-fix source those three fail (KeyError, `_start_engines` invoked, no re-arm) and the first-wait regression sees 0 probes where the fixed code sees 2. `test_first_wait_delays_checks_after_resume` is renamed to `test_first_wait_delays_the_initial_checks`: the grace period is armed at `start()` now, not on every resume. Signed-off-by: Serge Panev <spanev@nvidia.com>
…hread When the join times out, `stop` logged a warning and then cleared `_thread`, `_stop_event` and `_pause_event` regardless. The thread it failed to reap is still inside `_health_monitor_loop`, so its next iteration dereferenced the now-`None` event and the thread died with `AttributeError: 'NoneType' object has no attribute 'wait'` at the `self._stop_event.wait(self._check_interval)` at the bottom of the loop. The join budget is not generous enough to treat that path as unreachable. It is `timeout + interval + 5`, but a single 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 exactly when an engine is hung, which is the case the feature exists for. Clear `_is_checking_enabled` either way, since checking really has stopped, but leave the events in place when the thread outlived the join so it can observe the set stop event and exit on its own. The regression test drives `stop` against an in-flight health check that outlasts the join. Against the pre-fix source it fails on the first assertion and the monitor thread raises the AttributeError above; after the fix the thread exits cleanly and `threading.excepthook` records nothing. It cannot run faster than the join's fixed +5s floor. Signed-off-by: Serge Panev <spanev@nvidia.com>
The pinned SGLang server replaced disable_piecewise_cuda_graph with per-phase graph backends. The real fault-tolerance smoke already disables CUDA graphs entirely, so remove the stale argument that would otherwise fail ServerArgs construction. Signed-off-by: Serge Panev <spanev@nvidia.com>
The SGLang branch of sync_weights only called prepare_for_generation, which is resume_memory_occupation and nothing else, and then POSTed update_weights_from_tensor directly. At the pinned sglang rev that is rejected: the scheduler asserts on a session opened by begin_weight_update. end_weight_update is also what rebuilds quantized kernel layouts after the last bucket. Wrap the transfer in the same envelope the sibling SGLang synchronizer uses, so both transports drive the engine through one contract. The pause is load-bearing too: the buckets arrive as several update_weights_from_tensor calls, each taking the server's model update lock on its own, so a request admitted between buckets would run against a half-updated model. Also port the in_place pause rejection, which this path lacked, and drop the call-site config defaults the dp/pp guards were reading through. The existing tests could not see any of this -- they stub the worker's update_weights_from_tensor -- so add ordering and both failure paths. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Ray decides between a threaded and an asyncio actor with has_async_methods, which is inspect.getmembers over the whole MRO. The private coroutine on this mixin therefore turned SGLangGenerationWorker into an asyncio actor, and Ray runs an asyncio actor's sync methods on that event loop -- so the asyncio.run in the refit entry point raised 'asyncio.run() cannot be called from a running event loop' on the first refit. Making the entry point async would fix the crash but leave the flip in place, changing the concurrency semantics of every pre-existing worker RPC. Move the loop to a module-level coroutine instead: the mixin defines no coroutine members, so the actor stays threaded. The tests missed it because they drove the private coroutine on a plain object; cover the public wrapper, and assert the no-coroutine-member invariant, which is the thing that actually regresses. While here: a refit that received zero tensors now raises instead of bumping the weight version and reporting success, and the post-transfer invalidate_kv_cache is dropped -- the synchronizer now flushes before the buckets land. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
…ad code refit_cfg was NotRequired[Any] for a block normalized by exactly the vLLM schema, so give it the same type. The refit_transport comments and the grpo.py raise claimed null always means colocated CUDA-IPC, which is wrong for non-colocated SGLang, where it selects the NCCL weight-update group. get_rollout_engine_urls has no callers -- a leftover from the earlier HTTP transport design. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
With one one-GPU engine the payload-index -> TP-rank scatter never runs: one stream has nothing to align, _validate_rank_batches compares nothing, and update_weights_from_tensor gets a one-payload list. A transposed payload list would load every shard onto the wrong GPU while still succeeding, and only tp_size>1 can catch that. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Self-audit of the review fixes turned up four gaps. The no-asyncio-actor guard restated Ray's rule instead of asking Ray. Ray flips on iscoroutinefunction OR isasyncgenfunction, over the ACTOR class; the test checked only iscoroutinefunction, only on the mixin. An async generator moved back onto the mixin -- the exact shape at hand, since the receive loop is one -- left it green. It now imports Ray's own has_async_methods and asks about SGLangGenerationWorker. The KV pool was re-acquired on the success path while the pause was undone on every path, so a failed refit resumed the engines with the pool still released. Both now live in the same finally, KV first so requests are not readmitted before the memory is back. Nothing tested the two invariants the comments assert: that end_weight_update never closes a session begin_weight_update failed to open, and that a pause which itself raises still resumes. Both are covered now, as are the engine-reported-failure branch and the fact that the in_place guard is SGLang-only. The transport wording promised more than the code delivers: SGLang's NCCL weight-update group is Megatron-policy only (factory.py:165), so a DTensor user following the message would hit a second NotImplementedError. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Engine fault tolerance restarts a dead SGLang engine as a fresh Ray actor, but the checkpoint-engine refit fabric never noticed: the replacement has no NIXL receivers, its paired policy sender still binds the dead agent, and the synchronizer stays 'ready' forever. This wires recovery into the checkpoint-engine path as a rebind (never a teardown): - sync_weights probes recover_updatable_engines() every refit when use_fault_tolerance is on; any new engines invalidate readiness and re-run init_communicator, which reuses surviving engine objects, constructs the missing receivers on replacements, and reconnects pairs through NIXL's disconnect-before-connect init. - num_new_engines is consumed only after a successful communicator setup (the startup fleet reports through the same counter, so the first refit must not be misclassified as crash recovery). - Failure contracts are split three ways: recovery failure rolls the whole attempted cohort back (kill actors, restore None slots and the pre-attempt count) and stays retryable, escalating to terminal only if rollback fails (new RecoveryRollbackError); a rebind failure is terminal (NIXL prepare/add_remote_agent are not transactional, so retrying over partial state is unsafe); a failure after the transfer starts is terminal with an enforced latch - later syncs raise immediately with zero RPCs, serving and the health monitor stay paused, and shutdown() no longer runs in the unconditional finally on that path. Pre-transfer failures keep the safe cleanup/resume behavior; cleanup errors never mask the primary exception. - init_checkpoint_engine publishes checkpoint_engines only after every receiver constructs, so a mid-loop failure cannot arm the idempotency early-return and silently no-op the retry a rebind depends on. - Setup guards: use_fault_tolerance requires the built-in nixl backend (custom checkpoint engines have no reconnect contract) and the bf16 scheme (the sender streams raw BF16 tensors). The transfer-failure divergence from the NCCL sibling (no serving resume after a started transfer fails) is deliberate: an interrupted NCCL broadcast can be redone next refit; interrupted one-sided NIXL work cannot. Signed-off-by: Serge Panev <spanev@nvidia.com>
Two real PolicyCheckpointEngineMixin senders (NIXL, weights streamed from the HF checkpoint) plus two real SGLang engines with fault tolerance on, driven by the real CheckpointEngineWeightSynchronizer. Crashes one engine, waits for the health monitor, recovers, garbages the replacement's weights with the weights_checker reset_tensors control, runs the recovery refit, and proves via weights_checker compare that real current bytes reached the replacement — plus that only the replaced engine's NIXL agent_name changed and that a request routed directly to the replacement's server generates. Needs 4 GPUs (train_world_size >= rollout_world_size forces 2 senders for 2 engines). Signed-off-by: Serge Panev <spanev@nvidia.com>
…aths Address the five failure-path defects from the design re-review: 1. Move _start_engines inside the atomic recovery attempt: it publishes replacement actors and rewrites num_new_engines while it runs, so a synchronous mid-start failure already leaves partial state visible. 2. Extend the same attempt over the post-init work (count assert, health monitor arm, needs_offload release/resume): a failure there previously escaped with the cohort published, so the next recovery saw no dead slot and rebound a partially transitioned engine. 3. Roll back through the health monitor's kill path: a bounded best-effort graceful shutdown (router deregistration + spawned server process tree) before ray.kill, tolerating replacements whose init never reached self.process. The bound comes from the monitor's configured per-RPC timeout, exposed via a new check_timeout property. 4. Drop _checkpoint_engine_ready on any rolled-back (retryable) recovery failure instead of only latching on RecoveryRollbackError: readiness must not survive restored None slots. 5. Honor the terminal latch in the public shutdown(): after a terminal rebind/transfer failure the NIXL state is unknown and finalize_checkpoint_engine is not idempotent over it, so controller teardown becomes a zero-RPC no-op. The recovery catches use BaseException, matching the synchronizer's existing convention: a KeyboardInterrupt during the init wait must still roll the published cohort back (and drop readiness) before propagating. Tests: mid-start and offload-transition rollback (nonzero pre-attempt count; pinned graceful-shutdown timeout), recovery-rebind failure through sync_weights with zero RPCs on later syncs, terminal-latched shutdown as a zero-RPC no-op, factory negative tests for the fault-tolerance/quantization guards, and a 2x2 fake-NIXL survivor-peer reconnect test. Signed-off-by: Serge Panev <spanev@nvidia.com>
Two reviewer-requested tightenings: the recovery-rebind failure test now lets receiver init and the prepare gather succeed and fails the third ray.get — the non-transactional process-group rebind itself — asserting no session/transfer/count-clear RPC escaped the first call; and a mutation test bypasses schema validation via the scheme accessor to prove the factory's own bf16 guard refuses non-bf16 schemes. Signed-off-by: Serge Panev <spanev@nvidia.com>
…re test begin/end_weight_update (and the pause/invalidate/continue envelope) are called directly on the generation object, not through the run_checkpoint_engine_method wrapper the previous assertion inspected. Assert the direct lifecycle mocks are untouched after the failing rebind, and that the terminal retry records zero generation method calls at all. Signed-off-by: Serge Panev <spanev@nvidia.com>
github-actions
Bot
added
the
Documentation
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
3 participants
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.
What
When SGLang's fault tolerance (#3613) restarts a crashed rollout engine, the checkpoint-engine refit transport (#3519) must not keep streaming weights to the dead peer. This PR makes the two features compose: the weight synchronizer detects recovered engines every refit, restarts dead ones atomically, and rebinds the paired NIXL fabric to only the replaced peers — survivors' connections are reused, not rebuilt.
Failure contracts
_recovertreats everything from_start_enginesthrough the offload transition as one attempt: on failure the replacement cohort is gracefully shut down (router deregistration + server process tree, bounded by the health monitor's timeout), slots return toNone, and the unconsumed engine count is restored. A failed rollback escalates toRecoveryRollbackError(terminal).prepare()/add_remote_agent()are not transactional; retrying over partial state silently skips or double-registers. A failed (re)bind latches_terminal_error, and every later sync raises before issuing a single RPC.end_weight_updatefailure after the first bucket moved → terminal, no resume. This deliberately diverges from the NCCL sibling's finally-resume: an interrupted NCCL broadcast can be redone next refit, interrupted one-sided NIXL work cannot, so serving is not resumed over a half-updated model. Publicshutdown()also honors the latch (zero finalize RPCs over unknown transport state).Setup-time guards: fault tolerance requires the built-in
nixlbackend, and the transport is BF16-only (quantized schemes are rejected before any transfer).Validation
test_checkpoint_engine_recovery_real.py: baseline refit →_simulate_crash→ autonomous health-monitor detection → recovery → rebind → SGLangweights_checkersnapshot/reset/compare proves the replacement holds current bytes (311 tensors, 1.40 GiB) → replacementagent_namechanged, survivor unchanged → direct/generateon the replacement + fleet generation. Passed in 18:41.The design doc with the full contract rationale and review trail is available on request.