switch to a context manager

Author: Matt711

python/rapidsmpf/rapidsmpf/memory/buffer.pxd

@@ -3,6 +3,9 @@
 
 from libc.stddef cimport size_t
 from libcpp.memory cimport unique_ptr
+from rmm.pylibrmm.stream cimport Stream
+
+from rapidsmpf.memory.buffer_resource cimport BufferResource
 
 
 cdef extern from "<rapidsmpf/memory/buffer.hpp>" namespace "rapidsmpf" nogil:
@@ -13,14 +16,21 @@ cdef extern from "<rapidsmpf/memory/buffer.hpp>" namespace "rapidsmpf" nogil:
 
     cdef cppclass cpp_Buffer "rapidsmpf::Buffer":
         size_t size
-        # data() actually returns const std::byte*, declared as const void* here
-        # because const std::byte* -> const void* is an implicit C++ conversion.
-        const void* data() except +
+        # exclusive_data_access() returns std::byte*; void* works because
+        # non-const std::byte* -> void* is an implicit C++ conversion.
+        void* exclusive_data_access() except +
+        void unlock() noexcept
         MemoryType mem_type() noexcept
 
 
+cdef class BufferHostView:
+    cdef Buffer _buf
+
+
 cdef class Buffer:
     cdef unique_ptr[cpp_Buffer] _handle
+    cdef BufferResource _br
+    cdef Stream _stream
 
     @staticmethod
-    cdef Buffer from_handle(unique_ptr[cpp_Buffer] handle)
+    cdef Buffer from_handle(unique_ptr[cpp_Buffer] handle, BufferResource br, Stream stream)

python/rapidsmpf/rapidsmpf/memory/buffer.pyi

@@ -1,16 +1,21 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from contextlib import AbstractContextManager
 from enum import IntEnum
 
 class MemoryType(IntEnum):
     DEVICE = ...
     PINNED_HOST = ...
     HOST = ...
 
+class BufferHostView(AbstractContextManager[memoryview]):
+    def __enter__(self) -> memoryview: ...
+    def __exit__(self, *args: object) -> None: ...
+
 class Buffer:
     @property
     def size(self) -> int: ...
     @property
     def mem_type(self) -> MemoryType: ...
-    def __buffer__(self, flags: int, /) -> memoryview: ...
+    def host_view(self) -> BufferHostView: ...

python/rapidsmpf/rapidsmpf/memory/buffer.pyx

@@ -1,9 +1,32 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from cpython.buffer cimport PyBuffer_FillInfo
+from cpython.buffer cimport PyBUF_WRITE
+from cpython.memoryview cimport PyMemoryView_FromMemory
 from cython.operator cimport dereference as deref
 from libcpp.utility cimport move
+from rmm.pylibrmm.stream cimport Stream
+
+from rapidsmpf.memory.buffer_resource cimport BufferResource
+
+
+cdef class BufferHostView:
+    def __cinit__(self, Buffer buf):
+        self._buf = buf
+
+    def __enter__(self):
+        cdef void* ptr = deref(self._buf._handle).exclusive_data_access()
+        try:
+            return PyMemoryView_FromMemory(
+                <char*>ptr, <Py_ssize_t>deref(self._buf._handle).size, PyBUF_WRITE
+            )
+        except BaseException:
+            deref(self._buf._handle).unlock()
+            raise
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        deref(self._buf._handle).unlock()
+        return False
 
 
 cdef class Buffer:
@@ -22,9 +45,11 @@ cdef class Buffer:
             self._handle.reset()
 
     @staticmethod
-    cdef Buffer from_handle(unique_ptr[cpp_Buffer] handle):
+    cdef Buffer from_handle(unique_ptr[cpp_Buffer] handle, BufferResource br, Stream stream):
         cdef Buffer self = Buffer.__new__(Buffer)
         self._handle = move(handle)
+        self._br = br
+        self._stream = stream
         return self
 
     @property
@@ -37,16 +62,37 @@ cdef class Buffer:
         """Memory type of this buffer."""
         return deref(self._handle).mem_type()
 
