Handle host frames in serialization

CHANGELOG.md

@@ -11,6 +11,7 @@
 - PR #5288 Drop `auto_pickle` decorator #5288
 - PR #5231 Type `Buffer` as `uint8`
 - PR #5308 Coerce frames to `Buffer`s in deserialization
+- PR #5309 Handle host frames in serialization
 - PR #5312 Test serializing `Series` after `slice`
 - PR #5248 Support interleave_columns for string types
 

python/cudf/cudf/core/abc.py

@@ -21,27 +21,47 @@ def deserialize(cls, header, frames):
 
     def device_serialize(self):
         header, frames = self.serialize()
-        assert all((type(f) is cudf.core.buffer.Buffer) for f in frames)
+        assert all(
+            (type(f) in [cudf.core.buffer.Buffer, memoryview]) for f in frames
+        )
         header["type-serialized"] = pickle.dumps(type(self))
+        header["is-cuda"] = [
+            hasattr(f, "__cuda_array_interface__") for f in frames
+        ]
         header["lengths"] = [f.nbytes for f in frames]
         return header, frames
 
     @classmethod
     def device_deserialize(cls, header, frames):
         typ = pickle.loads(header["type-serialized"])
-        frames = [cudf.core.buffer.Buffer(f) for f in frames]
+        frames = [
+            cudf.core.buffer.Buffer(f) if c else memoryview(f)
+            for c, f in zip(header["is-cuda"], frames)
+        ]
+        assert all(
+            (type(f._owner) is rmm.DeviceBuffer)
+            if c
+            else (type(f) is memoryview)
+            for c, f in zip(header["is-cuda"], frames)
+        )
         obj = typ.deserialize(header, frames)
 
         return obj
 
     def host_serialize(self):
         header, frames = self.device_serialize()
-        frames = [f.to_host_array().data for f in frames]
+        frames = [
+            f.to_host_array().data if c else memoryview(f)
+            for c, f in zip(header["is-cuda"], frames)
+        ]
         return header, frames
 
     @classmethod
     def host_deserialize(cls, header, frames):
-        frames = [rmm.DeviceBuffer.to_device(memoryview(f)) for f in frames]
+        frames = [
+            rmm.DeviceBuffer.to_device(f) if c else f
+            for c, f in zip(header["is-cuda"], map(memoryview, frames))
+        ]
         obj = cls.device_deserialize(header, frames)
         return obj