[Core] Pluggable sleep-mode backend abstraction (RFC #34303)

tests/v1/worker/test_sleep_mode_backend.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU-only unit tests for the sleep-mode backend abstraction (RFC #34303).
+
+These cover the registry/factory contract and capability flags. They do not
+touch CUDA - the ``cumem`` suspend/resume path is exercised end-to-end on GPU
+in ``tests/basic_correctness/test_cumem.py``.
+"""
+
+import pytest
+
+from vllm.device_allocator.sleep_mode_backend import (
+    CuMemBackend,
+    SleepModeBackend,
+    SleepModeBackendFactory,
+)
+
+
+def test_cumem_is_the_default_registered_backend():
+    backend_cls = SleepModeBackendFactory.get_backend_class("cumem")
+    assert backend_cls is CuMemBackend
+    assert issubclass(backend_cls, SleepModeBackend)
+
+
+def test_cumem_capability_flags():
+    # cumem leaves NCCL untouched but does not preserve compiled artifacts,
+    # graphs, or durable state - these flags are what the executor and /health
+    # introspect to decide reinit / persistence behavior.
+    assert CuMemBackend.is_supported() is True
+    assert CuMemBackend.preserves_nccl() is True
+    assert CuMemBackend.preserves_compiled_artifacts() is False
+    assert CuMemBackend.preserves_graphs_with_nccl() is False
+    assert CuMemBackend.supports_durable_storage() is False
+
+
+def test_new_backend_starts_in_running_state():
+    # Constructing a backend must not touch the GPU; only suspend/resume do.
+    assert CuMemBackend().state() == "RUNNING"
+
+
+def test_unknown_backend_raises():
+    with pytest.raises(ValueError, match="Unsupported sleep-mode backend"):
+        SleepModeBackendFactory.get_backend_class("does-not-exist")
+
+
+def test_duplicate_registration_raises():
+    with pytest.raises(ValueError, match="already registered"):
+        SleepModeBackendFactory.register_backend(
+            "cumem",
+            "vllm.device_allocator.sleep_mode_backend",
+            "CuMemBackend",
+        )
+
+
+def test_third_party_backend_registration_and_resolution():
+    """A plugin registers a backend by name; the factory resolves it lazily."""
+    name = "_pytest_dummy_backend"
+    try:
+        SleepModeBackendFactory.register_backend(
+            name,
+            "tests.v1.worker.test_sleep_mode_backend",
+            "DummyBackend",
+        )
+        resolved = SleepModeBackendFactory.get_backend_class(name)
+        assert resolved is DummyBackend
+        assert resolved.supports_durable_storage() is True
+    finally:
+        SleepModeBackendFactory._registry.pop(name, None)
+
+
+def test_suspend_resume_state_transitions():
+    """Lifecycle state advances RUNNING -> SUSPENDED -> RUNNING without GPU."""
+    backend = DummyBackend()
+    assert backend.state() == "RUNNING"
+    backend.suspend(level=1)
+    assert backend.state() == "SUSPENDED"
+    backend.resume()
+    assert backend.state() == "RUNNING"
+
+
+class DummyBackend(SleepModeBackend):
+    """A no-GPU backend used to exercise lifecycle + registration in CPU tests."""
+
+    def suspend(self, level: int = 1) -> None:
+        self._state = "SUSPENDED"
+
+    def resume(self, tags: list[str] | None = None) -> None:
+        self._state = "RUNNING"
+
+    @classmethod
+    def supports_durable_storage(cls) -> bool:
+        return True

vllm/config/model.py

@@ -289,6 +289,11 @@ class ModelConfig:
     enable_sleep_mode: bool = False
     """Enable sleep mode for the engine (only cuda and
     hip platforms are supported)."""
+    sleep_mode_backend: str = "cumem"
+    """Mechanism used to free and restore GPU state for sleep mode. ``"cumem"``
+    (default) uses the built-in ``CuMemAllocator`` and is behavior-compatible
+    with prior releases. Additional backends (CUDA checkpoint, CRIU, durable
+    snapshot) may be registered in-tree or by plugins (RFC #34303)."""
     enable_cumem_allocator: bool = False
     """Enable the custom cumem allocator to leverage advanced GPU memory
     allocation features such as multi-node NVLink support.

vllm/device_allocator/sleep_mode_backend.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pluggable sleep-mode backends (RFC #34303).
+
+vLLM's sleep/wake-up today is hard-wired to ``CuMemAllocator``: the GPU worker
+calls ``allocator.sleep(...)`` / ``allocator.wake_up(...)`` directly. RFC #34303
+proposes additional mechanisms for freeing and restoring GPU state - CUDA
+process checkpoint, CRIU, durable snapshot/restore - that share the *dispatch*
+(``/sleep`` endpoint -> engine -> executor -> worker) but differ in *mechanism*
+and in which resources they preserve (NCCL communicators, compiled kernels,
+CUDA graphs, survival across process restart).
+
+This module introduces a thin backend abstraction so those mechanisms can be
+selected by name without changing the public API. The default ``cumem`` backend
+wraps today's ``CuMemAllocator`` path 1:1, so existing users see no behavior
+change. The factory mirrors ``KVConnectorFactory`` and lets third-party
+backends register through a ``vllm.general_plugins`` entry point at import time.
+"""
+
+from __future__ import annotations
+
+import importlib
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Literal
+
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config.model import ModelConfig
+
+logger = init_logger(__name__)
+
+SleepModeState = Literal["RUNNING", "SUSPENDED", "RESUMING"]
+
+
+class SleepModeBackend(ABC):
+    """Interface for a mechanism that frees and restores GPU state.
+
+    A backend owns the *mechanism* of suspend/resume. The dispatch path
+    (``/sleep`` endpoint -> engine -> executor -> worker) is shared across all
+    backends and lives outside this class.
+
+    Capability flags are ``@classmethod`` so callers (executor, ``/health``,
+    AUTO selection) can introspect a backend without instantiating it, matching
+    the capability-flag convention used by attention backends.
+    """
+
+    def __init__(self) -> None:
+        self._state: SleepModeState = "RUNNING"
+
+    @abstractmethod
+    def suspend(self, level: int = 1) -> None:
+        """Free GPU state.
+
+        ``level`` follows existing sleep-mode semantics: level 1 offloads
+        weights to host RAM (restorable in-process); level 2 discards weights
+        (reloaded from the model source on resume).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def resume(self, tags: list[str] | None = None) -> None:
+        """Restore previously-suspended GPU state.
+
+        ``tags`` optionally limits which tagged allocations are restored
+        (e.g. ``["weights"]`` or ``["kv_cache"]``).
+        """
+        raise NotImplementedError
+
+    def state(self) -> SleepModeState:
+        """Current lifecycle state. Lets ``/health`` distinguish a healthy-idle
+        (suspended) engine from a healthy-serving one (see RFC #34303)."""
+        return self._state
+
+    # -- Capability introspection (no instance required) --
+
+    @classmethod
+    def is_supported(cls) -> bool:
+        """Whether this backend can run on the current platform/driver."""
+        return True
+
+    @classmethod
+    def preserves_nccl(cls) -> bool:
+        """If False, NCCL communicators are destroyed by ``suspend`` and the
+        executor must re-initialize them on ``resume``."""
+        return False
+
+    @classmethod
+    def preserves_compiled_artifacts(cls) -> bool:
+        """If True, torch.compile / JIT kernels survive suspend/resume and need
+        not be recompiled on resume."""
+        return False
+
+    @classmethod
+    def preserves_graphs_with_nccl(cls) -> bool:
+        """If True, CUDA graphs containing NCCL collectives stay valid after
+        resume. False when NCCL is rebuilt (embedded comm handles go stale)."""
+        return False
+
+    @classmethod
+    def supports_durable_storage(cls) -> bool:
+        """If True, suspended state can be persisted beyond the process
+        lifetime (disk or object storage) and restored in a new process."""
+        return False
+
+
+class CuMemBackend(SleepModeBackend):
+    """Default backend.
+
+    Wraps the platform sleep-mode allocator exactly as the GPU worker did
+    before this abstraction existed, so behavior is identical to vLLM's current
+    sleep/wake-up. ``get_mem_allocator_instance()`` resolves to
+    ``CuMemAllocator`` on CUDA and ``XpuMemAllocator`` on XPU; suspend offloads
+    per-allocation between GPU and host, with NCCL buffers left untouched (they
+    are allocated outside the allocator pool).
+    """
+
+    def suspend(self, level: int = 1) -> None:
+        from vllm.device_allocator import get_mem_allocator_instance
+
+        self._state = "SUSPENDED"
+        allocator = get_mem_allocator_instance()
+        allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
+
+    def resume(self, tags: list[str] | None = None) -> None:
+        from vllm.device_allocator import get_mem_allocator_instance
+
+        self._state = "RESUMING"
+        allocator = get_mem_allocator_instance()
+        allocator.wake_up(tags)
+        self._state = "RUNNING"
+
+    @classmethod
+    def preserves_nccl(cls) -> bool:
+        # NCCL buffers live outside CuMemAllocator's pool, so an allocator-level
+        # sleep leaves the communicators intact (no reinit needed on resume).
+        return True
+
+
+class SleepModeBackendFactory:
+    """Registry and resolver for sleep-mode backends.
+
+    Mirrors ``KVConnectorFactory``: lazy module/class registration and a
+    built-in registry populated at import time. Third-party backends register
+    the same way from a ``vllm.general_plugins`` entry point.
+    """
+
+    _registry: dict[str, Callable[[], type[SleepModeBackend]]] = {}
+
+    @classmethod
+    def register_backend(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a backend with a lazy-loading module and class name."""
+        if name in cls._registry:
+            raise ValueError(f"Sleep-mode backend '{name}' is already registered.")
+
+        def loader() -> type[SleepModeBackend]:
+            module = importlib.import_module(module_path)
+            return getattr(module, class_name)
+
+        cls._registry[name] = loader
+
+    @classmethod
+    def get_backend_class(cls, name: str) -> type[SleepModeBackend]:
+        """Resolve a registered backend class by name."""
+        if name not in cls._registry:
+            available = ", ".join(sorted(cls._registry)) or "<none>"
+            raise ValueError(
+                f"Unsupported sleep-mode backend '{name}'. "
+                f"Registered backends: {available}."
+            )
+        return cls._registry[name]()
+
+    @classmethod
+    def create_backend(cls, model_config: ModelConfig) -> SleepModeBackend:
+        """Instantiate the backend selected by ``model_config``."""
+        name = model_config.sleep_mode_backend
+        backend_cls = cls.get_backend_class(name)
+        if not backend_cls.is_supported():
+            raise ValueError(
+                f"Sleep-mode backend '{name}' is not supported on this platform."
+            )
+        logger.info("Using sleep-mode backend: %s", name)
+        return backend_cls()
+
+
+# Register built-in backends here. Registration is lazy: only the module for the
+# selected backend is imported. Third-party backends (CUDA checkpoint, CRIU,
+# durable snapshot) register the same way through a vllm.general_plugins entry
+# point, without changes to vLLM core.
+SleepModeBackendFactory.register_backend(
+    "cumem",
+    "vllm.device_allocator.sleep_mode_backend",
+    "CuMemBackend",
+)

vllm/v1/worker/gpu_worker.py

@@ -86,6 +86,7 @@
 logger = init_logger(__name__)
 
 if TYPE_CHECKING:
+    from vllm.device_allocator.sleep_mode_backend import SleepModeBackend
     from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
@@ -170,6 +171,20 @@ def __init__(
         # pending non-blocking PP send work from the previous iteration
         self._pp_send_work: list[Handle] = []
 
+        # Resolved lazily on first sleep/wake; persists worker-process state.
+        self._sleep_mode_backend: SleepModeBackend | None = None
+
+    def _get_sleep_mode_backend(self) -> "SleepModeBackend":
+        if self._sleep_mode_backend is None:
+            from vllm.device_allocator.sleep_mode_backend import (
+                SleepModeBackendFactory,
+            )
+
+            self._sleep_mode_backend = SleepModeBackendFactory.create_backend(
+                self.vllm_config.model_config
+            )
+        return self._sleep_mode_backend
+
     def sleep(self, level: int = 1) -> None:
         torch.accelerator.synchronize()
         free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]
@@ -181,8 +196,7 @@ def sleep(self, level: int = 1) -> None:
                 name: buffer.cpu().clone() for name, buffer in model.named_buffers()
             }
 
-        allocator = get_mem_allocator_instance()
-        allocator.sleep(offload_tags=("weights",) if level == 1 else tuple())
+        self._get_sleep_mode_backend().suspend(level)
 
         torch.accelerator.synchronize()
         deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
@@ -202,8 +216,7 @@ def sleep(self, level: int = 1) -> None:
         )
 
     def wake_up(self, tags: list[str] | None = None) -> None:
-        allocator = get_mem_allocator_instance()
-        allocator.wake_up(tags)
+        self._get_sleep_mode_backend().resume(tags)
 
         # Restore the buffers after level 2 sleep
         if len(self._sleep_saved_buffers):