diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index b3d2e3fe373..34763135892 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -34,6 +34,7 @@ namespace cuda_core { decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain = nullptr; decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease = nullptr; decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent = nullptr; +decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent = nullptr; decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate = nullptr; decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy = nullptr; decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx = nullptr; @@ -188,6 +189,41 @@ class GILAcquireGuard { bool acquired_; }; +// Temporarily make a retained context current, restoring the caller's exact +// prior context (including no current context) on scope exit. +class ScopedCurrentContext { +public: + explicit ScopedCurrentContext(ContextHandle h_context) noexcept + : h_context_(std::move(h_context)) { + CUcontext target = as_cu(h_context_); + if (!target) { + return; + } + + GILReleaseGuard gil; + if (p_cuCtxGetCurrent(&previous_) != CUDA_SUCCESS || + previous_ == target) { + return; + } + changed_ = p_cuCtxSetCurrent(target) == CUDA_SUCCESS; + } + + ~ScopedCurrentContext() { + if (changed_) { + GILReleaseGuard gil; + p_cuCtxSetCurrent(previous_); + } + } + + ScopedCurrentContext(const ScopedCurrentContext&) = delete; + ScopedCurrentContext& operator=(const ScopedCurrentContext&) = delete; + +private: + ContextHandle h_context_; + CUcontext previous_ = nullptr; + bool changed_ = false; +}; + } // namespace // ============================================================================ @@ -906,6 +942,8 @@ struct DevicePtrBox { // through a const DevicePtrHandle. The stream can be changed after // allocation (e.g., to synchronize deallocation with a different stream). mutable StreamHandle h_stream; + // Used by MR-backed pointer deleters; native deleters ignore it. + mutable ContextHandle h_context; }; } // namespace @@ -925,8 +963,15 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept { return get_box(h)->h_stream; } -void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept { - get_box(h)->h_stream = h_stream; +void set_deallocation_stream( + const DevicePtrHandle& h, + const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept { + DevicePtrBox* box = get_box(h); + box->h_stream = h_stream; + if (!box->h_context) { + box->h_context = h_context; + } } DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) { @@ -937,7 +982,7 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, {}}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); @@ -955,7 +1000,7 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, {}}, [](DevicePtrBox* b) { GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); @@ -973,7 +1018,7 @@ DevicePtrHandle deviceptr_alloc(size_t size) { } auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, {}, {}}, [](DevicePtrBox* b) { GILReleaseGuard gil; p_cuMemFree(b->resource); @@ -991,7 +1036,7 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) { } auto box = std::shared_ptr( - new DevicePtrBox{reinterpret_cast(ptr), StreamHandle{}}, + new DevicePtrBox{reinterpret_cast(ptr), {}, {}}, [](DevicePtrBox* b) { GILReleaseGuard gil; p_cuMemFreeHost(reinterpret_cast(b->resource)); @@ -1002,7 +1047,8 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) { } DevicePtrHandle deviceptr_create_ref(CUdeviceptr ptr) { - auto box = std::make_shared(DevicePtrBox{ptr, StreamHandle{}}); + auto box = std::make_shared( + DevicePtrBox{ptr, {}, {}}); return DevicePtrHandle(box, &box->resource); } @@ -1018,7 +1064,7 @@ DevicePtrHandle deviceptr_create_with_owner(CUdeviceptr ptr, PyObject* owner) { } Py_INCREF(owner); auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, {}, {}}, [owner](DevicePtrBox* b) { GILAcquireGuard gil; if (gil.acquired()) { @@ -1036,7 +1082,7 @@ DevicePtrHandle deviceptr_create_mapped_graphics( const StreamHandle& h_stream ) { auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, {}}, [h_resource](DevicePtrBox* b) { GILReleaseGuard gil; CUgraphicsResource resource = as_cu(h_resource); @@ -1057,7 +1103,11 @@ void register_mr_dealloc_callback(MRDeallocCallback cb) { mr_dealloc_cb = cb; } -DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* mr) { +DevicePtrHandle deviceptr_create_with_mr( + CUdeviceptr ptr, + size_t size, + PyObject* mr, + const ContextHandle& h_context) { if (!mr) { return deviceptr_create_ref(ptr); } @@ -1068,8 +1118,9 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* } Py_INCREF(mr); auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, {}, h_context}, [mr, size](DevicePtrBox* b) { + ScopedCurrentContext context(b->h_context); GILAcquireGuard gil; if (gil.acquired()) { if (mr_dealloc_cb) { @@ -1162,7 +1213,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, {}}, [h_pool, key](DevicePtrBox* b) { ipc_ptr_cache.unregister_handle(key); GILReleaseGuard gil; @@ -1182,7 +1233,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, {}}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); diff --git a/cuda_core/cuda/core/_cpp/resource_handles.hpp b/cuda_core/cuda/core/_cpp/resource_handles.hpp index 3a9d2d75cff..6bdbdc0b316 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.hpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.hpp @@ -67,6 +67,7 @@ void clear_last_error() noexcept; extern decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain; extern decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease; extern decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent; +extern decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent; extern decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate; extern decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy; extern decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx; @@ -402,10 +403,15 @@ using MRDeallocCallback = void (*)(PyObject* mr, CUdeviceptr ptr, void register_mr_dealloc_callback(MRDeallocCallback cb); // Create a device pointer handle whose destructor calls mr.deallocate() -// via the registered callback. The mr's refcount is incremented and -// decremented when the handle is released. +// via the registered callback. The supplied context is retained and made +// current for the callback. The mr's refcount is incremented and decremented +// when the handle is released. // If mr is nullptr, equivalent to deviceptr_create_ref. -DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* mr); +DevicePtrHandle deviceptr_create_with_mr( + CUdeviceptr ptr, + size_t size, + PyObject* mr, + const ContextHandle& h_context); // Import a device pointer from IPC via cuMemPoolImportPointer. // When the last reference is released, cuMemFreeAsync is called on the stored stream. @@ -421,7 +427,11 @@ DevicePtrHandle deviceptr_import_ipc( StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept; // Set the deallocation stream for a device pointer handle. -void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept; +// If the handle has no cleanup context, use the supplied stream context. +void set_deallocation_stream( + const DevicePtrHandle& h, + const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept; // ============================================================================ // Library handle functions diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index d5ecac6cc09..b5f12c2cf8e 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -14,10 +14,13 @@ from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource from cuda.core._memory._ipc cimport IPCBufferDescriptor, IPCDataForBuffer from cuda.core._memory cimport _ipc from cuda.core._resource_handles cimport ( + ContextHandle, DevicePtrHandle, StreamHandle, + create_context_handle_ref, deviceptr_create_with_owner, deviceptr_create_with_mr, + get_stream_context, register_mr_dealloc_callback, as_intptr, as_cu, @@ -77,6 +80,65 @@ register_mr_dealloc_callback(_mr_dealloc_callback) __all__ = ['Buffer', 'MemoryResource'] +# Context Resolution Helpers +# -------------------------- +cdef ContextHandle _current_context_handle(): + cdef ContextHandle h_context + cdef cydriver.CUcontext context = NULL + with nogil: + HANDLE_RETURN(cydriver.cuCtxGetCurrent(&context)) + if context != NULL: + h_context = create_context_handle_ref(context) + return h_context + + +cdef ContextHandle _pointer_context_handle(cydriver.CUdeviceptr ptr): + """Return the pointer's context, or the current context for an async allocation.""" + cdef ContextHandle h_context + cdef cydriver.CUcontext context = NULL + cdef cydriver.CUresult result + if ptr == 0: + return h_context + + with nogil: + result = cydriver.cuPointerGetAttribute( + &context, + cydriver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_CONTEXT, + ptr, + ) + if result == cydriver.CUDA_SUCCESS: + if context != NULL: + return create_context_handle_ref(context) + # Stream-ordered allocations are context-independent. Retain the + # wrapping thread's current context for default-stream deallocation. + return _current_context_handle() + if result in ( + cydriver.CUresult.CUDA_ERROR_INVALID_VALUE, + cydriver.CUresult.CUDA_ERROR_NOT_INITIALIZED, + ): + # A foreign host pointer has no CUDA context to retain. + return h_context + with nogil: + HANDLE_RETURN(result) + return h_context + + +cdef ContextHandle _stream_context_handle(Stream stream): + cdef ContextHandle h_context = stream._h_context + cdef cydriver.CUcontext context = NULL + if h_context: + return h_context + h_context = get_stream_context(stream._h_stream) + if h_context: + return h_context + + with nogil: + HANDLE_RETURN(cydriver.cuStreamGetCtx(as_cu(stream._h_stream), &context)) + if context != NULL: + h_context = create_context_handle_ref(context) + return h_context + + # Memory Attribute Query Helpers # ------------------------------ cdef inline int _query_memory_attrs( @@ -188,8 +250,11 @@ cdef class Buffer: raise ValueError("owner and memory resource cannot be both specified together") cdef Buffer self = Buffer.__new__(cls) cdef uintptr_t c_ptr = (int(ptr)) + cdef ContextHandle h_context if mr is not None: - self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr) + h_context = _pointer_context_handle(c_ptr) + self._h_ptr = deviceptr_create_with_mr( + c_ptr, size, mr, h_context) else: self._h_ptr = deviceptr_create_with_owner(c_ptr, owner) self._size = size @@ -619,12 +684,14 @@ cdef Buffer Buffer_from_deviceptr_handle( cdef inline void Buffer_close(Buffer self, object stream): """Close a buffer, freeing its memory.""" cdef Stream s + cdef ContextHandle h_context if not self._h_ptr: return # Update deallocation stream if provided if stream is not None: s = Stream_accept(stream) - set_deallocation_stream(self._h_ptr, s._h_stream) + h_context = _stream_context_handle(s) + set_deallocation_stream(self._h_ptr, s._h_stream, h_context) # Reset handle - RAII deleter will free the memory (and release owner ref in C++) self._h_ptr.reset() self._size = 0 diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyx b/cuda_core/cuda/core/_memory/_memory_pool.pyx index fbb320e02ff..acebb626bce 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyx +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyx @@ -347,11 +347,8 @@ cdef inline void _MP_deallocate( ) noexcept nogil: cdef cydriver.CUstream s = as_cu(stream._h_stream) cdef cydriver.CUdeviceptr devptr = ptr - cdef cydriver.CUresult r with nogil: - r = cydriver.cuMemFreeAsync(devptr, s) - if r != cydriver.CUDA_ERROR_INVALID_CONTEXT: - HANDLE_RETURN(r) + HANDLE_RETURN(cydriver.cuMemFreeAsync(devptr, s)) cdef inline _MP_close(_MemPool self): diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd index 2f481fee4f8..f925c5be66b 100644 --- a/cuda_core/cuda/core/_resource_handles.pxd +++ b/cuda_core/cuda/core/_resource_handles.pxd @@ -211,7 +211,8 @@ cdef DevicePtrHandle deviceptr_create_mapped_graphics( const GraphicsResourceHandle& h_resource, const StreamHandle& h_stream) except+ nogil cdef DevicePtrHandle deviceptr_create_with_mr( - cydriver.CUdeviceptr ptr, size_t size, object mr) except+ nogil + cydriver.CUdeviceptr ptr, size_t size, object mr, + const ContextHandle& h_context) except+ nogil # MR deallocation callback type and registration ctypedef void (*MRDeallocCallback)( @@ -222,7 +223,9 @@ 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 void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil +cdef void set_deallocation_stream( + const DevicePtrHandle& h, const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept nogil # Library handles cdef LibraryHandle create_library_handle_from_file(const char* path) except+ nogil diff --git a/cuda_core/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx index f6fb6ac4e20..e3dd7ec68d7 100644 --- a/cuda_core/cuda/core/_resource_handles.pyx +++ b/cuda_core/cuda/core/_resource_handles.pyx @@ -122,14 +122,16 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": void register_mr_dealloc_callback "cuda_core::register_mr_dealloc_callback" ( MRDeallocCallback cb) noexcept DevicePtrHandle deviceptr_create_with_mr "cuda_core::deviceptr_create_with_mr" ( - cydriver.CUdeviceptr ptr, size_t size, object mr) except+ nogil + cydriver.CUdeviceptr ptr, size_t size, object mr, + const ContextHandle& h_context) except+ nogil DevicePtrHandle deviceptr_import_ipc "cuda_core::deviceptr_import_ipc" ( const MemoryPoolHandle& h_pool, const void* export_data, const StreamHandle& h_stream) except+ nogil StreamHandle deallocation_stream "cuda_core::deallocation_stream" ( const DevicePtrHandle& h) noexcept nogil void set_deallocation_stream "cuda_core::set_deallocation_stream" ( - const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil + const DevicePtrHandle& h, const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept nogil # Library handles LibraryHandle create_library_handle_from_file "cuda_core::create_library_handle_from_file" ( @@ -275,6 +277,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": void* p_cuDevicePrimaryCtxRetain "reinterpret_cast(cuda_core::p_cuDevicePrimaryCtxRetain)" void* p_cuDevicePrimaryCtxRelease "reinterpret_cast(cuda_core::p_cuDevicePrimaryCtxRelease)" void* p_cuCtxGetCurrent "reinterpret_cast(cuda_core::p_cuCtxGetCurrent)" + void* p_cuCtxSetCurrent "reinterpret_cast(cuda_core::p_cuCtxSetCurrent)" void* p_cuGreenCtxCreate "reinterpret_cast(cuda_core::p_cuGreenCtxCreate)" void* p_cuGreenCtxDestroy "reinterpret_cast(cuda_core::p_cuGreenCtxDestroy)" void* p_cuCtxFromGreenCtx "reinterpret_cast(cuda_core::p_cuCtxFromGreenCtx)" @@ -376,7 +379,8 @@ cdef void* _get_optional_driver_fn(str name): cdef void _init_driver_fn_pointers() noexcept: - global p_cuDevicePrimaryCtxRetain, p_cuDevicePrimaryCtxRelease, p_cuCtxGetCurrent + global p_cuDevicePrimaryCtxRetain, p_cuDevicePrimaryCtxRelease + global p_cuCtxGetCurrent, p_cuCtxSetCurrent global p_cuGreenCtxCreate, p_cuGreenCtxDestroy, p_cuCtxFromGreenCtx global p_cuDevResourceGenerateDesc, p_cuGreenCtxStreamCreate global p_cuStreamCreateWithPriority, p_cuStreamDestroy @@ -404,6 +408,7 @@ cdef void _init_driver_fn_pointers() noexcept: p_cuDevicePrimaryCtxRetain = _get_driver_fn("cuDevicePrimaryCtxRetain") p_cuDevicePrimaryCtxRelease = _get_driver_fn("cuDevicePrimaryCtxRelease") p_cuCtxGetCurrent = _get_driver_fn("cuCtxGetCurrent") + p_cuCtxSetCurrent = _get_driver_fn("cuCtxSetCurrent") p_cuGreenCtxCreate = _get_optional_driver_fn("cuGreenCtxCreate") p_cuGreenCtxDestroy = _get_optional_driver_fn("cuGreenCtxDestroy") p_cuCtxFromGreenCtx = _get_optional_driver_fn("cuCtxFromGreenCtx") diff --git a/cuda_core/tests/graph/test_graph_definition_lifetime.py b/cuda_core/tests/graph/test_graph_definition_lifetime.py index 804cccc1923..8ad5e23f76c 100644 --- a/cuda_core/tests/graph/test_graph_definition_lifetime.py +++ b/cuda_core/tests/graph/test_graph_definition_lifetime.py @@ -76,7 +76,13 @@ def _wait_until(predicate, timeout=None, interval=0.02): raise AssertionError(f"condition not satisfied within {timeout}s") -from cuda.core import Device, DeviceMemoryResource, EventOptions, Kernel, LaunchConfig +from cuda.core import ( + Device, + DeviceMemoryResource, + EventOptions, + Kernel, + LaunchConfig, +) from cuda.core.graph import ( ChildGraphNode, ConditionalNode, @@ -610,6 +616,68 @@ def test_user_object_cleanup_is_coalesced_on_python_thread(init_cuda): assert set(finalized_threads) == {main_thread} +@pytest.mark.agent_authored(model="gpt-5.6") +def test_deferred_buffer_cleanup_restores_no_current_context(tmp_path): + """Graph-deferred MR cleanup activates its retained allocation context.""" + code = f"timeout = {_FINALIZE_TIMEOUT!r}\n" + textwrap.dedent( + """ + import gc + import time + + from cuda.core import Device, LegacyPinnedMemoryResource + from cuda.core.graph import GraphDefinition + from cuda.core._utils.cuda_utils import driver, handle_return + + class TrackingPinnedMR(LegacyPinnedMemoryResource): + def __init__(self): + self.deallocations = [] + + def deallocate(self, ptr, size, *, stream=None): + super().deallocate(ptr, size, stream=stream) + self.deallocations.append((int(ptr), size)) + + Device(0).set_current() + mr = TrackingPinnedMR() + src = mr.allocate(4) + dst = mr.allocate(4) + graph_def = GraphDefinition() + graph_def.memcpy(dst, src, 4) + + src.close() + dst.close() + assert mr.deallocations == [] + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + del graph_def + + deadline = time.monotonic() + timeout + while len(mr.deallocations) != 2: + gc.collect() + if time.monotonic() >= deadline: + raise AssertionError("deferred buffer cleanup timed out") + time.sleep(0) + time.sleep(0.02) + + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + finally: + handle_return(driver.cuCtxSetCurrent(previous)) + """ + ) + result = subprocess.run( # noqa: S603 - controlled interpreter probe + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=60, + # Avoid shadowing the installed package with cuda_core/cuda/core. + cwd=tmp_path, + ) + assert result.returncode == 0, result.stderr + assert "mr.deallocate() failed" not in result.stderr + + @pytest.mark.agent_authored(model="gpt-5.6") def test_pending_call_queue_saturation_preserves_cleanup(tmp_path): """A full CPython queue neither strands nor mis-threads cleanup.""" diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 4427b899765..1e7bc825f06 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -495,6 +495,29 @@ def test_mr_deallocate_called_on_gc(): assert len(mr.active) == 0 +@pytest.mark.agent_authored(model="gpt-5.6") +@pytest.mark.parametrize("replace_stream", [False, True]) +def test_from_handle_mr_deallocation_restores_no_current_context(init_cuda, capsys, replace_stream): + """MR teardown uses retained pointer or replacement-stream context.""" + mr = TrackingMR() + buf = mr.allocate(1024) + stream = init_cuda.create_stream() if replace_stream else None + assert len(mr.active) == 1 + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + + buf.close(stream) + + assert len(mr.active) == 0 + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert "mr.deallocate() failed" not in capsys.readouterr().err + finally: + handle_return(driver.cuCtxSetCurrent(previous)) + + def test_mr_deallocate_receives_stream(): """Buffer.close(stream) forwards the stream to mr.deallocate() (issue #1619).""" device = Device() @@ -563,6 +586,33 @@ def deallocate(self, ptr, size, *, stream): assert captured["stream"].handle == default_stream().handle +@pytest.mark.agent_authored(model="gpt-5.6") +def test_from_handle_pool_pointer_deallocates_without_current_context(mempool_device, capsys): + """A context-independent pool pointer uses its deallocation stream context.""" + dev = mempool_device + stream = dev.create_stream() + mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + size = 256 + ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, mr.handle, stream.handle)) + stream.sync() + used_after_alloc = mr.attributes.used_mem_current + buf = Buffer.from_handle(int(ptr), size, mr=mr) + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + + buf.close(stream) + stream.sync() + + assert mr.attributes.used_mem_current < used_after_alloc + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert "mr.deallocate() failed" not in capsys.readouterr().err + finally: + handle_return(driver.cuCtxSetCurrent(previous)) + + def test_memory_resource_and_owner_disallowed(): with pytest.raises(ValueError, match="cannot be both specified together"): a = (ctypes.c_byte * 20)()