-    def __getbuffer__(self, Py_buffer* view, int flags):
+    def host_view(self):
+        """Context manager providing exclusive writable host access to the buffer.
+
+        Acquires an exclusive lock on the buffer for the duration of the ``with``
+        block, preventing concurrent stream-ordered operations on the C++ side.
+        The lock is released (and the returned ``memoryview`` must not be used)
+        once the block exits.
+
+        Returns
+        -------
+        BufferHostView
+            A context manager that yields a writable ``memoryview`` of the buffer.
+
+        Raises
+        ------
+        TypeError
+            If the buffer is not a host buffer (``HOST`` or ``PINNED_HOST``).
+        std::logic_error
+            If the buffer is already locked or a stream-ordered write is still
+            in flight (``is_latest_write_done() == False``).
+
+        Examples
+        --------
+        >>> with buf.host_view() as mv:
+        ...     mv[:] = b"\\x00" * buf.size
+        """
         if deref(self._handle).mem_type() not in {
             MemoryType.HOST, MemoryType.PINNED_HOST
         }:
             raise TypeError(
-                "buffer protocol is only supported for host buffers "
+                "host_view() is only supported for host buffers "
                 "(MemoryType.HOST or MemoryType.PINNED_HOST)"
             )
-        cdef void* ptr = <void*><const void*>deref(self._handle).data()
-        PyBuffer_FillInfo(view, self, ptr, deref(self._handle).size, False, flags)
-
-    def __releasebuffer__(self, Py_buffer* view):
-        pass
+        return BufferHostView(self)

python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx

@@ -595,7 +595,7 @@ cdef class BufferResource:
                     size, stream.view(), deref(reservation._handle)
                 )
             )
-        return Buffer.from_handle(move(handle))
+        return Buffer.from_handle(move(handle), self, stream)
 
     @property
     def statistics(self):

python/rapidsmpf/rapidsmpf/tests/test_buffer.py

@@ -63,37 +63,41 @@ def test_make_buffer(mem_type: MemoryType) -> None:
 
 
 @skip_if_no_pinned
-def test_buffer_protocol_pinned_host(pinned_br: BufferResource) -> None:
+def test_host_view_pinned_host(pinned_br: BufferResource) -> None:
     size = 256
     stream = Stream()
     reservation, _ = pinned_br.reserve(
         MemoryType.PINNED_HOST, size, allow_overbooking=False
     )
     buf = pinned_br.make_buffer(size, stream, reservation)
 
-    mv = memoryview(buf)
-    assert len(mv) == size
-    assert mv.format == "B"
-
     data = np.arange(size, dtype=np.uint8)
-    np.frombuffer(mv, dtype=np.uint8)[:] = data
-    assert np.array_equal(np.frombuffer(memoryview(buf), dtype=np.uint8), data)
+    with buf.host_view() as mv:
+        assert len(mv) == size
+        np.frombuffer(mv, dtype=np.uint8)[:] = data
+
+    with buf.host_view() as mv:
+        assert np.array_equal(np.frombuffer(mv, dtype=np.uint8), data)
 
 
-def test_buffer_protocol_host() -> None:
+def test_host_view_host() -> None:
     size = 256
     mr = rmm.mr.CudaMemoryResource()
     br = BufferResource(mr, memory_limits={MemoryType.HOST: size * 4})
     stream = Stream()
     reservation, _ = br.reserve(MemoryType.HOST, size, allow_overbooking=False)
     buf = br.make_buffer(size, stream, reservation)
 
-    mv = memoryview(buf)
-    assert len(mv) == size
-    assert mv.format == "B"
+    data = np.arange(size, dtype=np.uint8)
+    with buf.host_view() as mv:
+        assert len(mv) == size
+        np.frombuffer(mv, dtype=np.uint8)[:] = data
+
+    with buf.host_view() as mv:
+        assert np.array_equal(np.frombuffer(mv, dtype=np.uint8), data)
 
 
-def test_buffer_protocol_rejects_device() -> None:
+def test_host_view_rejects_device() -> None:
     size = 1024
     mr = rmm.mr.CudaMemoryResource()
     br = BufferResource(mr, memory_limits={MemoryType.DEVICE: size * 4})
@@ -102,4 +106,19 @@ def test_buffer_protocol_rejects_device() -> None:
     buf = br.make_buffer(size, stream, reservation)
 
     with pytest.raises(TypeError, match="host buffers"):
-        memoryview(buf)
+        buf.host_view()
+
+
+def test_host_view_lock_released_on_error() -> None:
+    size = 256
+    mr = rmm.mr.CudaMemoryResource()
+    br = BufferResource(mr, memory_limits={MemoryType.HOST: size * 4})
+    stream = Stream()
+    reservation, _ = br.reserve(MemoryType.HOST, size, allow_overbooking=False)
+    buf = br.make_buffer(size, stream, reservation)
+
+    with pytest.raises(RuntimeError), buf.host_view():
+        raise RuntimeError("intentional")
+
+    with buf.host_view():
+        pass