test(cuda.core): address review feedback on deallocation-stream PR

- Parametrize test_from_handle_mr_records_default_stream,
  test_from_handle_mr_records_explicit_stream, and
  test_from_handle_stream_requires_mr with [Buffer, ManagedBuffer] to
  cover the ManagedBuffer.from_handle entry point directly.
- Add test_close_with_default_stream_requires_context covering the
  _require_deallocation_stream_context guard in Buffer_close.
- Lift Stream_accept and default_stream to module-level imports.
- Replace _require_deallocation_stream_context (a pre-flight that
  duplicated make_deallocation_stream's context check) with
  _apply_deallocation_stream, which calls set_deallocation_stream once
  and translates CUDA_ERROR_INVALID_CONTEXT into a descriptive
  RuntimeError. Removes the redundant cuCtxGetCurrent call on the
  default-stream success path.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Andy-Jost

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -23,7 +23,6 @@ from cuda.core._resource_handles cimport (
     as_intptr,
     as_cu,
     get_current_context,
-    get_stream_context,
     set_deallocation_stream,
 )
 from cuda.core.typing import DevicePointerType
@@ -75,19 +74,21 @@ cdef void _mr_dealloc_callback(
 register_mr_dealloc_callback(_mr_dealloc_callback)
 
 
-cdef inline void _require_deallocation_stream_context(Stream s) except *:
-    """Default-stream tokens need a current context to pin into the free recipe."""
-    cdef ContextHandle h_ctx
-    if get_stream_context(s._h_stream):
-        return
-    h_ctx = get_current_context()
-    if h_ctx:
-        return
-    raise RuntimeError(
-        "Cannot record a default deallocation stream when no CUDA context is "
-        "current. Call Device.set_current() first, or pass stream= with a "
-        "non-default Stream."
-    )
+cdef inline void _apply_deallocation_stream(
+        const DevicePtrHandle& h_ptr, const StreamHandle& h_stream) except *:
+    """Record h_stream as the deallocation stream for h_ptr.
+
+    Translates CUDA_ERROR_INVALID_CONTEXT (default-stream token with no current
+    context) into a descriptive RuntimeError instead of a raw CUDAError.
+    """
+    cdef cydriver.CUresult status = set_deallocation_stream(h_ptr, h_stream)
+    if status == cydriver.CUresult.CUDA_ERROR_INVALID_CONTEXT:
+        raise RuntimeError(
+            "Cannot record a default deallocation stream when no CUDA context is "
+            "current. Call Device.set_current() first, or pass stream= with a "
+            "non-default Stream."
+        )
+    HANDLE_RETURN(status)
 
 
 __all__ = ['Buffer', 'MemoryResource']
@@ -214,9 +215,8 @@ cdef class Buffer:
         cdef Stream s
         if mr is not None:
             s = Stream_accept(default_stream() if stream is None else stream)
-            _require_deallocation_stream_context(s)
             self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr)
-            HANDLE_RETURN(set_deallocation_stream(self._h_ptr, s._h_stream))
+            _apply_deallocation_stream(self._h_ptr, s._h_stream)
         else:
             self._h_ptr = deviceptr_create_with_owner(c_ptr, owner)
         self._size = size
@@ -675,8 +675,7 @@ cdef inline void Buffer_close(Buffer self, object stream):
     # Update deallocation stream if provided
     if stream is not None:
         s = Stream_accept(stream)
-        _require_deallocation_stream_context(s)
-        HANDLE_RETURN(set_deallocation_stream(self._h_ptr, s._h_stream))
+        _apply_deallocation_stream(self._h_ptr, s._h_stream)
     # Reset handle - RAII deleter will free the memory (and release owner ref in C++)
     self._h_ptr.reset()
     self._size = 0

cuda_core/tests/test_memory.py

@@ -45,6 +45,7 @@
 )
 from cuda.core._dlpack import DLDeviceType
 from cuda.core._memory._ipc import IPCBufferDescriptor
