Repository: dask/distributed

Issue:
Deserialization of compressed data is sluggish and causes memory flares
If lz4 is installed and a chunk can be compressed efficiently, it is compressed automatically.
This happens both in network comms (configurable) and upon spilling (not configurable).

While working on https://github.com/coiled/coiled-runtime/pull/629, I found out that deserializing a compressed frame is abysmally slow:

![image](https://user-images.githubusercontent.com/6213168/209113228-9ca8f43e-1944-4c9b-965d-8c47b884dc2a.png)


When deserialization is triggered by network comms, it is offloaded to a separate thread for any message larger than 10 MiB (`distributed.comm.offload`). Note that a message will contain multiple small keys if possible.
I did not test if the GIL is released while this happens.

When deserialization is triggered by unspilling, the whole event loop is temporarily blocked.

### Microbenchmark
Serializing a 128 MiB compressible numpy array takes 9.96ms; there's a modest overhead on top of the raw call to lz4, which takes 9.75ms.
Deserializing the same takes a whopping 144ms. The raw lz4 decompress  takes an already huge 83ms.

At the very least we should be able to trim down the overhead. Also the raw decompression speed is suspiciously slow - typically, compression takes much longer than decompression so I suspect that something must be going on upstream.

```python
import io
import lz4.frame
import numpy
from distributed.protocol import deserialize_bytes, serialize_bytelist

a1 = numpy.random.random((4096, 4096))  # 128 MiB, uncompressible
a2 = numpy.ones((4096, 4096))  # 128 MiB, highly compressible

buf1 = io.BytesIO()
buf2 = io.BytesIO()
buf1.writelines(serialize_bytelist(a1))
buf2.writelines(serialize_bytelist(a2))
b1 = buf1.getvalue()
b2 = buf2.getvalue()
b3 = lz4.frame.compress(a2)

print("uncompressible size", len(b1))
print("compressible size", len(b2))
print("raw lz4 size", len(b3))
print()

print("serialize uncompressible")
%timeit serialize_bytelist(a1)
print("serialize compressible")
%timeit serialize_bytelist(a2)
print("raw lz4 compress")
%timeit lz4.frame.compress(a2)
print()

print("deserialize uncompressible")
%timeit deserialize_bytes(b1)
print("deserialize compressible")
%timeit deserialize_bytes(b2)
print("raw lz4 decompress")
%timeit lz4.frame.decompress(b3)
```
Output:
```
uncompressible size 134217969
compressible size 526635
raw lz4 size 566291

serialize uncompressible
89.9 µs ± 15.6 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
serialize compressible
9.96 ms ± 78.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
raw lz4 compress
9.75 ms ± 115 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

deserialize uncompressible
19.2 ms ± 138 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
deserialize compressible
143 ms ± 568 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
raw lz4 decompress
88.3 ms ± 113 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

`serialize_bytelist`  and `deserialize_bytes` in the code above are exactly what `distributed.spill` uses. However, I measured the timings in `distributed.comm.tcp` and I got the same number (134ms).


Discussion:
Notably, a lot of the tests in coiled-runtime start from `da.random.random()`, which produces uncompressible data, and will not show this issue.
Replicated on 

Python: 3.11.0
macOS: 13.1
arch : arm64
lz4 : 4.0.2
distributed : 2022.12.1
numpy : 1.24.0
msgpack : 1.0.4

```
uncompressible size 134217969
compressible size 526635
raw lz4 size 566291

serialize uncompressible
39.5 µs ± 3.99 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
serialize compressible
8.35 ms ± 48.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
raw lz4 compress
7.78 ms ± 49.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

deserialize uncompressible
11.3 ms ± 577 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
deserialize compressible
70.4 ms ± 659 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
raw lz4 decompress
40.6 ms ± 479 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
TL,DR; It's death by deep-copy. We're spending 85ms doing actual decompression and 49ms doing deep-copies.

```
   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.134    0.134 distributed/protocol/serialize.py:662(deserialize_bytes)
        1    0.000    0.000    0.098    0.098 distributed/protocol/compression.py:193(decompress)
        1    0.000    0.000    0.098    0.098 distributed/protocol/compression.py:195(<listcomp>)
        2    0.098    0.049    0.098    0.049 {built-in method lz4.block._block.decompress}
        1    0.000    0.000    0.036    0.036 distributed/protocol/serialize.py:473(merge_and_deserialize)
        1    0.036    0.036    0.036    0.036 {method 'join' of 'bytearray' objects}
```
The 128 MiB input buffer is split in two frames (`distributed.block.shard: 64 MiB`).
The two calls to `lz4.block._block.decompress`:
1. write the output to a local raw buffer `dest`:
https://github.com/python-lz4/python-lz4/blob/79370987909663d4e6ef743762768ebf970a2383/lz4/block/_block.c#L349-L359
3. then, they deep-copy to the output bytes object:
https://github.com/python-lz4/python-lz4/blob/79370987909663d4e6ef743762768ebf970a2383/lz4/block/_block.c#L389
4. The two bytes objects are not contiguous. Hence, they are deep-copied again into their final destination by `merge_and_deserialize`

Fixing this will require an upstream PR that lets us pass an `out=` parameter.

If I increase `distributed.block.shard` to 128 MiB, the profile changes - but not in a very good way:
```
   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.113    0.113 distributed/protocol/serialize.py:662(deserialize_bytes)
        1    0.000    0.000    0.094    0.094 distributed/protocol/compression.py:193(decompress)
        1    0.000    0.000    0.094    0.094 distributed/protocol/compression.py:195(<listcomp>)
        1    0.094    0.094    0.094    0.094 {built-in method lz4.block._block.decompress}
        1    0.000    0.000    0.019    0.019 distributed/protocol/serialize.py:473(merge_and_deserialize)
        1    0.000    0.000    0.019    0.019 distributed/protocol/serialize.py:371(deserialize)
        1    0.000    0.000    0.019    0.019 distributed/protocol/serialize.py:56(dask_loads)
        1    0.000    0.000    0.019    0.019 distributed/utils.py:746(wrapper)
        1    0.000    0.000    0.019    0.019 distributed/protocol/numpy.py:114(deserialize_numpy_ndarray)
        1    0.000    0.000    0.019    0.019 numpy/core/_asarray.py:22(require)
        1    0.019    0.019    0.019    0.019 {method 'copy' of 'numpy.ndarray' objects}
```
We can see that numpy is performing a deep copy - which should not happen. So even here we should be able to shave off 19ms off in numpy deserialization and 9ms in PyBytes_FromStringAndSize (9ms is the duratio of a raw deepcopy on my computer).

Worth noting that blosc (removed in #5269) does offer deserialization to already existing memory (albeit in an unpythonic and unsafe way).
This works, thanks to #5208:

```python
import dask
import distributed
import numpy

c = distributed.Client(n_workers=2, memory_limit="8 GiB")
x = c.submit(numpy.random.random, 4 * 2**30 // 8, workers=[0])
y = c.submit(lambda x: None, x, workers=[1])
```

Replace `numpy.random.random` with `numpy.ones`, and the receiving worker will be killed by the nanny.

> TL,DR; It's death by deep-copy. We're spending 85ms doing actual decompression and 49ms doing deep-copies.

Another important note:
The actual decompression part releases the GIL. So, as long as serialization/deserialization iss offloaded to another thread (#4424), it's fine.
To my understanding, the deep copies all hold the GIL while they are performed.

I would be curious to see if turning off compression entirely affects the benchmark suite.  That might give us a better holistic understanding of how much impact this has.  Is that easy to test?
> I would be curious to see if turning off compression entirely affects the benchmark suite. That might give us a better holistic understanding of how much impact this has. Is that easy to test?

Yes but a lot of data in the benchmark suite is created with numpy.random.random - which is uncompressible - so it's not necessarily (I'd dare say, unlikely) representative of real world data.

I think the whole thing can be dropped from 143ms to 18ms. I think we were too hasty to drop support for blosc - the performance benefit is huge.

```python
import pickle
from time import time
import blosc


def compress(frame):
    if isinstance(frame, pickle.PickleBuffer):
        mv = memoryview(frame)
        return blosc.compress(mv, clevel=5, cname="blosclz", typesize=mv.itemsize)
    else:
        return blosc.compress(frame, clevel=5, cname="blosclz")

# a = numpy.linspace(0, 1, 128 * 2**20 // 8)
a = numpy.ones(128 * 2**20 // 8)

t0 = time()
frames = []
pik = pickle.dumps(a, protocol=5, buffer_callback=frames.append)
t1 = time()
frames.insert(0, pik)
print([len(memoryview(frame)) for frame in frames])
cframes = [compress(frame) for frame in frames]
print([len(frame) for frame in cframes])
t2 = time()
frames2 = [blosc.decompress(frame) for frame in cframes]
t3 = time()
a2 = pickle.loads(frames2[0], buffers=frames2[1:])
t4 = time()
numpy.testing.assert_array_equal(a, a2)

print("dumps", t1 - t0)
print("compress", t2 - t1)
print("decompress", t3 - t2)
print("loads", t4 - t3)
```
```
[120, 16777216]
[136, 542224]
dumps 0.00045180320739746094
compress 0.011289596557617188
decompress 0.012519598007202148
loads 0.0037240982055664062
```
Same as above, but with `lz4.block.compress` / `lz4.block.decompress`:
```
[120, 16777216]
[121, 526366]
dumps 0.0005145072937011719
compress 0.010733366012573242
decompress 0.09740924835205078
loads 0.006391048431396484
```

linspace does not compress at all with lz4, while it compresses to 6% with blosc.

---

On a separate topic - what's the purpose of sharding? I can't see how it makes us avoid any memory spikes anyway, as we always have all the sharded frames in memory before we do anything with them?

While it's true that we won't be able to see benefits of compressor we will
be able to see the costs pretty clearly.  If we turn off compression and
nothing changes then we gain the option of ignoring this for a while.

Sharding is important as data sizes grow to 2GB, when many algorithms break
down.  Below that we're guessing that it's nice (smoother maybe).

On Fri, Dec 23, 2022, 5:04 PM crusaderky ***@***.***> wrote:

> I think the whole thing can be dropped from 143ms to 18ms. I think we were
> too hasty to drop support for blosc - the performance benefit is huge.
>
> import picklefrom time import timeimport blosc
>
> def compress(frame):
>     if isinstance(frame, pickle.PickleBuffer):
>         mv = memoryview(frame)
>         return blosc.compress(mv, clevel=5, cname="blosclz", typesize=mv.itemsize)
>     else:
>         return blosc.compress(frame, clevel=5, cname="blosclz")
> a = numpy.linspace(0, 1, 128 * 2**20 // 8)# a = numpy.ones(128 * 2**20 // 8)
> t0 = time()frames = []pik = pickle.dumps(a, protocol=5, buffer_callback=frames.append)t1 = time()frames.insert(0, pik)print([len(memoryview(frame)) for frame in frames])cframes = [compress(frame) for frame in frames]print([len(frame) for frame in cframes])t2 = time()frames2 = [blosc.decompress(frame, as_bytearray=True) for frame in cframes]t3 = time()a2 = pickle.loads(frames2[0], buffers=frames2[1:])t4 = time()numpy.testing.assert_array_equal(a, a2)
> print("dumps", t1 - t0)print("compress", t2 - t1)print("decompress", t3 - t2)print("loads", t4 - t3)
>
>
> [120, 16777216]
> [136, 1069658]
> dumps 0.0005013942718505859
> compress 0.006549835205078125
> decompress 0.012061119079589844
> loads 0.0056269168853759766
>
> On a separate topic - what's the purpose of sharding? I can't see how it
> makes us avoid any memory spikes anyway, as we always have all the sharded
> frames in memory before we do anything with them?
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/distributed/issues/7433#issuecomment-1364415978>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTFJQDZDUOPN3EKV5BTWOY4XBANCNFSM6AAAAAATHEYDAI>
> .
> You are receiving this because you commented.Message ID:
> ***@***.***>
>


Resolving patch:
diff --git a/distributed/protocol/compression.py b/distributed/protocol/compression.py
--- a/distributed/protocol/compression.py
+++ b/distributed/protocol/compression.py
@@ -8,6 +8,7 @@
 import logging
 from collections.abc import Callable
 from contextlib import suppress
+from functools import partial
 from random import randint
 from typing import Literal
 
@@ -64,12 +65,16 @@
     if parse_version(lz4.__version__) < parse_version("0.23.1"):
         raise ImportError("Need lz4 >= 0.23.1")
 
-    from lz4.block import compress as lz4_compress
-    from lz4.block import decompress as lz4_decompress
+    import lz4.block
 
     compressions["lz4"] = {
-        "compress": lz4_compress,
-        "decompress": lz4_decompress,
+        "compress": lz4.block.compress,
+        # Avoid expensive deep copies when deserializing writeable numpy arrays
+        # See distributed.protocol.numpy.deserialize_numpy_ndarray
+        # Note that this is only useful for buffers smaller than distributed.comm.shard;
+        # larger ones are deep-copied between decompression and serialization anyway in
+        # order to merge them.
+        "decompress": partial(lz4.block.decompress, return_bytearray=True),
     }
     default_compression = "lz4"
 
diff --git a/distributed/protocol/numpy.py b/distributed/protocol/numpy.py
--- a/distributed/protocol/numpy.py
+++ b/distributed/protocol/numpy.py
@@ -138,7 +138,12 @@ def deserialize_numpy_ndarray(header, frames):
         # This should exclusively happen when the underlying buffer is read-only, e.g.
         # a read-only mmap.mmap or a bytes object.
         # Specifically, these are the known use cases:
-        # 1. decompressed output of a buffer that was not sharded
+        # 1. decompression with a library that does not support output to bytearray
+        #    (lz4 does; snappy, zlib, and zstd don't).
+        #    Note that this only applies to buffers whose uncompressed size was small
+        #    enough that they weren't sharded (distributed.comm.shard); for larger
+        #    buffers the decompressed output is deep-copied beforehand into a bytearray
+        #    in order to merge it.
         # 2. unspill with zict <2.3.0 (https://github.com/dask/zict/pull/74)
         x = np.require(x, requirements=["W"])
 


Test patch:
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
@@ -17,7 +17,7 @@
     serialize,
     to_serialize,
 )
-from distributed.protocol.compression import maybe_compress
+from distributed.protocol.compression import default_compression, maybe_compress
 from distributed.protocol.numpy import itemsize
 from distributed.protocol.utils import BIG_BYTES_SHARD_SIZE
 from distributed.system import MEMORY_LIMIT
@@ -216,8 +216,8 @@ def test_itemsize(dt, size):
     assert itemsize(np.dtype(dt)) == size
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_compress_numpy():
-    pytest.importorskip("lz4")
     x = np.ones(10000000, dtype="i4")
     frames = dumps({"x": to_serialize(x)})
     assert sum(map(nbytes, frames)) < x.nbytes
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
@@ -31,6 +31,25 @@ def test_protocol():
         assert loads(dumps(msg)) == msg
 
 
+def test_default_compression():
+    """Test that the default compression algorithm is lz4 -> snappy -> None.
+    If neither is installed, test that we don't fall back to the very slow zlib.
+    """
+    try:
+        import lz4  # noqa: F401
+
+        assert default_compression == "lz4"
+        return
+    except ImportError:
+        pass
+    try:
+        import snappy  # noqa: F401
+
+        assert default_compression == "snappy"
+    except ImportError:
+        assert default_compression is None
+
+
 @pytest.mark.parametrize(
     "config,default",
     [
@@ -49,8 +68,8 @@ def test_compression_config(config, default):
             assert get_default_compression() == default
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_compression_1():
-    pytest.importorskip("lz4")
     np = pytest.importorskip("numpy")
     x = np.ones(1000000)
     b = x.tobytes()
@@ -60,8 +79,8 @@ def test_compression_1():
     assert {"x": b} == y
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_compression_2():
-    pytest.importorskip("lz4")
     np = pytest.importorskip("numpy")
     x = np.random.random(10000)
     msg = dumps(to_serialize(x.data))
@@ -69,8 +88,8 @@ def test_compression_2():
     assert all(c is None for c in compression)
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_compression_3():
-    pytest.importorskip("lz4")
     np = pytest.importorskip("numpy")
     x = np.ones(1000000)
     frames = dumps({"x": Serialize(x.data)})
@@ -79,8 +98,8 @@ def test_compression_3():
     assert {"x": x.data} == y
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_compression_without_deserialization():
-    pytest.importorskip("lz4")
     np = pytest.importorskip("numpy")
     x = np.ones(1000000)
 
@@ -91,6 +110,18 @@ def test_compression_without_deserialization():
     assert all(len(frame) < 1000000 for frame in msg["x"].frames)
 
 
+def test_lz4_decompression_avoids_deep_copy():
+    """Test that lz4 output is a bytearray, not bytes, so that numpy deserialization is
+    not forced to perform a deep copy to obtain a writeable array.
+    Note that zlib, zstandard, and snappy don't have this option.
+    """
+    pytest.importorskip("lz4")
+    a = bytearray(1_000_000)
+    b = compressions["lz4"]["compress"](a)
+    c = compressions["lz4"]["decompress"](b)
+    assert isinstance(c, bytearray)
+
+
 def test_small():
     assert sum(map(nbytes, dumps(b""))) < 10
     assert sum(map(nbytes, dumps(1))) < 10
@@ -106,7 +137,13 @@ def test_small_and_big():
 
 @pytest.mark.parametrize(
     "lib,compression",
-    [(None, None), ("zlib", "zlib"), ("lz4", "lz4"), ("zstandard", "zstd")],
+    [
+        (None, None),
+        ("zlib", "zlib"),
+        ("lz4", "lz4"),
+        ("snappy", "snappy"),
+        ("zstandard", "zstd"),
+    ],
 )
 def test_maybe_compress(lib, compression):
     if lib:
@@ -126,7 +163,13 @@ def test_maybe_compress(lib, compression):
 
 @pytest.mark.parametrize(
     "lib,compression",
-    [(None, None), ("zlib", "zlib"), ("lz4", "lz4"), ("zstandard", "zstd")],
+    [
+        (None, None),
+        ("zlib", "zlib"),
+        ("lz4", "lz4"),
+        ("snappy", "snappy"),
+        ("zstandard", "zstd"),
+    ],
 )
 def test_compression_thread_safety(lib, compression):
     if lib:
@@ -164,7 +207,13 @@ def test_compress_decompress(fn):
 
 @pytest.mark.parametrize(
     "lib,compression",
-    [(None, None), ("zlib", "zlib"), ("lz4", "lz4"), ("zstandard", "zstd")],
+    [
+        (None, None),
+        ("zlib", "zlib"),
+        ("lz4", "lz4"),
+        ("snappy", "snappy"),
+        ("zstandard", "zstd"),
+    ],
 )
 def test_maybe_compress_config_default(lib, compression):
     if lib:
@@ -183,9 +232,9 @@ def test_maybe_compress_config_default(lib, compression):
             assert compressions[rc]["decompress"](rd) == payload
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_maybe_compress_sample():
     np = pytest.importorskip("numpy")
-    lz4 = pytest.importorskip("lz4")
     payload = np.random.randint(0, 255, size=10000).astype("u1").tobytes()
     fmt, compressed = maybe_compress(payload)
     assert fmt is None
@@ -202,10 +251,9 @@ def test_large_bytes():
         assert len(frames[1]) < 1000
 
 
-@pytest.mark.slow
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_large_messages():
     np = pytest.importorskip("numpy")
-    pytest.importorskip("lz4")
     if MEMORY_LIMIT < 8e9:
         pytest.skip("insufficient memory")
 
@@ -248,8 +296,8 @@ def test_loads_deserialize_False():
     assert result == 123
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_loads_without_deserialization_avoids_compression():
-    pytest.importorskip("lz4")
     b = b"0" * 100000
 
     msg = {"x": 1, "data": to_serialize(b)}
@@ -311,12 +359,12 @@ def test_dumps_loads_Serialized():
     assert result == result3
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 def test_maybe_compress_memoryviews():
     np = pytest.importorskip("numpy")
-    pytest.importorskip("lz4")
     x = np.arange(1000000, dtype="int64")
     compression, payload = maybe_compress(x.data)
-    assert compression == "lz4"
+    assert compression == default_compression
     assert len(payload) < x.nbytes * 0.75
 
 
diff --git a/distributed/protocol/tests/test_serialize.py b/distributed/protocol/tests/test_serialize.py
--- a/distributed/protocol/tests/test_serialize.py
+++ b/distributed/protocol/tests/test_serialize.py
@@ -33,6 +33,7 @@
     serialize_bytes,
     to_serialize,
 )
+from distributed.protocol.compression import default_compression
 from distributed.protocol.serialize import check_dask_serializable
 from distributed.utils import ensure_memoryview, nbytes
 from distributed.utils_test import NO_AMM, gen_test, inc
@@ -279,9 +280,9 @@ def test_serialize_bytes(kwargs):
         assert str(x) == str(y)
 
 
+@pytest.mark.skipif(default_compression is None, reason="requires lz4 or snappy")
 @pytest.mark.skipif(np is None, reason="Test needs numpy")
 def test_serialize_list_compress():
-    pytest.importorskip("lz4")
     x = np.ones(1000000)
     L = serialize_bytelist(x)
     assert sum(map(nbytes, L)) < x.nbytes / 2