Support out-of-band buffers in Python pickling

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## New Features
 
+- PR #375 Support out-of-band buffers in Python pickling
+
 ## Improvements
 
 ## Bug Fixes

python/rmm/_lib/device_buffer.pyx

@@ -18,6 +18,8 @@
 # cython: language_level = 3
 
 
+import pickle
+
 import numpy as np
 
 from libcpp.memory cimport unique_ptr
@@ -109,12 +111,11 @@ cdef class DeviceBuffer:
     def size(self):
         return int(self.c_size())
 
-    def __getstate__(self):
-        return self.tobytes()
-
-    def __setstate__(self, state):
-        cdef DeviceBuffer other = DeviceBuffer.c_to_device(state)
-        self.c_obj = move(other.c_obj)
+    def __reduce_ex__(self, protocol):
+        host_data = self.tobytes()
+        if protocol >= 5:
+            host_data = pickle.PickleBuffer(host_data)
+        return to_device, (host_data,)
 
     @property
     def __cuda_array_interface__(self):

python/rmm/tests/test_rmm.py

@@ -250,6 +250,18 @@ def test_rmm_device_buffer_pickle_roundtrip(hb):
     db2 = pickle.loads(pb)
     hb2 = db2.tobytes()
     assert hb == hb2
+    # out-of-band
+    if pickle.HIGHEST_PROTOCOL >= 5:
+        db = rmm.DeviceBuffer.to_device(hb)
+        buffers = []
+        pb2 = pickle.dumps(db, protocol=5, buffer_callback=buffers.append)
+        del db
+        assert len(buffers) == 1
+        assert isinstance(buffers[0], pickle.PickleBuffer)
+        assert bytes(buffers[0]) == hb
+        db3 = pickle.loads(pb2, buffers=buffers)
+        hb3 = db3.tobytes()
+        assert hb3 == hb
 
 
 def test_rmm_cupy_allocator():