def deserialize(cls, header, frames):
"""Called when dask.distributed is performing a deserialization for
data of this class.
Do not use this directly. It is invoked by dask.distributed.
Parameters
----------
deserialize : callable
Used to deserialize data that needs further deserialization .
header, frames : dict
See custom serialization documentation in dask.distributed.
Returns
-------
obj : Buffer
Returns an instance of Buffer.
"""
assert len(frames) == 1, "Use Buffer.serialize() for serialization"
iface = header["cuda_array_interface"]
if (
len(iface["shape"]) == 0 or 0 in iface["shape"]
): # The array is empty
arr = rmm.device_array(
iface["shape"], dtype=np.dtype(iface["typestr"])
)
else:
# Updating the data pointer to the frame data.
iface["data"] = frames[0].__cuda_array_interface__["data"]
# We need an object that exposes __cuda_array_interface__
_dummy_iface = types.SimpleNamespace()
_dummy_iface.__cuda_array_interface__ = iface
# Allocating a new RMM CUDA array with shape and dtype as specified
# in `iface` and copy the frame data to it, which makes the memory
# survive until the new Buffer goes out of scope.
# TODO: when DASK supports RMM, we can simply return it as is.
arr, _ = rmm.auto_device(_dummy_iface)
return Buffer(arr)
def to_device(ary):
dary, _ = rmm.auto_device(ary)
return dary