From 44d3245a6dbdb172172d3e6e3cb620f4ae537ce8 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Mon, 17 Aug 2026 19:14:07 +0800 Subject: [PATCH] [Diffusion] Clear handled CUDA registration errors Signed-off-by: Hongsheng Liu --- .../offloader/test_cuda_host_registration.py | 37 +++++++++++++++++ .../offloader/cuda_host_registration.py | 40 +++++++++++++++++-- 2 files changed, 74 insertions(+), 3 deletions(-) diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py index 01a1f763e..262ce920c 100644 --- a/tests/diffusion/offloader/test_cuda_host_registration.py +++ b/tests/diffusion/offloader/test_cuda_host_registration.py @@ -130,6 +130,8 @@ def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) runtime = _FakeRuntime([0, 7]) monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + consumed_errors: list[int] = [] + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append) with pytest.raises(CudaHostRegistrationError, match="error-7"): CudaHostRegistration.create( @@ -141,6 +143,7 @@ def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch ) assert runtime.unregistered == [0x1000] + assert consumed_errors == [7] def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None: @@ -164,6 +167,8 @@ def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.Monke monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) runtime = _FakeRuntime([0, 7], unregister_results=[9]) monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + consumed_errors: list[int] = [] + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append) with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"): CudaHostRegistration.create( @@ -174,6 +179,35 @@ def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.Monke max_bytes=8192, ) + assert consumed_errors == [7, 9] + + +@pytest.mark.parametrize("pending_error", [0, 801]) +def test_consume_last_cuda_error_accepts_cleared_or_matching_state( + monkeypatch: pytest.MonkeyPatch, + pending_error: int, +) -> None: + class _GetLastError: + def __call__(self) -> int: + return pending_error + + runtime = SimpleNamespace(cudaGetLastError=_GetLastError()) + monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime) + + registration_module._consume_last_cuda_error(801) + + +def test_consume_last_cuda_error_rejects_unrelated_pending_error(monkeypatch: pytest.MonkeyPatch) -> None: + class _GetLastError: + def __call__(self) -> int: + return 700 + + runtime = SimpleNamespace(cudaGetLastError=_GetLastError()) + monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: runtime) + + with pytest.raises(CudaHostRegistrationCleanupError, match="cudaGetLastError reported 700"): + registration_module._consume_last_cuda_error(801) + def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) @@ -219,6 +253,8 @@ def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatc monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) runtime = _FakeRuntime([0], unregister_results=[9, 0]) monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + consumed_errors: list[int] = [] + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", consumed_errors.append) registration = CudaHostRegistration.create( {"weights": [_FakeTensor(0x1003, 1)]}, # type: ignore[list-item] max_bytes=4096, @@ -227,3 +263,4 @@ def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatc assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"] assert registration.close() == [] assert runtime.unregistered == [0x1000, 0x1000] + assert consumed_errors == [9] diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py index 8b9e6314e..6ee54b02d 100644 --- a/vllm_omni/diffusion/offloader/cuda_host_registration.py +++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py @@ -4,6 +4,7 @@ from __future__ import annotations +import ctypes import mmap from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -99,6 +100,39 @@ def _error_message(runtime: _CudaRuntime, error: int) -> str: return str(message) +def _consume_last_cuda_error(expected_error: int) -> None: + """Clear a handled CUDA Runtime error before returning to PyTorch. + + PyTorch's ``torch.cuda.cudart()`` binding exposes host registration but not + ``cudaGetLastError``. Resolve the already-loaded process symbol so this call + clears the same CUDA Runtime instance that produced the registration error. + A different pending error is not safe to hide behind the staging fallback. + """ + try: + get_last_error = ctypes.CDLL(None).cudaGetLastError + get_last_error.argtypes = [] + get_last_error.restype = ctypes.c_int + pending_error = int(get_last_error()) + except (AttributeError, OSError) as exc: + raise HostRegistrationCleanupError("cannot clear CUDA's pending error after host-registration failure") from exc + + if pending_error not in (0, expected_error): + raise HostRegistrationCleanupError( + f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}" + ) + + +def _handled_error_message(runtime: _CudaRuntime, error: int) -> str: + """Format and consume one nonzero CUDA Runtime return code.""" + error_code = int(error) + message = _error_message(runtime, error) + try: + _consume_last_cuda_error(error_code) + except HostRegistrationCleanupError as exc: + raise HostRegistrationCleanupError(f"{message}; {exc}") from exc + return message + + class CudaHostRegistration: """Own CUDA registrations for already-existing file-backed CPU tensors.""" @@ -147,7 +181,7 @@ class CudaHostRegistration: if int(error) != 0: raise HostRegistrationError( "cudaHostRegister(read-only) failed for " - f"[{region.start:#x}, {region.end:#x}): {_error_message(runtime, error)}" + f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}" ) registered.append(region) @@ -167,7 +201,7 @@ class CudaHostRegistration: error = runtime.cudaHostUnregister(region.start) if int(error) != 0: rollback_errors.append( - f"cudaHostUnregister({region.start:#x}) failed: {_error_message(runtime, error)}" + f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}" ) except Exception as rollback_exc: rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}") @@ -200,7 +234,7 @@ class CudaHostRegistration: error = self._runtime.cudaHostUnregister(region.start) if int(error) != 0: errors.append( - f"cudaHostUnregister({region.start:#x}) failed: {_error_message(self._runtime, error)}" + f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}" ) failed.append(region) except Exception as exc: