diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 2329d51a149e..5d4986db2254 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -123,6 +123,21 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, CUmemGenericAllocationHandle** p_memHandle, unsigned long long* chunk_sizes, size_t num_chunks) { #endif + // Reset the global error_code at entry so a sticky error left over from a + // prior cumem call (e.g. a silently-failed unmap on a previous sleep cycle) + // does not cause the early-return guards below to short-circuit a fresh + // mapping attempt. Without this reset, a wake_up that follows a failed + // sleep cycle can return CUDA_SUCCESS-equivalent behaviour from cuMemCreate + // but then bail out at the first `if (error_code != 0) return;` check + // before cuMemMap is ever attempted, leaving a half-initialised allocation. + // + // THREADING: the global error_code (and error_msg) are intentionally NOT + // synchronized. CuMemAllocator is a documented singleton and its sleep()/ + // wake_up()/allocate paths are driven serially by a single worker thread + // (no concurrent create_and_map / my_malloc), so resetting error_code here + // cannot race a peer thread's error state. If that single-threaded invariant + // is ever relaxed, error_code/error_msg must become thread_local. + error_code = no_error; ensure_context(device); // Define memory allocation properties CUmemAllocationProp prop = {}; @@ -166,8 +181,33 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, if (error_code != 0) { return; } - CUDA_CHECK(cuMemMap(d_mem, size, 0, *p_memHandle, 0)); + // Recovery path for wake-time cuMemMap failures (most commonly + // CUDA_ERROR_INVALID_VALUE / "invalid argument" at this line). The + // failure mode this guards against: + // - A prior sleep cycle's cuMemUnmap silently failed (sticky error_code + // was masked) leaving `d_mem` with an existing mapping; cuMemCreate + // then succeeded for the wake (it just allocates a fresh handle) but + // cuMemMap fails because `d_mem` already maps to the previous handle. + // - A late free callback ran during sleep and the VA was reused by an + // overlapping reservation in the meantime. + // Either way, an idempotent cuMemUnmap of the target VA range followed by + // a single cuMemMap retry is safe: if no overlapping mapping exists, + // cuMemUnmap returns CUDA_ERROR_INVALID_VALUE which we ignore; if one + // exists, the unmap clears it and the retry succeeds. + CUresult map_status = cuMemMap(d_mem, size, 0, *p_memHandle, 0); + if (map_status == CUDA_ERROR_INVALID_VALUE) { + // Best-effort unmap of any pre-existing mapping at this VA. Ignore the + // return value: if nothing is mapped, this fails harmlessly. + (void)cuMemUnmap(d_mem, size); + map_status = cuMemMap(d_mem, size, 0, *p_memHandle, 0); + } + CUDA_CHECK(map_status); if (error_code != 0) { + // Map ultimately failed. Release the handle we created above so we + // don't leak a CUmemGenericAllocationHandle on the recovery path; the + // caller's error path (my_malloc / python_create_and_map) frees the + // outer slot and VA reservation as usual. + (void)cuMemRelease(*p_memHandle); return; } #else diff --git a/tests/device_allocator/__init__.py b/tests/device_allocator/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/device_allocator/test_cumem_wake_up_recovery.py b/tests/device_allocator/test_cumem_wake_up_recovery.py new file mode 100644 index 000000000000..3eb85d6c7924 --- /dev/null +++ b/tests/device_allocator/test_cumem_wake_up_recovery.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for ``CuMemAllocator.wake_up`` recovery semantics. + +These tests exercise the Python-layer behaviour of the cumem allocator's +wake-up loop without requiring a GPU. They drive the loop with a mocked C +extension (``python_create_and_map``) so we can deterministically simulate +the failure modes observed on multi-process TP/PP workloads: + +* Transient wake-time ``cuMemMap`` failures that should be surfaced as a + structured ``WakeUpPartialFailure`` rather than a bare ``RuntimeError``. +* Persistent failures that should leave ``pointer_to_data`` consistent so a + subsequent retry can re-attempt only the failed entries. +* A failure on one entry must not silently abort the wake of the remaining + entries — the loop must complete and report every failure together. + +These tests would FAIL against the pre-fix wake_up loop because: + 1. A bare exception from the C call propagates out of the first failing + iteration, leaving later allocations un-mapped with no signal that they + were skipped. + 2. The exception type is the generic ``RuntimeError`` raised by the C + extension, with no structured list of failed pointers. +""" + +from __future__ import annotations + +import sys +from unittest import mock + +import pytest + + +# ----------------------------------------------------------------------------- +# Module loader: stub the C extension + cuda wrapper so cumem.py imports +# cleanly on non-CUDA test runners (CI, contributor laptops). Without these +# stubs, ``from vllm.cumem_allocator import ...`` raises ModuleNotFoundError +# and the module sets ``cumem_available = False`` — which makes +# ``CuMemAllocator.get_instance()`` assert out before we can exercise it. +# ----------------------------------------------------------------------------- + + +@pytest.fixture +def cumem_module(monkeypatch): + """Import (or re-import) ``vllm.device_allocator.cumem`` with the C + extension and CUDA wrapper stubbed out so the import succeeds without a + GPU. Returns the freshly-imported module so tests can patch attributes + on it directly.""" + + # Stub C extension before import. + fake_ext = mock.MagicMock(name="vllm.cumem_allocator") + monkeypatch.setitem(sys.modules, "vllm.cumem_allocator", fake_ext) + + # Stub the CUDA wrapper module that cumem imports at top-level. + fake_wrapper_mod = mock.MagicMock( + name="vllm.distributed.device_communicators.cuda_wrapper" + ) + fake_wrapper_mod.CudaRTLibrary = mock.MagicMock(name="CudaRTLibrary") + monkeypatch.setitem( + sys.modules, + "vllm.distributed.device_communicators.cuda_wrapper", + fake_wrapper_mod, + ) + + # Stub find_loaded_library so it returns a fake path without scanning + # /proc/self/maps for an actual loaded library. + fake_sys_utils = mock.MagicMock(name="vllm.utils.system_utils") + fake_sys_utils.find_loaded_library = mock.MagicMock( + return_value="/fake/path/cumem_allocator.so" + ) + monkeypatch.setitem(sys.modules, "vllm.utils.system_utils", fake_sys_utils) + + # Force a fresh import so the top-level try/except sees our stubs. + sys.modules.pop("vllm.device_allocator.cumem", None) + import vllm.device_allocator.cumem as cumem # noqa: E402 + + # Reset the singleton between tests so state doesn't leak. + cumem.CuMemAllocator.instance = None + + yield cumem + + # Cleanup after the test. + cumem.CuMemAllocator.instance = None + sys.modules.pop("vllm.device_allocator.cumem", None) + + +def _make_handle(device: int, size: int, d_mem: int, p_handle: int) -> tuple: + """Construct a HandleType tuple matching the + ``(device, aligned_size, d_mem_ptr, p_memHandle_ptr)`` C-extension ABI.""" + return (device, size, d_mem, p_handle) + + +def _seed_allocator_with_handles(cumem, allocator, handles): + """Populate ``allocator.pointer_to_data`` with ``AllocationData`` entries + that look like a post-sleep state ready to be woken.""" + from vllm.device_allocator import AllocationData + + for handle in handles: + ptr = handle[2] + allocator.pointer_to_data[ptr] = AllocationData( + handle=handle, + tag="weights", + cpu_backup_tensor=None, + ) + + +# ----------------------------------------------------------------------------- +# Test 1: per-allocation failures don't abort the loop and don't corrupt state. +# ----------------------------------------------------------------------------- + + +def test_wake_up_propagates_per_allocation_failure_without_corrupting_state( + cumem_module, +): + """A failure on the 2nd of 3 allocations must (a) be captured rather + than silently abort, (b) NOT skip the 3rd allocation, and (c) leave + ``pointer_to_data`` intact so a retry can re-attempt only the failed + entry. Pre-fix: the first raise from ``create_and_map`` propagates and + the 3rd entry never gets mapped, so its restore is skipped silently.""" + + allocator = cumem_module.CuMemAllocator.get_instance() + handles = [ + _make_handle(0, 4096, 0x1000, 0xA000), + _make_handle(0, 4096, 0x2000, 0xA100), + _make_handle(0, 4096, 0x3000, 0xA200), + ] + _seed_allocator_with_handles(cumem_module, allocator, handles) + + call_log: list[int] = [] + + def fake_create_and_map(handle): + call_log.append(handle[2]) + if handle[2] == 0x2000: + raise RuntimeError( + "CUDA Error: invalid argument at csrc/cumem_allocator.cpp:169" + ) + + with mock.patch.object( + cumem_module, "create_and_map", side_effect=fake_create_and_map + ): + with pytest.raises(cumem_module.WakeUpPartialFailure) as excinfo: + allocator.wake_up() + + # 1. The loop visited ALL three entries — failures don't abort iteration. + assert call_log == [0x1000, 0x2000, 0x3000] + + # 2. The structured exception lists exactly the failed pointer(s). + assert excinfo.value.failed_pointers == [0x2000] + + # 3. State is intact: the entry still lives in pointer_to_data so the + # caller can retry just that one (or fall back to a cold restart). + assert 0x2000 in allocator.pointer_to_data + assert len(allocator.pointer_to_data) == 3 + + +# ----------------------------------------------------------------------------- +# Test 2: structured exception type, not bare RuntimeError. +# ----------------------------------------------------------------------------- + + +def test_wake_up_raises_structured_exception_on_persistent_failure(cumem_module): + """When ``create_and_map`` always fails, ``wake_up`` must raise the + structured ``WakeUpPartialFailure`` (a ``RuntimeError`` subclass) + carrying the list of failed pointers and the first underlying + exception. Pre-fix: a generic ``RuntimeError`` from the C extension is + re-raised verbatim with no structured payload.""" + + allocator = cumem_module.CuMemAllocator.get_instance() + handles = [ + _make_handle(0, 4096, 0x1000, 0xA000), + _make_handle(0, 4096, 0x2000, 0xA100), + ] + _seed_allocator_with_handles(cumem_module, allocator, handles) + + original = RuntimeError( + "CUDA Error: invalid argument at csrc/cumem_allocator.cpp:169" + ) + + def fake_create_and_map(handle): + raise original + + with mock.patch.object( + cumem_module, "create_and_map", side_effect=fake_create_and_map + ): + with pytest.raises(cumem_module.WakeUpPartialFailure) as excinfo: + allocator.wake_up() + + err = excinfo.value + # Subclass relationship lets existing ``except RuntimeError`` callers + # still catch it, but new callers can `isinstance(e, WakeUpPartialFailure)` + # for structured handling. + assert isinstance(err, RuntimeError) + assert isinstance(err, cumem_module.WakeUpPartialFailure) + assert set(err.failed_pointers) == {0x1000, 0x2000} + assert err.first_exception is original + + +# ----------------------------------------------------------------------------- +# Test 3: structured failure includes ALL failed pointers, not just the first. +# ----------------------------------------------------------------------------- + + +def test_wake_up_collects_all_failed_pointers(cumem_module): + """The structured exception must record every failed pointer the loop + encountered, not just the first one. This is what lets a calling + executor decide between per-allocation retry and a worker-wide cold + restart based on the breadth of the failure.""" + + allocator = cumem_module.CuMemAllocator.get_instance() + handles = [ + _make_handle(0, 4096, 0x1000, 0xA000), + _make_handle(0, 4096, 0x2000, 0xA100), + _make_handle(0, 4096, 0x3000, 0xA200), + _make_handle(0, 4096, 0x4000, 0xA300), + ] + _seed_allocator_with_handles(cumem_module, allocator, handles) + + def fake_create_and_map(handle): + # Fail on the middle two — succeed on the outer two. + if handle[2] in (0x2000, 0x3000): + raise RuntimeError("CUDA Error: invalid argument at .../cumem:169") + + with mock.patch.object( + cumem_module, "create_and_map", side_effect=fake_create_and_map + ): + with pytest.raises(cumem_module.WakeUpPartialFailure) as excinfo: + allocator.wake_up() + + # All failed pointers reported in iteration order. + assert excinfo.value.failed_pointers == [0x2000, 0x3000] + + +# ----------------------------------------------------------------------------- +# Test 4: success path is unchanged — no exception, no skips. +# ----------------------------------------------------------------------------- + + +def test_wake_up_success_path_unchanged(cumem_module): + """Sanity check: when every ``create_and_map`` succeeds, ``wake_up`` + must not raise, must visit every entry, and must not leak the new + ``WakeUpPartialFailure`` type into the success path.""" + + allocator = cumem_module.CuMemAllocator.get_instance() + handles = [ + _make_handle(0, 4096, 0x1000, 0xA000), + _make_handle(0, 4096, 0x2000, 0xA100), + ] + _seed_allocator_with_handles(cumem_module, allocator, handles) + + call_log: list[int] = [] + + def fake_create_and_map(handle): + call_log.append(handle[2]) + + with mock.patch.object( + cumem_module, "create_and_map", side_effect=fake_create_and_map + ): + allocator.wake_up() # must not raise + + assert call_log == [0x1000, 0x2000] diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 7c4fedd34a3b..259dbfc1c7f7 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -55,6 +55,38 @@ def unmap_and_release(allocation_handle: HandleType) -> None: python_unmap_and_release(*allocation_handle) +class WakeUpPartialFailure(RuntimeError): + """ + Raised by ``CuMemAllocator.wake_up`` when one or more per-allocation + ``create_and_map`` calls fail. Unlike a bare ``RuntimeError`` propagated + from the C extension, this exception type carries structured information + about which allocations failed so that callers (worker / executor / + engine-core liveness probes) can distinguish a partial wake from a clean + one and recover deterministically rather than returning HTTP 200 while + background workers are wedged. + + Attributes: + failed_pointers: device-pointer addresses (``ptr``) of the + allocations whose wake-time mapping failed. + first_exception: the first underlying exception raised by the C + extension, preserved so the original CUDA error string remains + visible in tracebacks. + """ + + def __init__( + self, + failed_pointers: list[int], + first_exception: BaseException, + ) -> None: + self.failed_pointers = list(failed_pointers) + self.first_exception = first_exception + super().__init__( + f"CuMemAllocator.wake_up failed for " + f"{len(self.failed_pointers)} allocation(s); " + f"first error: {first_exception!r}" + ) + + def get_pluggable_allocator( python_malloc_fn: Callable[[HandleType], None], python_free_func: Callable[[int], HandleType], @@ -290,11 +322,45 @@ def wake_up(self, tags: list[str] | None = None) -> None: tags: The tags of the memory allocation that will be loaded back to GPU memory. If None, all memory allocation will be loaded back to GPU memory. + + Raises: + WakeUpPartialFailure: if one or more per-allocation + ``create_and_map`` calls fail. The exception carries the list + of failed device pointers and the first underlying exception + so callers can decide how to recover (e.g. mark the worker + as needing a cold restart). Iteration continues through all + allocations even after a failure so the post-wake state is + deterministic — every allocation has either been remapped or + recorded as failed. """ + failed_pointers: list[int] = [] + first_exception: BaseException | None = None for ptr, data in self.pointer_to_data.items(): if tags is None or data.tag in tags: handle = data.handle - create_and_map(handle) + try: + create_and_map(handle) + except Exception as exc: # noqa: BLE001 — preserve and re-raise + failed_pointers.append(ptr) + if first_exception is None: + first_exception = exc + logger.error( + "CuMemAllocator.wake_up: create_and_map failed for " + "ptr=0x%x tag=%s size=%s: %r", + ptr, + data.tag, + handle[1], + exc, + ) + # Skip the cudaMemcpy restore for this allocation; the + # backing tensor is still pinned in `data.cpu_backup_tensor` + # and will be retried on a subsequent wake_up or freed on + # a cold restart, but copying into an unmapped VA would + # crash the process. + continue + # Only mark awake once the VA is successfully remapped; on the + # failure path above we `continue`, so it stays flagged asleep + # and is reported via WakeUpPartialFailure. data.is_asleep = False if data.cpu_backup_tensor is not None: cpu_backup_tensor = data.cpu_backup_tensor @@ -306,6 +372,10 @@ def wake_up(self, tags: list[str] | None = None) -> None: libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None + if failed_pointers: + assert first_exception is not None + raise WakeUpPartialFailure(failed_pointers, first_exception) + @contextmanager def use_memory_pool(self, tag: str | None = None): """