Serialization family to preserve headers of the underlying dumps functions.

[madsbk]

Currently, the dask and cuda serialization families overwrite the header of the underlying dumps functions. For instance, the cuda_dumps overwrite the "type-serialized", which is a problem when Dask uses pickle5 and the underlying loads function use pickle protocol=4.
Beside, overwriting an underlying protocol's header is bad style.

This PR fixes this issue

• Closes CI fails: "ValueError: unsupported pickle protocol: 5" rapidsai/dask-cuda#746
• Tests added / passed
• Passes pre-commit run --all-files

[madsbk]

Notice, I have changed the test test_compression_numpy_list() to better reflect what (I think) we want to support.

[jrbourbeau]

Thanks @madsbk! It looks like there's another spot we need to insert ["sub-headers"]

____________________________ test_deserialize_grad _____________________________

    def test_deserialize_grad():
        a = np.random.rand(8, 1)
        t = torch.tensor(a, requires_grad=True, dtype=torch.float)
>       t2 = deserialize(*serialize(t))

distributed/protocol/tests/test_torch.py:48: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
distributed/protocol/serialize.py:410: in deserialize
    return loads(header, frames)
distributed/protocol/serialize.py:50: in dask_loads
    return loads(header["sub-header"], frames)
distributed/protocol/torch.py:36: in deserialize_torch_Tensor
    x = dask_deserialize.dispatch(np.ndarray)(header, frames)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

header = {'device': 'cpu', 'requires_grad': True, 'serializer': 'dask', 'sub-header': {'dtype': (0, '<f4'), 'shape': (8, 1), 'strides': (4, 4), 'writeable': [True]}, ...}
frames = [<memory at 0x13abbdb80>]

    @dask_deserialize.register(np.ndarray)
    def deserialize_numpy_ndarray(header, frames):
        with log_errors():
            if header.get("pickle"):
                return pickle.loads(frames[0], buffers=frames[1:])
    
            (frame,) = frames
>           (writeable,) = header["writeable"]
E           KeyError: 'writeable'

distributed/protocol/numpy.py:116: KeyError

@jakirkham would you have time to take a look at this PR?

[madsbk]

Thanks @madsbk! It looks like there's another spot we need to insert ["sub-headers"]

Thanks, the torch protocol now also uses a sub-header

[jakirkham]

@jakirkham would you have time to take a look at this PR?

@jrbourbeau, were you just wanting feedback on that test failure or was there something else you were looking for?

[madsbk]

Hmm, distributed/protocol/tests/test_torch.py::test_resnet is still failing. Will look at it tomorrow.

[madsbk]

As far as I can see, the CI errors isn't related to this PR.

FAILED distributed/tests/test_stress.py::test_stress_creation_and_deletion - ...
FAILED distributed/comm/tests/test_ws.py::test_collections - tornado.util.Tim..
FAILED distributed/deploy/tests/test_adaptive.py::test_adaptive_local_cluster_multi_workers
FAILED distributed/tests/test_asyncprocess.py::test_exit_callback 

[madsbk]

@jrbourbeau, I think this is ready to be merged.

[jakirkham]

@jrbourbeau was there anything else we needed to do here or is this good to go?

[madsbk]

@jrbourbeau it would be good to get this merged, it is blocking Dask-CUDA: rapidsai/dask-cuda#746

[jakirkham]

Going to merge to unblock Dask-CUDA. The PyTorch issue noted originally has been resolved.

Most of the failures here seem to be related to Distributed's concurrent.future support, which are unrelated. There is one failure with test_spill_to_disk, but it only happens in one job. We might want to keep an eye on that though.