Repository: dask/distributed

Issue:
Large bytestrings making their way into msgpack
Msgpack does not handle large bytestrings well.  Usually we evict large bytestrings from messages and send them along separately (see http://matthewrocklin.com/blog/work/2016/04/14/dask-distributed-optimizing-protocol).  

From https://github.com/dask/dask/issues/1411#issuecomment-235655955 it appears that this isn't happening during a `scatter` operation.  This may be because the large bytestrings aren't at the top-level of the message, but instead in a `data=` key.  

If this is the case then there are a few options:
1.  Dump scattered keys into the top-level keys of the message (and watch out for key conflicts)
2.  Teach the protocol to look more deeply into sub dictionaries (but be careful that this doesn't affect performance in the many-small message case)
3.  Special case a few keys, like `data` and `payload`

This might be a good learning issue.  Pinging @jcrist 



Discussion:
I've tried repeatedly to reproduce the issue, and am failing. I've tried:
- 1 single large numpy array to send to one worker
- Many single large numpy arrays to send to one worker.

Both cases have passed fine. The code is below:

``` python
import numpy as np
from distributed import Executor, LocalCluster
import pickle
from pandas import msgpack

c = LocalCluster(n_workers=1, threads_per_worker=1)
exc = Executor(c)

# -- Single large array --
# Shape picked so that pickle.dumps(data, protocol=2) results in a bytestring
# longer than 2**31. You can't pickle numpy arrays where the data bytes are longer
# than 2**31 - 6.
data = np.zeros(2**31 - 6, dtype='i1')
exc.scatter([data])

# -- Many Slices --
# Shape taken from original issue, div by 2 so it fits on my machine twice
shape = (145, 174, 145, 288 / 2)
data = np.zeros(shape=shape)
slices = [data[..., i] for i in range(data.shape[-1])]
remote = exc.scatter(slices)
```

@akcheung, do you have a reproducible example for this?

Try the code pasted in the beginning of [#1411](https://github.com/dask/dask/issues/1411), uncommenting out `data = np.random.randint(0, 1000, size=(145, 174, 145, 288))`

Regardless of whether or not we can get this to fail it remains a performance bottleneck.  

This should produce unfriendly results:

```
from distributed.protocol import dumps

msg = {'x': {'y': b'large-random-byte-string'}}
%prun dumps(msg)
```

@jcrist I suspect that your example did not fail due the data being highly compressible.  Try switching out for random data for a better (worse) result.

I've been thinking about this a bit more and I think that the following could work:

Currently messages are broken into two parts
1.  Small msgpack friendly bits
2.  key-value pairs of large bytestrings in the top-level dictionary

I propose that instead we swap out part 2 for key-sequences and large byte strings.  Consider the following example

```
{'x': 1, 'y': b'big-byte-string-1', 'z': {'a': 2, 'b': b'big-byte-string-2'}}
```

As currently implemented the protocol would decompose this into the following two pieces

```
msg = {'x': 1, 'z': {'a': 2, 'b': b'big-byte-string-2'}} 
keys = ['y']
values = [b'big-byte-string-1']
```

Instead I propose the following:

```
msg = {'x': 1, 'z': {'a': 2}} 
keys = [('y',), ('z', 'b')]
values = [b'big-byte-string-1', b'big-byte-string-2']
```

Keys now contains paths from which large bytestrings originated.  These paths are similar to what you find in `toolz.update_in`. 

This doesn't solve the performance problem; we need to verify that traversing through every message doesn't cause undue overhead.  We should benchmark this a bit.  If it is a problem then we can probably get around it with a combination of opt-in (we know which messages are likely to contain large bytestrings) and exception handling (if the protocol fails to serialize something, then try traversing through.)


Resolving patch:
diff --git a/distributed/protocol.py b/distributed/protocol.py
--- a/distributed/protocol.py
+++ b/distributed/protocol.py
@@ -27,17 +27,16 @@
 from __future__ import print_function, division, absolute_import
 
 import random
-import struct
+from copy import deepcopy
 
 try:
     import pandas.msgpack as msgpack
 except ImportError:
     import msgpack
 
-from toolz import first, keymap, identity, merge
+from toolz import identity, get_in
 
 from .utils import ignoring
-from .compatibility import unicode
 
 
 compressions = {None: {'compress': identity,
@@ -64,35 +63,93 @@
     default_compression = 'lz4'
 
 
+BIG_BYTES_SIZE = 2**20
+BIG_BYTES_SHARD_SIZE = 2**28
+
+
+def extract_big_bytes(x):
+    big = {}
+    _extract_big_bytes(x, big)
+    if big:
+        x = deepcopy(x)
+        for path in big:
+            t = get_in(path[:-1], x)
+            if isinstance(t, dict):
+                del t[path[-1]]
+            else:
+                t[path[-1]] = None
+    return x, big
+
+
+def _extract_big_bytes(x, big, path=()):
+    if type(x) is dict:
+        for k, v in x.items():
+            if isinstance(v, (list, dict)):
+                _extract_big_bytes(v, big, path + (k,))
+            elif type(v) is bytes and len(v) >= BIG_BYTES_SIZE:
+                big[path + (k,)] = v
+    elif type(x) is list:
+        for k, v in enumerate(x):
+            if isinstance(v, (list, dict)):
+                _extract_big_bytes(v, big, path + (k,))
+            elif type(v) is bytes and len(v) >= BIG_BYTES_SIZE:
+                big[path + (k,)] = v
+
+
 def dumps(msg):
     """ Transform Python value to bytestream suitable for communication """
-    small_header = {}
-
+    big = {}
+    # Only dicts can contain big values
     if isinstance(msg, dict):
-        big = {k: v for k, v in msg.items()
-                    if isinstance(v, bytes) and len(v) > 1e6}
-    else:
-        big = False
-    if big:
-        small = {k: v for k, v in msg.items() if k not in big}
-    else:
-        small = msg
+        msg, big = extract_big_bytes(msg)
+    small_header, small_payload = dumps_msgpack(msg)
+    if not big:
+        return small_header, small_payload
+
+    # Shard the big segments
+    shards = []
+    res = {}
+    for k, v in list(big.items()):
+        L = []
+        for i, j in enumerate(range(0, len(v), BIG_BYTES_SHARD_SIZE)):
+            key = '.shard-%d-%s' % (i, k)
+            res[key] = v[j: j + BIG_BYTES_SHARD_SIZE]
+            L.append(key)
+        shards.append((k, L))
+
+    keys, values = zip(*res.items())
 
-    frames = dumps_msgpack(small)
-    if big:
-        frames += dumps_big_byte_dict(big)
+    compression = []
+    values2 = []
+    for v in values:
+        fmt, vv = maybe_compress(v)
+        compression.append(fmt)
+        values2.append(vv)
 
-    return frames
+    header = {'encoding': 'big-byte-dict',
+              'keys': keys,
+              'compression': compression,
+              'shards': shards}
+
+    return [small_header, small_payload,
+            msgpack.dumps(header, use_bin_type=True)] + values2
 
 
 def loads(frames):
     """ Transform bytestream back into Python value """
-    header, payload, frames = frames[0], frames[1], frames[2:]
-    msg = loads_msgpack(header, payload)
+    small_header, small_payload, frames = frames[0], frames[1], frames[2:]
+    msg = loads_msgpack(small_header, small_payload)
 
     if frames:
-        big = loads_big_byte_dict(*frames)
-        msg.update(big)
+        header = msgpack.loads(frames[0], encoding='utf8')
+
+        values2 = [compressions[c]['decompress'](v)
+                   for c, v in zip(header['compression'], frames[1:])]
+        lk = dict(zip(header['keys'], values2))
+
+        for k, keys in header['shards']:
+            v = b''.join(lk.pop(kk) for kk in keys)
+            get_in(k[:-1], msg)[k[-1]] = v
 
     return msg
 
@@ -109,7 +166,7 @@ def byte_sample(b, size, n):
 
 
 def maybe_compress(payload, compression=default_compression, min_size=1e4,
-        sample_size=1e4, nsamples=5):
+                   sample_size=1e4, nsamples=5):
     """ Maybe compress payload
 
     1.  We don't compress small messages
@@ -183,65 +240,6 @@ def loads_msgpack(header, payload):
             payload = decompress(payload)
         except KeyError:
             raise ValueError("Data is compressed as %s but we don't have this"
-                    " installed" % header['compression'].decode())
+                             " installed" % header['compression'].decode())
 
     return msgpack.loads(payload, encoding='utf8')
-
-
-def dumps_big_byte_dict(d):
-    """ Serialize large byte dictionary to sequence of frames
-
-    The input must be a dictionary and all values of that dictionary must be
-    bytestrings.  These should probably be large.
-
-    Returns a sequence of frames, one header followed by each of the values
-
-    See Also:
-        loads_big_byte_dict
-    """
-    assert isinstance(d, dict) and all(isinstance(v, bytes) for v in d.values())
-    shards = {}
-    for k, v in list(d.items()):
-        if len(v) >= 2**31:
-            L = []
-            for i, j in enumerate(range(0, len(v), 2**30)):
-                key = '.shard-%d-%s' % (i, k)
-                d[key] = v[j: j + 2**30]
-                L.append(key)
-            del d[k]
-            shards[k] = L
-
-    keys, values = zip(*d.items())
-
-    compress = compressions[default_compression]['compress']
-    compression = []
-    values2 = []
-    for v in values:
-        fmt, vv = maybe_compress(v)
-        compression.append(fmt)
-        values2.append(vv)
-
-    header = {'encoding': 'big-byte-dict',
-              'keys': keys,
-              'compression': compression}
-    if shards:
-        header['shards'] = shards
-
-    return [msgpack.dumps(header, use_bin_type=True)] + values2
-
-
-def loads_big_byte_dict(header, *values):
-    """ Deserialize big-byte frames to large byte dictionary
-
-    See Also:
-        dumps_big_byte_dict
-    """
-    header = msgpack.loads(header, encoding='utf8')
-
-    values2 = [compressions[c]['decompress'](v)
-               for c, v in zip(header['compression'], values)]
-    result = dict(zip(header['keys'], values2))
-
-    for k, keys in header.get('shards', {}).items():
-        result[k] = b''.join(result.pop(kk) for kk in keys)
-    return result


Test patch:
diff --git a/distributed/tests/test_protocol.py b/distributed/tests/test_protocol.py
--- a/distributed/tests/test_protocol.py
+++ b/distributed/tests/test_protocol.py
@@ -1,15 +1,15 @@
 from __future__ import print_function, division, absolute_import
 
-from distributed.protocol import (loads, dumps, dumps_msgpack, loads_msgpack,
-        dumps_big_byte_dict, loads_big_byte_dict, msgpack, maybe_compress)
+from distributed.protocol import loads, dumps, msgpack, maybe_compress
 import pytest
 
+
 def test_protocol():
     for msg in [1, 'a', b'a', {'x': 1}, {b'x': 1}, {'x': b''}, {}]:
         assert loads(dumps(msg)) == msg
 
 
-def test_compression():
+def test_compression_1():
     pytest.importorskip('lz4')
     np = pytest.importorskip('numpy')
     x = np.ones(1000000)
@@ -19,7 +19,7 @@ def test_compression():
     assert x.tobytes() == y
 
 
-def test_compression():
+def test_compression_2():
     pytest.importorskip('lz4')
     np = pytest.importorskip('numpy')
     x = np.random.random(10000)
@@ -41,22 +41,29 @@ def test_small_and_big():
     # assert loads([big_header, big]) == {'y': d['y']}
 
 
-def test_big_bytes():
-    d = {'x': b'123', b'y': b'4567'}
-    L = dumps_big_byte_dict(d)
-    assert isinstance(L, (tuple, list))
-    assert all(isinstance(item, bytes) for item in L)
-    dd = loads_big_byte_dict(*L)
+def test_big_bytes_protocol():
+    np = pytest.importorskip('numpy')
+    data = np.random.randint(0, 255, dtype='u1', size=2**21).tobytes()
+
+    d = {'x': data, 'y': 'foo'}
+    frames = dumps(d)
+    assert len(frames) == 4     # Only `data` is extracted
+    assert data is frames[3]    # `data` isn't sharded as it's too short
+    dd = loads(frames)
     assert dd == d
 
+    d = {'x': [data], 'y': 'foo'}
+    frames = dumps(d)
+    assert len(frames) == 4
+    assert data is frames[3]
+    dd = loads(frames)
+    assert dd == d
 
-def test_big_bytes_protocol():
-    np = pytest.importorskip('numpy')
-    data = np.random.randint(0, 255, dtype='u1', size=2000000).tobytes()
-    d = {'x': data, 'y': b'1' * 2000000}
-    L = dumps(d)
-    assert d['x'] in L
-    dd = loads(L)
+    d = {'x': {'z': [data, 'small_data']}, 'y': 'foo'}
+    frames = dumps(d)
+    assert len(frames) == 4
+    assert data is frames[3]
+    dd = loads(frames)
     assert dd == d
 
 
@@ -86,8 +93,16 @@ def test_large_messages():
     pytest.importorskip('lz4')
     if psutil.virtual_memory().total < 8e9:
         return
-    msg = {'x': b'0' * (2**31 + 10)}
+
+    big_bytes = b'0' * (2**31 + 10)
+    msg = {'x': [big_bytes, b'small_bytes'],
+           'y': {'a': big_bytes, 'b': b'small_bytes'}}
 
     b = dumps(msg)
     msg2 = loads(b)
     assert msg == msg2
+
+    assert len(b) >= 2
+    big_header = msgpack.loads(b[2], encoding='utf8')
+    assert len(big_header['shards']) == 2
+    assert len(big_header['keys']) + 2 + 1 == len(b)