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 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