How to use the distributed.protocol.serialize.dask_deserialize.register function in distributed

To help you get started, we’ve selected a few distributed examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github dask / distributed / distributed / protocol / h5py.py View on Github external
@dask_deserialize.register(h5py.File)
def deserialize_h5py_file(header, frames):
    import h5py

    return h5py.File(header["filename"], mode="r")
github dask / distributed / distributed / protocol / arrow.py View on Github external
@dask_deserialize.register(pyarrow.RecordBatch)
def deserialize_batch(header, frames):
    blob = frames[0]
    reader = pyarrow.RecordBatchStreamReader(pyarrow.BufferReader(blob))
    return reader.read_next_batch()
github dask / distributed / distributed / protocol / numpy.py View on Github external
@dask_deserialize.register(np.ma.core.MaskedConstant)
def deserialize_numpy_ma_masked(header, frames):
    return np.ma.masked
github dask / distributed / distributed / protocol / numba.py View on Github external
@dask_deserialize.register(numba.cuda.devicearray.DeviceNDArray)
def dask_deserialize_numba_array(header, frames):
    if dask_deserialize_rmm_device_buffer:
        frames = [dask_deserialize_rmm_device_buffer(header, frames)]
    else:
        frames = [numba.cuda.to_device(np.asarray(memoryview(f))) for f in frames]
        for f in frames:
            weakref.finalize(f, numba.cuda.current_context)

    arr = cuda_deserialize_numba_ndarray(header, frames)
    return arr
github dask / distributed / distributed / protocol / cupy.py View on Github external
    @dask_deserialize.register(MatDescriptor)
    def deserialize_cupy_matdescriptor(header, frames):
        return MatDescriptor.create()
github dask / distributed / distributed / protocol / torch.py View on Github external
@dask_deserialize.register(torch.nn.Parameter)
def deserialize_torch_Parameters(header, frames):
    t = dask_deserialize.dispatch(torch.Tensor)(header, frames)
    return torch.nn.Parameter(data=t, requires_grad=header["requires_grad"])
github dask / distributed / distributed / protocol / sparse.py View on Github external
@dask_deserialize.register(sparse.COO)
def deserialize_sparse(header, frames):

    coords_frames = frames[: header["nframes"][0]]
    data_frames = frames[header["nframes"][0] :]

    coords = deserialize(header["coords-header"], coords_frames)
    data = deserialize(header["data-header"], data_frames)

    shape = header["shape"]

    return sparse.COO(coords, data, shape=shape)
github dask / distributed / distributed / protocol / netcdf4.py View on Github external
@dask_deserialize.register(netCDF4.Variable)
def deserialize_netcdf4_variable(header, frames):
    header["type"] = header["parent-type"]
    header["type-serialized"] = header["parent-type-serialized"]
    parent = deserialize(header, frames)
    return parent.variables[header["name"]]
github dask / distributed / distributed / protocol / netcdf4.py View on Github external
@dask_deserialize.register(netCDF4.Group)
def deserialize_netcdf4_group(header, frames):
    file = deserialize_netcdf4_dataset(header, frames)
    return file[header["path"]]