| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Pluggable sleep-mode backends (RFC #34303). |
| 4 | |
| 5 | vLLM's sleep/wake-up today is hard-wired to ``CuMemAllocator``: the GPU worker |
| 6 | calls ``allocator.sleep(...)`` / ``allocator.wake_up(...)`` directly. RFC #34303 |
| 7 | proposes additional mechanisms for freeing and restoring GPU state - CUDA |
| 8 | process checkpoint, CRIU, durable snapshot/restore - that share the *dispatch* |
| 9 | (``/sleep`` endpoint -> engine -> executor -> worker) but differ in *mechanism* |
| 10 | and in which resources they preserve (NCCL communicators, compiled kernels, |
| 11 | CUDA graphs, survival across process restart). |
| 12 | |
| 13 | This module introduces a thin backend abstraction so those mechanisms can be |
| 14 | selected by name without changing the public API. The default ``cumem`` backend |
| 15 | wraps today's ``CuMemAllocator`` path 1:1, so existing users see no behavior |
| 16 | change. The factory mirrors ``KVConnectorFactory`` and lets third-party |
| 17 | backends register through a ``vllm.general_plugins`` entry point at import time. |
| 18 | """ |
| 19 | |
| 20 | from __future__ import annotations |
| 21 | |
| 22 | import importlib |
| 23 | from abc import ABC, abstractmethod |
| 24 | from collections.abc import Callable |
| 25 | from typing import TYPE_CHECKING, Literal |
| 26 | |
| 27 | from vllm.logger import init_logger |
| 28 | |
| 29 | if TYPE_CHECKING: |
| 30 | from vllm.config.model import ModelConfig |
| 31 | |
| 32 | logger = init_logger(__name__) |
| 33 | |
| 34 | SleepModeState = Literal["RUNNING", "SUSPENDED", "RESUMING"] |
| 35 | |
| 36 | |
| 37 | class SleepModeBackend(ABC): |
| 38 | """Interface for a mechanism that frees and restores GPU state. |
| 39 | |
| 40 | A backend owns the *mechanism* of suspend/resume. The dispatch path |
| 41 | (``/sleep`` endpoint -> engine -> executor -> worker) is shared across all |
| 42 | backends and lives outside this class. |
| 43 | |
| 44 | Capability flags are ``@classmethod`` so callers (executor, ``/health``, |
| 45 | AUTO selection) can introspect a backend without instantiating it, matching |
| 46 | the capability-flag convention used by attention backends. |
| 47 | """ |
| 48 | |
| 49 | def __init__(self) -> None: |
| 50 | self._state: SleepModeState = "RUNNING" |
| 51 | |
| 52 | @abstractmethod |
| 53 | def suspend(self, level: int = 1) -> None: |
| 54 | """Free GPU state. |
| 55 | |
| 56 | ``level`` follows existing sleep-mode semantics: level 1 offloads |
| 57 | weights to host RAM (restorable in-process); level 2 discards weights |
| 58 | (reloaded from the model source on resume). |
| 59 | """ |
| 60 | raise NotImplementedError |
| 61 | |
| 62 | @abstractmethod |
| 63 | def resume(self, tags: list[str] | None = None) -> None: |
| 64 | """Restore previously-suspended GPU state. |
| 65 | |
| 66 | ``tags`` optionally limits which tagged allocations are restored |
| 67 | (e.g. ``["weights"]`` or ``["kv_cache"]``). |
| 68 | """ |
| 69 | raise NotImplementedError |
| 70 | |
| 71 | def state(self) -> SleepModeState: |
| 72 | """Current lifecycle state. Lets ``/health`` distinguish a healthy-idle |
| 73 | (suspended) engine from a healthy-serving one (see RFC #34303).""" |
| 74 | return self._state |
| 75 | |
| 76 | def discard(self, tags: tuple[str, ...]) -> None: |
| 77 | """Discard selected tagged allocations without preserving contents.""" |
| 78 | raise RuntimeError( |
| 79 | f"Sleep-mode backend {type(self).__name__!r} does not support " |
| 80 | "selective discard." |
| 81 | ) |
| 82 | |
| 83 | # -- Capability introspection (no instance required) -- |
| 84 | |
| 85 | @classmethod |
| 86 | def is_supported(cls) -> bool: |
| 87 | """Whether this backend can run on the current platform/driver.""" |
| 88 | return True |
| 89 | |
| 90 | @classmethod |
| 91 | def preserves_communicators(cls) -> bool: |
| 92 | """If False, collective communicators (e.g. NCCL) are destroyed by |
| 93 | ``suspend`` and the executor must re-initialize them on ``resume``.""" |
| 94 | return False |
| 95 | |
| 96 | @classmethod |
| 97 | def preserves_compiled_artifacts(cls) -> bool: |
| 98 | """If True, torch.compile / JIT kernels survive suspend/resume and need |
| 99 | not be recompiled on resume.""" |
| 100 | return False |
| 101 | |
| 102 | @classmethod |
| 103 | def preserves_graphs_with_communicators(cls) -> bool: |
| 104 | """If True, CUDA graphs containing collective communicators (e.g. NCCL) |
| 105 | stay valid after resume. False when communicators are rebuilt (embedded |
| 106 | comm handles go stale).""" |
| 107 | return False |
| 108 | |
| 109 | @classmethod |
| 110 | def supports_durable_storage(cls) -> bool: |
| 111 | """If True, suspended state can be persisted beyond the process |
| 112 | lifetime (disk or object storage) and restored in a new process.""" |
| 113 | return False |
| 114 | |
| 115 | |
| 116 | class CuMemBackend(SleepModeBackend): |
| 117 | """Default backend. |
| 118 | |
| 119 | Wraps the platform sleep-mode allocator exactly as the GPU worker did |
| 120 | before this abstraction existed, so behavior is identical to vLLM's current |
| 121 | sleep/wake-up. ``get_mem_allocator_instance()`` resolves to |
| 122 | ``CuMemAllocator`` on CUDA and ``XpuMemAllocator`` on XPU; suspend offloads |
| 123 | per-allocation between GPU and host, with NCCL buffers left untouched (they |
| 124 | are allocated outside the allocator pool). |
| 125 | """ |
| 126 | |
| 127 | def suspend(self, level: int = 1) -> None: |
| 128 | from vllm.device_allocator import get_mem_allocator_instance |
| 129 | |
| 130 | self._state = "SUSPENDED" |
| 131 | allocator = get_mem_allocator_instance() |
| 132 | # Runtime state is kept at every level; weights only at level 1. |
| 133 | offload_tags = ("weights", "runtime") if level == 1 else ("runtime",) |
| 134 | allocator.sleep(offload_tags=offload_tags) |
| 135 | |
| 136 | def resume(self, tags: list[str] | None = None) -> None: |
| 137 | from vllm.device_allocator import get_mem_allocator_instance |
| 138 | |
| 139 | self._state = "RESUMING" |
| 140 | allocator = get_mem_allocator_instance() |
| 141 | allocator.wake_up(tags) |
| 142 | self._state = "RUNNING" |
| 143 | |
| 144 | def discard(self, tags: tuple[str, ...]) -> None: |
| 145 | from vllm.device_allocator import get_mem_allocator_instance |
| 146 | |
| 147 | allocator = get_mem_allocator_instance() |
| 148 | allocator.discard(tags) |
| 149 | self._state = "SUSPENDED" |
| 150 | |
| 151 | @classmethod |
| 152 | def preserves_communicators(cls) -> bool: |
| 153 | # Comm identity/topology survive memory suspension; no reinit needed. |
| 154 | return True |
| 155 | |
| 156 | |
| 157 | class SleepModeBackendFactory: |
| 158 | """Registry and resolver for sleep-mode backends. |
| 159 | |
| 160 | Mirrors ``KVConnectorFactory``: lazy module/class registration and a |
| 161 | built-in registry populated at import time. Third-party backends register |
| 162 | the same way from a ``vllm.general_plugins`` entry point. |
| 163 | """ |
| 164 | |
| 165 | _registry: dict[str, Callable[[], type[SleepModeBackend]]] = {} |
| 166 | |
| 167 | @classmethod |
| 168 | def register_backend(cls, name: str, module_path: str, class_name: str) -> None: |
| 169 | """Register a backend with a lazy-loading module and class name.""" |
| 170 | if name in cls._registry: |
| 171 | raise ValueError(f"Sleep-mode backend '{name}' is already registered.") |
| 172 | |
| 173 | def loader() -> type[SleepModeBackend]: |
| 174 | module = importlib.import_module(module_path) |
| 175 | return getattr(module, class_name) |
| 176 | |
| 177 | cls._registry[name] = loader |
| 178 | |
| 179 | @classmethod |
| 180 | def get_backend_class(cls, name: str) -> type[SleepModeBackend]: |
| 181 | """Resolve a registered backend class by name.""" |
| 182 | if name not in cls._registry: |
| 183 | available = ", ".join(sorted(cls._registry)) or "<none>" |
| 184 | raise ValueError( |
| 185 | f"Unsupported sleep-mode backend '{name}'. " |
| 186 | f"Registered backends: {available}." |
| 187 | ) |
| 188 | return cls._registry[name]() |
| 189 | |
| 190 | @classmethod |
| 191 | def create_backend(cls, model_config: ModelConfig) -> SleepModeBackend: |
| 192 | """Instantiate the backend selected by ``model_config``.""" |
| 193 | name = model_config.sleep_mode_backend |
| 194 | backend_cls = cls.get_backend_class(name) |
| 195 | if not backend_cls.is_supported(): |
| 196 | raise ValueError( |
| 197 | f"Sleep-mode backend '{name}' is not supported on this platform." |
| 198 | ) |
| 199 | logger.info("Using sleep-mode backend: %s", name) |
| 200 | return backend_cls() |
| 201 | |
| 202 | |
| 203 | # Register built-in backends here. Registration is lazy: only the module for the |
| 204 | # selected backend is imported. Third-party backends (CUDA checkpoint, CRIU, |
| 205 | # durable snapshot) register the same way through a vllm.general_plugins entry |
| 206 | # point, without changes to vLLM core. |
| 207 | SleepModeBackendFactory.register_backend( |
| 208 | "cumem", |
| 209 | "vllm.device_allocator.sleep_mode_backend", |
| 210 | "CuMemBackend", |
| 211 | ) |