From a0a49d46c9299237434e3eeaf630cee2f8c9517f Mon Sep 17 00:00:00 2001 From: Nils Matteson Date: Sat, 30 May 2026 12:45:01 -0600 Subject: [PATCH] [Core] Pluggable sleep-mode backend abstraction (RFC #34303) Introduce a thin backend abstraction in front of the sleep/wake-up GPU path so alternative mechanisms (CUDA checkpoint, CRIU, durable snapshot) proposed in RFC #34303 can be selected by name without changing the public API. - New vllm/device_allocator/sleep_mode_backend.py: - SleepModeBackend ABC (suspend/resume + capability classmethods + RUNNING/SUSPENDED/RESUMING state). - CuMemBackend default - wraps CuMemAllocator 1:1. - SleepModeBackendFactory mirroring KVConnectorFactory (lazy registry, plugin-registerable via vllm.general_plugins). - ModelConfig.sleep_mode_backend: str = "cumem" (new field, default preserves current behavior; auto-exposed as --sleep-mode-backend). - GPUWorker.sleep()/wake_up() dispatch through the factory. The cumem backend issues the identical allocator calls, so behavior is unchanged for every existing user. - CPU-only unit tests for the registry/factory contract and capability flags (GPU suspend/resume stays covered by test_cumem.py). Refs #34303 Signed-off-by: Nils Matteson --- tests/v1/worker/test_sleep_mode_backend.py | 92 +++++++++ vllm/config/model.py | 5 + vllm/device_allocator/sleep_mode_backend.py | 195 ++++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 21 ++- 4 files changed, 309 insertions(+), 4 deletions(-) create mode 100644 tests/v1/worker/test_sleep_mode_backend.py create mode 100644 vllm/device_allocator/sleep_mode_backend.py diff --git a/tests/v1/worker/test_sleep_mode_backend.py b/tests/v1/worker/test_sleep_mode_backend.py new file mode 100644 index 000000000000..684ff87837e2 --- /dev/null +++ b/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 diff --git a/vllm/config/model.py b/vllm/config/model.py index b12639d51604..e19c9408914e 100644 --- a/vllm/config/model.py +++ b/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. diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py new file mode 100644 index 000000000000..fe0b31a60d2b --- /dev/null +++ b/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 "" + 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", +) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 0c5512d5e15f..554ffbaea60f 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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):