Repository: dask/distributed

Issue:
`distributed.protocol.dumps` does not handle objects not serializable via `msgpack.dumps` as suggested by docstring
When using `distributed.protocol.dumps` to serialize arbitrary data, it fails and raises a `TypeError`. The docstring suggests it should be able to handle this.

**Minimal Complete Verifiable Example**:

```python
import numpy as np
from distributed.protocol import dumps

dumps(np.int64(123))
```

raises

```python
2024-01-09 19:55:04,572 - distributed.protocol.core - CRITICAL - Failed to Serialize
Traceback (most recent call last):
  File "/Users/hendrikmakait/projects/dask/distributed/distributed/protocol/core.py", line 108, in dumps
    frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/dask-distributed/lib/python3.12/site-packages/msgpack/__init__.py", line 36, in packb
    return Packer(**kwargs).pack(o)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "msgpack/_packer.pyx", line 294, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 300, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 297, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 291, in msgpack._cmsgpack.Packer._pack
TypeError: can not serialize 'numpy.int64' object
Traceback (most recent call last):
  File "/Users/hendrikmakait/projects/dask/distributed/sandbox.py", line 5, in <module>
    dumps(np.int64(123))
  File "/Users/hendrikmakait/projects/dask/distributed/distributed/protocol/core.py", line 108, in dumps
    frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/dask-distributed/lib/python3.12/site-packages/msgpack/__init__.py", line 36, in packb
    return Packer(**kwargs).pack(o)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "msgpack/_packer.pyx", line 294, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 300, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 297, in msgpack._cmsgpack.Packer.pack
  File "msgpack/_packer.pyx", line 291, in msgpack._cmsgpack.Packer._pack
TypeError: can not serialize 'numpy.int64' object
```

Relevant section of `dumps`'s docstring: https://github.com/dask/distributed/blob/7562f9c566978de4f3f5b73920a24ea1813d6e28/distributed/protocol/core.py#L30-L35


Discussion:

Resolving patch:
diff --git a/distributed/protocol/core.py b/distributed/protocol/core.py
--- a/distributed/protocol/core.py
+++ b/distributed/protocol/core.py
@@ -11,6 +11,7 @@
     Serialize,
     Serialized,
     ToPickle,
+    _is_msgpack_serializable,
     merge_and_deserialize,
     msgpack_decode_default,
     msgpack_encode_default,
@@ -104,7 +105,31 @@ def _encode_default(obj):
             else:
                 return msgpack_encode_default(obj)
 
-        frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
+        try:
+            frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
+        except TypeError as e:
+            logger.info(
+                f"Failed to serialize ({e}); falling back to pickle. "
+                "Be aware that this may degrade performance."
+            )
+
+            def _encode_default_safe(obj):
+                encoded = _encode_default(obj)
+                if encoded is not obj or _is_msgpack_serializable(obj):
+                    return encoded
+
+                obj = ToPickle(obj)
+                offset = len(frames)
+                frames.extend(create_pickled_sub_frames(obj))
+                return {"__Pickled__": offset}
+
+            # If possible, we want to avoid the performance penalty from the checks
+            # implemented in _encode_default_safe to fall back to pickle, so we
+            # try to serialize the data without the fallback first assuming that
+            # this succeeds in the overwhelming majority of cases.
+            frames[0] = msgpack.dumps(
+                msg, default=_encode_default_safe, use_bin_type=True
+            )
         return frames
 
     except Exception:


Test patch:
diff --git a/distributed/protocol/tests/test_protocol.py b/distributed/protocol/tests/test_protocol.py
--- a/distributed/protocol/tests/test_protocol.py
+++ b/distributed/protocol/tests/test_protocol.py
@@ -20,6 +20,7 @@
     serialize,
 )
 from distributed.utils import nbytes
+from distributed.utils_test import captured_logger
 
 
 def test_protocol():
@@ -174,3 +175,36 @@ def gen_deeply_nested(depth):
 
     msg = gen_deeply_nested(sys.getrecursionlimit() // 4)
     assert isinstance(serialize(msg), tuple)
+
+
+def test_fallback_to_pickle():
+    np = pytest.importorskip("numpy")
+
+    d = 1
+    L = dumps(d)
+    assert b"__Pickled__" not in L[0]
+    assert loads(L) == d
+
+    d = np.int64(1)
+    with captured_logger("distributed.protocol.core") as logger:
+        L = dumps(d)
+    assert "can not serialize 'numpy.int64'" in logger.getvalue()
+    assert L[0].count(b"__Pickled__") == 1
+    assert loads(L) == d
+
+    d = {np.int64(1): {np.int64(2): "a"}, 3: ("b", "c"), 4: "d"}
+    with captured_logger("distributed.protocol.core") as logger:
+        L = dumps(d)
+    assert "can not serialize 'numpy.int64'" in logger.getvalue()
+    # Make sure that we pickle the individual ints, not the entire dict
+    assert L[0].count(b"__Pickled__") == 2
+    assert loads(L) == d
+
+    d = {np.int64(1): {Serialize(2): "a"}, 3: ("b", "c"), 4: "d"}
+    with captured_logger("distributed.protocol.core") as logger:
+        L = dumps(d)
+    assert "can not serialize 'numpy.int64'" in logger.getvalue()
+    # Make sure that we still serialize and don't pickle indiscriminately
+    assert L[0].count(b"__Pickled__") == 1
+    assert L[0].count(b"__Serialized__") == 1
+    assert loads(L) == {np.int64(1): {2: "a"}, 3: ("b", "c"), 4: "d"}