Repository: dask/distributed

Issue:
Client spews errors in JupyterLab during `compute`
In recent versions of distributed, during a compute, tons of errors sometimes start spewing out in JupyterLab like this:

![Screen Shot 2021-10-27 at 3 53 08 PM](https://user-images.githubusercontent.com/3309802/139152831-2fe9f2be-0ab5-4600-a83a-51732429d620.png)

I've heard other people complain about this too.

For searchability, here are some of the logs:
<details>

```
tornado.application - ERROR - Exception in callback functools.partial(<bound method IOLoop._discard_future_result of <zmq.eventloop.ioloop.ZMQIOLoop object at 0x1334d13d0>>, <Task finished name='Task-337' coro=<Cluster._sync_cluster_info() done, defined at /Users/gabe/dev/distributed/distributed/deploy/cluster.py:104> exception=OSError('Timed out trying to connect to tls://54.212.201.147:8786 after 5 s')>)
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/comm/tcp.py", line 398, in connect
    stream = await self.client.connect(
  File "/Users/gabe/dev/dask-playground/env/lib/python3.9/site-packages/tornado/tcpclient.py", line 288, in connect
    stream = await stream.start_tls(
asyncio.exceptions.CancelledError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/gabe/.pyenv/versions/3.9.1/lib/python3.9/asyncio/tasks.py", line 489, in wait_for
    fut.result()
asyncio.exceptions.CancelledError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/comm/core.py", line 284, in connect
    comm = await asyncio.wait_for(
  File "/Users/gabe/.pyenv/versions/3.9.1/lib/python3.9/asyncio/tasks.py", line 491, in wait_for
    raise exceptions.TimeoutError() from exc
asyncio.exceptions.TimeoutError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/gabe/dev/dask-playground/env/lib/python3.9/site-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/Users/gabe/dev/dask-playground/env/lib/python3.9/site-packages/tornado/ioloop.py", line 765, in _discard_future_result
    future.result()
  File "/Users/gabe/dev/distributed/distributed/deploy/cluster.py", line 105, in _sync_cluster_info
    await self.scheduler_comm.set_metadata(
  File "/Users/gabe/dev/distributed/distributed/core.py", line 785, in send_recv_from_rpc
    comm = await self.live_comm()
  File "/Users/gabe/dev/distributed/distributed/core.py", line 742, in live_comm
    comm = await connect(
  File "/Users/gabe/dev/distributed/distributed/comm/core.py", line 308, in connect
    raise OSError(
OSError: Timed out trying to connect to tls://54.212.201.147:8786 after 5 s
tornado.application - ERROR - Exception in callback functools.partial(<bound method IOLoop._discard_future_result of <zmq.eventloop.ioloop.ZMQIOLoop object at 0x1334d13d0>>, <Task finished name='Task-340' coro=<Cluster._sync_cluster_info() done, defined at /Users/gabe/dev/distributed/distributed/deploy/cluster.py:104> exception=OSError('Timed out trying to connect to tls://54.212.201.147:8786 after 5 s')>)
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/comm/tcp.py", line 398, in connect
    stream = await self.client.connect(
  File "/Users/gabe/dev/dask-playground/env/lib/python3.9/site-packages/tornado/tcpclient.py", line 288, in connect
    stream = await stream.start_tls(
asyncio.exceptions.CancelledError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/gabe/.pyenv/versions/3.9.1/lib/python3.9/asyncio/tasks.py", line 489, in wait_for
    fut.result()
asyncio.exceptions.CancelledError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/comm/core.py", line 284, in connect
    comm = await asyncio.wait_for(
  File "/Users/gabe/.pyenv/versions/3.9.1/lib/python3.9/asyncio/tasks.py", line 491, in wait_for
    raise exceptions.TimeoutError() from exc
asyncio.exceptions.TimeoutError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/gabe/dev/dask-playground/env/lib/python3.9/site-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/Users/gabe/dev/dask-playground/env/lib/python3.9/site-packages/tornado/ioloop.py", line 765, in _discard_future_result
    future.result()
  File "/Users/gabe/dev/distributed/distributed/deploy/cluster.py", line 105, in _sync_cluster_info
    await self.scheduler_comm.set_metadata(
  File "/Users/gabe/dev/distributed/distributed/core.py", line 785, in send_recv_from_rpc
    comm = await self.live_comm()
  File "/Users/gabe/dev/distributed/distributed/core.py", line 742, in live_comm
    comm = await connect(
  File "/Users/gabe/dev/distributed/distributed/comm/core.py", line 308, in connect
    raise OSError(
OSError: Timed out trying to connect to tls://54.212.201.147:8786 after 5 s
tornado.application - ERROR - Exception in callback functools.partial(<bound method IOLoop._discard_future_result of <zmq.eventloop.ioloop.ZMQIOLoop object at 0x1334d13d0>>, <Task finished name='Task-349' coro=<Cluster._sync_cluster_info() done, defined at /Users/gabe/dev/distributed/distributed/deploy/cluster.py:104> exception=OSError('Timed out trying to connect to tls://54.212.201.147:8786 after 5 s')>)
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/comm/tcp.py", line 398, in connect
    stream = await self.client.connect(
  File "/Users/gabe/dev/dask-playground/env/lib/python3.9/site-packages/tornado/tcpclient.py", line 288, in connect
    stream = await stream.start_tls(
asyncio.exceptions.CancelledError
```

</details>

From reading the traceback, it appears to be unhandled exceptions in the new cluster<->scheduler synced dict from https://github.com/dask/distributed/pull/5033. I'd guess that when the scheduler gets overwhelmed from any of the myriad of things that can block its event loop, something breaks. I'm not sure why it stays broken and the comm doesn't reconnect, but after one failure it seems you'll get this message once per second.
1. There should be error handling in the cluster info syncing; at a minimum, any error here should be handled and logged, but not allowed to propagate up
2. Why is this error happening, and why does it not seem to recover?

cc @jacobtomlinson 

Discussion:
jcrist: Test failures are unrelated.

Patch:
diff --git a/README.rst b/README.rst
index 4a539443..119ba559 100644
--- a/README.rst
+++ b/README.rst
@@ -1,7 +1,7 @@
 Distributed
 ===========
 
-|Test Status| |Coverage| |Doc Status| |Gitter| |Version Status| |NumFOCUS|
+|Test Status| |Coverage| |Doc Status| |Discourse| |Version Status| |NumFOCUS|
 
 A library for distributed computation.  See documentation_ for more details.
 
@@ -14,9 +14,9 @@ A library for distributed computation.  See documentation_ for more details.
 .. |Coverage| image:: https://codecov.io/gh/dask/distributed/branch/main/graph/badge.svg
    :target: https://codecov.io/gh/dask/distributed/branch/main
    :alt: Coverage status
-.. |Gitter| image:: https://badges.gitter.im/Join%20Chat.svg
-   :alt: Join the chat at https://gitter.im/dask/dask
-   :target: https://gitter.im/dask/dask?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge
+.. |Discourse| image:: https://img.shields.io/discourse/users?logo=discourse&server=https%3A%2F%2Fdask.discourse.group
+   :alt: Discuss Dask-related things and ask for help
+   :target: https://dask.discourse.group
 .. |Version Status| image:: https://img.shields.io/pypi/v/distributed.svg
    :target: https://pypi.python.org/pypi/distributed/
 .. |NumFOCUS| image:: https://img.shields.io/badge/powered%20by-NumFOCUS-orange.svg?style=flat&colorA=E1523D&colorB=007D8A
diff --git a/distributed/comm/ucx.py b/distributed/comm/ucx.py
index e47b61a4..61c6c0dc 100644
--- a/distributed/comm/ucx.py
+++ b/distributed/comm/ucx.py
@@ -5,6 +5,7 @@ See :ref:`communications` for more.
 
 .. _UCX: https://github.com/openucx/ucx
 """
+import asyncio
 import functools
 import logging
 import os
@@ -278,16 +279,6 @@ class UCX(Comm):
 
                 # Send meta data
 
-                # Send close flag and number of frames (_Bool, int64)
-                await self.ep.send(struct.pack("?Q", False, nframes))
-                # Send which frames are CUDA (bool) and
-                # how large each frame is (uint64)
-                await self.ep.send(
-                    struct.pack(nframes * "?" + nframes * "Q", *cuda_frames, *sizes)
-                )
-
-                # Send frames
-
                 # It is necessary to first synchronize the default stream before start
                 # sending We synchronize the default stream because UCX is not
                 # stream-ordered and syncing the default stream will wait for other
@@ -296,8 +287,22 @@ class UCX(Comm):
                 if any(cuda_send_frames):
                     synchronize_stream(0)
 
+                tasks = []
+
+                # Send close flag and number of frames (_Bool, int64)
+                tasks.append(self.ep.send(struct.pack("?Q", False, nframes)))
+                # Send which frames are CUDA (bool) and
+                # how large each frame is (uint64)
+                tasks.append(
+                    self.ep.send(
+                        struct.pack(nframes * "?" + nframes * "Q", *cuda_frames, *sizes)
+                    )
+                )
+
+                # Send frames
                 for each_frame in send_frames:
-                    await self.ep.send(each_frame)
+                    tasks.append(self.ep.send(each_frame))
+                await asyncio.gather(*tasks)
                 return sum(sizes)
             except (ucp.exceptions.UCXBaseException):
                 self.abort()
@@ -354,8 +359,10 @@ class UCX(Comm):
                 if any(cuda_recv_frames):
                     synchronize_stream(0)
 
+                tasks = []
                 for each_frame in recv_frames:
-                    await self.ep.recv(each_frame)
+                    tasks.append(self.ep.recv(each_frame))
+                await asyncio.gather(*tasks)
                 msg = await from_frames(
                     frames,
                     deserialize=self.deserialize,
diff --git a/distributed/deploy/cluster.py b/distributed/deploy/cluster.py
index 20a0a990..fb83808f 100644
--- a/distributed/deploy/cluster.py
+++ b/distributed/deploy/cluster.py
@@ -1,5 +1,4 @@
 import asyncio
-import copy
 import datetime
 import logging
 import threading
@@ -60,6 +59,7 @@ class Cluster:
         self._sync_interval = parse_timedelta(
             scheduler_sync_interval, default="seconds"
         )
+        self._sync_cluster_info_task = None
 
         if name is None:
             name = str(uuid.uuid4())[:8]
@@ -94,23 +94,56 @@ class Cluster:
         )
         self._cluster_info.update(info)
 
-        self.periodic_callbacks["sync-cluster-info"] = PeriodicCallback(
-            self._sync_cluster_info, self._sync_interval * 1000
-        )
+        # Start a background task for syncing cluster info with the scheduler
+        self._sync_cluster_info_task = asyncio.ensure_future(self._sync_cluster_info())
+
         for pc in self.periodic_callbacks.values():
             pc.start()
         self.status = Status.running
 
     async def _sync_cluster_info(self):
-        await self.scheduler_comm.set_metadata(
-            keys=["cluster-manager-info"],
-            value=copy.copy(self._cluster_info),
-        )
+        err_count = 0
+        warn_at = 5
+        max_interval = 10 * self._sync_interval
+        # Loop until the cluster is shutting down. We shouldn't really need
+        # this check (the `CancelledError` should be enough), but something
+        # deep in the comms code is silencing `CancelledError`s _some_ of the
+        # time, resulting in a cancellation not always bubbling back up to
+        # here. Relying on the status is fine though, not worth changing.
+        while self.status == Status.running:
+            try:
+                await self.scheduler_comm.set_metadata(
+                    keys=["cluster-manager-info"],
+                    value=self._cluster_info.copy(),
+                )
+                err_count = 0
+            except asyncio.CancelledError:
+                # Task is being closed. When we drop Python < 3.8 we can drop
+                # this check (since CancelledError is not a subclass of
+                # Exception then).
+                break
+            except Exception:
+                err_count += 1
+                # Only warn if multiple subsequent attempts fail, and only once
+                # per set of subsequent failed attempts. This way we're not
+                # excessively noisy during a connection blip, but we also don't
+                # silently fail.
+                if err_count == warn_at:
+                    logger.warning(
+                        "Failed to sync cluster info multiple times - perhaps "
+                        "there's a connection issue? Error:",
+                        exc_info=True,
+                    )
+            # Sleep, with error backoff
+            interval = min(max_interval, self._sync_interval * 1.5 ** err_count)
+            await asyncio.sleep(interval)
 
     async def _close(self):
         if self.status == Status.closed:
             return
 
+        self.status = Status.closing
+
         with suppress(AttributeError):
             self._adaptive.stop()
 
@@ -119,6 +152,11 @@ class Cluster:
         if self._watch_worker_status_task:
             await self._watch_worker_status_task
 
+        if self._sync_cluster_info_task:
+            self._sync_cluster_info_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._sync_cluster_info_task
+
         if self.scheduler_comm:
             await self.scheduler_comm.close_rpc()
 
diff --git a/distributed/protocol/cuda.py b/distributed/protocol/cuda.py
index 572b6354..e21ea5f7 100644
--- a/distributed/protocol/cuda.py
+++ b/distributed/protocol/cuda.py
@@ -15,17 +15,20 @@ def cuda_dumps(x):
     except TypeError:
         raise NotImplementedError(type_name)
 
-    header, frames = dumps(x)
-    header["type-serialized"] = pickle.dumps(type(x))
-    header["serializer"] = "cuda"
-    header["compression"] = (False,) * len(frames)  # no compression for gpu data
+    sub_header, frames = dumps(x)
+    header = {
+        "sub-header": sub_header,
+        "type-serialized": pickle.dumps(type(x)),
+        "serializer": "cuda",
+        "compression": (False,) * len(frames),  # no compression for gpu data
+    }
     return header, frames
 
 
 def cuda_loads(header, frames):
     typ = pickle.loads(header["type-serialized"])
     loads = cuda_deserialize.dispatch(typ)
-    return loads(header, frames)
+    return loads(header["sub-header"], frames)
 
 
 register_serialization_family("cuda", cuda_dumps, cuda_loads)
diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py
index 2dbac0a4..1a957f90 100644
--- a/distributed/protocol/serialize.py
+++ b/distributed/protocol/serialize.py
@@ -32,20 +32,23 @@ def dask_dumps(x, context=None):
     except TypeError:
         raise NotImplementedError(type_name)
     if has_keyword(dumps, "context"):
-        header, frames = dumps(x, context=context)
+        sub_header, frames = dumps(x, context=context)
     else:
-        header, frames = dumps(x)
+        sub_header, frames = dumps(x)
 
-    header["type"] = type_name
-    header["type-serialized"] = pickle.dumps(type(x), protocol=4)
-    header["serializer"] = "dask"
+    header = {
+        "sub-header": sub_header,
+        "type": type_name,
+        "type-serialized": pickle.dumps(type(x), protocol=4),
+        "serializer": "dask",
+    }
     return header, frames
 
 
 def dask_loads(header, frames):
     typ = pickle.loads(header["type-serialized"])
     loads = dask_deserialize.dispatch(typ)
-    return loads(header, frames)
+    return loads(header["sub-header"], frames)
 
 
 def pickle_dumps(x, context=None):
@@ -775,6 +778,7 @@ class ObjectDictSerializer:
             else:
                 if isinstance(v, dict):
                     h, f = self.serialize(v)
+                    h = {"nested-dict": h}
                 else:
                     h, f = serialize(v, serializers=(self.serializer, "pickle"))
                 header["complex"][k] = {
@@ -796,7 +800,11 @@ class ObjectDictSerializer:
         for k, d in header["complex"].items():
             h = d["header"]
             f = frames[d["start"] : d["stop"]]
-            v = deserialize(h, f)
+            nested_dict = h.get("nested-dict")
+            if nested_dict:
+                v = self.deserialize(nested_dict, f)
+            else:
+                v = deserialize(h, f)
             dd[k] = v
 
         return obj
diff --git a/distributed/protocol/torch.py b/distributed/protocol/torch.py
index f8b6acb1..f6315f69 100644
--- a/distributed/protocol/torch.py
+++ b/distributed/protocol/torch.py
@@ -1,7 +1,12 @@
-import numpy as np
 import torch
 
-from .serialize import dask_deserialize, dask_serialize, register_generic, serialize
+from .serialize import (
+    dask_deserialize,
+    dask_serialize,
+    deserialize,
+    register_generic,
+    serialize,
+)
 
 
 @dask_serialize.register(torch.Tensor)
@@ -9,10 +14,11 @@ def serialize_torch_Tensor(t):
     requires_grad_ = t.requires_grad
 
     if requires_grad_:
-        header, frames = serialize(t.detach().numpy())
+        sub_header, frames = serialize(t.detach().numpy())
     else:
-        header, frames = serialize(t.numpy())
+        sub_header, frames = serialize(t.numpy())
 
+    header = {"sub-header": sub_header}
     if t.grad is not None:
         grad_header, grad_frames = serialize(t.grad.numpy())
         header["grad"] = {"header": grad_header, "start": len(frames)}
@@ -27,13 +33,11 @@ def deserialize_torch_Tensor(header, frames):
     if header.get("grad", False):
         i = header["grad"]["start"]
         frames, grad_frames = frames[:i], frames[i:]
-        grad = dask_deserialize.dispatch(np.ndarray)(
-            header["grad"]["header"], grad_frames
-        )
+        grad = deserialize(header["grad"]["header"], grad_frames)
     else:
         grad = None
 
-    x = dask_deserialize.dispatch(np.ndarray)(header, frames)
+    x = deserialize(header["sub-header"], frames)
     if header["device"] == "cpu":
         t = torch.from_numpy(x)
         if header["requires_grad"]:
@@ -49,14 +53,15 @@ def deserialize_torch_Tensor(header, frames):
 
 @dask_serialize.register(torch.nn.Parameter)
 def serialize_torch_Parameters(p):
-    header, frames = serialize(p.detach())
+    sub_header, frames = serialize(p.detach())
+    header = {"sub-header": sub_header}
     header["requires_grad"] = p.requires_grad
     return header, frames
 
 
 @dask_deserialize.register(torch.nn.Parameter)
 def deserialize_torch_Parameters(header, frames):
-    t = dask_deserialize.dispatch(torch.Tensor)(header, frames)
+    t = deserialize(header["sub-header"], frames)
     return torch.nn.Parameter(data=t, requires_grad=header["requires_grad"])
 
 
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 185ecacc..bf384f34 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -7078,17 +7078,12 @@ class Scheduler(SchedulerState, ServerNode):
 
     def set_metadata(self, comm=None, keys=None, value=None):
         parent: SchedulerState = cast(SchedulerState, self)
-        try:
-            metadata = parent._task_metadata
-            for key in keys[:-1]:
-                if key not in metadata or not isinstance(metadata[key], (dict, list)):
-                    metadata[key] = {}
-                metadata = metadata[key]
-            metadata[keys[-1]] = value
-        except Exception:
-            import pdb
-
-            pdb.set_trace()
+        metadata = parent._task_metadata
+        for key in keys[:-1]:
+            if key not in metadata or not isinstance(metadata[key], (dict, list)):
+                metadata[key] = {}
+            metadata = metadata[key]
+        metadata[keys[-1]] = value
 
     def get_metadata(self, comm=None, keys=None, default=no_default):
         parent: SchedulerState = cast(SchedulerState, self)


Test patch:
diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py
index a8876e57..22c453b7 100644
--- a/distributed/deploy/tests/test_local.py
+++ b/distributed/deploy/tests/test_local.py
@@ -1103,6 +1103,42 @@ async def test_cluster_info_sync():
         assert info["foo"] == "bar"
 
 
+@pytest.mark.asyncio
+async def test_cluster_info_sync_is_robust_to_network_blips(monkeypatch):
+    async with LocalCluster(
+        processes=False, asynchronous=True, scheduler_sync_interval="1ms"
+    ) as cluster:
+        assert cluster._cluster_info["name"] == cluster.name
+
+        error_called = False
+
+        async def error(*args, **kwargs):
+            nonlocal error_called
+            await asyncio.sleep(0.001)
+            error_called = True
+            raise OSError
+
+        # Temporarily patch the `set_metadata` RPC to error
+        with monkeypatch.context() as patch:
+            patch.setattr(cluster.scheduler_comm, "set_metadata", error)
+
+            # Set a new cluster_info value
+            cluster._cluster_info["foo"] = "bar"
+
+            # Wait for the bad method to be called at least once
+            while not error_called:
+                await asyncio.sleep(0.01)
+
+        # Check that cluster_info is resynced after the error condition is fixed
+        while "foo" not in cluster.scheduler.get_metadata(
+            keys=["cluster-manager-info"]
+        ):
+            await asyncio.sleep(0.01)
+
+        info = cluster.scheduler.get_metadata(keys=["cluster-manager-info"])
+        assert info["foo"] == "bar"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("host", [None, "127.0.0.1"])
 @pytest.mark.parametrize("use_nanny", [True, False])
diff --git a/distributed/protocol/tests/test_protocol.py b/distributed/protocol/tests/test_protocol.py
index 411fcc0c..5c6b67f7 100644
--- a/distributed/protocol/tests/test_protocol.py
+++ b/distributed/protocol/tests/test_protocol.py
@@ -2,7 +2,15 @@ import pytest
 
 from distributed.protocol import dumps, loads, maybe_compress, msgpack, to_serialize
 from distributed.protocol.compression import compressions
-from distributed.protocol.serialize import Serialize, Serialized, deserialize, serialize
+from distributed.protocol.cuda import cuda_deserialize, cuda_serialize
+from distributed.protocol.serialize import (
+    Serialize,
+    Serialized,
+    dask_deserialize,
+    dask_serialize,
+    deserialize,
+    serialize,
+)
 from distributed.system import MEMORY_LIMIT
 from distributed.utils import nbytes
 
@@ -218,3 +226,30 @@ def test_maybe_compress_memoryviews():
     else:
         assert compression == "blosc"
         assert len(payload) < x.nbytes / 10
+
+
+@pytest.mark.parametrize("serializers", [("dask",), ("cuda",)])
+def test_preserve_header(serializers):
+    """
+    Test that a serialization family doesn't overwrite the headers
+    of the underlying registered dumps/loads functions.
+    """
+
+    class MyObj:
+        pass
+
+    @cuda_serialize.register(MyObj)
+    @dask_serialize.register(MyObj)
+    def _(x):
+        return {}, []
+
+    @cuda_deserialize.register(MyObj)
+    @dask_deserialize.register(MyObj)
+    def _(header, frames):
+        assert header == {}
+        assert frames == []
+        return MyObj()
+
+    header, frames = serialize(MyObj(), serializers=serializers)
+    o = deserialize(header, frames)
+    assert isinstance(o, MyObj)
diff --git a/distributed/protocol/tests/test_rmm.py b/distributed/protocol/tests/test_rmm.py
index 528b9460..1e9b7e0e 100644
--- a/distributed/protocol/tests/test_rmm.py
+++ b/distributed/protocol/tests/test_rmm.py
@@ -24,10 +24,10 @@ def test_serialize_rmm_device_buffer(size, serializers):
     y_np = y.copy_to_host()
 
     if serializers[0] == "cuda":
-        assert header["strides"] == (1,)
+        assert header["sub-header"]["strides"] == (1,)
         assert all(hasattr(f, "__cuda_array_interface__") for f in frames)
     elif serializers[0] == "dask":
-        assert header["strides"] == (1,)
+        assert header["sub-header"]["strides"] == (1,)
         assert all(isinstance(f, memoryview) for f in frames)
 
     assert (x_np == y_np).all()
diff --git a/distributed/protocol/tests/test_serialize.py b/distributed/protocol/tests/test_serialize.py
index 74b5d2d2..f6bf653b 100644
--- a/distributed/protocol/tests/test_serialize.py
+++ b/distributed/protocol/tests/test_serialize.py
@@ -430,18 +430,45 @@ async def test_profile_nested_sizeof():
     frames = await to_frames(msg)
 
 
-def test_compression_numpy_list():
-    class MyObj:
+def test_different_compression_families():
+    """Test serialization of a collection of items that use different compression
+
+    This scenario happens for instance when serializing collections of
+    cupy and numpy arrays.
+    """
+
+    class MyObjWithCompression:
         pass
 
-    @dask_serialize.register(MyObj)
-    def _(x):
+    class MyObjWithNoCompression:
+        pass
+
+    def my_dumps_compression(obj, context=None):
+        if not isinstance(obj, MyObjWithCompression):
+            raise NotImplementedError()
+        header = {"compression": [True]}
+        return header, [bytes(2 ** 20)]
+
+    def my_dumps_no_compression(obj, context=None):
+        if not isinstance(obj, MyObjWithNoCompression):
+            raise NotImplementedError()
+
         header = {"compression": [False]}
-        frames = [b""]
-        return header, frames
+        return header, [bytes(2 ** 20)]
+
+    def my_loads(header, frames):
+        return pickle.loads(frames[0])
+
+    register_serialization_family("with-compression", my_dumps_compression, my_loads)
+    register_serialization_family("no-compression", my_dumps_no_compression, my_loads)
 
-    header, frames = serialize([MyObj(), MyObj()])
-    assert header["compression"] == [False, False]
+    header, _ = serialize(
+        [MyObjWithCompression(), MyObjWithNoCompression()],
+        serializers=("with-compression", "no-compression"),
+        on_error="raise",
+        iterate_collection=True,
+    )
+    assert header["compression"] == [True, False]
 
 
 @gen_test()