[Core] Consume SleepModeBackend capability flags in the worker suspend/resume path#47500
Draft
matteso1 wants to merge 1 commit into
Draft
[Core] Consume SleepModeBackend capability flags in the worker suspend/resume path#47500matteso1 wants to merge 1 commit into
matteso1 wants to merge 1 commit into
Conversation
galletas1712
reviewed
Jul 6, 2026
| if self.mq_broadcaster is not None: | ||
| self.mq_broadcaster = None | ||
|
|
||
| def prepare_for_suspend(self) -> None: |
Contributor
There was a problem hiding this comment.
Let's introduce these APIs in device_communicators as well, since there are other comms aside from NCCL (flashinfer, custom allreduce, etc) that need to be suspended.
3 tasks
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
ae55672 to
2fc395b
Compare
The capability flags merged in vllm-project#44074 (preserves_communicators, preserves_graphs_with_communicators) document a contract nothing in-tree enforces: Worker.sleep/wake_up assume the cumem world where communicators and captured CUDA graphs survive suspend untouched. Checkpoint-style backends (RFC vllm-project#34303 phase 2) break both assumptions - cuda-checkpoint cannot restore a TP>1 vLLM server today because comm state does not survive (NVIDIA/cuda-checkpoint#27). This wires the vLLM-side contract: - GroupCoordinator.prepare_for_suspend()/reinit_after_resume(): communicator-owned teardown/rebuild (each communicator owns destroy/recreate inside suspend/resume; the worker only sequences the calls). Single-member groups no-op, so the TP=1 checkpoint path works today. Multi-member groups fail loudly pending checkpoint support in the comms libraries (NCCL contrib/nccl_checkpoint, FlashInfer vllm-project#3791). - parallel_state.prepare_communicators_for_suspend()/ reinit_communicators_after_resume(): sequencing (subgroups before world on suspend, world first on resume), mirroring destroy_model_parallel. - Worker.sleep/wake_up consult the flags. A backend that rebuilds communicators without guaranteeing graph validity fails loudly at wake instead of replaying graphs that launch into destroyed comm state. No behavior change for existing users: the default cumem backend reports preserves_communicators() == True, so every new branch is unreachable on today's configurations. Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>
2fc395b to
ac77c38
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
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.
Prior art / provenance
SleepModeBackendwith capability flags (preserves_communicators,preserves_graphs_with_communicators,preserves_compiled_artifacts,supports_durable_storage) — but the flags currently have zero in-tree consumers:Worker.sleep/wake_upstill assume the cumem world where communicators and captured CUDA graphs survive suspend untouched. Thepreserves_communicatorsdocstring even says "the executor must re-initialize them on resume" — nothing does.cuda-checkpointcannot restore 2-GPU vLLM server process NVIDIA/cuda-checkpoint#27) because comm state does not survive suspend.contrib/nccl_checkpointand FlashInfer RFC [RFC] Checkpoint lifecycle APIs for CUDA-graph-stable communication workspaces flashinfer-ai/flashinfer#3791 (@galletas1712) keep CUDA graphs valid across checkpoint while the uncheckpointable parts of comm state are destroyed/recreated.What this adds
GroupCoordinator.prepare_for_suspend()/reinit_after_resume()— the communicator-owned lifecycle hooks. Single-member groups hold no peer state and no-op (so a TP=1 checkpoint backend works today). Multi-member groups fail loudly with a pointer to the required comms-library support, instead of resuming into corrupt comm state.prepare_communicators_for_suspend()/reinit_communicators_after_resume()inparallel_state— sequencing only (subgroups before world on suspend, world first on resume), mirroringdestroy_model_parallel's group handling.Worker.sleep/wake_upconsume the flags — a backend reportingpreserves_communicators() == Falsegets teardown beforesuspend()and rebuild afterresume(); if it also reportspreserves_graphs_with_communicators() == False, wake fails loudly rather than replaying captured graphs that launch into destroyed comm state (discard+recapture is follow-up work; a suspended engine is not serving, so re-arming is an init-time-style operation consistent with the capture-at-init-only constraint).What this deliberately does not do
CudaCheckpointBackend— that lands separately on the [Core] Pluggable sleep-mode backend abstraction (RFC #34303) #44074 factory (being driven by @elizabetht; this PR is the seam it plugs into).Behavior
No change for existing users: the default
cumembackend reportspreserves_communicators() == True, so every new branch is unreachable on today's configurations. Unit tests cover the orchestration order (teardown before suspend, rebuild after resume, graph invalidation after rebuild), the untouched cumem path, and single- vs multi-member hook behavior.Open questions for reviewers
preserves_graphs_with_communicatorsis a classmethod, but graph validity is workload-dependent (a graph survives checkpoint iff nothing comm-device-side was captured into it — FlashInfer-path capture sets pass today, DeepEP does not until NCCL checkpoint learns graphs). Should the flag become instance/capture-aware, or is conservative-False per backend acceptable until then?destroy_model_parallel's explicit globals (_TP/_DCP/_PCP/_PP/_DP/_EP/_EPLB+_WORLD). Should it iterate the_groupsregistry instead to cover ad-hoc groups?compile_or_warm_up_modelsubset) or as a reset API on the graph-capture manager?