From 2fdc5a08ee9a9d3c31a086f8c32c454b0959327f Mon Sep 17 00:00:00 2001 From: Simon Mo Date: Tue, 30 Jun 2026 18:05:43 -0700 Subject: [PATCH 1/2] Address sleep mode backend review comments Co-authored-by: GPT-5.5 --- tests/test_sleep_mode_backend.py | 92 --------------------- tests/v1/worker/test_gpu_worker.py | 66 +++++++++++++++ vllm/device_allocator/sleep_mode_backend.py | 49 ++--------- 3 files changed, 72 insertions(+), 135 deletions(-) delete mode 100644 tests/test_sleep_mode_backend.py diff --git a/tests/test_sleep_mode_backend.py b/tests/test_sleep_mode_backend.py deleted file mode 100644 index a5bb2d315b2c..000000000000 --- a/tests/test_sleep_mode_backend.py +++ /dev/null @@ -1,92 +0,0 @@ -# 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.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/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 31be4a8402f8..fcf6cede1c83 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -6,6 +6,11 @@ import pytest import vllm.v1.worker.gpu_worker as gpu_worker_module +from vllm.device_allocator.sleep_mode_backend import ( + CuMemBackend, + SleepModeBackend, + SleepModeBackendFactory, +) from vllm.multimodal.video import ( PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, @@ -46,6 +51,67 @@ def _pynvvideocodec_decoder_budget(api_process_count: int = 1) -> int: ) +def test_sleep_mode_cumem_is_the_default_registered_backend(): + backend_cls = SleepModeBackendFactory.get_backend_class("cumem") + assert backend_cls is CuMemBackend + assert issubclass(backend_cls, SleepModeBackend) + assert backend_cls.is_supported() is True + + +def test_sleep_mode_new_backend_starts_in_running_state(): + # Constructing a backend must not touch the GPU; only suspend/resume do. + assert CuMemBackend().state() == "RUNNING" + + +def test_sleep_mode_unknown_backend_raises(): + with pytest.raises(ValueError, match="Unsupported sleep-mode backend"): + SleepModeBackendFactory.get_backend_class("does-not-exist") + + +def test_sleep_mode_duplicate_registration_raises(): + with pytest.raises(ValueError, match="already registered"): + SleepModeBackendFactory.register_backend( + "cumem", + "vllm.device_allocator.sleep_mode_backend", + "CuMemBackend", + ) + + +def test_sleep_mode_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, + __name__, + "DummySleepModeBackend", + ) + resolved = SleepModeBackendFactory.get_backend_class(name) + assert resolved is DummySleepModeBackend + finally: + SleepModeBackendFactory._registry.pop(name, None) + + +def test_sleep_mode_suspend_resume_state_transitions(): + """Lifecycle state advances RUNNING -> SUSPENDED -> RUNNING without GPU.""" + backend = DummySleepModeBackend() + assert backend.state() == "RUNNING" + backend.suspend(level=1) + assert backend.state() == "SUSPENDED" + backend.resume() + assert backend.state() == "RUNNING" + + +class DummySleepModeBackend(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" + + @pytest.mark.parametrize("video_backend", [None, "opencv"]) def test_reserve_mm_ipc_gpu_memory_raw_frame_budget_only( monkeypatch: pytest.MonkeyPatch, diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py index 81fccf5325ec..80f78acfe8df 100644 --- a/vllm/device_allocator/sleep_mode_backend.py +++ b/vllm/device_allocator/sleep_mode_backend.py @@ -4,11 +4,9 @@ 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). +proposes additional mechanisms for freeing and restoring GPU state that share +the *dispatch* (``/sleep`` endpoint -> engine -> executor -> worker) but differ +in *mechanism*. This module introduces a thin backend abstraction so those mechanisms can be selected by name without changing the public API. The default ``cumem`` backend @@ -41,9 +39,8 @@ class SleepModeBackend(ABC): (``/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. + Backend classes expose support checks without requiring instantiation so + callers can validate platform availability before creating a backend. """ def __init__(self) -> None: @@ -73,37 +70,11 @@ def state(self) -> SleepModeState: (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. @@ -112,8 +83,7 @@ class CuMemBackend(SleepModeBackend): 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). + per-allocation between GPU and host. """ def suspend(self, level: int = 1) -> None: @@ -131,13 +101,6 @@ def resume(self, tags: list[str] | None = None) -> None: 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. From 4ff7f02daaec45f66b5a86f8cc66830e4955303f Mon Sep 17 00:00:00 2001 From: Simon Mo Date: Tue, 30 Jun 2026 18:06:47 -0700 Subject: [PATCH 2/2] Apply sleep backend lint fix Co-authored-by: GPT-5.5 --- vllm/device_allocator/sleep_mode_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py index 80f78acfe8df..72d7d8d6464f 100644 --- a/vllm/device_allocator/sleep_mode_backend.py +++ b/vllm/device_allocator/sleep_mode_backend.py @@ -135,7 +135,7 @@ def get_backend_class(cls, name: str) -> type[SleepModeBackend]: return cls._registry[name]() @classmethod - def create_backend(cls, model_config: "ModelConfig") -> SleepModeBackend: + 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)