+from cuda.core._stream import Stream_accept, default_stream
 from cuda.core._utils.cuda_utils import CUDAError, handle_return
 from cuda.core.typing import (
     ManagedMemoryLocationType,
@@ -514,16 +515,15 @@ def deallocate(self, ptr, size, *, stream=None):
     assert received["stream"].handle == stream.handle
 
 
-def test_from_handle_mr_records_default_stream():
-    """When a Buffer is minted via :meth:`Buffer.from_handle` with ``mr`` but
-    without an explicit ``stream=``, the deallocation stream is recorded at
+@pytest.mark.parametrize("buffer_type", [Buffer, ManagedBuffer])
+def test_from_handle_mr_records_default_stream(buffer_type):
+    """When a Buffer/ManagedBuffer is minted via :meth:`from_handle` with ``mr``
+    but without an explicit ``stream=``, the deallocation stream is recorded at
     creation as ``default_stream()`` (not chosen later in the destructor).
     See `#2497`.
     """
     import gc
 
-    from cuda.core._stream import Stream_accept, default_stream
-
     device = Device()
     device.set_current()
     captured = {}
@@ -551,7 +551,7 @@ def deallocate(self, ptr, size, *, stream):
 
     mr = StrictCapturingMR()
     # ptr=1 is fine because StrictCapturingMR.deallocate does not free.
-    buf = Buffer.from_handle(1, 1024, mr=mr)
+    buf = buffer_type.from_handle(1, 1024, mr=mr)
     del buf
     gc.collect()
 
@@ -560,12 +560,11 @@ def deallocate(self, ptr, size, *, stream):
 
 
 @pytest.mark.agent_authored(model="cursor-grok-4.5")
-def test_from_handle_mr_records_explicit_stream():
-    """Buffer.from_handle(..., mr=mr, stream=s) stores s for teardown."""
+@pytest.mark.parametrize("buffer_type", [Buffer, ManagedBuffer])
+def test_from_handle_mr_records_explicit_stream(buffer_type):
+    """Buffer/ManagedBuffer.from_handle(..., mr=mr, stream=s) stores s for teardown."""
     import gc
 
-    from cuda.core._stream import Stream_accept
-
     device = Device()
     device.set_current()
     stream = device.create_stream()
@@ -591,20 +590,68 @@ def deallocate(self, ptr, size, *, stream):
             captured["stream"] = Stream_accept(stream)
 
     mr = StrictCapturingMR()
-    buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream)
+    buf = buffer_type.from_handle(1, 1024, mr=mr, stream=stream)
     del buf
     gc.collect()
 
     assert captured["stream"].handle == stream.handle
 
 
 @pytest.mark.agent_authored(model="cursor-grok-4.5")
-def test_from_handle_stream_requires_mr():
+@pytest.mark.parametrize("buffer_type", [Buffer, ManagedBuffer])
+def test_from_handle_stream_requires_mr(buffer_type):
     device = Device()
     device.set_current()
     stream = device.create_stream()
     with pytest.raises(ValueError, match="stream requires a memory resource"):
-        Buffer.from_handle(1, 1024, stream=stream)
+        buffer_type.from_handle(1, 1024, stream=stream)
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-4-6")
+def test_close_with_default_stream_requires_context():
+    """Buffer.close(stream=default_stream()) raises when no context is current.
+
+    ``default_stream()`` has no bound context, so the close path must find
+    a current context to anchor the free.  Without one it should raise rather
+    than silently record an unusable stream handle.
+    """
+    device = Device()
+    device.set_current()
+    stream = device.create_stream()
+
+    class NoopMR(MemoryResource):
+        @property
+        def is_device_accessible(self):
+            return True
+
+        @property
+        def is_host_accessible(self):
+            return False
+
+        @property
+        def device_id(self):
+            return device.device_id
+
+        def allocate(self, size, *, stream):
+            raise NotImplementedError
+
+        def deallocate(self, ptr, size, *, stream):
+            pass
+
+    mr = NoopMR()
+    # Use a real stream at creation so _init succeeds without a current context later.
+    buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream)
+
+    previous = handle_return(driver.cuCtxPopCurrent())
+    assert int(previous) != 0
+    try:
+        assert int(handle_return(driver.cuCtxGetCurrent())) == 0
+        with pytest.raises(RuntimeError, match="no CUDA context is current"):
+            buf.close(stream=default_stream())
+    finally:
+        handle_return(driver.cuCtxSetCurrent(previous))
+
+    buf.close()  # clean up using the recorded stream (which carries a context)
 
 
 @pytest.mark.agent_authored(model="cursor-grok-4.5")
@@ -648,8 +695,6 @@ def deallocate(self, ptr, size, *, stream):
 @pytest.mark.parametrize("buffer_type", [Buffer, ManagedBuffer])
 def test_from_handle_mr_explicit_stream_without_current_context(buffer_type):
     """A context-bound stream makes owning from_handle context-independent."""
-    from cuda.core._stream import Stream_accept
-
     device = Device()
     device.set_current()
     stream = device.create_stream()