Repository: dask/distributed

Issue:
Handling custom serialization with MsgPack directly
Today we using things like `extract_serialize` to pull out objects MsgPack can't handle and serialize alongside it. In benchmarks we have done our extra handling code takes about 3x more time than MsgPack alone. An interesting idea to follow up on would be to see if we can add an `ExtType` or something to default encoding/decoding to handle out-of-band buffers and merely track where to insert them later. This would be analogous to how pickle works with out-of-band buffers. Though it may speed up serializing and deserializing by doing fewer passes over the data by leveraging MsgPack's own passes. In theory we could get up to a 4x speed up in serialization by following this strategy.


Discussion:
cc @mrocklin @quasiben
If there is an easy system within msgpack that we can hijack to avoid our
traversing the object then great, I'm all for it.

Aside: I like all of these micro-improvements, but I think that we should
also take a step back, and figure out if they're going to get us where we
want to go.  I am curious if we're going to get an eventual 5-10x speedup
if we continue along the many-small-improvements path.  If not, then we
might want to bail and reconsider things more broadly.


On Tue, Jan 5, 2021 at 1:21 PM jakirkham <notifications@github.com> wrote:

> cc @mrocklin <https://github.com/mrocklin> @quasiben
> <https://github.com/quasiben>
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/4379#issuecomment-754906380>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTHRJHKXG5VFYOPPKEDSYN7FFANCNFSM4VBYVWQQ>
> .
>

This came up again today when we were revisiting one of the case where the idea originated that the scheduler is slow. In that case about 50% of the time on the scheduler is spent evenly receiving/deserializing data (admittedly this before the HLG work) and performing transitions. Additionally another ~20% of time on the scheduler is spent sending data most of which is spent serializing it. As a result it seems reasonable to conclude improvements in serialization/deserialization are well worth our time.
@jakirkham maybe spending time on Cythonization is the wrong move if this can be done easily.
Yeah mostly was curious if the `extract_serialize` Cythonization would be easy to test/show some notable improvement. Agree if that's not easy, we can just ignore it.

Looking at `transitions` atm.

