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

tests/v1/worker/test_gpu_worker.py

@@ -7,6 +7,7 @@
 import pytest
 
 import vllm.v1.worker.gpu_worker as gpu_worker_module
+from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.multimodal.video import (
     PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES,
     PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES,
@@ -186,3 +187,146 @@ def test_startup_plan_apply_gate(plan_env):
     explicit = _plan_worker(kv_bytes=7 * GiB_bytes)
     maybe_apply_startup_plan(explicit)
     assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes
+
+
+# Suspend/resume orchestration: Worker.sleep/wake_up consume the
+# SleepModeBackend capability flags (vllm/device_allocator/sleep_mode_backend.py).
+
+
+class _FlaggedBackend:
+    """Minimal stand-in for a SleepModeBackend with configurable flags."""
+
+    _preserves_communicators = True
+    _preserves_graphs = True
+
+    def __init__(self):
+        self.calls: list[tuple] = []
+
+    def suspend(self, level: int = 1) -> None:
+        self.calls.append(("suspend", level))
+
+    def resume(self, tags: list[str] | None = None) -> None:
+        self.calls.append(("resume", tags))
+
+    @classmethod
+    def preserves_communicators(cls) -> bool:
+        return cls._preserves_communicators
+
+    @classmethod
+    def preserves_graphs_with_communicators(cls) -> bool:
+        return cls._preserves_graphs
+
+
+class _PreservingBackend(_FlaggedBackend):
+    """cumem-shaped: communicators survive suspend untouched."""
+
+
+class _RebuildingBackend(_FlaggedBackend):
+    """checkpoint-shaped: communicators are destroyed by suspend and must be
+    rebuilt on resume; captured graphs do not survive the rebuild."""
+
+    _preserves_communicators = False
+    _preserves_graphs = False
+
+
+def _sleep_worker(backend: _FlaggedBackend) -> Worker:
+    worker = object.__new__(Worker)
+    worker._sleep_mode_backend = backend
+    worker._sleep_saved_buffers = {}
+    worker.model_runner = SimpleNamespace(
+        post_kv_cache_wake_up=lambda: None,
+    )
+    return worker
+
+
+def _record_comm_hooks(monkeypatch: pytest.MonkeyPatch, events: list[str]) -> None:
+    monkeypatch.setattr(
+        gpu_worker_module,
+        "prepare_communicators_for_suspend",
+        lambda: events.append("prepare_comms"),
+    )
+    monkeypatch.setattr(
+        gpu_worker_module,
+        "reinit_communicators_after_resume",
+        lambda: events.append("reinit_comms"),
+    )
+
+
+def _stub_accelerator(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        gpu_worker_module.torch.accelerator, "synchronize", lambda: None
+    )
+    monkeypatch.setattr(
+        gpu_worker_module.torch.accelerator, "get_memory_info", lambda: (0, 0)
+    )
+
+
+def test_sleep_prepares_communicators_before_suspend(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    events: list[str] = []
+    _record_comm_hooks(monkeypatch, events)
+    _stub_accelerator(monkeypatch)
+    backend = _RebuildingBackend()
+    worker = _sleep_worker(backend)
+
+    worker.sleep(level=1)
+
+    # Teardown must land before the backend snapshots/frees GPU state.
+    assert events == ["prepare_comms"]
+    assert backend.calls == [("suspend", 1)]
+
+
+def test_wake_up_rebuilds_communicators_then_invalidates_graphs(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    events: list[str] = []
+    _record_comm_hooks(monkeypatch, events)
+    backend = _RebuildingBackend()
+    worker = _sleep_worker(backend)
+
+    # Graphs embedding comm handles are stale after the rebuild; the draft
+    # contract fails loudly until discard+recapture exists.
+    with pytest.raises(NotImplementedError, match="invalidates captured"):
+        worker.wake_up()
+
+    assert backend.calls == [("resume", None)]
+    assert events == ["reinit_comms"]
+
+
+def test_preserving_backend_skips_comm_orchestration(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    events: list[str] = []
+    _record_comm_hooks(monkeypatch, events)
+    _stub_accelerator(monkeypatch)
+    backend = _PreservingBackend()
+    worker = _sleep_worker(backend)
+
+    worker.sleep(level=1)
+    worker.wake_up()
+
+    # cumem-shaped backends keep today's exact path: no hook is invoked.
+    assert events == []
+    assert backend.calls == [("suspend", 1), ("resume", None)]
+
+
+def _group_with_world_size(world_size: int) -> GroupCoordinator:
+    group = object.__new__(GroupCoordinator)
+    group.world_size = world_size
+    group.unique_name = f"test_group_ws{world_size}"
+    return group
+
+
+def test_group_suspend_hooks_noop_for_single_member_groups():
+    group = _group_with_world_size(1)
+    group.prepare_for_suspend()
+    group.reinit_after_resume()
+
+
+def test_group_suspend_hooks_fail_loudly_for_multi_member_groups():
+    group = _group_with_world_size(2)
+    with pytest.raises(NotImplementedError, match="nccl_checkpoint"):
+        group.prepare_for_suspend()
+    with pytest.raises(NotImplementedError, match="nccl_checkpoint"):
+        group.reinit_after_resume()

vllm/distributed/parallel_state.py

@@ -1213,6 +1213,46 @@ def destroy(self):
         if self.mq_broadcaster is not None:
             self.mq_broadcaster = None
 
+    def prepare_for_suspend(self) -> None:
+        """Tear down communicator state that cannot survive a process
+        suspend (peer mappings, RDMA connections, IPC handles).
+
+        Called by the GPU worker before ``SleepModeBackend.suspend`` when
+        the selected backend reports ``preserves_communicators() == False``
+        (see ``vllm/device_allocator/sleep_mode_backend.py``). Each
+        communicator owns its own teardown and rebuild; the worker only
+        sequences the calls.
+
+        Single-member groups hold no peer state and no-op. Multi-member
+        teardown/rebuild requires checkpoint support inside the comms
+        library itself (e.g. NCCL's ``contrib/nccl_checkpoint``), so until
+        a communicator implements it, suspending a multi-member group
+        fails loudly here instead of resuming into corrupt comm state.
+        """
+        if self.world_size == 1:
+            return
+        raise NotImplementedError(
+            f"Communicator group '{self.unique_name}' (world_size="
+            f"{self.world_size}) cannot be suspended yet: multi-member "
+            "communicator teardown/rebuild requires checkpoint support in "
+            "the comms library (RFC #34303; NCCL contrib/nccl_checkpoint)."
+        )
+
+    def reinit_after_resume(self) -> None:
+        """Recreate communicator state after ``SleepModeBackend.resume``.
+
+        Mirror of ``prepare_for_suspend``; see its docstring for the
+        contract and the single- vs multi-member behavior.
+        """
+        if self.world_size == 1:
+            return
+        raise NotImplementedError(
+            f"Communicator group '{self.unique_name}' (world_size="
+            f"{self.world_size}) cannot be rebuilt after resume yet: "
+            "requires checkpoint support in the comms library "
+            "(RFC #34303; NCCL contrib/nccl_checkpoint)."
+        )
+
     def prepare_communication_buffer_for_model(self, model: torch.nn.Module):
         if self.device_communicator is not None:
             self.device_communicator.prepare_communication_buffer_for_model(model)
@@ -2083,6 +2123,39 @@ def destroy_model_parallel():
     _EPLB = None
 
 
+def _live_parallel_groups() -> list[GroupCoordinator]:
+    """The model-parallel subgroups currently initialized, in the order
+    ``destroy_model_parallel`` handles them. Excludes the world group."""
+    return [g for g in (_TP, _DCP, _PCP, _PP, _DP, _EP, _EPLB) if g is not None]
+
+
+def prepare_communicators_for_suspend() -> None:
+    """Prepare every live communicator group for a process suspend.
+
+    Consumed by the GPU worker when the selected sleep-mode backend
+    reports ``preserves_communicators() == False``: teardown is owned by
+    each communicator (``GroupCoordinator.prepare_for_suspend``), and this
+    function only sequences it — subgroups first, world group last.
+    ``reinit_communicators_after_resume`` rebuilds in the opposite order.
+    """
+    for group in _live_parallel_groups():
+        group.prepare_for_suspend()
+    if _WORLD is not None:
+        _WORLD.prepare_for_suspend()
+
+
+def reinit_communicators_after_resume() -> None:
+    """Rebuild every live communicator group after a process resume.
+
+    Counterpart of ``prepare_communicators_for_suspend``: world group
+    first, then the model-parallel subgroups.
+    """
+    if _WORLD is not None:
+        _WORLD.reinit_after_resume()
+    for group in _live_parallel_groups():
+        group.reinit_after_resume()
+
+
 def destroy_distributed_environment():
     global _WORLD, _NODE_COUNT
     if _WORLD:

vllm/v1/worker/gpu_worker.py

@@ -43,6 +43,8 @@
     Handle,
     get_pp_group,
     get_tp_group,
+    prepare_communicators_for_suspend,
+    reinit_communicators_after_resume,
 )
 from vllm.distributed.weight_transfer import (
     WeightTransferEngine,
@@ -199,7 +201,13 @@ def sleep(self, level: int = 1) -> None:
                 name: buffer.cpu().clone() for name, buffer in model.named_buffers()
             }
 
-        self._get_sleep_mode_backend().suspend(level)
+        backend = self._get_sleep_mode_backend()
+        if not backend.preserves_communicators():
+            # Communicator-owned teardown must complete before the backend
+            # snapshots/frees GPU state (peer mappings and RDMA connections
+            # cannot survive a suspend). See prepare_communicators_for_suspend.
+            prepare_communicators_for_suspend()
+        backend.suspend(level)
 
         torch.accelerator.synchronize()
         deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
@@ -219,7 +227,15 @@ def sleep(self, level: int = 1) -> None:
         )
 
     def wake_up(self, tags: list[str] | None = None) -> None:
-        self._get_sleep_mode_backend().resume(tags)
+        backend = self._get_sleep_mode_backend()
+        backend.resume(tags)
+
+        if not backend.preserves_communicators():
+            # Mirror of the suspend-side teardown: rebuild is owned by each
+            # communicator, sequenced world-group-first.
+            reinit_communicators_after_resume()
+            if not backend.preserves_graphs_with_communicators():
+                self._invalidate_captured_graphs()
 
         # Restore the buffers after level 2 sleep
         if len(self._sleep_saved_buffers):
@@ -232,6 +248,21 @@ def wake_up(self, tags: list[str] | None = None) -> None:
         if tags is None or "kv_cache" in tags:
             self.model_runner.post_kv_cache_wake_up()
 
+    def _invalidate_captured_graphs(self) -> None:
+        """Captured CUDA graphs embed communicator handles; once communicators
+        are rebuilt on resume (``preserves_graphs_with_communicators() ==
+        False``), replaying a captured graph would launch into destroyed comm
+        state. The safe sequel is discard + recapture — an init-time-style
+        operation, since a suspended engine is not serving traffic — but the
+        discard half has no in-tree mechanism yet, so this fails loudly
+        instead of resuming with stale graphs (RFC #34303).
+        """
+        raise NotImplementedError(
+            "The sleep-mode backend rebuilt communicators on resume, which "
+            "invalidates captured CUDA graphs; discard+recapture after resume "
+            "is not implemented yet (RFC #34303)."
+        )
+
     def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
         if (
             current_platform.is_cuda_alike()