Track mutable frames

distributed/protocol/core.py

@@ -7,7 +7,7 @@
 from .compression import compressions, maybe_compress, decompress
 from .serialize import serialize, deserialize, Serialize, Serialized, extract_serialize
 from .utils import frame_split_size, merge_frames, msgpack_opts
-from ..utils import nbytes
+from ..utils import is_writeable, nbytes
 
 _deserialize = deserialize
 
@@ -46,6 +46,8 @@ def dumps(msg, serializers=None, on_error="message", context=None):
         out_frames = []
 
         for key, (head, frames) in data.items():
+            if "writeable" not in head:
+                head["writeable"] = tuple(map(is_writeable, frames))
             if "lengths" not in head:
                 head["lengths"] = tuple(map(nbytes, frames))
 
@@ -71,6 +73,8 @@ def dumps(msg, serializers=None, on_error="message", context=None):
             out_frames.extend(_out_frames)
 
         for key, (head, frames) in pre.items():
+            if "writeable" not in head:
+                head["writeable"] = tuple(map(is_writeable, frames))
             if "lengths" not in head:
                 head["lengths"] = tuple(map(nbytes, frames))
             head["count"] = len(frames)

distributed/protocol/cupy.py

@@ -47,6 +47,7 @@ def cuda_deserialize_cupy_ndarray(header, frames):
 @dask_serialize.register(cupy.ndarray)
 def dask_serialize_cupy_ndarray(x):
     header, frames = cuda_serialize_cupy_ndarray(x)
+    header["writeable"] = (None,) * len(frames)
     frames = [memoryview(cupy.asnumpy(f)) for f in frames]
     return header, frames
 

distributed/protocol/numba.py

@@ -51,6 +51,7 @@ def cuda_deserialize_numba_ndarray(header, frames):
 @dask_serialize.register(numba.cuda.devicearray.DeviceNDArray)
 def dask_serialize_numba_ndarray(x):
     header, frames = cuda_serialize_numba_ndarray(x)
+    header["writeable"] = (None,) * len(frames)
     frames = [memoryview(f.copy_to_host()) for f in frames]
     return header, frames
 

distributed/protocol/rmm.py

@@ -30,6 +30,7 @@ def cuda_deserialize_rmm_device_buffer(header, frames):
     @dask_serialize.register(rmm.DeviceBuffer)
     def dask_serialize_rmm_device_buffer(x):
         header, frames = cuda_serialize_rmm_device_buffer(x)
+        header["writeable"] = (None,) * len(frames)
         frames = [numba.cuda.as_cuda_array(f).copy_to_host().data for f in frames]
         return header, frames
 

distributed/protocol/serialize.py

@@ -9,7 +9,7 @@
 import msgpack
 
 from . import pickle