@madsbk, if you have some time on Monday, maybe we can chat about this? 🙂
Just chatted with @madsbk about this, I think we will start with trying to remove `extract_serialize` by using MsgPack to do this work ( https://github.com/dask/distributed/issues/4379 ). Our hope is this already addresses performance for the status message use case ( https://github.com/dask/distributed/issues/4376 ).

Will continue pushing on optimizing transitions and moving communication out of there ( https://github.com/dask/distributed/issues/4454 ) ( https://github.com/dask/distributed/pull/4451 ).

cc @quasiben (for vis)
As a first step we are adding a fast path for things that can be handled by MsgPack alone ( https://github.com/dask/distributed/pull/4480 ). Though we are still interested in improving the serialization workflow overall, which may be handled in later work.

Resolving patch:
diff --git a/distributed/protocol/core.py b/distributed/protocol/core.py
--- a/distributed/protocol/core.py
+++ b/distributed/protocol/core.py
@@ -1,14 +1,10 @@
-from functools import reduce
 import logging
-import operator
-
 import msgpack
 
 from .compression import compressions, maybe_compress, decompress
 from .serialize import (
     Serialize,
     Serialized,
-    extract_serialize,
     msgpack_decode_default,
     msgpack_encode_default,
     merge_and_deserialize,
@@ -20,78 +16,59 @@
 logger = logging.getLogger(__name__)
 
 
-def dumps(msg, serializers=None, on_error="message", context=None):
-    """ Transform Python message to bytestream suitable for communication """
-    try:
-        data = {}
+def dumps(msg, serializers=None, on_error="message", context=None) -> list:
+    """Transform Python message to bytestream suitable for communication
 
+    Developer Notes
+    ---------------
+    The approach here is to use `msgpack.dumps()` to serialize `msg` and
+    write the result to the first output frame. If `msgpack.dumps()`
+    encounters an object it cannot serialize like a NumPy array, it is handled
+    out-of-band by `_encode_default()` and appended to the output frame list.
+    """
+    try:
         if context and "compression" in context:
             compress_opts = {"compression": context["compression"]}
         else:
             compress_opts = {}
 
-        # Only lists and dicts can contain serialized values
-        if isinstance(msg, (list, dict)):
-            msg, data, bytestrings = extract_serialize(msg)
-        small_header, small_payload = dumps_msgpack(msg, **compress_opts)
-
-        if not data:  # fast path without serialized data
-            return small_header, small_payload
+        def _inplace_compress_frames(header, frames):
+            compression = list(header.get("compression", [None] * len(frames)))
+
+            for i in range(len(frames)):
+                if compression[i] is None:
+                    compression[i], frames[i] = maybe_compress(
+                        frames[i], **compress_opts
+                    )
+
+            header["compression"] = tuple(compression)
+
+        frames = [None]
+
+        def _encode_default(obj):
+            typ = type(obj)
+            if typ is Serialize or typ is Serialized:
+                if typ is Serialize:
+                    obj = obj.data
+                offset = len(frames)
+                sub_header, sub_frames = serialize_and_split(
+                    obj, serializers=serializers, on_error=on_error, context=context
+                )
+                sub_header["num-sub-frames"] = len(sub_frames)
+                _inplace_compress_frames(sub_header, sub_frames)
+                frames.append(
+                    msgpack.dumps(
+                        sub_header, default=msgpack_encode_default, use_bin_type=True
+                    )
+                )
+                frames.extend(sub_frames)
+                return {"__Serialized__": offset}
+            else:
+                return msgpack_encode_default(obj)
 
-        pre = {
-            key: (value.header, value.frames)
-            for key, value in data.items()
-            if type(value) is Serialized
-        }
+        frames[0] = msgpack.dumps(msg, default=_encode_default, use_bin_type=True)
+        return frames
 
-        data = {
-            key: serialize_and_split(
-                value.data, serializers=serializers, on_error=on_error, context=context
-            )
-            for key, value in data.items()
-            if type(value) is Serialize
-        }
-
-        header = {"headers": {}, "keys": [], "bytestrings": list(bytestrings)}
-
-        out_frames = []
-
-        for key, (head, frames) in data.items():
-            # Compress frames that are not yet compressed
-            out_compression = []
-            for frame, compression in zip(
-                frames, head.get("compression") or [None] * len(frames)
-            ):
-                if compression is None:
-                    compression, frame = maybe_compress(frame, **compress_opts)
-
-                out_compression.append(compression)
-                out_frames.append(frame)
-
-            head["compression"] = out_compression
-            head["count"] = len(frames)
-            header["headers"][key] = head
-            header["keys"].append(key)
-
-        for key, (head, frames) in pre.items():
-            head["count"] = len(frames)
-            header["headers"][key] = head
-            header["keys"].append(key)
-            out_frames.extend(frames)
-
-        for i, frame in enumerate(out_frames):
-            if type(frame) is memoryview and frame.strides != (1,):
-                try:
-                    frame = frame.cast("B")
-                except TypeError:
-                    frame = frame.tobytes()
-                out_frames[i] = frame
-
-        return [
-            small_header,
-            small_payload,
-            msgpack.dumps(header, use_bin_type=True),
-        ] + out_frames
     except Exception:
         logger.critical("Failed to Serialize", exc_info=True)
         raise
@@ -99,53 +76,35 @@ def dumps(msg, serializers=None, on_error="message", context=None):
 
 def loads(frames, deserialize=True, deserializers=None):
     """ Transform bytestream back into Python value """
-    frames = frames[::-1]  # reverse order to improve pop efficiency
-    if not isinstance(frames, list):
-        frames = list(frames)
+
     try:
-        small_header = frames.pop()
-        small_payload = frames.pop()
-        msg = loads_msgpack(small_header, small_payload)
-        if not frames:
-            return msg
-
-        header = frames.pop()
-        header = msgpack.loads(header, use_list=False, **msgpack_opts)
-        keys = header["keys"]
-        headers = header["headers"]
-        bytestrings = set(header["bytestrings"])
-
-        for key in keys:
-            head = headers[key]
-            count = head["count"]
-            if count:
-                fs = frames[-count::][::-1]
-                del frames[-count:]
-            else:
-                fs = []
 
-            if deserialize or key in bytestrings:
-                if "compression" in head:
-                    fs = decompress(head, fs)
-                value = merge_and_deserialize(head, fs, deserializers=deserializers)
-            else:
-                value = Serialized(head, fs)
-
-            def put_in(keys, coll, val):
-                """Inverse of get_in, but does type promotion in the case of lists"""
-                if keys:
-                    holder = reduce(operator.getitem, keys[:-1], coll)
-                    if isinstance(holder, tuple):
-                        holder = list(holder)
-                        coll = put_in(keys[:-1], coll, holder)
-                    holder[keys[-1]] = val
+        def _decode_default(obj):
+            offset = obj.get("__Serialized__", 0)
+            if offset > 0:
+                sub_header = msgpack.loads(
+                    frames[offset],
+                    object_hook=msgpack_decode_default,
+                    use_list=False,
+                    **msgpack_opts
+                )
+                offset += 1
+                sub_frames = frames[offset : offset + sub_header["num-sub-frames"]]
+                if deserialize:
+                    if "compression" in sub_header:
+                        sub_frames = decompress(sub_header, sub_frames)
+                    return merge_and_deserialize(
+                        sub_header, sub_frames, deserializers=deserializers
+                    )
                 else:
-                    coll = val
-                return coll
+                    return Serialized(sub_header, sub_frames)
+            else:
+                return msgpack_decode_default(obj)
 
-            msg = put_in(key, msg, value)
+        return msgpack.loads(
+            frames[0], object_hook=_decode_default, use_list=False, **msgpack_opts
+        )
 
-        return msg
     except Exception:
         logger.critical("Failed to deserialize", exc_info=True)
         raise
diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py
--- a/distributed/protocol/serialize.py
+++ b/distributed/protocol/serialize.py
@@ -507,70 +507,6 @@ def __ne__(self, other):
         return not (self == other)
 
 
-def extract_serialize(x) -> tuple:
-    """Pull out Serialize objects from message
-
-    This also remove large bytestrings from the message into a second
-    dictionary.
-
-    Examples
-    --------
-    >>> from distributed.protocol import to_serialize
-    >>> msg = {'op': 'update', 'data': to_serialize(123)}
-    >>> extract_serialize(msg)
-    ({'op': 'update'}, {('data',): <Serialize: 123>}, set())
-    """
-    typ_x: type = type(x)
-    if typ_x is dict:
-        x_d: dict = x
-        x_items = x_d.items()
-        x2 = {}
-    elif typ_x is list:
-        x_l: list = x
-        x_items = enumerate(x_l)
-        x2 = len(x_l) * [None]
-
-    ser = {}
-    bytestrings = set()
-    path = ()
-    _extract_serialize(x_items, x2, ser, bytestrings, path)
-    return x2, ser, bytestrings
-
-
-def _extract_serialize(x_items, x2, ser: dict, bytestrings: set, path: tuple) -> None:
-    for k, v in x_items:
-        path_k = path + (k,)
-        typ_v: type = type(v)
-        if typ_v is dict:
-            v_d: dict = v
-            v_items = v_d.items()
-            x2[k] = v2 = {}
-            _extract_serialize(v_items, v2, ser, bytestrings, path_k)
-        elif typ_v is list:
-            v_l: list = v
-            v_items = enumerate(v_l)
-            x2[k] = v2 = len(v_l) * [None]
-            _extract_serialize(v_items, v2, ser, bytestrings, path_k)
-        elif typ_v is Serialize or typ_v is Serialized:
-            ser[path_k] = v
-        elif typ_v is bytes:
-            v_b: bytes = v
-            if len(v_b) > 2 ** 16:
-                ser[path_k] = to_serialize(v_b)
-                bytestrings.add(path_k)
-            else:
-                x2[k] = v_b
-        elif typ_v is bytearray:
-            v_ba: bytearray = v
-            if len(v_ba) > 2 ** 16:
-                ser[path_k] = to_serialize(v_ba)
-                bytestrings.add(path_k)
-            else:
-                x2[k] = v_ba
-        else:
-            x2[k] = v
-
-
 def nested_deserialize(x):
     """
     Replace all Serialize and Serialized values nested in *x*


Test patch:
diff --git a/distributed/comm/tests/test_comms.py b/distributed/comm/tests/test_comms.py
--- a/distributed/comm/tests/test_comms.py
+++ b/distributed/comm/tests/test_comms.py
@@ -1037,7 +1037,7 @@ def check_out_false(out_value):
         assert isinstance(ser, Serialized)
         assert deserialize(ser.header, ser.frames) == 456
 
-        assert isinstance(to_ser, list)
+        assert isinstance(to_ser, (tuple, list)) and len(to_ser) == 1
         (to_ser,) = to_ser
         # The to_serialize() value could have been actually serialized
         # or not (it's a transport-specific optimization)
@@ -1051,6 +1051,8 @@ def check_out_true(out_value):
         expected_msg = msg.copy()
         expected_msg["ser"] = 456
         expected_msg["to_ser"] = [123]
+        # Notice, we allow "to_ser" to be a tuple or a list
+        assert list(out_value.pop("to_ser")) == expected_msg.pop("to_ser")
         assert out_value == expected_msg
 
     await check_listener_deserialize(addr, False, msg, check_out_false)
@@ -1061,13 +1063,14 @@ def check_out_true(out_value):
 
     # Test with long bytestrings, large enough to be transferred
     # as a separate payload
+    # TODO: currently bytestrings are not transferred as a separate payload
 
     _uncompressible = os.urandom(1024 ** 2) * 4  # end size: 8 MB
 
     msg = {
         "op": "update",
         "x": _uncompressible,
-        "to_ser": [to_serialize(_uncompressible)],
+        "to_ser": (to_serialize(_uncompressible),),
         "ser": Serialized(*serialize(_uncompressible)),
     }
     msg_orig = msg.copy()
@@ -1089,7 +1092,7 @@ def check_out(deserialize_flag, out_value):
         else:
             assert isinstance(ser, Serialized)
             assert deserialize(ser.header, ser.frames) == _uncompressible
-            assert isinstance(to_ser, list)
+            assert isinstance(to_ser, tuple) and len(to_ser) == 1
             (to_ser,) = to_ser
             # The to_serialize() value could have been actually serialized
             # or not (it's a transport-specific optimization)
diff --git a/distributed/protocol/tests/test_collection.py b/distributed/protocol/tests/test_collection.py
--- a/distributed/protocol/tests/test_collection.py
+++ b/distributed/protocol/tests/test_collection.py
@@ -47,4 +47,4 @@ def test_nested_types():
     header, frames = serialize([[[x]]])
     assert "dask" in str(header)
     assert len(frames) == 1
-    assert x.data in frames
+    assert x.data == np.frombuffer(frames[0]).data
diff --git a/distributed/protocol/tests/test_numpy.py b/distributed/protocol/tests/test_numpy.py
--- a/distributed/protocol/tests/test_numpy.py
+++ b/distributed/protocol/tests/test_numpy.py
@@ -219,13 +219,13 @@ def test_compress_numpy():
     frames = dumps({"x": to_serialize(x)})
     assert sum(map(nbytes, frames)) < x.nbytes
 
-    header = msgpack.loads(frames[2], raw=False, use_list=False, strict_map_key=False)
+    header = msgpack.loads(frames[1], raw=False, use_list=False, strict_map_key=False)
     try:
         import blosc  # noqa: F401
     except ImportError:
         pass
     else:
-        assert all(c == "blosc" for c in header["headers"][("x",)]["compression"])
+        assert all(c == "blosc" for c in header["compression"])
 
 
 def test_compress_memoryview():
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
@@ -26,8 +26,9 @@ def test_compression_2():
     pytest.importorskip("lz4")
     np = pytest.importorskip("numpy")
     x = np.random.random(10000)
-    header, payload = dumps(x.tobytes())
-    assert not header or not msgpack.loads(header, encoding="utf8").get("compression")
+    msg = dumps(to_serialize(x.tobytes()))
+    compression = msgpack.loads(msg[1]).get("compression")
+    assert all(c is None for c in compression)
 
 
 def test_compression_without_deserialization():
@@ -87,14 +88,13 @@ def test_maybe_compress_sample():
 
 def test_large_bytes():
     for tp in (bytes, bytearray):
-        msg = {"x": tp(b"0" * 1000000), "y": 1}
+        msg = {"x": to_serialize(tp(b"0" * 1000000)), "y": 1}
         frames = dumps(msg)
+        msg["x"] = msg["x"].data
         assert loads(frames) == msg
         assert len(frames[0]) < 1000
         assert len(frames[1]) < 1000
 
-        assert loads(frames, deserialize=False) == msg
-
 
 @pytest.mark.slow
 def test_large_messages():