Address PR 44074 sleep backend review comments

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,

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.
 
@@ -172,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)