diff --git a/cuda_core/cuda/core/_cpp/rt/api.hpp b/cuda_core/cuda/core/_cpp/rt/api.hpp index cd39a27c91e..a7f752fbe8f 100644 --- a/cuda_core/cuda/core/_cpp/rt/api.hpp +++ b/cuda_core/cuda/core/_cpp/rt/api.hpp @@ -239,8 +239,9 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept; // Set the deallocation stream for a device pointer handle. // Returns CUDA_ERROR_INVALID_CONTEXT when a default-stream token cannot be // bound because no CUDA context is current. +// May throw std::bad_alloc. CUresult set_deallocation_stream( - const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept; + const DevicePtrHandle& h, const StreamHandle& h_stream); // ============================================================================ // Virtual memory management (VMM_DESIGN.md) diff --git a/cuda_core/cuda/core/_cpp/rt/internal.hpp b/cuda_core/cuda/core/_cpp/rt/internal.hpp index 39b3dae0886..9cd32d6c19d 100644 --- a/cuda_core/cuda/core/_cpp/rt/internal.hpp +++ b/cuda_core/cuda/core/_cpp/rt/internal.hpp @@ -59,8 +59,9 @@ struct DeallocationStream { // Implemented in stream.cpp ContextHandle deallocation_context(const DeallocationStream& stream) noexcept; -// Implemented in stream.cpp -bool make_deallocation_stream(const StreamHandle& h, DeallocationStream& out) noexcept; +// Implemented in stream.cpp. Not noexcept: binding a default-stream token +// allocates. May throw std::bad_alloc. +bool make_deallocation_stream(const StreamHandle& h, DeallocationStream& out); // Implemented in virtual_memory.cpp. Synchronize the stream a VMM buffer // recorded before its mappings are released. Skips the sync, with a report diff --git a/cuda_core/cuda/core/_cpp/rt/memory.cpp b/cuda_core/cuda/core/_cpp/rt/memory.cpp index 47bdd89e387..ce9e2513b06 100644 --- a/cuda_core/cuda/core/_cpp/rt/memory.cpp +++ b/cuda_core/cuda/core/_cpp/rt/memory.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -171,7 +172,7 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept { } // Replace the stream that orders a device pointer's deallocation. -CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept { +CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) { if (!h) { return CUDA_ERROR_INVALID_VALUE; } @@ -183,6 +184,32 @@ CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h return CUDA_SUCCESS; } +// Build the box for a freshly acquired device pointer. Binding the deallocation +// stream and allocating the box can fail: make_deallocation_stream returns +// false (err set) when no context is current, and either step can throw +// std::bad_alloc. On any failure run undo() to release the pointer and return +// nullptr; a thrown exception is left in `failure` for the caller to rethrow +// once it holds no lock (undo may not report under a lock, see +// deviceptr_import_ipc). Never throws. +template +static DevicePtrBox* make_deviceptr_box( + CUdeviceptr ptr, const StreamHandle& h_stream, Undo&& undo, + std::exception_ptr& failure) noexcept { + DevicePtrBox* box = nullptr; + try { + DeallocationStream ds; + if (make_deallocation_stream(h_stream, ds)) { + box = new DevicePtrBox{ptr, std::move(ds)}; + } + } catch (...) { + failure = std::current_exception(); + } + if (!box) { + undo(); + } + return box; +} + DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) { GILReleaseGuard gil; CUdeviceptr ptr; @@ -190,14 +217,18 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h return {}; } - DeallocationStream ds; - if (!make_deallocation_stream(h_stream, ds)) { - pw_cuMemFreeAsync(ptr, as_cu(h_stream)); + std::exception_ptr failure; + DevicePtrBox* raw = make_deviceptr_box( + ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure); + if (!raw) { + if (failure) { + std::rethrow_exception(failure); + } return {}; } auto box = std::shared_ptr( - new DevicePtrBox{ptr, std::move(ds)}, + raw, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; const DeallocationStream& stream = b->deallocation; @@ -220,14 +251,18 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) return {}; } - DeallocationStream ds; - if (!make_deallocation_stream(h_stream, ds)) { - pw_cuMemFreeAsync(ptr, as_cu(h_stream)); + std::exception_ptr failure; + DevicePtrBox* raw = make_deviceptr_box( + ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure); + if (!raw) { + if (failure) { + std::rethrow_exception(failure); + } return {}; } auto box = std::shared_ptr( - new DevicePtrBox{ptr, std::move(ds)}, + raw, [](DevicePtrBox* b) { GILReleaseGuard gil; const DeallocationStream& stream = b->deallocation; @@ -488,6 +523,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* GILReleaseGuard gil; CUresult discard_status = CUDA_SUCCESS; CUdeviceptr discarded = 0; + std::exception_ptr failure; { std::lock_guard lock(ipc_import_mutex); @@ -500,10 +536,19 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* return {}; } - DeallocationStream ds; - if (make_deallocation_stream(h_stream, ds)) { + // No deallocation stream could be recorded (or the box could not be + // allocated): discard the import with the raw call (a pw_ report would + // acquire the GIL under the mutex). + DevicePtrBox* raw = make_deviceptr_box( + ptr, h_stream, + [&]() noexcept { + discard_status = p_cuMemFreeAsync(ptr, as_cu(h_stream)); // raw: under ipc_import_mutex + discarded = ptr; + }, + failure); + if (raw) { auto box = std::shared_ptr( - new DevicePtrBox{ptr, std::move(ds)}, + raw, [h_pool, key](DevicePtrBox* b) { // Release the GIL first (the GIL is the outermost lock), then hold the // mutex across unregister + free: a concurrent import that finds this @@ -530,11 +575,6 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* ipc_ptr_cache.register_handle(key, h); return h; } - - // No deallocation stream could be recorded: discard the import with - // the raw call (a pw_ report would acquire the GIL under the mutex). - discard_status = p_cuMemFreeAsync(ptr, as_cu(h_stream)); // raw: under ipc_import_mutex - discarded = ptr; } if (discard_status != CUDA_SUCCESS) { char operation[160]; @@ -543,6 +583,9 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* "failed while discarding an IPC import that could not record a " "deallocation stream; the mapping leaked"); } + if (failure) { + std::rethrow_exception(failure); + } return {}; } else { @@ -552,14 +595,18 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* return {}; } - DeallocationStream ds; - if (!make_deallocation_stream(h_stream, ds)) { - pw_cuMemFreeAsync(ptr, as_cu(h_stream)); + std::exception_ptr failure; + DevicePtrBox* raw = make_deviceptr_box( + ptr, h_stream, [&]() noexcept { pw_cuMemFreeAsync(ptr, as_cu(h_stream)); }, failure); + if (!raw) { + if (failure) { + std::rethrow_exception(failure); + } return {}; } auto box = std::shared_ptr( - new DevicePtrBox{ptr, std::move(ds)}, + raw, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; const DeallocationStream& stream = b->deallocation; diff --git a/cuda_core/cuda/core/_cpp/rt/stream.cpp b/cuda_core/cuda/core/_cpp/rt/stream.cpp index 0497c0fabb5..c76a0ae2e9f 100644 --- a/cuda_core/cuda/core/_cpp/rt/stream.cpp +++ b/cuda_core/cuda/core/_cpp/rt/stream.cpp @@ -189,9 +189,11 @@ StreamHandle create_context_bound_legacy_stream(const ContextHandle& h_context) namespace detail { // Real streams are copied unchanged. Default-stream tokens without an embedded // context are bound to the current context. Returns false (and sets err) when a -// default-stream token cannot be bound because no context is current. +// default-stream token cannot be bound because no context is current. Binding +// a default-stream token allocates (the bound StreamBox and the current-context +// handle), so this may throw std::bad_alloc and must not be noexcept. bool make_deallocation_stream( - const StreamHandle& h, DeallocationStream& out) noexcept { + const StreamHandle& h, DeallocationStream& out) { out = {}; if (!h) { return true; diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 23ab10c06e9..52fba667886 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -316,7 +316,13 @@ cdef class Buffer: record_stream = mr.is_device_accessible self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr) if record_stream: - _ds_status = set_deallocation_stream(self._h_ptr, s._h_stream) + try: + _ds_status = set_deallocation_stream(self._h_ptr, s._h_stream) + except MemoryError: + # Reset for the same reason as below, inside the handler so the + # destructor cannot swallow the MemoryError that is propagating. + self._h_ptr.reset() + raise if _ds_status != cydriver.CUresult.CUDA_SUCCESS: # Reset before raising: the DevicePtrHandle destructor would otherwise # invoke _mr_dealloc_callback, which catches any inner exception and diff --git a/cuda_core/cuda/core/_rt.pxd b/cuda_core/cuda/core/_rt.pxd index 36fbd82acc4..aece0cf32b2 100644 --- a/cuda_core/cuda/core/_rt.pxd +++ b/cuda_core/cuda/core/_rt.pxd @@ -292,7 +292,7 @@ cdef void register_mr_dealloc_callback(MRDeallocCallback cb) noexcept cdef DevicePtrHandle deviceptr_import_ipc( const MemoryPoolHandle& h_pool, const void* export_data, const StreamHandle& h_stream) except+ nogil cdef StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept nogil -cdef cydriver.CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil +cdef cydriver.CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) except+ nogil # Virtual memory management (VMM_DESIGN.md) cdef MemAllocationHandle create_mem_allocation_handle( diff --git a/cuda_core/cuda/core/_rt.pyx b/cuda_core/cuda/core/_rt.pyx index 1ebe3a1b460..6a6a95a88cc 100644 --- a/cuda_core/cuda/core/_rt.pyx +++ b/cuda_core/cuda/core/_rt.pyx @@ -172,7 +172,7 @@ cdef extern from "_cpp/rt/rt.hpp" namespace "cuda_core::rt": StreamHandle deallocation_stream "cuda_core::rt::deallocation_stream" ( const DevicePtrHandle& h) noexcept nogil cydriver.CUresult set_deallocation_stream "cuda_core::rt::set_deallocation_stream" ( - const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil + const DevicePtrHandle& h, const StreamHandle& h_stream) except+ nogil # Virtual memory management (VMM_DESIGN.md) MemAllocationHandle create_mem_allocation_handle "cuda_core::rt::create_mem_allocation_handle" ( diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index f292bb209a6..dda0be2e5b5 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -104,6 +104,10 @@ New features Fixes and enhancements ---------------------- +- :meth:`Buffer.set_deallocation_stream` and :meth:`Buffer.from_handle` now raise + :class:`MemoryError` when the host runs out of memory while binding a default + stream, instead of aborting the process with ``std::terminate``. + - :attr:`system.Device.uuid_without_prefix` preserves UUIDs that NVML already reports without a prefix, such as on Orin. This also allows :meth:`system.Device.to_cuda_device` to match those devices correctly. diff --git a/cuda_core/tests/test_allocation_failure.py b/cuda_core/tests/test_allocation_failure.py new file mode 100644 index 00000000000..0629e65a9c2 --- /dev/null +++ b/cuda_core/tests/test_allocation_failure.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Host allocation failure in the C++ handle layer must raise, never terminate. + +``std::bad_alloc`` that escapes a ``noexcept`` function calls ``std::terminate`` +and aborts the interpreter. These tests make the Nth C++ allocation on a thread +throw ``std::bad_alloc`` by preloading a tiny ``operator new`` replacement into a +child process, so an abort shows up as a non-zero exit status of the child +instead of taking down the test run. +""" + +import os +import shutil +import subprocess +import sys +import textwrap + +import pytest + +pytestmark = pytest.mark.skipif( + not sys.platform.startswith("linux"), reason="the allocation shim is preloaded with LD_PRELOAD" +) + +# The Nth C++ allocation made by the arming thread throws std::bad_alloc. Python +# allocates through its own allocators, so only C++ allocations are affected. +SHIM_SOURCE = textwrap.dedent( + """ + #include + #include + + static thread_local long countdown = 0; + + extern "C" void arm_bad_alloc(long n) { countdown = n; } + + void* operator new(std::size_t size) { + if (countdown > 0 && --countdown == 0) { + throw std::bad_alloc(); + } + if (void* p = std::malloc(size ? size : 1)) { + return p; + } + throw std::bad_alloc(); + } + void operator delete(void* p) noexcept { std::free(p); } + void operator delete(void* p, std::size_t) noexcept { std::free(p); } + """ +) + +# Runs in the child. argv[1] is the shim path, argv[2] the scenario. +CHILD = textwrap.dedent( + """ + import ctypes + import sys + + shim = ctypes.CDLL(sys.argv[1]) + scenario = sys.argv[2] + + from cuda.core import Buffer, Device, MemoryResource + from cuda.core._stream import default_stream + from cuda.core._utils.cuda_utils import driver, handle_return + + + class DeviceMR(MemoryResource): + def allocate(self, size, *, stream=None): + raise NotImplementedError + + def deallocate(self, ptr, size, *, stream=None): + handle_return(driver.cuMemFree(ptr)) + + is_device_accessible = True + is_host_accessible = False + device_id = 0 + + + def arm(n): + shim.arm_bad_alloc(ctypes.c_long(n)) + + + Device(0).set_current() + mr = DeviceMR() + outcomes = set() + + if scenario == "set_deallocation_stream": + buf = Buffer.from_handle(handle_return(driver.cuMemAlloc(64)), 64, mr=mr) + token = default_stream() + arm(1) + try: + buf.set_deallocation_stream(token) + except MemoryError: + outcomes.add("MemoryError") + else: + outcomes.add("ok") + finally: + arm(0) + buf.close() + elif scenario == "from_handle": + # Fail each successive allocation on the way through Buffer.from_handle. + for n in range(1, 9): + ptr = handle_return(driver.cuMemAlloc(64)) + arm(n) + try: + buf = Buffer.from_handle(ptr, 64, mr=mr) + except MemoryError: + outcomes.add("MemoryError") + else: + outcomes.add("ok") + buf.close() + finally: + arm(0) + else: + raise SystemExit("unknown scenario " + scenario) + + print(",".join(sorted(outcomes))) + """ +) + + +@pytest.fixture(scope="module") +def alloc_shim(tmp_path_factory): + compiler = shutil.which("c++") or shutil.which("g++") or shutil.which("clang++") + if compiler is None: + pytest.skip("no C++ compiler available to build the allocation shim") + workdir = tmp_path_factory.mktemp("alloc_shim") + source = workdir / "shim.cpp" + source.write_text(SHIM_SOURCE) + shim = workdir / "libshim.so" + subprocess.run([compiler, "-shared", "-fPIC", "-O1", "-o", str(shim), str(source)], check=True) # noqa: S603 + return shim + + +def run_child(shim, scenario): + env = dict(os.environ) + preload = [str(shim)] + if env.get("LD_PRELOAD"): + preload.append(env["LD_PRELOAD"]) + env["LD_PRELOAD"] = " ".join(preload) + return subprocess.run( # noqa: S603 + [sys.executable, "-c", CHILD, str(shim), scenario], + env=env, + # Run outside the source tree: "-c" puts the cwd first on sys.path, so a + # cwd of cuda_core/ would import the unbuilt source package. + cwd=shim.parent, + capture_output=True, + text=True, + timeout=300, + check=False, + ) + + +@pytest.mark.agent_authored(model="claude-sonnet-5-5") +def test_set_deallocation_stream_raises_memory_error_instead_of_terminating(init_cuda, alloc_shim): + """Binding a default-stream token allocates; failure there must be a MemoryError. + + Before the fix, ``set_deallocation_stream`` and ``make_deallocation_stream`` + were ``noexcept``, so the ``std::bad_alloc`` called ``std::terminate`` and the + child died with SIGABRT. + """ + result = run_child(alloc_shim, "set_deallocation_stream") + assert result.returncode == 0, f"child exited with {result.returncode}\n{result.stderr}" + assert result.stdout.strip() == "MemoryError" + + +@pytest.mark.agent_authored(model="claude-sonnet-5-5") +def test_buffer_from_handle_survives_allocation_failure(init_cuda, alloc_shim): + """Failing each allocation in ``Buffer.from_handle`` raises MemoryError or succeeds, never aborts.""" + result = run_child(alloc_shim, "from_handle") + assert result.returncode == 0, f"child exited with {result.returncode}\n{result.stderr}" + assert "MemoryError" in result.stdout.strip().split(",")