Skip to content

[Core] Consume SleepModeBackend capability flags in the worker suspend/resume path#47500

Draft
matteso1 wants to merge 1 commit into
vllm-project:mainfrom
thaw-ai:suspend-resume-orchestration
Draft

[Core] Consume SleepModeBackend capability flags in the worker suspend/resume path#47500
matteso1 wants to merge 1 commit into
vllm-project:mainfrom
thaw-ai:suspend-resume-orchestration

Conversation

@matteso1

@matteso1 matteso1 commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Prior art / provenance

What this adds

  1. 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.
  2. prepare_communicators_for_suspend() / reinit_communicators_after_resume() in parallel_state — sequencing only (subgroups before world on suspend, world first on resume), mirroring destroy_model_parallel's group handling.
  3. Worker.sleep/wake_up consume the flags — a backend reporting preserves_communicators() == False gets teardown before suspend() and rebuild after resume(); if it also reports preserves_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

Behavior

No change for existing users: the default cumem backend reports preserves_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

  1. preserves_graphs_with_communicators is 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?
  2. Group enumeration mirrors destroy_model_parallel's explicit globals (_TP/_DCP/_PCP/_PP/_DP/_EP/_EPLB + _WORLD). Should it iterate the _groups registry instead to cover ad-hoc groups?
  3. Once discard+recapture exists, should it live in the worker (a compile_or_warm_up_model subset) or as a reset API on the graph-capture manager?

@mergify mergify Bot added the v1 label Jul 3, 2026
if self.mq_broadcaster is not None:
self.mq_broadcaster = None

def prepare_for_suspend(self) -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mergify

mergify Bot commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @matteso1.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Jul 8, 2026
@matteso1 matteso1 force-pushed the suspend-resume-orchestration branch from ae55672 to 2fc395b Compare July 8, 2026 12:33
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants