Repository: dask/distributed

Issue:
`tblib=3.0.0` breaks `test_fail_to_pickle_execute_1`
Upgrading `tblib` from `2.0.0` to `3.0.0` breaks `test_fail_to_pickle_execute_1`

Full traceback (https://github.com/dask/distributed/actions/runs/7177526736/job/19544154681#step:19:2575):

```python3
________________________ test_fail_to_pickle_execute_1 _________________________

c = <Client: No scheduler connected>
s = <Scheduler 'tcp://127.0.0.1:44071', workers: 0, cores: 0, tasks: 0>
a = <Worker 'tcp://127.0.0.1:46865', name: 0, status: closed, stored: 0, running: 0/1, ready: 0, comm: 0, waiting: 0>
b = <Worker 'tcp://127.0.0.1:35527', name: 1, status: closed, stored: 0, running: 0/2, ready: 0, comm: 0, waiting: 0>

    @gen_cluster(client=True)
    async def test_fail_to_pickle_execute_1(c, s, a, b):
        """Test failure to serialize triggered by computing a key which is individually
        larger than target. The data is lost and the task is marked as failed; the worker
        remains in usable condition.
    
        See also
        --------
        test_workerstate_fail_to_pickle_execute_1
        test_workerstate_fail_to_pickle_flight
        test_fail_to_pickle_execute_2
        test_fail_to_pickle_spill
        """
        x = c.submit(FailToPickle, reported_size=100e9, key="x")
        await wait(x)
    
        assert x.status == "error"
    
        with pytest.raises(TypeError, match="Could not serialize"):
>           await x

distributed/tests/test_worker_memory.py:178: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
distributed/client.py:330: in _result
    raise exc.with_traceback(tb)
distributed/spill.py:181: in __setitem__
    with self._capture_metrics(), self._handle_errors(key):
../../../miniconda3/envs/dask-distributed/lib/python3.11/contextlib.py:155: in __exit__
    self.gen.throw(typ, value, traceback)
distributed/spill.py:149: in _handle_errors
    raise orig_e
distributed/spill.py:322: in __setitem__
    pickled = self.dump(value)
distributed/protocol/serialize.py:671: in serialize_bytelist
    header, frames = serialize_and_split(x, **kwargs)
../../../miniconda3/envs/dask-distributed/lib/python3.11/contextlib.py:81: in inner
    return func(*args, **kwds)
distributed/protocol/serialize.py:458: in serialize_and_split
    header, frames = serialize(x, serializers, on_error, context)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

    from __future__ import annotations
    
    import codecs
    import importlib
    import traceback
    from array import array
    from enum import Enum
    from functools import partial
    from types import ModuleType
    from typing import Any, Generic, Literal, TypeVar
    
    import msgpack
    
    import dask
    from dask.base import normalize_token
    from dask.sizeof import sizeof
    from dask.utils import typename
    
    from distributed.metrics import context_meter
    from distributed.protocol import pickle
    from distributed.protocol.compression import decompress, maybe_compress
    from distributed.protocol.utils import (
        frame_split_size,
        host_array_from_buffers,
        merge_memoryviews,
        msgpack_opts,
        pack_frames_prelude,
        unpack_frames,
    )
    from distributed.utils import ensure_memoryview, has_keyword
    
    T = TypeVar("T")
    
    dask_serialize = dask.utils.Dispatch("dask_serialize")
    dask_deserialize = dask.utils.Dispatch("dask_deserialize")
    
    _cached_allowed_modules: dict[str, ModuleType] = {}
    
    
    def dask_dumps(x, context=None):
        """Serialize object using the class-based registry"""
        type_name = typename(type(x))
        try:
            dumps = dask_serialize.dispatch(type(x))
        except TypeError:
            raise NotImplementedError(type_name)
        if has_keyword(dumps, "context"):
            sub_header, frames = dumps(x, context=context)
        else:
            sub_header, frames = dumps(x)
    
        header = {
            "sub-header": sub_header,
            "type": type_name,
            "type-serialized": pickle.dumps(type(x)),
            "serializer": "dask",
        }
        return header, frames
    
    
    def dask_loads(header, frames):
        typ = pickle.loads(header["type-serialized"])
        loads = dask_deserialize.dispatch(typ)
        return loads(header["sub-header"], frames)
    
    
    def pickle_dumps(x, context=None):
        frames = [None]
        writeable = []
    
        def buffer_callback(f):
            f = memoryview(f)
            frames.append(f)
            writeable.append(not f.readonly)
    
        frames[0] = pickle.dumps(
            x,
            buffer_callback=buffer_callback,
            protocol=context.get("pickle-protocol", None) if context else None,
        )
        header = {
            "serializer": "pickle",
            "writeable": tuple(writeable),
        }
    
        return header, frames
    
    
    def pickle_loads(header, frames):
        x, buffers = frames[0], frames[1:]
    
        writeable = header.get("writeable")
        if not writeable:
            writeable = len(buffers) * (None,)
    
        buffers = [
            memoryview(bytearray(mv) if w else bytes(mv)) if w == mv.readonly else mv
            for w, mv in zip(writeable, map(ensure_memoryview, buffers))
        ]
    
        return pickle.loads(x, buffers=buffers)
    
    
    def import_allowed_module(name):
        if name in _cached_allowed_modules:
            return _cached_allowed_modules[name]
    
        # Check for non-ASCII characters
        name = name.encode("ascii").decode()
        # We only compare the root module
        root = name.split(".", 1)[0]
    
        # Note, if an empty string creeps into allowed-imports it is disallowed explicitly
        if root and root in dask.config.get("distributed.scheduler.allowed-imports"):
            _cached_allowed_modules[name] = importlib.import_module(name)
            return _cached_allowed_modules[name]
        else:
            raise RuntimeError(
                f"Importing {repr(name)} is not allowed, please add it to the list of "
                "allowed modules the scheduler can import via the "
                "distributed.scheduler.allowed-imports configuration setting."
            )
    
    
    def msgpack_decode_default(obj):
        """
        Custom packer/unpacker for msgpack
        """
        if "__Enum__" in obj:
            mod = import_allowed_module(obj["__module__"])
            typ = getattr(mod, obj["__name__"])
            return getattr(typ, obj["name"])
    
        if "__Set__" in obj:
            return set(obj["as-list"])
    
        if "__Serialized__" in obj:
            # Notice, the data here is marked a Serialized rather than deserialized. This
            # is because deserialization requires Pickle which the Scheduler cannot run
            # because of security reasons.
            # By marking it Serialized, the data is passed through to the workers that
            # eventually will deserialize it.
            return Serialized(*obj["data"])
    
        return obj
    
    
    def msgpack_encode_default(obj):
        """
        Custom packer/unpacker for msgpack
        """
    
        if isinstance(obj, Serialize):
            return {"__Serialized__": True, "data": serialize(obj.data)}
    
        if isinstance(obj, Enum):
            return {
                "__Enum__": True,
                "name": obj.name,
                "__module__": obj.__module__,
                "__name__": type(obj).__name__,
            }
    
        if isinstance(obj, set):
            return {"__Set__": True, "as-list": list(obj)}
    
        return obj
    
    
    def msgpack_dumps(x):
        try:
            frame = msgpack.dumps(x, use_bin_type=True)
        except Exception:
            raise NotImplementedError()
        else:
            return {"serializer": "msgpack"}, [frame]
    
    
    def msgpack_loads(header, frames):
        return msgpack.loads(b"".join(frames), use_list=False, **msgpack_opts)
    
    
    def serialization_error_loads(header, frames):
        msg = "\n".join([codecs.decode(frame, "utf8") for frame in frames])
        raise TypeError(msg)
    
    
    families = {}
    
    
    def register_serialization_family(name, dumps, loads):
        families[name] = (dumps, loads, dumps and has_keyword(dumps, "context"))
    
    
    register_serialization_family("dask", dask_dumps, dask_loads)
    register_serialization_family("pickle", pickle_dumps, pickle_loads)
    register_serialization_family("msgpack", msgpack_dumps, msgpack_loads)
    register_serialization_family("error", None, serialization_error_loads)
    
    
    def check_dask_serializable(x):
        if type(x) in (list, set, tuple) and len(x):
            return check_dask_serializable(next(iter(x)))
        elif type(x) is dict and len(x):
            return check_dask_serializable(next(iter(x.items()))[1])
        else:
            try:
                dask_serialize.dispatch(type(x))
                return True
            except TypeError:
                pass
        return False
    
    
    def serialize(  # type: ignore[no-untyped-def]
        x: object,
        serializers=None,
        on_error: Literal["message" | "raise"] = "message",
        context=None,
        iterate_collection: bool | None = None,
    ) -> tuple[dict[str, Any], list[bytes | memoryview]]:
        r"""
        Convert object to a header and list of bytestrings
    
        This takes in an arbitrary Python object and returns a msgpack serializable
        header and a list of bytes or memoryview objects.
    
        The serialization protocols to use are configurable: a list of names
        define the set of serializers to use, in order. These names are keys in
        the ``serializer_registry`` dict (e.g., 'pickle', 'msgpack'), which maps
        to the de/serialize functions. The name 'dask' is special, and will use the
        per-class serialization methods. ``None`` gives the default list
        ``['dask', 'pickle']``.
    
        Notes on the ``iterate_collection`` argument (only relevant when
        ``x`` is a collection):
        - ``iterate_collection=True``: Serialize collection elements separately.
        - ``iterate_collection=False``: Serialize collection elements together.
        - ``iterate_collection=None`` (default): Infer the best setting.
    
        Examples
        --------
        >>> serialize(1)
        ({}, [b'\x80\x04\x95\x03\x00\x00\x00\x00\x00\x00\x00K\x01.'])
    
        >>> serialize(b'123')  # some special types get custom treatment
        ({'type': 'builtins.bytes'}, [b'123'])
    
        >>> deserialize(*serialize(1))
        1
    
        Returns
        -------
        header: dictionary containing any msgpack-serializable metadata
        frames: list of bytes or memoryviews, commonly of length one
    
        See Also
        --------
        deserialize : Convert header and frames back to object
        to_serialize : Mark that data in a message should be serialized
        register_serialization : Register custom serialization functions
        """
        if serializers is None:
            serializers = ("dask", "pickle")  # TODO: get from configuration
    
        # Handle objects that are marked as `Serialize`, or that are
        # already `Serialized` objects (don't want to serialize them twice)
        if isinstance(x, Serialized):
            return x.header, x.frames
        if isinstance(x, Serialize):
            return serialize(
                x.data,
                serializers=serializers,
                on_error=on_error,
                context=context,
                iterate_collection=True,
            )
    
        # Note: don't use isinstance(), as it would match subclasses
        # (e.g. namedtuple, defaultdict) which however would revert to the base class on a
        # round-trip through msgpack
        if iterate_collection is None and type(x) in (list, set, tuple, dict):
            if type(x) is list and "msgpack" in serializers:
                # Note: "msgpack" will always convert lists to tuples
                #       (see GitHub #3716), so we should iterate
                #       through the list if "msgpack" comes before "pickle"
                #       in the list of serializers.
                iterate_collection = ("pickle" not in serializers) or (
                    serializers.index("pickle") > serializers.index("msgpack")
                )
            if not iterate_collection:
                # Check for "dask"-serializable data in dict/list/set
                iterate_collection = check_dask_serializable(x)
    
        # Determine whether keys are safe to be serialized with msgpack
        if type(x) is dict and iterate_collection:
            try:
                msgpack.dumps(list(x.keys()))
            except Exception:
                dict_safe = False
            else:
                dict_safe = True
    
        if (
            type(x) in (list, set, tuple)
            and iterate_collection
            or type(x) is dict
            and iterate_collection
            and dict_safe
        ):
            if isinstance(x, dict):
                headers_frames = []
                for k, v in x.items():
                    _header, _frames = serialize(
                        v, serializers=serializers, on_error=on_error, context=context
                    )
                    _header["key"] = k
                    headers_frames.append((_header, _frames))
            else:
                assert isinstance(x, (list, set, tuple))
                headers_frames = [
                    serialize(
                        obj, serializers=serializers, on_error=on_error, context=context
                    )
                    for obj in x
                ]
    
            frames = []
            lengths = []
            compressions: list[str | None] = []
            for _header, _frames in headers_frames:
                frames.extend(_frames)
                length = len(_frames)
                lengths.append(length)
                compressions.extend(_header.get("compression") or [None] * len(_frames))
    
            headers = {
                "sub-headers": [obj[0] for obj in headers_frames],
                "is-collection": True,
                "frame-lengths": lengths,
                "type-serialized": type(x).__name__,
            }
            if any(compression is not None for compression in compressions):
                headers["compression"] = compressions
            return headers, frames
    
        tb = ""
        exc = None
    
        for name in serializers:
            dumps, _, wants_context = families[name]
            try:
                header, frames = dumps(x, context=context) if wants_context else dumps(x)
                header["serializer"] = name
                return header, frames
            except NotImplementedError:
                continue
            except Exception as e:
                exc = e
                tb = traceback.format_exc()
                break
        type_x = type(x)
        if isinstance(x, (ToPickle, Serialize)):
            type_x = type(x.data)
        msg = f"Could not serialize object of type {type_x.__name__}"
        if on_error == "message":
            txt_frames = [msg]
            if tb:
                txt_frames.append(tb[:100000])
    
            frames = [frame.encode() for frame in txt_frames]
    
            return {"serializer": "error"}, frames
        elif on_error == "raise":
            try:
                str_x = str(x)[:10000]
            except Exception:
                raise TypeError(msg) from exc
>           raise TypeError(msg, str_x) from exc
E           Exception: TypeError('Could not serialize object of type FailToPickle', '<test_worker_memory.FailToPickle object at 0x7f75935d5f50>')

distributed/protocol/serialize.py:379: Exception
----------------------------- Captured stdout call -----------------------------
Dumped cluster state to test_cluster_dump/test_fail_to_pickle_execute_1.yaml
----------------------------- Captured stderr call -----------------------------
2023-12-12 06:27:59,142 - distributed.scheduler - INFO - State start
2023-12-12 06:27:59,146 - distributed.scheduler - INFO -   Scheduler at:     tcp://127.0.0.1:44071
2023-12-12 06:27:59,147 - distributed.scheduler - INFO -   dashboard at:  http://127.0.0.1:40863/status
2023-12-12 06:27:59,148 - distributed.scheduler - INFO - Registering Worker plugin shuffle
2023-12-12 06:27:59,155 - distributed.worker - INFO -       Start worker at:      tcp://127.0.0.1:46865
2023-12-12 06:27:59,155 - distributed.worker - INFO -          Listening to:      tcp://127.0.0.1:46865
2023-12-12 06:27:59,155 - distributed.worker - INFO -           Worker name:                          0
2023-12-12 06:27:59,155 - distributed.worker - INFO -          dashboard at:            127.0.0.1:33601
2023-12-12 06:27:59,156 - distributed.worker - INFO - Waiting to connect to:      tcp://127.0.0.1:44071
2023-12-12 06:27:59,156 - distributed.worker - INFO - -------------------------------------------------
2023-12-12 06:27:59,156 - distributed.worker - INFO -               Threads:                          1
2023-12-12 06:27:59,156 - distributed.worker - INFO -                Memory:                  15.61 GiB
2023-12-12 06:27:59,156 - distributed.worker - INFO -       Local Directory: /tmp/dask-scratch-space/worker-b1h25pds
2023-12-12 06:27:59,156 - distributed.worker - INFO - -------------------------------------------------
2023-12-12 06:27:59,157 - distributed.worker - INFO -       Start worker at:      tcp://127.0.0.1:35527
2023-12-12 06:27:59,157 - distributed.worker - INFO -          Listening to:      tcp://127.0.0.1:35527
2023-12-12 06:27:59,157 - distributed.worker - INFO -           Worker name:                          1
2023-12-12 06:27:59,158 - distributed.worker - INFO -          dashboard at:            127.0.0.1:44531
2023-12-12 06:27:59,158 - distributed.worker - INFO - Waiting to connect to:      tcp://127.0.0.1:44071
2023-12-12 06:27:59,158 - distributed.worker - INFO - -------------------------------------------------
2023-12-12 06:27:59,158 - distributed.worker - INFO -               Threads:                          2
2023-12-12 06:27:59,158 - distributed.worker - INFO -                Memory:                  15.61 GiB
2023-12-12 06:27:59,158 - distributed.worker - INFO -       Local Directory: /tmp/dask-scratch-space/worker-bq3s0idz
2023-12-12 06:27:59,158 - distributed.worker - INFO - -------------------------------------------------
2023-12-12 06:27:59,162 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://127.0.0.1:46865', name: 0, status: init, memory: 0, processing: 0>
2023-12-12 06:27:59,163 - distributed.scheduler - INFO - Starting worker compute stream, tcp://127.0.0.1:46865
2023-12-12 06:27:59,164 - distributed.core - INFO - Starting established connection to tcp://127.0.0.1:49854
2023-12-12 06:27:59,164 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://127.0.0.1:35527', name: 1, status: init, memory: 0, processing: 0>
2023-12-12 06:27:59,165 - distributed.scheduler - INFO - Starting worker compute stream, tcp://127.0.0.1:35527
2023-12-12 06:27:59,165 - distributed.core - INFO - Starting established connection to tcp://127.0.0.1:49866
2023-12-12 06:27:59,166 - distributed.worker - INFO - Starting Worker plugin shuffle
2023-12-12 06:27:59,166 - distributed.worker - INFO - Starting Worker plugin shuffle
2023-12-12 06:27:59,167 - distributed.worker - INFO -         Registered to:      tcp://127.0.0.1:44071
2023-12-12 06:27:59,167 - distributed.worker - INFO - -------------------------------------------------
2023-12-12 06:27:59,167 - distributed.worker - INFO -         Registered to:      tcp://127.0.0.1:44071
2023-12-12 06:27:59,167 - distributed.worker - INFO - -------------------------------------------------
2023-12-12 06:27:59,168 - distributed.core - INFO - Starting established connection to tcp://127.0.0.1:44071
2023-12-12 06:27:59,168 - distributed.core - INFO - Starting established connection to tcp://127.0.0.1:44071
2023-12-12 06:27:59,172 - distributed.scheduler - INFO - Receive client connection: Client-9800c358-98b7-11ee-8bc2-0022482c46a0
2023-12-12 06:27:59,173 - distributed.core - INFO - Starting established connection to tcp://127.0.0.1:49868
2023-12-12 06:27:59,189 - distributed.protocol.pickle - ERROR - Failed to serialize <test_worker_memory.FailToPickle object at 0x7f75935d5f50>.
Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 63, in dumps
    result = pickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/tests/test_worker_memory.py", line 147, in __getstate__
    raise CustomError()
test_worker_memory.CustomError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 68, in dumps
    pickler.dump(x)
  File "/home/runner/work/distributed/distributed/distributed/tests/test_worker_memory.py", line 147, in __getstate__
    raise CustomError()
test_worker_memory.CustomError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 81, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1479, in dumps
    cp.dump(obj)
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1245, in dump
    return super().dump(obj)
           ^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/tests/test_worker_memory.py", line 147, in __getstate__
    raise CustomError()
test_worker_memory.CustomError
2023-12-12 06:27:59,208 - distributed.protocol.pickle - ERROR - Failed to serialize ('Could not serialize object of type FailToPickle', '<test_worker_memory.FailToPickle object at 0x7f75935d5f50>').
Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 63, in dumps
    result = pickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/tests/test_worker_memory.py", line 147, in __getstate__
    raise CustomError()
test_worker_memory.CustomError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 68, in dumps
    pickler.dump(x)
  File "/home/runner/work/distributed/distributed/distributed/tests/test_worker_memory.py", line 147, in __getstate__
    raise CustomError()
test_worker_memory.CustomError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/serialize.py", line 353, in serialize
    header, frames = dumps(x, context=context) if wants_context else dumps(x)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/protocol/serialize.py", line 76, in pickle_dumps
    frames[0] = pickle.dumps(
                ^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 81, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1479, in dumps
    cp.dump(obj)
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1245, in dump
    return super().dump(obj)
           ^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/tests/test_worker_memory.py", line 147, in __getstate__
    raise CustomError()
test_worker_memory.CustomError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/worker_state_machine.py", line 2418, in _transition_to_memory
    self.data[ts.key] = value
    ~~~~~~~~~^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/spill.py", line 181, in __setitem__
    with self._capture_metrics(), self._handle_errors(key):
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/contextlib.py", line 155, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/home/runner/work/distributed/distributed/distributed/spill.py", line 149, in _handle_errors
    raise orig_e
  File "/home/runner/work/distributed/distributed/distributed/spill.py", line 322, in __setitem__
    pickled = self.dump(value)
              ^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/protocol/serialize.py", line 671, in serialize_bytelist
    header, frames = serialize_and_split(x, **kwargs)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/protocol/serialize.py", line 458, in serialize_and_split
    header, frames = serialize(x, serializers, on_error, context)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/protocol/serialize.py", line 379, in serialize
    raise TypeError(msg, str_x) from exc
TypeError: ('Could not serialize object of type FailToPickle', '<test_worker_memory.FailToPickle object at 0x7f75935d5f50>')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 63, in dumps
    result = pickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/tblib/pickling_support.py", line 46, in pickle_exception
    rv = obj.__reduce_ex__(3)
         ^^^^^^^^^^^^^^^^^^^^
RecursionError: maximum recursion depth exceeded while calling a Python object

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 68, in dumps
    pickler.dump(x)
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 24, in reducer_override
    if _always_use_pickle_for(obj):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 45, in _always_use_pickle_for
    return isinstance(x, (str, bytes))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
RecursionError: maximum recursion depth exceeded in __instancecheck__

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1245, in dump
    return super().dump(obj)
           ^^^^^^^^^^^^^^^^^
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1323, in reducer_override
    is_anyclass = issubclass(t, type)
                  ^^^^^^^^^^^^^^^^^^^
RecursionError: maximum recursion depth exceeded while calling a Python object

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/runner/work/distributed/distributed/distributed/protocol/pickle.py", line 81, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1479, in dumps
    cp.dump(obj)
  File "/home/runner/miniconda3/envs/dask-distributed/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1249, in dump
    raise pickle.PicklingError(msg) from e
_pickle.PicklingError: Could not pickle object as excessively deep recursion required.
2023-12-12 06:27:59,349 - distributed.scheduler - INFO - Remove client Client-9800c358-98b7-11ee-8bc2-0022482c46a0
2023-12-12 06:27:59,350 - distributed.core - INFO - Received 'close-stream' from tcp://127.0.0.1:49868; closing.
2023-12-12 06:27:59,351 - distributed.scheduler - INFO - Remove client Client-9800c358-98b7-11ee-8bc2-0022482c46a0
2023-12-12 06:27:59,352 - distributed.scheduler - INFO - Close client connection: Client-9800c358-98b7-11ee-8bc2-0022482c46a0
2023-12-12 06:27:59,353 - distributed.worker - INFO - Stopping worker at tcp://127.0.0.1:46865. Reason: worker-close
2023-12-12 06:27:59,354 - distributed.worker - INFO - Stopping worker at tcp://127.0.0.1:35527. Reason: worker-close
2023-12-12 06:27:59,356 - distributed.core - INFO - Connection to tcp://127.0.0.1:44071 has been closed.
2023-12-12 06:27:59,357 - distributed.core - INFO - Connection to tcp://127.0.0.1:44071 has been closed.
2023-12-12 06:27:59,357 - distributed.core - INFO - Received 'close-stream' from tcp://127.0.0.1:49854; closing.
2023-12-12 06:27:59,357 - distributed.core - INFO - Received 'close-stream' from tcp://127.0.0.1:49866; closing.
2023-12-12 06:27:59,357 - distributed.scheduler - INFO - Remove worker <WorkerState 'tcp://127.0.0.1:46865', name: 0, status: closing, memory: 0, processing: 0> (stimulus_id='handle-worker-cleanup-1702362479.3578002')
2023-12-12 06:27:59,358 - distributed.scheduler - INFO - Remove worker <WorkerState 'tcp://127.0.0.1:35527', name: 1, status: closing, memory: 0, processing: 0> (stimulus_id='handle-worker-cleanup-1702362479.3585954')
2023-12-12 06:27:59,359 - distributed.scheduler - INFO - Lost all workers
2023-12-12 06:27:59,563 - distributed.scheduler - INFO - Scheduler closing due to unknown reason...
2023-12-12 06:27:59,564 - distributed.scheduler - INFO - Scheduler closing all comms
- generated xml file: /home/runner/work/distributed/distributed/reports/pytest.xml -
```


Discussion:
From what I see, `tblib` switched from iteration to recursion in https://github.com/ionelmc/python-tblib/pull/72 (specifically, in https://github.com/ionelmc/python-tblib/commit/166bc55d1485544c259a246b24abc1b463381e6c#diff-fd132c82e8f7c2524f71c27f06d99ba7d8b6cef0dc8256c19cf4006b36ed9751R122).
This causes us to run into a `RecursionError`, which then blows up.
Turns out I was wrong, the problem is caused by `tblib` starting to roundtrip the `__context__` (https://github.com/ionelmc/python-tblib/commit/166bc55d1485544c259a246b24abc1b463381e6c#diff-fd132c82e8f7c2524f71c27f06d99ba7d8b6cef0dc8256c19cf4006b36ed9751R56).

Resolving patch:
diff --git a/distributed/spill.py b/distributed/spill.py
--- a/distributed/spill.py
+++ b/distributed/spill.py
@@ -136,25 +136,24 @@ def _handle_errors(self, key: Key | None) -> Iterator[None]:
             logger.error("Spill to disk failed; keeping data in memory", exc_info=True)
             raise HandledError()
         except PickleError as e:
-            key_e, orig_e = e.args
-            assert key_e in self.fast
-            assert key_e not in self.slow
-            if key_e == key:
+            assert e.key in self.fast
+            assert e.key not in self.slow
+            if e.key == key:
                 assert key is not None
                 # The key we just inserted failed to serialize.
                 # This happens only when the key is individually larger than target.
                 # The exception will be caught by Worker and logged; the status of
                 # the task will be set to error.
                 del self[key]
-                raise orig_e
+                raise
             else:
                 # The key we just inserted is smaller than target, but it caused
                 # another, unrelated key to be spilled out of the LRU, and that key
                 # failed to serialize. There's nothing wrong with the new key. The older
                 # key is still in memory.
-                if key_e not in self.logged_pickle_errors:
-                    logger.error("Failed to pickle %r", key_e, exc_info=True)
-                    self.logged_pickle_errors.add(key_e)
+                if e.key not in self.logged_pickle_errors:
+                    logger.error("Failed to pickle %r", e.key, exc_info=True)
+                    self.logged_pickle_errors.add(e.key)
                 raise HandledError()
 
     def __setitem__(self, key: Key, value: object) -> None:
@@ -267,8 +266,13 @@ class MaxSpillExceeded(Exception):
     pass
 
 
-class PickleError(Exception):
-    pass
+class PickleError(TypeError):
+    def __str__(self) -> str:
+        return f"Failed to pickle {self.key!r}"
+
+    @property
+    def key(self) -> Key:
+        return self.args[0]
 
 
 class HandledError(Exception):
@@ -324,7 +328,7 @@ def __setitem__(self, key: Key, value: object) -> None:
             # zict.LRU ensures that the key remains in fast if we raise.
             # Wrap the exception so that it's recognizable by SpillBuffer,
             # which will then unwrap it.
-            raise PickleError(key, e)
+            raise PickleError(key) from e
 
         # Thanks to Buffer.__setitem__, we never update existing
         # keys in slow, but always delete them and reinsert them.


Test patch:
diff --git a/distributed/tests/test_spill.py b/distributed/tests/test_spill.py
--- a/distributed/tests/test_spill.py
+++ b/distributed/tests/test_spill.py
@@ -219,9 +219,10 @@ def test_spillbuffer_fail_to_serialize(tmp_path):
     a = Bad(size=201)
 
     # Exception caught in the worker
-    with pytest.raises(TypeError, match="Could not serialize"):
+    with pytest.raises(TypeError, match="Failed to pickle 'a'") as e:
         with captured_logger("distributed.spill") as logs_bad_key:
             buf["a"] = a
+    assert isinstance(e.value.__cause__.__cause__, MyError)
 
     # spill.py must remain silent because we're already logging in worker.py
     assert not logs_bad_key.getvalue()
@@ -240,7 +241,7 @@ def test_spillbuffer_fail_to_serialize(tmp_path):
 
     # worker.py won't intercept the exception here, so spill.py must dump the traceback
     logs_value = logs_bad_key_mem.getvalue()
-    assert "Failed to pickle" in logs_value  # from distributed.spill
+    assert "Failed to pickle 'b'" in logs_value  # from distributed.spill
     assert "Traceback" in logs_value  # from distributed.spill
     assert_buf(buf, tmp_path, {"b": b, "c": c}, {})
 
diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
--- a/distributed/tests/test_worker_memory.py
+++ b/distributed/tests/test_worker_memory.py
@@ -174,8 +174,9 @@ async def test_fail_to_pickle_execute_1(c, s, a, b):
 
     assert x.status == "error"
 
-    with pytest.raises(TypeError, match="Could not serialize"):
+    with pytest.raises(TypeError, match="Failed to pickle 'x'") as e:
         await x
+    assert isinstance(e.value.__cause__.__cause__, CustomError)
 
     await assert_basic_futures(c)