add resource management to mempool

Author: ksimpson-work

cuda_core/cuda/core/experimental/_memory.py

@@ -213,7 +213,8 @@ def __release_buffer__(self, buffer: memoryview, /):
 class IPCBuffer(Buffer):
     """Buffer class to represent a buffer description which can be shared across processes.
     It is not a valid buffer containing data, but rather a description used by the importing
-    process to construct a valid buffer."""
+    process to construct a valid buffer. It's primary use is to provide a serialization
+    mechanism for passing exported buffers between processes."""
 
     def __init__(self, reserved: bytes, size):
         super().__init__(0, 0)
@@ -248,10 +249,36 @@ def _reconstruct(cls, reserved, size):
 
 
 class MemoryResource(abc.ABC):
-    __slots__ = ("_handle",)
+    """Base class for memory resources.
+
+    This class provides an abstract interface for memory resources and includes
+    an optional destruction mechanism through _MembersNeededForFinalize.
+    """
+
+    class _MembersNeededForFinalize:
+        __slots__ = ("handle", "is_initialized")
+
+        def __init__(self, mr_obj, handle=None):
+            self.handle = handle
+            self.is_initialized = False
+            weakref.finalize(mr_obj, self.close)
+
+        def close(self):
+            if self.is_initialized and self.handle:
+                # Specific cleanup can be implemented by derived classes
+                self.handle = None
+                self.is_initialized = False
+
+    __slots__ = ("__weakref__", "_mnff")
 
     @abc.abstractmethod
-    def __init__(self, *args, **kwargs): ...
+    def __init__(self, *args, **kwargs):
+        self._mnff = MemoryResource._MembersNeededForFinalize(self)
+
+    def close(self):
+        """Release any resources associated with this memory resource."""
+        if hasattr(self, "_mnff"):
+            self._mnff.close()
 
     @abc.abstractmethod
     def allocate(self, size, stream=None) -> Buffer: ...
@@ -327,7 +354,22 @@ class Mempool(MemoryResource):
     from_shared_handle : Import an existing memory pool from another process
     """
 
-    __slots__ = ("_dev_id", "_handle", "_ipc_enabled")
+    class _MembersNeededForFinalize:
+        __slots__ = ("handle", "is_initialized")
+
+        def __init__(self, mr_obj, handle=None):
+            self.handle = handle
+            weakref.finalize(mr_obj, self.close)
+
+        def close(self):
+            if self.is_initialized and self.handle:
+                handle_return(driver.cuMemPoolDestroy(self.handle))
+                self.handle = None
+
+    __slots__ = (
+        "__weakref__",
+        "_mnff",
+    )
 
     def __init__(self):
         """Direct instantiation is not supported.
@@ -360,8 +402,8 @@ def _init(dev_id: int, handle: int, ipc_enabled: bool) -> Mempool:
         """
         self = Mempool.__new__(Mempool)
         self._dev_id = dev_id
-        self._handle = handle
         self._ipc_enabled = ipc_enabled
+        self._mnff = Mempool._MembersNeededForFinalize(self, handle)
         return self
 
     @staticmethod
@@ -463,7 +505,7 @@ def get_shareable_handle(self) -> int:
         """
         if not self._ipc_enabled:
             raise RuntimeError("This memory pool was not created with IPC support enabled")
-        return handle_return(driver.cuMemPoolExportToShareableHandle(self._handle, _get_platform_handle_type(), 0))
+        return handle_return(driver.cuMemPoolExportToShareableHandle(self._mnff.handle, _get_platform_handle_type(), 0))
 
     def export_buffer(self, buffer: Buffer) -> IPCBuffer:
         """Export a buffer allocated from this pool for sharing between processes.
@@ -528,7 +570,9 @@ def import_buffer(self, ipc_buffer: IPCBuffer) -> Buffer:
             raise RuntimeError("This memory pool was not created with IPC support enabled")
         share_data = driver.CUmemPoolPtrExportData()
         share_data.reserved = ipc_buffer.reserved
-        return Buffer(handle_return(driver.cuMemPoolImportPointer(self._handle, share_data)), ipc_buffer._size, self)
+        return Buffer(
+            handle_return(driver.cuMemPoolImportPointer(self._mnff.handle, share_data)), ipc_buffer._size, self
+        )
 
     def allocate(self, size: int, stream=None) -> Buffer:
         """Allocate memory from the pool.
@@ -552,7 +596,7 @@ def allocate(self, size: int, stream=None) -> Buffer:
         """
         if stream is None:
             stream = default_stream()
-        ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, self._handle, stream.handle))
+        ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, self._mnff.handle, stream.handle))
         return Buffer(ptr, size, self)
 
     def deallocate(self, ptr: int, size: int, stream=None) -> None:
@@ -615,7 +659,7 @@ def reuse_follow_event_dependencies(self) -> bool:
         return bool(
             handle_return(
                 driver.cuMemPoolGetAttribute(
-                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES
                 )
             )
         )
@@ -626,7 +670,7 @@ def reuse_allow_opportunistic(self) -> bool:
         return bool(
             handle_return(
                 driver.cuMemPoolGetAttribute(
-                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC
                 )
             )
         )
@@ -637,7 +681,7 @@ def reuse_allow_internal_dependencies(self) -> bool:
         return bool(
             handle_return(
                 driver.cuMemPoolGetAttribute(
-                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES
                 )
             )
         )
@@ -647,7 +691,9 @@ def release_threshold(self) -> int:
         """Amount of reserved memory to hold before OS release."""
         return int(
             handle_return(
-                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD)
+                driver.cuMemPoolGetAttribute(
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD
+                )
             )
         )
 
@@ -657,7 +703,7 @@ def reserved_mem_current(self) -> int:
         return int(
             handle_return(
                 driver.cuMemPoolGetAttribute(
-                    self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT
                 )
             )
         )
@@ -667,7 +713,9 @@ def reserved_mem_high(self) -> int:
         """High watermark of backing memory allocated."""
         return int(
             handle_return(
-                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH)
+                driver.cuMemPoolGetAttribute(
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH
+                )
             )
         )
 
@@ -676,7 +724,9 @@ def used_mem_current(self) -> int:
         """Current amount of memory in use."""
         return int(
             handle_return(
-                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_CURRENT)
+                driver.cuMemPoolGetAttribute(
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_CURRENT
+                )
             )
         )
 
@@ -685,7 +735,9 @@ def used_mem_high(self) -> int:
         """High watermark of memory in use."""
         return int(
             handle_return(
-                driver.cuMemPoolGetAttribute(self._handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_HIGH)
+                driver.cuMemPoolGetAttribute(
+                    self._mnff.handle, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_HIGH
+                )
             )
         )