Moves AllocationHandle serialization to a registration with multiprocessing, since it depends on DupFd.

Author: Andy-Jost

cuda_core/cuda/core/experimental/_device.py

@@ -1337,15 +1337,15 @@ def create_graph_builder(self) -> GraphBuilder:
         return GraphBuilder._init(stream=self.create_stream(), is_stream_owner=True)
 
 
+def _reduce_device(device):
+    return _reconstruct_device, (device.device_id,)
+
+
 def _reconstruct_device(device_id):
     device = Device(device_id)
     if not device._has_inited:
         device.set_current()
     return device
 
 
-def _reduce_device(device):
-    return _reconstruct_device, (device.device_id,)
-
-
 multiprocessing.reduction.register(Device, _reduce_device)

cuda_core/cuda/core/experimental/_memory.pyx

@@ -16,6 +16,7 @@ import abc
 import array
 import contextlib
 import cython
+import multiprocessing
 import os
 import platform
 import sys
@@ -421,16 +422,6 @@ cdef class IPCAllocationHandle:
         """Close the handle."""
         self.close()
 
-    def __reduce__(self):
-        import multiprocessing
-        multiprocessing.context.assert_spawning(self)
-        df = multiprocessing.reduction.DupFd(self.handle)
-        return self._reconstruct, (df, self._uuid)
-
-    @classmethod
-    def _reconstruct(cls, df, uuid):
-        return cls._init(df.detach(), uuid)
-
     def __int__(self) -> int:
         if self._handle < 0:
             raise ValueError(
@@ -447,6 +438,17 @@ cdef class IPCAllocationHandle:
         return self._uuid
 
 
+def _reduce_allocation_handle(alloc_handle):
+    df = multiprocessing.reduction.DupFd(alloc_handle.handle)
+    return _reconstruct_allocation_handle, (type(alloc_handle), df, alloc_handle.uuid)
+
+def _reconstruct_allocation_handle(cls, df, uuid):
+    return cls._init(df.detach(), uuid)
+
+
+multiprocessing.reduction.register(IPCAllocationHandle, _reduce_allocation_handle)
+
+
 @dataclass
 cdef class DeviceMemoryResourceOptions:
     """Customizable :obj:`~_memory.DeviceMemoryResource` options.