diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index cdaa644b62ec..67e3a431fa60 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -7,6 +7,7 @@ import pytest import vllm.v1.worker.gpu_worker as gpu_worker_module +from vllm.distributed.parallel_state import GroupCoordinator from vllm.multimodal.video import ( PYNVVIDEOCODEC_CUDA_CONTEXT_BYTES, PYNVVIDEOCODEC_DECODER_GPU_MEMORY_BYTES, @@ -186,3 +187,146 @@ def test_startup_plan_apply_gate(plan_env): explicit = _plan_worker(kv_bytes=7 * GiB_bytes) maybe_apply_startup_plan(explicit) assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes + + +# Suspend/resume orchestration: Worker.sleep/wake_up consume the +# SleepModeBackend capability flags (vllm/device_allocator/sleep_mode_backend.py). + + +class _FlaggedBackend: + """Minimal stand-in for a SleepModeBackend with configurable flags.""" + + _preserves_communicators = True + _preserves_graphs = True + + def __init__(self): + self.calls: list[tuple] = [] + + def suspend(self, level: int = 1) -> None: + self.calls.append(("suspend", level)) + + def resume(self, tags: list[str] | None = None) -> None: + self.calls.append(("resume", tags)) + + @classmethod + def preserves_communicators(cls) -> bool: + return cls._preserves_communicators + + @classmethod + def preserves_graphs_with_communicators(cls) -> bool: + return cls._preserves_graphs + + +class _PreservingBackend(_FlaggedBackend): + """cumem-shaped: communicators survive suspend untouched.""" + + +class _RebuildingBackend(_FlaggedBackend): + """checkpoint-shaped: communicators are destroyed by suspend and must be + rebuilt on resume; captured graphs do not survive the rebuild.""" + + _preserves_communicators = False + _preserves_graphs = False + + +def _sleep_worker(backend: _FlaggedBackend) -> Worker: + worker = object.__new__(Worker) + worker._sleep_mode_backend = backend + worker._sleep_saved_buffers = {} + worker.model_runner = SimpleNamespace( + post_kv_cache_wake_up=lambda: None, + ) + return worker + + +def _record_comm_hooks(monkeypatch: pytest.MonkeyPatch, events: list[str]) -> None: + monkeypatch.setattr( + gpu_worker_module, + "prepare_communicators_for_suspend", + lambda: events.append("prepare_comms"), + ) + monkeypatch.setattr( + gpu_worker_module, + "reinit_communicators_after_resume", + lambda: events.append("reinit_comms"), + ) + + +def _stub_accelerator(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + gpu_worker_module.torch.accelerator, "synchronize", lambda: None + ) + monkeypatch.setattr( + gpu_worker_module.torch.accelerator, "get_memory_info", lambda: (0, 0) + ) + + +def test_sleep_prepares_communicators_before_suspend( + monkeypatch: pytest.MonkeyPatch, +): + events: list[str] = [] + _record_comm_hooks(monkeypatch, events) + _stub_accelerator(monkeypatch) + backend = _RebuildingBackend() + worker = _sleep_worker(backend) + + worker.sleep(level=1) + + # Teardown must land before the backend snapshots/frees GPU state. + assert events == ["prepare_comms"] + assert backend.calls == [("suspend", 1)] + + +def test_wake_up_rebuilds_communicators_then_invalidates_graphs( + monkeypatch: pytest.MonkeyPatch, +): + events: list[str] = [] + _record_comm_hooks(monkeypatch, events) + backend = _RebuildingBackend() + worker = _sleep_worker(backend) + + # Graphs embedding comm handles are stale after the rebuild; the draft + # contract fails loudly until discard+recapture exists. + with pytest.raises(NotImplementedError, match="invalidates captured"): + worker.wake_up() + + assert backend.calls == [("resume", None)] + assert events == ["reinit_comms"] + + +def test_preserving_backend_skips_comm_orchestration( + monkeypatch: pytest.MonkeyPatch, +): + events: list[str] = [] + _record_comm_hooks(monkeypatch, events) + _stub_accelerator(monkeypatch) + backend = _PreservingBackend() + worker = _sleep_worker(backend) + + worker.sleep(level=1) + worker.wake_up() + + # cumem-shaped backends keep today's exact path: no hook is invoked. + assert events == [] + assert backend.calls == [("suspend", 1), ("resume", None)] + + +def _group_with_world_size(world_size: int) -> GroupCoordinator: + group = object.__new__(GroupCoordinator) + group.world_size = world_size + group.unique_name = f"test_group_ws{world_size}" + return group + + +def test_group_suspend_hooks_noop_for_single_member_groups(): + group = _group_with_world_size(1) + group.prepare_for_suspend() + group.reinit_after_resume() + + +def test_group_suspend_hooks_fail_loudly_for_multi_member_groups(): + group = _group_with_world_size(2) + with pytest.raises(NotImplementedError, match="nccl_checkpoint"): + group.prepare_for_suspend() + with pytest.raises(NotImplementedError, match="nccl_checkpoint"): + group.reinit_after_resume() diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 162ed03d23b2..0dbdc64e7e41 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1213,6 +1213,46 @@ def destroy(self): if self.mq_broadcaster is not None: self.mq_broadcaster = None + def prepare_for_suspend(self) -> None: + """Tear down communicator state that cannot survive a process + suspend (peer mappings, RDMA connections, IPC handles). + + Called by the GPU worker before ``SleepModeBackend.suspend`` when + the selected backend reports ``preserves_communicators() == False`` + (see ``vllm/device_allocator/sleep_mode_backend.py``). Each + communicator owns its own teardown and rebuild; the worker only + sequences the calls. + + Single-member groups hold no peer state and no-op. Multi-member + teardown/rebuild requires checkpoint support inside the comms + library itself (e.g. NCCL's ``contrib/nccl_checkpoint``), so until + a communicator implements it, suspending a multi-member group + fails loudly here instead of resuming into corrupt comm state. + """ + if self.world_size == 1: + return + raise NotImplementedError( + f"Communicator group '{self.unique_name}' (world_size=" + f"{self.world_size}) cannot be suspended yet: multi-member " + "communicator teardown/rebuild requires checkpoint support in " + "the comms library (RFC #34303; NCCL contrib/nccl_checkpoint)." + ) + + def reinit_after_resume(self) -> None: + """Recreate communicator state after ``SleepModeBackend.resume``. + + Mirror of ``prepare_for_suspend``; see its docstring for the + contract and the single- vs multi-member behavior. + """ + if self.world_size == 1: + return + raise NotImplementedError( + f"Communicator group '{self.unique_name}' (world_size=" + f"{self.world_size}) cannot be rebuilt after resume yet: " + "requires checkpoint support in the comms library " + "(RFC #34303; NCCL contrib/nccl_checkpoint)." + ) + def prepare_communication_buffer_for_model(self, model: torch.nn.Module): if self.device_communicator is not None: self.device_communicator.prepare_communication_buffer_for_model(model) @@ -2083,6 +2123,39 @@ def destroy_model_parallel(): _EPLB = None +def _live_parallel_groups() -> list[GroupCoordinator]: + """The model-parallel subgroups currently initialized, in the order + ``destroy_model_parallel`` handles them. Excludes the world group.""" + return [g for g in (_TP, _DCP, _PCP, _PP, _DP, _EP, _EPLB) if g is not None] + + +def prepare_communicators_for_suspend() -> None: + """Prepare every live communicator group for a process suspend. + + Consumed by the GPU worker when the selected sleep-mode backend + reports ``preserves_communicators() == False``: teardown is owned by + each communicator (``GroupCoordinator.prepare_for_suspend``), and this + function only sequences it — subgroups first, world group last. + ``reinit_communicators_after_resume`` rebuilds in the opposite order. + """ + for group in _live_parallel_groups(): + group.prepare_for_suspend() + if _WORLD is not None: + _WORLD.prepare_for_suspend() + + +def reinit_communicators_after_resume() -> None: + """Rebuild every live communicator group after a process resume. + + Counterpart of ``prepare_communicators_for_suspend``: world group + first, then the model-parallel subgroups. + """ + if _WORLD is not None: + _WORLD.reinit_after_resume() + for group in _live_parallel_groups(): + group.reinit_after_resume() + + def destroy_distributed_environment(): global _WORLD, _NODE_COUNT if _WORLD: diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 182476e25336..95b2f43fbd81 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -43,6 +43,8 @@ Handle, get_pp_group, get_tp_group, + prepare_communicators_for_suspend, + reinit_communicators_after_resume, ) from vllm.distributed.weight_transfer import ( WeightTransferEngine, @@ -199,7 +201,13 @@ def sleep(self, level: int = 1) -> None: name: buffer.cpu().clone() for name, buffer in model.named_buffers() } - self._get_sleep_mode_backend().suspend(level) + backend = self._get_sleep_mode_backend() + if not backend.preserves_communicators(): + # Communicator-owned teardown must complete before the backend + # snapshots/frees GPU state (peer mappings and RDMA connections + # cannot survive a suspend). See prepare_communicators_for_suspend. + prepare_communicators_for_suspend() + backend.suspend(level) torch.accelerator.synchronize() deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0) @@ -219,7 +227,15 @@ def sleep(self, level: int = 1) -> None: ) def wake_up(self, tags: list[str] | None = None) -> None: - self._get_sleep_mode_backend().resume(tags) + backend = self._get_sleep_mode_backend() + backend.resume(tags) + + if not backend.preserves_communicators(): + # Mirror of the suspend-side teardown: rebuild is owned by each + # communicator, sequenced world-group-first. + reinit_communicators_after_resume() + if not backend.preserves_graphs_with_communicators(): + self._invalidate_captured_graphs() # Restore the buffers after level 2 sleep if len(self._sleep_saved_buffers): @@ -232,6 +248,21 @@ def wake_up(self, tags: list[str] | None = None) -> None: if tags is None or "kv_cache" in tags: self.model_runner.post_kv_cache_wake_up() + def _invalidate_captured_graphs(self) -> None: + """Captured CUDA graphs embed communicator handles; once communicators + are rebuilt on resume (``preserves_graphs_with_communicators() == + False``), replaying a captured graph would launch into destroyed comm + state. The safe sequel is discard + recapture — an init-time-style + operation, since a suspended engine is not serving traffic — but the + discard half has no in-tree mechanism yet, so this fails loudly + instead of resuming with stale graphs (RFC #34303). + """ + raise NotImplementedError( + "The sleep-mode backend rebuilt communicators on resume, which " + "invalidates captured CUDA graphs; discard+recapture after resume " + "is not implemented yet (RFC #34303)." + ) + def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: if ( current_platform.is_cuda_alike()