| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """CPU-only unit tests for the sleep-mode backend abstraction (RFC #34303). |
| 4 | |
| 5 | These cover the registry/factory contract and capability flags. They do not |
| 6 | touch CUDA - the ``cumem`` suspend/resume path is exercised end-to-end on GPU |
| 7 | in ``tests/basic_correctness/memory/``. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | |
| 12 | from vllm.device_allocator.sleep_mode_backend import ( |
| 13 | CuMemBackend, |
| 14 | SleepModeBackend, |
| 15 | SleepModeBackendFactory, |
| 16 | ) |
| 17 | |
| 18 | |
| 19 | def test_cumem_is_the_default_registered_backend(): |
| 20 | backend_cls = SleepModeBackendFactory.get_backend_class("cumem") |
| 21 | assert backend_cls is CuMemBackend |
| 22 | assert issubclass(backend_cls, SleepModeBackend) |
| 23 | |
| 24 | |
| 25 | def test_cumem_capability_flags(): |
| 26 | # cumem leaves communicators untouched but does not preserve compiled |
| 27 | # artifacts, graphs, or durable state - these flags are what the executor and |
| 28 | # /health introspect to decide reinit / persistence behavior. |
| 29 | assert CuMemBackend.is_supported() is True |
| 30 | assert CuMemBackend.preserves_communicators() is True |
| 31 | assert CuMemBackend.preserves_compiled_artifacts() is False |
| 32 | assert CuMemBackend.preserves_graphs_with_communicators() is False |
| 33 | assert CuMemBackend.supports_durable_storage() is False |
| 34 | |
| 35 | |
| 36 | def test_new_backend_starts_in_running_state(): |
| 37 | # Constructing a backend must not touch the GPU; only suspend/resume do. |
| 38 | assert CuMemBackend().state() == "RUNNING" |
| 39 | |
| 40 | |
| 41 | @pytest.mark.parametrize("enable_nccl_comm_suspend", [True, False]) |
| 42 | def test_worker_drives_communicator_suspension(monkeypatch, enable_nccl_comm_suspend): |
| 43 | """Comm walkers run around sleep/wake only when explicitly enabled.""" |
| 44 | from types import SimpleNamespace |
| 45 | |
| 46 | from vllm.v1.worker.gpu_worker import Worker |
| 47 | |
| 48 | calls: list[tuple[str, object]] = [] |
| 49 | |
| 50 | class Backend: |
| 51 | def suspend(self, level: int = 1) -> None: |
| 52 | calls.append(("backend.suspend", level)) |
| 53 | |
| 54 | def resume(self, tags: list[str] | None = None) -> None: |
| 55 | calls.append(("backend.resume", tuple(tags) if tags else None)) |
| 56 | |
| 57 | worker = object.__new__(Worker) |
| 58 | worker._sleep_mode_backend = Backend() |
| 59 | worker._sleep_saved_parameters = {} |
| 60 | worker._sleep_saved_buffers = {} |
| 61 | worker._sleep_saved_draft_buffers = {} |
| 62 | worker.vllm_config = SimpleNamespace( |
| 63 | model_config=SimpleNamespace(enable_nccl_comm_suspend=enable_nccl_comm_suspend) |
| 64 | ) |
| 65 | |
| 66 | monkeypatch.setattr("torch.accelerator.synchronize", lambda: None) |
| 67 | monkeypatch.setattr("torch.accelerator.get_memory_info", lambda: (0, 0)) |
| 68 | monkeypatch.setattr( |
| 69 | "vllm.v1.worker.gpu_worker.suspend_device_comms", |
| 70 | lambda: calls.append(("comms.suspend", None)), |
| 71 | ) |
| 72 | monkeypatch.setattr( |
| 73 | "vllm.v1.worker.gpu_worker.resume_device_comms", |
| 74 | lambda: calls.append(("comms.resume", None)), |
| 75 | ) |
| 76 | |
| 77 | worker.sleep(level=1) |
| 78 | worker.wake_up(tags=["weights"]) |
| 79 | |
| 80 | expected = [ |
| 81 | ("backend.suspend", 1), |
| 82 | ("comms.suspend", None), |
| 83 | ("backend.resume", ("weights",)), |
| 84 | ("comms.resume", None), |
| 85 | ] |
| 86 | if not enable_nccl_comm_suspend: |
| 87 | expected = [c for c in expected if not c[0].startswith("comms.")] |
| 88 | assert calls == expected |
| 89 | |
| 90 | |
| 91 | def test_unknown_backend_raises(): |
| 92 | with pytest.raises(ValueError, match="Unsupported sleep-mode backend"): |
| 93 | SleepModeBackendFactory.get_backend_class("does-not-exist") |
| 94 | |
| 95 | |
| 96 | def test_duplicate_registration_raises(): |
| 97 | with pytest.raises(ValueError, match="already registered"): |
| 98 | SleepModeBackendFactory.register_backend( |
| 99 | "cumem", |
| 100 | "vllm.device_allocator.sleep_mode_backend", |
| 101 | "CuMemBackend", |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | def test_third_party_backend_registration_and_resolution(): |
| 106 | """A plugin registers a backend by name; the factory resolves it lazily.""" |
| 107 | name = "_pytest_dummy_backend" |
| 108 | try: |
| 109 | SleepModeBackendFactory.register_backend( |
| 110 | name, |
| 111 | "tests.v1.worker.test_sleep_mode_backend", |
| 112 | "DummyBackend", |
| 113 | ) |
| 114 | resolved = SleepModeBackendFactory.get_backend_class(name) |
| 115 | assert resolved is DummyBackend |
| 116 | assert resolved.supports_durable_storage() is True |
| 117 | finally: |
| 118 | SleepModeBackendFactory._registry.pop(name, None) |
| 119 | |
| 120 | |
| 121 | def test_suspend_resume_state_transitions(): |
| 122 | """Lifecycle state advances RUNNING -> SUSPENDED -> RUNNING without GPU.""" |
| 123 | backend = DummyBackend() |
| 124 | assert backend.state() == "RUNNING" |
| 125 | backend.suspend(level=1) |
| 126 | assert backend.state() == "SUSPENDED" |
| 127 | backend.resume() |
| 128 | assert backend.state() == "RUNNING" |
| 129 | |
| 130 | |
| 131 | class DummyBackend(SleepModeBackend): |
| 132 | """A no-GPU backend used to exercise lifecycle + registration in CPU tests.""" |
| 133 | |
| 134 | def suspend(self, level: int = 1) -> None: |
| 135 | self._state = "SUSPENDED" |
| 136 | |
| 137 | def resume(self, tags: list[str] | None = None) -> None: |
| 138 | self._state = "RUNNING" |
| 139 | |
| 140 | @classmethod |
| 141 | def supports_durable_storage(cls) -> bool: |
| 142 | return True |