def test_feed_setup_teardown(s, a, b): def setup(scheduler): return 1 def func(scheduler, state): assert state == 1 return 'OK' def teardown(scheduler, state): scheduler.flag = 'done' stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': dumps(func), 'setup': dumps(setup), 'teardown': dumps(teardown), 'interval': 0.01}) for i in range(5): response = yield read(stream) assert response == 'OK' close(stream) start = time() while not hasattr(s, 'flag'): yield gen.sleep(0.01) assert time() - start < 5
def test_pickle_numpy(): np = pytest.importorskip('numpy') x = np.ones(5) assert (loads(dumps(x)) == x).all() x = np.ones(5000) assert (loads(dumps(x)) == x).all()
def test_pickle_functions(): value = 1 def f(x): # closure return x + value for func in [f, lambda x: x + 1, partial(add, 1)]: assert loads(dumps(func))(1) == func(1)
def dont_test_workers_update_center(s, a, b): aa = rpc(ip=a.ip, port=a.port) response = yield aa.update_data(data={'x': dumps(1), 'y': dumps(2)}) assert response['status'] == 'OK' assert response['nbytes'] == {'x': sizeof(1), 'y': sizeof(2)} assert a.data == {'x': 1, 'y': 2} assert s.who_has == {'x': {a.address}, 'y': {a.address}} assert s.has_what[a.address] == {'x', 'y'} yield aa.delete_data(keys=['x'], close=True) assert not s.who_has['x'] assert all('x' not in s for s in c.has_what.values()) aa.close_rpc()
def test_run(s): pytest.importorskip('psutil') n = Nanny(s.ip, s.port, ncores=2, ip='127.0.0.1', loop=s.loop) yield n._start() with rpc(n.address) as nn: response = yield nn.run(function=dumps(lambda: 1)) assert response['status'] == 'OK' assert response['result'] == 1 yield n._close()
def processing(): with log_errors(): from distributed.diagnostics.scheduler import processing stream = yield connect(ip=options["host"], port=options["tcp-port"]) yield write(stream, {"op": "feed", "function": dumps(processing), "interval": 0.200}) while True: try: msg = yield read(stream) except StreamClosedError: break else: messages["processing"] = msg
def test_feed(s, a, b): def func(scheduler): return dumps(scheduler.processing) stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': dumps(func), 'interval': 0.01}) for i in range(5): response = yield read(stream) expected = s.processing assert cloudpickle.loads(response) == expected close(stream)
def test_feed(s, a, b): def func(scheduler): return dumps(dict(scheduler.worker_info)) comm = yield connect(s.address) yield comm.write({'op': 'feed', 'function': dumps(func), 'interval': 0.01}) for i in range(5): response = yield comm.read() expected = dict(s.worker_info) assert cloudpickle.loads(response) == expected yield comm.close()
def test_feed_large_bytestring(s, a, b): np = pytest.importorskip('numpy') x = np.ones(10000000) def func(scheduler): y = x return True comm = yield connect(s.address) yield comm.write({'op': 'feed', 'function': dumps(func), 'interval': 0.05}) for i in range(5): response = yield comm.read() assert response is True yield comm.close()
def test_feed_large_bytestring(s, a, b): np = pytest.importorskip('numpy') x = np.ones(10000000) def func(scheduler): y = x return True stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': dumps(func), 'interval': 0.01}) for i in range(5): response = yield read(stream) assert response == True close(stream)
def test_pickle_functions(): def make_closure(): value = 1 def f(x): # closure return x + value return f def funcs(): yield make_closure() yield (lambda x: x + 1) yield partial(add, 1) for func in funcs(): wr = weakref.ref(func) func2 = loads(dumps(func)) wr2 = weakref.ref(func2) assert func2(1) == func(1) del func, func2 assert wr() is None assert wr2() is None
def serialize_numpy_ndarray(x, context=None): if x.dtype.hasobject or (x.dtype.flags & np.core.multiarray.LIST_PICKLE): header = {"pickle": True} frames = [None] buffer_callback = lambda f: frames.append(memoryview(f)) frames[0] = pickle.dumps( x, buffer_callback=buffer_callback, protocol=(context or {}).get("pickle-protocol", None), ) return header, frames # We cannot blindly pickle the dtype as some may fail pickling, # so we have a mixture of strategies. if x.dtype.kind == "V": # Preserving all the information works best when pickling try: # Only use stdlib pickle as cloudpickle is slow when failing # (microseconds instead of nanoseconds) dt = ( 1, pickle.pickle.dumps(x.dtype, protocol=(context or {}).get( "pickle-protocol", None)), ) pickle.loads(dt[1]) # does it unpickle fine? except Exception: # dtype fails pickling => fall back on the descr if reasonable. if x.dtype.type is not np.void or x.dtype.alignment != 1: raise else: dt = (0, x.dtype.descr) else: dt = (0, x.dtype.str) # Only serialize broadcastable data for arrays with zero strided axes broadcast_to = None if 0 in x.strides: broadcast_to = x.shape strides = x.strides writeable = x.flags.writeable x = x[tuple(slice(None) if s != 0 else slice(1) for s in strides)] if not x.flags.c_contiguous and not x.flags.f_contiguous: # Broadcasting can only be done with contiguous arrays x = np.ascontiguousarray(x) x = np.lib.stride_tricks.as_strided( x, strides=[ j if i != 0 else i for i, j in zip(strides, x.strides) ], writeable=writeable, ) if not x.shape: # 0d array strides = x.strides data = x.ravel() elif x.flags.c_contiguous or x.flags.f_contiguous: # Avoid a copy and respect order when unserializing strides = x.strides data = x.ravel(order="K") else: x = np.ascontiguousarray(x) strides = x.strides data = x.ravel() if data.dtype.fields or data.dtype.itemsize > 8: data = data.view("u%d" % math.gcd(x.dtype.itemsize, 8)) try: data = data.data except ValueError: # "ValueError: cannot include dtype 'M' in a buffer" data = data.view("u%d" % math.gcd(x.dtype.itemsize, 8)).data header = { "dtype": dt, "shape": x.shape, "strides": strides, "writeable": [x.flags.writeable], } if broadcast_to is not None: header["broadcast_to"] = broadcast_to frames = [data] return header, frames
def func(scheduler): return dumps(scheduler.processing)
def func(scheduler): return dumps(dict(scheduler.worker_info))
def test_worker_bad_args(c, a, b): aa = rpc(ip=a.ip, port=a.port) bb = rpc(ip=b.ip, port=b.port) class NoReprObj(object): """ This object cannot be properly represented as a string. """ def __str__(self): raise ValueError("I have no str representation.") def __repr__(self): raise ValueError("I have no repr representation.") response = yield aa.compute(key='x', function=dumps(NoReprObj), args=dumps(()), who_has={}) assert not a.active assert response['status'] == 'OK' assert a.data['x'] assert isinstance(response['compute_start'], float) assert isinstance(response['compute_stop'], float) assert isinstance(response['thread'], Integral) def bad_func(*args, **kwargs): 1 / 0 class MockLoggingHandler(logging.Handler): """Mock logging handler to check for expected logs.""" def __init__(self, *args, **kwargs): self.reset() logging.Handler.__init__(self, *args, **kwargs) def emit(self, record): self.messages[record.levelname.lower()].append(record.getMessage()) def reset(self): self.messages = { 'debug': [], 'info': [], 'warning': [], 'error': [], 'critical': [], } hdlr = MockLoggingHandler() old_level = logger.level logger.setLevel(logging.DEBUG) logger.addHandler(hdlr) response = yield bb.compute(key='y', function=dumps(bad_func), args=dumps(['x']), kwargs=dumps({'k': 'x'}), who_has={'x': [a.address]}) assert not b.active assert response['status'] == 'error' # Make sure job died because of bad func and not because of bad # argument. assert isinstance(loads(response['exception']), ZeroDivisionError) if sys.version_info[0] >= 3: assert any('1 / 0' in line for line in pluck(3, traceback.extract_tb( loads(response['traceback']))) if line) assert hdlr.messages['warning'][0] == " Compute Failed\n" \ "Function: bad_func\n" \ "args: (< could not convert arg to str >)\n" \ "kwargs: {'k': < could not convert arg to str >}\n" assert re.match(r"^Send compute response to scheduler: y, " \ "\{.*'args': \(< could not convert arg to str >\), .*" \ "'kwargs': \{'k': < could not convert arg to str >\}.*\}", hdlr.messages['debug'][0]) or \ re.match("^Send compute response to scheduler: y, " \ "\{.*'kwargs': \{'k': < could not convert arg to str >\}, .*" \ "'args': \(< could not convert arg to str >\).*\}", hdlr.messages['debug'][0]) logger.setLevel(old_level) # Now we check that both workers are still alive. assert not a.active response = yield aa.compute(key='z', function=dumps(add), args=dumps([1, 2]), who_has={}, close=True) assert not a.active assert response['status'] == 'OK' assert a.data['z'] == 3 assert isinstance(response['compute_start'], float) assert isinstance(response['compute_stop'], float) assert isinstance(response['thread'], Integral) assert not b.active response = yield bb.compute(key='w', function=dumps(add), args=dumps([1, 2]), who_has={}, close=True) assert not b.active assert response['status'] == 'OK' assert b.data['w'] == 3 assert isinstance(response['compute_start'], float) assert isinstance(response['compute_stop'], float) assert isinstance(response['thread'], Integral) aa.close_rpc() bb.close_rpc()
def test_pickle_data(): data = [1, b'123', '123', [123], {}, set()] for d in data: assert loads(dumps(d)) == d
def test_pickle_data(): data = [1, b"123", "123", [123], {}, set()] for d in data: assert loads(dumps(d)) == d assert deserialize(*serialize(d, serializers=("pickle", ))) == d