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:
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

    27 files  ± 0      27 suites  ±0   9h 37m 42s :stopwatch: - 22m 5s
 3 952 tests + 1   3 840 :white_check_mark:  -  2    109 :zzz: ±0  2 :x: +2  1 :fire: +1 
49 710 runs  +14  47 419 :white_check_mark: +12  2 288 :zzz:  - 1  2 :x: +2  1 :fire: +1 

For more details on these failures and errors, see [this check](https://github.com/dask/distributed/runs/20348680595).

Results for commit d88c407b. ± Comparison against base commit 5c481ddc.

[test-results]:data:application/gzip;base64,H4sIAEK1nmUC/02Myw6DIBBFf8Ww7gLGUYb+TAOICamPBmFl+u+FtqLLc+7N2dnoJ7exewPy1rAt+XjSkIKOfl0yt9j3kFVeY9lb1VV8bMna4gj56Z7+lZ3gqqpR+6mkq3AhrKGcsglpKV1UUvA/Hl2UKFR1vy4A0aEu3S9fs3adZx8zsIHIIpeG0NnOaiFAG2lAg+FKtTTwsVPgjGbvD0Oj2R4UAQAA


Patch:
diff --git a/distributed/protocol/core.py b/distributed/protocol/core.py
index e698335a..fbeaabc3 100644
--- a/distributed/protocol/core.py
+++ b/distributed/protocol/core.py
@@ -11,6 +11,7 @@ from distributed.protocol.serialize import (
     Serialize,
     Serialized,
     ToPickle,
+    _is_msgpack_serializable,
     merge_and_deserialize,
     msgpack_decode_default,
     msgpack_encode_default,
@@ -104,7 +105,31 @@ def dumps(  # type: ignore[no-untyped-def]
             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
index 0b765806..630239c1 100644
--- a/distributed/protocol/tests/test_protocol.py
+++ b/distributed/protocol/tests/test_protocol.py
@@ -20,6 +20,7 @@ from distributed.protocol.serialize import (
     serialize,
 )
 from distributed.utils import nbytes
+from distributed.utils_test import captured_logger
 
 
 def test_protocol():
@@ -174,3 +175,36 @@ def test_deeply_nested_structures():
 
     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"}