def test_serializers():
with echo_server() as e:
comm = yield connect(e.address)
b = BatchedSend(interval='10ms', serializers=['msgpack'])
b.start(comm)
b.send({'x': to_serialize(123)})
b.send({'x': to_serialize('hello')})
yield gen.sleep(0.100)
b.send({'x': to_serialize(lambda x: x + 1)})
with captured_logger('distributed.protocol') as sio:
yield gen.sleep(0.100)
value = sio.getvalue()
assert 'serialize' in value
assert 'type' in value
assert 'function' in value
msg = yield comm.read()
assert list(msg) == [{'x': 123}, {'x': 'hello'}]
with pytest.raises(gen.TimeoutError):
msg = yield gen.with_timeout(timedelta(milliseconds=100), comm.read())
def f():
server = Server({'echo': echo_serialize})
server.listen('tcp://')
with rpc(server.address, serializers=['msgpack']) as r:
with pytest.raises(TypeError):
yield r.echo(x=to_serialize(inc))
with rpc(server.address, serializers=['msgpack', 'pickle']) as r:
result = yield r.echo(x=to_serialize(inc))
assert result == {'result': inc}
server.stop()
def test_nested_deserialize():
x = {'op': 'update',
'x': [to_serialize(123), to_serialize(456), 789],
'y': {'a': ['abc', Serialized(*serialize('def'))],
'b': b'ghi'}
}
x_orig = copy.deepcopy(x)
assert nested_deserialize(x) == {'op': 'update',
'x': [123, 456, 789],
'y': {'a': ['abc', 'def'],
'b': b'ghi'}
}
assert x == x_orig # x wasn't mutated
def check_deserialize_roundtrip(addr):
"""
Sanity check round-tripping with "deserialize" on and off.
"""
# Test with long bytestrings, large enough to be transferred
# as a separate payload
_uncompressible = os.urandom(1024 ** 2) * 4 # end size: 4 MB
msg = {'op': 'update',
'x': _uncompressible,
'to_ser': [to_serialize(_uncompressible)],
'ser': Serialized(*serialize(_uncompressible)),
}
for should_deserialize in (True, False):
a, b = yield get_comm_pair(addr, deserialize=should_deserialize)
yield a.write(msg)
got = yield b.read()
yield b.write(got)
got = yield a.read()
assert sorted(got) == sorted(msg)
for k in ('op', 'x'):
assert got[k] == msg[k]
if should_deserialize:
assert isinstance(got['to_ser'][0], (bytes, bytearray))
assert isinstance(got['ser'], (bytes, bytearray))
else:
assert isinstance(got['to_ser'][0], (to_serialize, Serialized))
assert isinstance(got['ser'], Serialized)
def check_out(deserialize_flag, out_value):
# Check output with deserialize=False
assert sorted(out_value) == sorted(msg_orig)
out_value = out_value.copy() # in case transport passed the object as-is
to_ser = out_value.pop('to_ser')
ser = out_value.pop('ser')
expected_msg = msg_orig.copy()
del expected_msg['ser']
del expected_msg['to_ser']
assert out_value == expected_msg
if deserialize_flag:
assert isinstance(ser, (bytes, bytearray))
assert bytes(ser) == _uncompressible
else:
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == _uncompressible
assert isinstance(to_ser, list)
to_ser, = to_ser
# The to_serialize() value could have been actually serialized
# or not (it's a transport-specific optimization)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == _uncompressible
else:
assert to_ser == to_serialize(_uncompressible)
def run_traffic_jam(nsends, nbytes):
# This test eats `nsends * nbytes` bytes in RAM
np = pytest.importorskip('numpy')
from distributed.protocol import to_serialize
data = bytes(np.random.randint(0, 255, size=(nbytes,)).astype('u1').data)
with echo_server() as e:
comm = yield connect(e.address)
b = BatchedSend(interval=0.01)
b.start(comm)
msg = {'x': to_serialize(data)}
for i in range(nsends):
b.send(assoc(msg, 'i', i))
if np.random.random() > 0.5:
yield gen.sleep(0.001)
results = []
count = 0
while len(results) < nsends:
# If this times out then I think it's a backpressure issue
# Somehow we're able to flood the socket so that the receiving end
# loses some of our messages
L = yield gen.with_timeout(timedelta(seconds=5), comm.read())
count += 1
results.extend(r['i'] for r in L)
assert count == b.batch_count == e.count
assert b.message_count == nsends
assert results == list(range(nsends))
comm.close() # external closing
yield b.close()
def check_rpc_message_lifetime(*listen_args):
# Issue #956: rpc arguments and result shouldn't be kept alive longer
# than necessary
server = Server({'echo': echo_serialize})
server.listen(*listen_args)
# Sanity check
obj = CountedObject()
assert CountedObject.n_instances == 1
del obj
assert CountedObject.n_instances == 0
with rpc(server.address) as remote:
obj = CountedObject()
res = yield remote.echo(x=to_serialize(obj))
assert isinstance(res['result'], CountedObject)
# Make sure resource cleanup code in coroutines runs
yield gen.sleep(0.05)
w1 = weakref.ref(obj)
w2 = weakref.ref(res['result'])
del obj, res
assert w1() is None
assert w2() is None
# If additional instances were created, they were deleted as well
assert CountedObject.n_instances == 0
server.stop()
def test_err_on_bad_deserializer():
frames = yield to_frames({'x': to_serialize(1234)},
serializers=['pickle'])
result = yield from_frames(frames, deserializers=['pickle', 'foo'])
assert result == {'x': 1234}
with pytest.raises(TypeError) as info:
yield from_frames(frames, deserializers=['msgpack'])
def f():
server = Server({'echo': serialize})
server.listen('tcp://')
with rpc(server.address) as r:
data = b'1' * 1000000
result = yield r.echo(x=to_serialize(data))
assert result == {'result': data}
server.stop()
def test_errors():
msg = {'data': {'foo': to_serialize(inc)}}
header, frames = serialize(msg, serializers=['msgpack', 'pickle'])
assert header['serializer'] == 'pickle'
header, frames = serialize(msg, serializers=['msgpack'])
assert header['serializer'] == 'error'
with pytest.raises(TypeError):
serialize(msg, serializers=['msgpack'], on_error='raise')
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
header = msgpack.loads(frames[2], raw=False, use_list=False)
try:
import blosc # noqa: F401
except ImportError:
pass
else:
assert all(c == 'blosc' for c in
header['headers'][('x',)]['compression'])
def test_dumps_serialize_numpy_large():
psutil = pytest.importorskip('psutil')
if psutil.virtual_memory().total < 2e9:
return
x = np.random.random(size=int(BIG_BYTES_SHARD_SIZE * 2 // 8)).view('u1')
assert x.nbytes == BIG_BYTES_SHARD_SIZE * 2
frames = dumps([to_serialize(x)])
dtype, shape = x.dtype, x.shape
checksum = crc32(x)
del x
[y] = loads(frames)
assert (y.dtype, y.shape) == (dtype, shape)
assert crc32(y) == checksum, "Arrays are unequal"
def test_serialize_deserialize_model():
model = keras.models.Sequential()
model.add(keras.layers.Dense(5, input_dim=3))
model.add(keras.layers.Dense(2))
model.compile(optimizer='sgd', loss='mse')
x = np.random.random((1, 3))
y = np.random.random((1, 2))
model.train_on_batch(x, y)
loaded = deserialize(*serialize(model))
assert_allclose(loaded.predict(x), model.predict(x))
data = {'model': to_serialize(model)}
frames = dumps(data)
result = loads(frames)
assert_allclose(result['model'].predict(x), model.predict(x))
def test_loads_without_deserialization_avoids_compression():
pytest.importorskip('lz4')
b = b'0' * 100000
msg = {'x': 1, 'data': to_serialize(b)}
frames = dumps(msg)
assert sum(map(nbytes, frames)) < 10000
msg2 = loads(frames, deserialize=False)
assert sum(map(nbytes, msg2['data'].frames)) < 10000
msg3 = dumps(msg2)
msg4 = loads(msg3)
assert msg4 == {'x': 1, 'data': b'0' * 100000}
def check_out_false(out_value):
# Check output with deserialize=False
out_value = out_value.copy() # in case transport passed the object as-is
to_ser = out_value.pop('to_ser')
ser = out_value.pop('ser')
expected_msg = msg_orig.copy()
del expected_msg['ser']
del expected_msg['to_ser']
assert out_value == expected_msg
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == 456
assert isinstance(to_ser, list)
to_ser, = to_ser
# The to_serialize() value could have been actually serialized
# or not (it's a transport-specific optimization)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == 123
else:
assert to_ser == to_serialize(123)
def handle_comm(comm):
yield comm.write({'data': to_serialize(_EOFRaising())})
with pytest.raises(CommClosedError):
yield comm.read()
async def handle_comm(comm):
await comm.write({"data": to_serialize(_EOFRaising())})
with pytest.raises(CommClosedError):
await comm.read()
def test_worker_task(s, a, b):
with rpc(ip=a.ip, port=a.port) as aa:
yield aa.compute(task=to_serialize((inc, 1)), key='x', report=False)
assert a.data['x'] == 2
def echo_serialize(comm, x):
return {'result': to_serialize(x)}
def echo_serialize(comm, x):
return {"result": to_serialize(x)}
def test_empty_loads():
from distributed.protocol import loads, dumps
e = Empty()
e2 = loads(dumps([to_serialize(e)]))
assert isinstance(e2[0], Empty)
def test_serialize_numpy_ma_masked_array(x):
y, = loads(dumps([to_serialize(x)]))
assert x.data.dtype == y.data.dtype
np.testing.assert_equal(x.data, y.data)
np.testing.assert_equal(x.mask, y.mask)
np.testing.assert_equal(x.fill_value, y.fill_value)
def check_deserialize(addr):
# Test with Serialize and Serialized objects
msg = {
'op': 'update',
'x': b'abc',
'to_ser': [to_serialize(123)],
'ser': Serialized(*serialize(456)),
}
msg_orig = msg.copy()
def check_out_false(out_value):
# Check output with deserialize=False
out_value = out_value.copy(
) # in case transport passed the object as-is
to_ser = out_value.pop('to_ser')
ser = out_value.pop('ser')
expected_msg = msg_orig.copy()
del expected_msg['ser']
del expected_msg['to_ser']
assert out_value == expected_msg
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == 456
assert isinstance(to_ser, list)
to_ser, = to_ser
# The to_serialize() value could have been actually serialized
# or not (it's a transport-specific optimization)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == 123
else:
assert to_ser == to_serialize(123)
def check_out_true(out_value):
# Check output with deserialize=True
expected_msg = msg.copy()
expected_msg['ser'] = 456
expected_msg['to_ser'] = [123]
assert out_value == expected_msg
yield check_listener_deserialize(addr, False, msg, check_out_false)
yield check_connector_deserialize(addr, False, msg, check_out_false)
yield check_listener_deserialize(addr, True, msg, check_out_true)
yield check_connector_deserialize(addr, True, msg, check_out_true)
# Test with a long bytestring
msg = {
'op': 'update',
'x': b'abc',
'y': b'def\n' * (3 * 1024**2), # end size: 12 MB
}
msg_orig = msg.copy()
def check_out(out_value):
assert out_value == msg_orig
yield check_listener_deserialize(addr, False, msg, check_out)
yield check_connector_deserialize(addr, False, msg, check_out)
yield check_listener_deserialize(addr, True, msg, check_out)
yield check_connector_deserialize(addr, True, msg, check_out)
def test_empty_loads_deep():
from distributed.protocol import loads, dumps
e = Empty()
e2 = loads(dumps([[[to_serialize(e)]]]))
assert isinstance(e2[0][0][0], Empty)
def check_deserialize(addr):
"""
Check the "deserialize" flag on connect() and listen().
"""
# Test with Serialize and Serialized objects
msg = {'op': 'update',
'x': b'abc',
'to_ser': [to_serialize(123)],
'ser': Serialized(*serialize(456)),
}
msg_orig = msg.copy()
def check_out_false(out_value):
# Check output with deserialize=False
out_value = out_value.copy() # in case transport passed the object as-is
to_ser = out_value.pop('to_ser')
ser = out_value.pop('ser')
expected_msg = msg_orig.copy()
del expected_msg['ser']
del expected_msg['to_ser']
assert out_value == expected_msg
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == 456
assert isinstance(to_ser, list)
to_ser, = to_ser
# The to_serialize() value could have been actually serialized
# or not (it's a transport-specific optimization)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == 123
else:
assert to_ser == to_serialize(123)
def check_out_true(out_value):
# Check output with deserialize=True
expected_msg = msg.copy()
expected_msg['ser'] = 456
expected_msg['to_ser'] = [123]
assert out_value == expected_msg
yield check_listener_deserialize(addr, False, msg, check_out_false)
yield check_connector_deserialize(addr, False, msg, check_out_false)
yield check_listener_deserialize(addr, True, msg, check_out_true)
yield check_connector_deserialize(addr, True, msg, check_out_true)
# Test with long bytestrings, large enough to be 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)],
'ser': Serialized(*serialize(_uncompressible)),
}
msg_orig = msg.copy()
def check_out(deserialize_flag, out_value):
# Check output with deserialize=False
assert sorted(out_value) == sorted(msg_orig)
out_value = out_value.copy() # in case transport passed the object as-is
to_ser = out_value.pop('to_ser')
ser = out_value.pop('ser')
expected_msg = msg_orig.copy()
del expected_msg['ser']
del expected_msg['to_ser']
assert out_value == expected_msg
if deserialize_flag:
assert isinstance(ser, (bytes, bytearray))
assert bytes(ser) == _uncompressible
else:
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == _uncompressible
assert isinstance(to_ser, list)
to_ser, = to_ser
# The to_serialize() value could have been actually serialized
# or not (it's a transport-specific optimization)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == _uncompressible
else:
assert to_ser == to_serialize(_uncompressible)
yield check_listener_deserialize(addr, False, msg, partial(check_out, False))
yield check_connector_deserialize(addr, False, msg, partial(check_out, False))
yield check_listener_deserialize(addr, True, msg, partial(check_out, True))
yield check_connector_deserialize(addr, True, msg, partial(check_out, True))
def echo_serialize(comm, x):
return {'result': to_serialize(x)}
def handle_comm(comm):
yield comm.write({'data': to_serialize(_EOFRaising())})
with pytest.raises(CommClosedError):
yield comm.read()
def check_deserialize(addr):
"""
Check the "deserialize" flag on connect() and listen().
"""
# Test with Serialize and Serialized objects
msg = {'op': 'update',
'x': b'abc',
'to_ser': [to_serialize(123)],
'ser': Serialized(*serialize(456)),
}
msg_orig = msg.copy()
def check_out_false(out_value):
# Check output with deserialize=False
out_value = out_value.copy() # in case transport passed the object as-is
to_ser = out_value.pop('to_ser')
ser = out_value.pop('ser')
expected_msg = msg_orig.copy()
del expected_msg['ser']
del expected_msg['to_ser']
assert out_value == expected_msg
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == 456
assert isinstance(to_ser, list)
to_ser, = to_ser
# The to_serialize() value could have been actually serialized
# or not (it's a transport-specific optimization)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == 123
else:
assert to_ser == to_serialize(123)
def check_out_true(out_value):
# Check output with deserialize=True
expected_msg = msg.copy()
expected_msg['ser'] = 456
expected_msg['to_ser'] = [123]
assert out_value == expected_msg
yield check_listener_deserialize(addr, False, msg, check_out_false)
yield check_connector_deserialize(addr, False, msg, check_out_false)
yield check_listener_deserialize(addr, True, msg, check_out_true)
yield check_connector_deserialize(addr, True, msg, check_out_true)
# Test with long bytestrings, large enough to be 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)],
'ser': Serialized(*serialize(_uncompressible)),
}
msg_orig = msg.copy()
def check_out(deserialize_flag, out_value):
# Check output with deserialize=False
assert sorted(out_value) == sorted(msg_orig)
out_value = out_value.copy() # in case transport passed the object as-is
to_ser = out_value.pop('to_ser')
ser = out_value.pop('ser')
expected_msg = msg_orig.copy()
del expected_msg['ser']
del expected_msg['to_ser']
assert out_value == expected_msg
if deserialize_flag:
assert isinstance(ser, (bytes, bytearray))
assert bytes(ser) == _uncompressible
else:
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == _uncompressible
assert isinstance(to_ser, list)
to_ser, = to_ser
# The to_serialize() value could have been actually serialized
# or not (it's a transport-specific optimization)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == _uncompressible
else:
assert to_ser == to_serialize(_uncompressible)
yield check_listener_deserialize(addr, False, msg, partial(check_out, False))
yield check_connector_deserialize(addr, False, msg, partial(check_out, False))
yield check_listener_deserialize(addr, True, msg, partial(check_out, True))
yield check_connector_deserialize(addr, True, msg, partial(check_out, True))
def test_serialize_numpy_ma_masked():
y, = loads(dumps([to_serialize(np.ma.masked)]))
assert y is np.ma.masked
def test_serialize_numpy_ma_masked_array(x):
(y,) = loads(dumps([to_serialize(x)]))
assert x.data.dtype == y.data.dtype
np.testing.assert_equal(x.data, y.data)
np.testing.assert_equal(x.mask, y.mask)
np.testing.assert_equal(x.fill_value, y.fill_value)
async def check_deserialize(addr):
"""
Check the "deserialize" flag on connect() and listen().
"""
# Test with Serialize and Serialized objects
msg = {
"op": "update",
"x": b"abc",
"to_ser": [to_serialize(123)],
"ser": Serialized(*serialize(456)),
}
msg_orig = msg.copy()
def check_out_false(out_value):
# Check output with deserialize=False
out_value = out_value.copy(
) # in case transport passed the object as-is
to_ser = out_value.pop("to_ser")
ser = out_value.pop("ser")
expected_msg = msg_orig.copy()
del expected_msg["ser"]
del expected_msg["to_ser"]
assert out_value == expected_msg
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == 456
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)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header, to_ser.frames) == 123
else:
assert to_ser == to_serialize(123)
def check_out_true(out_value):
# Check output with deserialize=True
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)
await check_connector_deserialize(addr, False, msg, check_out_false)
await check_listener_deserialize(addr, True, msg, check_out_true)
await check_connector_deserialize(addr, True, msg, check_out_true)
# 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), ),
"ser": Serialized(*serialize(_uncompressible)),
}
msg_orig = msg.copy()
def check_out(deserialize_flag, out_value):
# Check output with deserialize=False
assert sorted(out_value) == sorted(msg_orig)
out_value = out_value.copy(
) # in case transport passed the object as-is
to_ser = out_value.pop("to_ser")
ser = out_value.pop("ser")
expected_msg = msg_orig.copy()
del expected_msg["ser"]
del expected_msg["to_ser"]
assert out_value == expected_msg
if deserialize_flag:
assert isinstance(ser, (bytes, bytearray))
assert bytes(ser) == _uncompressible
else:
assert isinstance(ser, Serialized)
assert deserialize(ser.header, ser.frames) == _uncompressible
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)
if isinstance(to_ser, Serialized):
assert deserialize(to_ser.header,
to_ser.frames) == _uncompressible
else:
assert to_ser == to_serialize(_uncompressible)
await check_listener_deserialize(addr, False, msg,
partial(check_out, False))
await check_connector_deserialize(addr, False, msg,
partial(check_out, False))
await check_listener_deserialize(addr, True, msg, partial(check_out, True))
await check_connector_deserialize(addr, True, msg,
partial(check_out, True))
def test_serialize_numpy_ma_masked():
(y,) = loads(dumps([to_serialize(np.ma.masked)]))
assert y is np.ma.masked
