Skip to content

Use pickle protocol 5 with NumPy object arrays #3871

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mrocklin merged 8 commits into from
Jun 9, 2020
Merged

Use pickle protocol 5 with NumPy object arrays #3871

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
d3c1d62
Test pickle serialization with NumPy object arrays
jakirkham Jun 9, 2020
77427ef
Test equality after checking array strides
jakirkham Jun 9, 2020
6529237
Check array shape and type as well
jakirkham Jun 9, 2020
9e82ce1
Handle object array tests specially
jakirkham Jun 9, 2020
3759929
Test pickling a ragged NumPy array
jakirkham Jun 9, 2020
1487612
Use pickle protocol 5 with NumPy object arrays
jakirkham Jun 9, 2020
406dd3b
Assert ragged array components extracted as frames
jakirkham Jun 9, 2020
e07d26b
Assert additional frames when pickling
jakirkham Jun 9, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions distributed/protocol/numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from .serialize import dask_serialize, dask_deserialize
from . import pickle

from ..utils import log_errors
from ..utils import log_errors, nbytes


def itemsize(dt):
Expand All @@ -22,7 +22,10 @@ def itemsize(dt):
def serialize_numpy_ndarray(x):
if x.dtype.hasobject:
header = {"pickle": True}
frames = [pickle.dumps(x)]
frames = [None]
buffer_callback = lambda f: frames.append(memoryview(f))
Copy link
Member

@mrocklin mrocklin Jun 9, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are we converting into a memoryview here?

Copy link
Member Author

@jakirkham jakirkham Jun 9, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The buffer_callback is called on values of type PickleBuffer. Here's an example:

In [1]: import pickle                                                           

In [2]: import numpy as np                                                      

In [3]: a = np.array([np.arange(3), np.arange(4, 6)], dtype=object)             

In [4]: l = []; d = pickle.dumps(a, protocol=5, buffer_callback=l.append)       

In [5]: l                                                                       
Out[5]: [<pickle.PickleBuffer at 0x115a64740>, <pickle.PickleBuffer at 0x115a647c0>]

Though PickleBuffer is a wrapper that expects to be passed some bytes-like object, which it wraps zero-copy. In the example above these are the nested NumPy arrays.

However as the rest of the serialization code expects frames to be bytes or memoryviews. We would either need to retool it to support PickleBuffer (with special casing when it is not available) or we can simply convert it to a memoryview, which is also zero-copy.

We've gone for the latter path of converting to a memoryview similar to what was done in PR ( #3784 ). Note this is done without loss of generality.

frames[0] = pickle.dumps(x, buffer_callback=buffer_callback)
header["lengths"] = tuple(map(nbytes, frames))
return header, frames

# We cannot blindly pickle the dtype as some may fail pickling,
Expand Down Expand Up @@ -96,10 +99,10 @@ def serialize_numpy_ndarray(x):
@dask_deserialize.register(np.ndarray)
def deserialize_numpy_ndarray(header, frames):
with log_errors():
(frame,) = frames

if header.get("pickle"):
return pickle.loads(frame)
return pickle.loads(frames[0], buffers=frames[1:])

(frame,) = frames

is_custom, dt = header["dtype"]
if is_custom:
Expand Down
16 changes: 15 additions & 1 deletion distributed/protocol/tests/test_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
)
from distributed.protocol.utils import BIG_BYTES_SHARD_SIZE
from distributed.protocol.numpy import itemsize
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
Expand Down Expand Up @@ -57,6 +58,7 @@ def test_serialize():
np.array(["abc"], dtype="S3"),
np.array(["abc"], dtype="U3"),
np.array(["abc"], dtype=object),
np.array([np.arange(3), np.arange(4, 6)], dtype=object),
np.ones(shape=(5,), dtype=("f8", 32)),
np.ones(shape=(5,), dtype=[("x", "f8", 32)]),
np.ones(shape=(5,), dtype=np.dtype([("a", "i1"), ("b", "f8")], align=False)),
Expand All @@ -79,12 +81,24 @@ def test_dumps_serialize_numpy(x):
frames = decompress(header, frames)
for frame in frames:
assert isinstance(frame, (bytes, memoryview))
if x.dtype.char == "O" and any(isinstance(e, np.ndarray) for e in x.flat):
if HIGHEST_PROTOCOL >= 5:
assert len(frames) > 1
else:
assert len(frames) == 1
y = deserialize(header, frames)

np.testing.assert_equal(x, y)
assert x.shape == y.shape
assert x.dtype == y.dtype
if x.flags.c_contiguous or x.flags.f_contiguous:
assert x.strides == y.strides

if x.dtype.char == "O":
for e_x, e_y in zip(x.flat, y.flat):
np.testing.assert_equal(e_x, e_y)
else:
np.testing.assert_equal(x, y)


@pytest.mark.parametrize(
"x",
Expand Down
19 changes: 19 additions & 0 deletions distributed/protocol/tests/test_pickle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,25 @@ def test_pickle_numpy():
assert (loads(dumps(x)) == x).all()
assert (deserialize(*serialize(x, serializers=("pickle",))) == x).all()

x = np.array([np.arange(3), np.arange(4, 6)], dtype=object)
x2 = loads(dumps(x))
assert x.shape == x2.shape
assert x.dtype == x2.dtype
assert x.strides == x2.strides
for e_x, e_x2 in zip(x.flat, x2.flat):
np.testing.assert_equal(e_x, e_x2)
h, f = serialize(x, serializers=("pickle",))
if HIGHEST_PROTOCOL >= 5:
assert len(f) == 3
else:
assert len(f) == 1
x3 = deserialize(h, f)
assert x.shape == x3.shape
assert x.dtype == x3.dtype
assert x.strides == x3.strides
for e_x, e_x3 in zip(x.flat, x3.flat):
np.testing.assert_equal(e_x, e_x3)

if HIGHEST_PROTOCOL >= 5:
x = np.ones(5000)

Expand Down