-from ..utils import has_keyword, nbytes, typename, ensure_bytes
+from ..utils import has_keyword, nbytes, typename, ensure_bytes, is_writeable
 from .compression import maybe_compress, decompress
 from .utils import (
     unpack_frames,
@@ -473,6 +473,8 @@ def replace_inner(x):
 
 def serialize_bytelist(x, **kwargs):
     header, frames = serialize(x, **kwargs)
+    if "writeable" not in header:
+        header["writeable"] = tuple(map(is_writeable, frames))
     if "lengths" not in header:
         header["lengths"] = tuple(map(nbytes, frames))
     if frames:
@@ -501,8 +503,7 @@ def deserialize_bytes(b):
     else:
         header = {}
     frames = decompress(header, frames)
-    if not any(hasattr(f, "__cuda_array_interface__") for f in frames):
-        frames = merge_frames(header, frames)
+    frames = merge_frames(header, frames)
     return deserialize(header, frames)
 
 

distributed/protocol/tests/test_numpy.py

@@ -17,7 +17,7 @@
 from distributed.protocol.pickle import HIGHEST_PROTOCOL
 from distributed.protocol.compression import maybe_compress
 from distributed.system import MEMORY_LIMIT
-from distributed.utils import tmpfile, nbytes
+from distributed.utils import ensure_bytes, tmpfile, nbytes
 from distributed.utils_test import gen_cluster
 
 
@@ -100,12 +100,16 @@ def test_dumps_serialize_numpy(x):
         np.testing.assert_equal(x, y)
 
 
-def test_dumps_numpy_writable():
+@pytest.mark.parametrize("writeable", [True, False])
+def test_dumps_numpy_writable(writeable):
     a1 = np.arange(1000)
-    a1.flags.writeable = False
-    (a2,) = loads(dumps([to_serialize(a1)]))
+    a1.flags.writeable = writeable
+    fs = dumps([to_serialize(a1)])
+    # Make all frames read-only
+    fs = list(map(ensure_bytes, fs))
+    (a2,) = loads(fs)
     assert (a1 == a2).all()
-    assert a2.flags.writeable
+    assert a2.flags.writeable == a1.flags.writeable
 
 
 @pytest.mark.parametrize(

distributed/protocol/tests/test_pandas.py

@@ -4,7 +4,15 @@
 
 from dask.dataframe.utils import assert_eq
 
-from distributed.protocol import serialize, deserialize, decompress
+from distributed.protocol import (
+    serialize,
+    deserialize,
+    decompress,
+    dumps,
+    loads,
+    to_serialize,
+)
+from distributed.utils import ensure_bytes
 
 
 dfs = [
@@ -63,10 +71,21 @@
 
 
 @pytest.mark.parametrize("df", dfs)
-def test_dumps_serialize_numpy(df):
+def test_dumps_serialize_pandas(df):
     header, frames = serialize(df)
     if "compression" in header:
         frames = decompress(header, frames)
     df2 = deserialize(header, frames)
 
     assert_eq(df, df2)
+
+
+def test_dumps_pandas_writable():
+    a1 = np.arange(1000)
+    s1 = pd.Series(a1)
+    fs = dumps([to_serialize(s1)])
+    # Make all frames read-only
+    fs = list(map(ensure_bytes, fs))
+    (s2,) = loads(fs)
+    assert (s1 == s2).all()
+    s2[...] = 0

distributed/protocol/tests/test_protocol_utils.py

@@ -1,21 +1,30 @@
 import pytest
 
 from distributed.protocol.utils import merge_frames, pack_frames, unpack_frames
-from distributed.utils import ensure_bytes
+from distributed.utils import ensure_bytes, is_writeable
 
 
 @pytest.mark.parametrize(
-    "lengths,frames",
+    "lengths,writeable,frames",
     [
-        ([3], [b"123"]),
-        ([3, 3], [b"123", b"456"]),
-        ([2, 3, 2], [b"12345", b"67"]),
-        ([5, 2], [b"123", b"45", b"67"]),
-        ([3, 4], [b"12", b"34", b"567"]),
+        ([3], [False], [b"123"]),
+        ([3], [True], [b"123"]),
+        ([3], [None], [b"123"]),
+        ([3], [False], [bytearray(b"123")]),
+        ([3], [True], [bytearray(b"123")]),
+        ([3], [None], [bytearray(b"123")]),
+        ([3, 3], [False, False], [b"123", b"456"]),
+        ([2, 3, 2], [False, True, None], [b"12345", b"67"]),
+        ([2, 3, 2], [False, True, None], [bytearray(b"12345"), bytearray(b"67")]),
+        ([5, 2], [False, True], [b"123", b"45", b"67"]),
+        ([3, 4], [None, False], [b"12", b"34", b"567"]),
     ],
 )
-def test_merge_frames(lengths, frames):
-    header = {"lengths": lengths}
+def test_merge_frames(lengths, writeable, frames):
+    header = {
+        "lengths": lengths,
+        "writeable": writeable,
+    }
     result = merge_frames(header, frames)
 
     data = b"".join(frames)
@@ -24,8 +33,8 @@ def test_merge_frames(lengths, frames):
         expected.append(data[:i])
         data = data[i:]
 
-    assert all(isinstance(f, memoryview) for f in result)
-    assert all(not f.readonly for f in result)
+    writeables = list(map(is_writeable, result))
+    assert (r == e for r, e in zip(writeables, header["writeable"]) if e is not None)
     assert list(map(ensure_bytes, result)) == expected
 
 

distributed/protocol/utils.py

@@ -51,36 +51,44 @@ def merge_frames(header, frames):
     [b'123456']
     """
     lengths = list(header["lengths"])
-    frames = list(map(memoryview, frames))
+    writeables = list(header["writeable"])
 
+    assert len(lengths) == len(writeables)
     assert sum(lengths) == sum(map(nbytes, frames))
 
-    if not all(len(f) == l for f, l in zip(frames, lengths)):
-        frames = frames[::-1]
-        lengths = lengths[::-1]
-
-        out = []
-        while lengths:
-            l = lengths.pop()
-            L = []
-            while l:
-                frame = frames.pop()
-                if nbytes(frame) <= l:
-                    L.append(frame)
-                    l -= nbytes(frame)
-                else:
-                    L.append(frame[:l])
-                    frames.append(frame[l:])
-                    l = 0
-            if len(L) == 1:  # no work necessary
-                out.append(L[0])
-            else:
-                out.append(memoryview(bytearray().join(L)))
-        frames = out
-
-    frames = [memoryview(bytearray(f)) if f.readonly else f for f in frames]
+    if all(len(f) == l for f, l in zip(frames, lengths)):
+        return [
+            (bytearray(f) if w else bytes(f)) if w == memoryview(f).readonly else f
+            for w, f in zip(header["writeable"], frames)
+        ]
 
-    return frames
+    frames = frames[::-1]
+    lengths = lengths[::-1]
+    writeables = writeables[::-1]
+
+    out = []
+    while lengths:
+        l = lengths.pop()
+        w = writeables.pop()
+        L = []
+        while l:
+            frame = frames.pop()
+            if nbytes(frame) <= l:
+                L.append(frame)
+                l -= nbytes(frame)
+            else:
+                frame = memoryview(frame)
+                L.append(frame[:l])
+                frames.append(frame[l:])
+                l = 0
+        if len(L) == 1 and w != memoryview(L[0]).readonly:  # no work necessary
+            out.extend(L)
+        elif w:
+            out.append(bytearray().join(L))
+        else:
+            out.append(bytes().join(L))
+
+    return out
 
 
 def pack_frames_prelude(frames):

distributed/utils.py

@@ -1107,6 +1107,19 @@ def nbytes(frame, _bytes_like=(bytes, bytearray)):
             return len(frame)
 
 
+def is_writeable(frame):
+    """
+    Check whether frame is writeable
+
+    Will return ``True`` if writeable, ``False`` if readonly, and
+    ``None`` if undetermined.
+    """
+    try:
+        return not memoryview(frame).readonly
+    except TypeError:
+        return None
+
+
 @contextmanager
 def time_warn(duration, text):
     start = time()