Renames Buffer import_/export methods.

Author: Andy-Jost

cuda_core/cuda/core/experimental/_device.py

@@ -1343,8 +1343,9 @@ def _reconstruct_device(device_id):
         device.set_current()
     return device
 
+
 def _reduce_device(device):
     return _reconstruct_device, (device.device_id,)
 
-multiprocessing.reduction.register(Device, _reduce_device)
 
+multiprocessing.reduction.register(Device, _reduce_device)

cuda_core/cuda/core/experimental/_memory.pyx

@@ -29,7 +29,6 @@ if platform.system() == "Linux":
 
 if TYPE_CHECKING:
     from ._device import Device
-    import cuda.bindings.driver
     import uuid
 
 # TODO: define a memory property mixin class and make Buffer and
@@ -75,7 +74,7 @@ cdef class Buffer:
         self.close()
 
     def __reduce__(self):
-        return Buffer.import_, (self.memory_resource, self.export())
+        return Buffer.from_ipc_descriptor, (self.memory_resource, self.get_ipc_descriptor())
 
     cpdef close(self, stream: Stream = None):
         """Deallocate this buffer asynchronously on the given stream.
@@ -137,7 +136,7 @@ cdef class Buffer:
             return self._mr.device_id
         raise NotImplementedError("WIP: Currently this property only supports buffers with associated MemoryResource")
 
-    def export(self) -> IPCBufferDescriptor:
+    def get_ipc_descriptor(self) -> IPCBufferDescriptor:
         """Export a buffer allocated for sharing between processes."""
         if not self._mr.is_ipc_enabled:
             raise RuntimeError("Memory resource is not IPC-enabled")
@@ -146,7 +145,7 @@ cdef class Buffer:
         return IPCBufferDescriptor._init(ptr.reserved, self.size)
 
     @classmethod
-    def import_(cls, mr: MemoryResource, ipc_buffer: IPCBufferDescriptor) -> Buffer:
+    def from_ipc_descriptor(cls, mr: MemoryResource, ipc_buffer: IPCBufferDescriptor) -> Buffer:
         """Import a buffer that was exported from another process."""
         if not mr.is_ipc_enabled:
             raise RuntimeError("Memory resource is not IPC-enabled")
@@ -621,7 +620,7 @@ class DeviceMemoryResource(MemoryResource):
             raise_if_driver_error(err)
 
             if opts.ipc_enabled:
-                self.get_allocation_handle() # enables Buffer.export, sets uuid
+                self.get_allocation_handle() # enables Buffer.get_ipc_descriptor, sets uuid
 
     def __del__(self):
         self.close()
@@ -675,8 +674,9 @@ class DeviceMemoryResource(MemoryResource):
         return self
 
     def unregister(self):
-        with contextlib.suppress(KeyError):
-            del _ipc_registry[self.uuid]
+        if _ipc_registry is not None:
+            with contextlib.suppress(KeyError):
+                del _ipc_registry[self.uuid]
 
     @property
     def uuid(self):

cuda_core/tests/memory_ipc/test_errors.py

@@ -73,11 +73,11 @@ class TestImportWrongMR(ChildErrorHarness):
     def PARENT_ACTION(self, queue):
         mr2 = DeviceMemoryResource(self.device, dict(max_size=POOL_SIZE, ipc_enabled=True))
         buffer = mr2.allocate(NBYTES)
-        queue.put([self.mr, buffer.export()])  # Note: mr does not own this buffer
+        queue.put([self.mr, buffer.get_ipc_descriptor()])  # Note: mr does not own this buffer
 
     def CHILD_ACTION(self, queue):
         mr, buffer_desc = queue.get(timeout=CHILD_TIMEOUT_SEC)
-        Buffer.import_(mr, buffer_desc)
+        Buffer.from_ipc_descriptor(mr, buffer_desc)
 
     def ASSERT(self, exc_type, exc_msg):
         assert exc_type is CUDAError
@@ -110,7 +110,7 @@ def PARENT_ACTION(self, queue):
 
     def CHILD_ACTION(self, queue):
         buffer = queue.get(timeout=CHILD_TIMEOUT_SEC)
-        Buffer.import_(self.mr, buffer)
+        Buffer.from_ipc_descriptor(self.mr, buffer)
 
     def ASSERT(self, exc_type, exc_msg):
         assert exc_type is TypeError

cuda_core/tests/memory_ipc/test_leaks.py

@@ -80,7 +80,7 @@ def __reduce__(self):
         lambda mr: mr.get_allocation_handle(),
         lambda mr: mr,
         lambda mr: mr.allocate(NBYTES),
-        lambda mr: mr.allocate(NBYTES).export(),
+        lambda mr: mr.allocate(NBYTES).get_ipc_descriptor(),
     ],
     ids=["alloc_handle", "mr", "buffer", "buffer_desc"],
 )

cuda_core/tests/memory_ipc/test_memory_ipc.py

@@ -109,8 +109,8 @@ def test_main(self, device, ipc_memory_resource):
         # Allocate and share memory.
         buf1 = mr.allocate(NBYTES)
         buf2 = mr.allocate(NBYTES)
-        q1.put(buf1.export())
-        q2.put(buf2.export())
+        q1.put(buf1.get_ipc_descriptor())
+        q2.put(buf2.get_ipc_descriptor())
 
         # Wait for children.
         p1.join(timeout=CHILD_TIMEOUT_SEC)
@@ -130,7 +130,7 @@ def child_main(self, alloc_handle, idx, queue):
         device.set_current()
         mr = DeviceMemoryResource.from_allocation_handle(device, alloc_handle)
         buffer_descriptor = queue.get(timeout=CHILD_TIMEOUT_SEC)
-        buffer = Buffer.import_(mr, buffer_descriptor)
+        buffer = Buffer.from_ipc_descriptor(mr, buffer_descriptor)
         IPCBufferTestHelper(device, buffer).fill_buffer(starting_from=idx)
 
 

cuda_core/tests/memory_ipc/test_serialize.py

@@ -39,7 +39,7 @@ def test_main(self, device, ipc_memory_resource):
         parent_conn.send(buffer1)  # directly
 
         buffer2 = mr.allocate(NBYTES)
-        parent_conn.send(buffer2.export())  # by descriptor
+        parent_conn.send(buffer2.get_ipc_descriptor())  # by descriptor
 
         # Wait for the child process.
         process.join(timeout=CHILD_TIMEOUT_SEC)
@@ -64,7 +64,7 @@ def child_main(self, conn):
         # Receive the buffers.
         buffer1 = conn.recv()  # directly
         buffer_desc = conn.recv()
-        buffer2 = Buffer.import_(mr, buffer_desc)  # by descriptor
+        buffer2 = Buffer.from_ipc_descriptor(mr, buffer_desc)  # by descriptor
 
         # Modify the buffers.
         IPCBufferTestHelper(device, buffer1).fill_buffer(flipped=True)
@@ -126,7 +126,7 @@ def test_object_passing(device, ipc_memory_resource):
     mr = ipc_memory_resource
     alloc_handle = mr.get_allocation_handle()
     buffer = mr.allocate(NBYTES)
-    buffer_desc = buffer.export()
+    buffer_desc = buffer.get_ipc_descriptor()
 
     helper = IPCBufferTestHelper(device, buffer)
     helper.fill_buffer(flipped=False)
@@ -145,8 +145,8 @@ def child_main(device, alloc_handle, mr1, buffer_desc, buffer1):
 
     # OK to build the buffer from either mr and the descriptor.
     # All buffer* objects point to the same memory.
-    buffer2 = Buffer.import_(mr1, buffer_desc)
-    buffer3 = Buffer.import_(mr2, buffer_desc)
+    buffer2 = Buffer.from_ipc_descriptor(mr1, buffer_desc)
+    buffer3 = Buffer.from_ipc_descriptor(mr2, buffer_desc)
 
     helper1 = IPCBufferTestHelper(device, buffer1)
     helper2 = IPCBufferTestHelper(device, buffer2)

cuda_core/tests/memory_ipc/test_workerpool.py

@@ -36,7 +36,7 @@ def test_ipc_workerpool(self, device, ipc_memory_resource):
         mr = ipc_memory_resource
         buffers = [mr.allocate(NBYTES) for _ in range(NTASKS)]
         with multiprocessing.Pool(processes=NWORKERS, initializer=self.init_worker, initargs=([mr],)) as pool:
-            pool.starmap(self.process_buffer, [(0, buffer.export()) for buffer in buffers])
+            pool.starmap(self.process_buffer, [(0, buffer.get_ipc_descriptor()) for buffer in buffers])
 
         for buffer in buffers:
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
@@ -49,15 +49,16 @@ def test_ipc_workerpool_multi_mr(self, device, ipc_memory_resource):
         buffers = [mr.allocate(NBYTES) for mr, _ in zip(cycle(mrs), range(NTASKS))]
         with multiprocessing.Pool(processes=NWORKERS, initializer=self.init_worker, initargs=(mrs,)) as pool:
             pool.starmap(
-                self.process_buffer, [(mrs.index(buffer.memory_resource), buffer.export()) for buffer in buffers]
+                self.process_buffer,
+                [(mrs.index(buffer.memory_resource), buffer.get_ipc_descriptor()) for buffer in buffers],
             )
 
         for buffer in buffers:
             IPCBufferTestHelper(device, buffer).verify_buffer(flipped=True)
 
     def process_buffer(self, mr_idx, buffer_desc):
         device = Device()
-        buffer = Buffer.import_(g_mrs[mr_idx], buffer_desc)
+        buffer = Buffer.from_ipc_descriptor(g_mrs[mr_idx], buffer_desc)
         IPCBufferTestHelper(device, buffer).fill_buffer(flipped=True)
 
 

cuda_core/tests/test_memory.py

@@ -350,11 +350,11 @@ def test_mempool(mempool_device):
         mr.get_allocation_handle()
 
     with pytest.raises(RuntimeError, match=ipc_error_msg):
-        buffer.export()
+        buffer.get_ipc_descriptor()
 
     with pytest.raises(RuntimeError, match=ipc_error_msg):
         handle = IPCBufferDescriptor._init(b"", 0)
-        Buffer.import_(mr, handle)
+        Buffer.from_ipc_descriptor(mr, handle)
 
     buffer.close()