def f(): s = Scheduler() s.listen(0) x = Worker(s.ip, s.port, ip='127.0.0.1') y = Worker(s.ip, s.port, ip='127.0.0.1') z = Worker(s.ip, s.port, ip='127.0.0.1') x.data['a'] = 1 y.data['a'] = 2 yield [x._start(), y._start(), z._start()] zz = rpc(ip=z.ip, port=z.port) yield zz.compute(function=dumps(inc), args=dumps(('a',)), who_has={'a': [x.address]}, key='b') assert z.data['b'] == 2 if 'a' in z.data: del z.data['a'] yield zz.compute(function=dumps(inc), args=dumps(('a',)), who_has={'a': [y.address]}, key='c') assert z.data['c'] == 3 yield [x._close(), y._close(), z._close()] zz.close_streams()
def test_update_state_supports_recomputing_released_results(s): s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={'x': 1, 'y': (inc, 'x'), 'z': (inc, 'x')}, keys=['z'], dependencies={'y': {'x'}, 'x': set(), 'z': {'y'}}, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int), compute_start=10, compute_stop=11) s.ensure_occupied(alice) s.mark_task_finished('y', alice, nbytes=10, type=dumps(int), compute_start=10, compute_stop=11) s.ensure_occupied(alice) s.mark_task_finished('z', alice, nbytes=10, type=dumps(int), compute_start=10, compute_stop=11) s.ensure_occupied(alice) assert not s.waiting assert not s.ready assert s.waiting_data == {'z': set()} assert s.who_has == {'z': {alice}} s.update_graph(tasks={'x': 1, 'y': (inc, 'x')}, keys=['y'], dependencies={'y': {'x'}}, client='client') assert s.waiting == {'y': {'x'}} assert s.waiting_data == {'x': {'y'}, 'y': set(), 'z': set()} assert s.who_wants == {'z': {'client'}, 'y': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} assert set(s.processing[alice]) == {'x'}
def test_upload_file(s, a, b): assert not os.path.exists(os.path.join(a.local_dir, 'foobar.py')) assert not os.path.exists(os.path.join(b.local_dir, 'foobar.py')) assert a.local_dir != b.local_dir aa = rpc(ip=a.ip, port=a.port) bb = rpc(ip=b.ip, port=b.port) yield [aa.upload_file(filename='foobar.py', data=b'x = 123'), bb.upload_file(filename='foobar.py', data='x = 123')] assert os.path.exists(os.path.join(a.local_dir, 'foobar.py')) assert os.path.exists(os.path.join(b.local_dir, 'foobar.py')) def g(): import foobar return foobar.x yield aa.compute(function=dumps(g), key='x') result = yield aa.get_data(keys=['x']) assert result == {'x': dumps(123)} yield a._close() yield b._close() aa.close_streams() bb.close_streams() assert not os.path.exists(os.path.join(a.local_dir, 'foobar.py'))
def test_upload_egg(s, a, b): eggname = 'mytestegg-1.0.0-py3.4.egg' local_file = __file__.replace('test_worker.py', eggname) assert not os.path.exists(os.path.join(a.local_dir, eggname)) assert not os.path.exists(os.path.join(b.local_dir, eggname)) assert a.local_dir != b.local_dir aa = rpc(ip=a.ip, port=a.port) bb = rpc(ip=b.ip, port=b.port) with open(local_file, 'rb') as f: payload = f.read() yield [aa.upload_file(filename=eggname, data=payload), bb.upload_file(filename=eggname, data=payload)] assert os.path.exists(os.path.join(a.local_dir, eggname)) assert os.path.exists(os.path.join(b.local_dir, eggname)) def g(x): import testegg return testegg.inc(x) yield aa.compute(function=dumps(g), key='x', args=dumps((10,))) result = yield aa.get_data(keys=['x']) assert result == {'x': dumps(10 + 1)} yield a._close() yield b._close() aa.close_streams() bb.close_streams() assert not os.path.exists(os.path.join(a.local_dir, eggname))
def test_update_state_respects_data_in_memory(loop): s = Scheduler() s.start(0) s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={'x': 1, 'y': (inc, 'x')}, keys=['y'], dependencies={'y': {'x'}, 'x': set()}, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int), compute_start=10, compute_stop=11) s.ensure_occupied(alice) s.mark_task_finished('y', alice, nbytes=10, type=dumps(int), compute_start=11, compute_stop=12) s.ensure_occupied(alice) assert s.released == {'x'} assert s.who_has == {'y': {alice}} s.update_graph(tasks={'x': 1, 'y': (inc, 'x'), 'z': (add, 'y', 'x')}, keys=['z'], dependencies={'y': {'x'}, 'z': {'y', 'x'}}, client='client') assert s.released == set() assert s.waiting == {'z': {'x'}} assert s.processing[alice] == {'x'} # x was released, need to recompute assert s.waiting_data == {'x': {'z'}, 'y': {'z'}, 'z': set()} assert s.who_wants == {'y': {'client'}, 'z': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} s.stop()
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' stream.close() start = time() while not hasattr(s, 'flag'): yield gen.sleep(0.01) assert time() - start < 5
def test_worker_task_bytes(s, a, b): aa = rpc(ip=a.ip, port=a.port) yield aa.compute(task=dumps((inc, 1)), key='x', serialized=True) assert a.data['x'] == 2 yield aa.compute(function=dumps(inc), args=dumps((10,)), key='y', serialized=True) assert a.data['y'] == 11
def test_scatter_no_workers(s): with pytest.raises(gen.TimeoutError): yield gen.with_timeout(timedelta(seconds=0.1), s.scatter(data={'x': dumps(1)}, client='alice')) w = Worker(s.ip, s.port, ncores=3, ip='127.0.0.1') yield [s.scatter(data={'x': dumps(1)}, client='alice'), w._start()] assert w.data['x'] == 1
def test_update_state_respects_data_in_memory(loop): s = Scheduler() s.start(0) s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={ 'x': 1, 'y': (inc, 'x') }, keys=['y'], dependencies={ 'y': {'x'}, 'x': set() }, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int), compute_start=10, compute_stop=11) s.ensure_occupied(alice) s.mark_task_finished('y', alice, nbytes=10, type=dumps(int), compute_start=11, compute_stop=12) s.ensure_occupied(alice) assert s.released == {'x'} assert s.who_has == {'y': {alice}} s.update_graph(tasks={ 'x': 1, 'y': (inc, 'x'), 'z': (add, 'y', 'x') }, keys=['z'], dependencies={ 'y': {'x'}, 'z': {'y', 'x'} }, client='client') assert s.released == set() assert s.waiting == {'z': {'x'}} assert set(s.processing[alice]) == {'x' } # x was released need to recompute assert set(s.rprocessing['x']) == {alice } # x was released need to recompute assert s.waiting_data == {'x': {'z'}, 'y': {'z'}, 'z': set()} assert s.who_wants == {'y': {'client'}, 'z': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} s.stop()
def test_worker_task_bytes(s, a, b): with rpc(ip=a.ip, port=a.port) as aa: yield aa.compute(task=dumps((inc, 1)), key='x', report=False) assert a.data['x'] == 2 yield aa.compute(function=dumps(inc), args=dumps((10, )), key='y', report=False) assert a.data['y'] == 11
def test_worker_task_bytes(s, a, b): aa = rpc(ip=a.ip, port=a.port) yield aa.compute(task=dumps((inc, 1)), key='x', serialized=True) assert a.data['x'] == 2 yield aa.compute(function=dumps(inc), args=dumps((10, )), key='y', serialized=True) assert a.data['y'] == 11
def test_update_state_with_processing(loop): s = Scheduler() s.start(0) s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={'x': 1, 'y': (inc, 'x'), 'z': (inc, 'y')}, keys=['z'], dependencies={'y': {'x'}, 'x': set(), 'z': {'y'}}, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int)) assert s.waiting == {'z': {'y'}} assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'z': set()} assert list(s.ready) == [] assert s.who_wants == {'z': {'client'}} assert s.wants_what == {'client': {'z'}} assert s.who_has == {'x': {alice}} s.update_graph(tasks={'a': (inc, 'x'), 'b': (add,'a','y'), 'c': (inc, 'z')}, keys=['b', 'c'], dependencies={'a': {'x'}, 'b': {'a', 'y'}, 'c': {'z'}}, client='client') assert s.waiting == {'z': {'y'}, 'b': {'a', 'y'}, 'c': {'z'}} assert s.stacks[alice] == ['a'] assert not s.ready assert s.waiting_data == {'x': {'y', 'a'}, 'y': {'z', 'b'}, 'z': {'c'}, 'a': {'b'}, 'b': set(), 'c': set()} assert s.who_wants == {'b': {'client'}, 'c': {'client'}, 'z': {'client'}} assert s.wants_what == {'client': {'b', 'c', 'z'}} s.stop()
def test_update_state(loop): s = Scheduler() s.start(0) s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={'x': 1, 'y': (inc, 'x')}, keys=['y'], dependencies={'y': 'x', 'x': set()}, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int)) assert s.processing[alice] == {'y'} assert not s.ready assert s.who_wants == {'y': {'client'}} assert s.wants_what == {'client': {'y'}} s.update_graph(tasks={'a': 1, 'z': (add, 'y', 'a')}, keys=['z'], dependencies={'z': {'y', 'a'}}, client='client') assert s.tasks == {'x': 1, 'y': (inc, 'x'), 'a': 1, 'z': (add, 'y', 'a')} assert s.dependencies == {'x': set(), 'a': set(), 'y': {'x'}, 'z': {'a', 'y'}} assert s.dependents == {'z': set(), 'y': {'z'}, 'a': {'z'}, 'x': {'y'}} assert s.waiting == {'z': {'a', 'y'}} assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'a': {'z'}, 'z': set()} assert s.who_wants == {'z': {'client'}, 'y': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} assert list(s.ready) == ['a'] s.stop()
def test_update_state_with_processing(s): s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={'x': 1, 'y': (inc, 'x'), 'z': (inc, 'y')}, keys=['z'], dependencies={'y': {'x'}, 'x': set(), 'z': {'y'}}, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int), compute_start=10, compute_stop=11) s.ensure_occupied(alice) assert s.waiting == {'z': {'y'}} assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'z': set()} assert list(s.ready) == [] assert s.who_wants == {'z': {'client'}} assert s.wants_what == {'client': {'z'}} assert s.who_has == {'x': {alice}} s.update_graph(tasks={'a': (inc, 'x'), 'b': (add,'a','y'), 'c': (inc, 'z')}, keys=['b', 'c'], dependencies={'a': {'x'}, 'b': {'a', 'y'}, 'c': {'z'}}, client='client') assert s.waiting == {'z': {'y'}, 'b': {'a', 'y'}, 'c': {'z'}} assert 'a' in s.stacks[alice] or 'a' in s.processing[alice] assert not s.ready assert s.waiting_data == {'x': {'y', 'a'}, 'y': {'z', 'b'}, 'z': {'c'}, 'a': {'b'}, 'b': set(), 'c': set()} assert s.who_wants == {'b': {'client'}, 'c': {'client'}, 'z': {'client'}} assert s.wants_what == {'client': {'b', 'c', 'z'}}
def f(c, 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 c.who_has == {'x': {a.address}, 'y': {a.address}} assert c.has_what[a.address] == {'x', 'y'} yield aa.delete_data(keys=['x'], close=True) assert not c.who_has['x'] assert all('x' not in s for s in c.has_what.values()) aa.close_streams()
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_streams()
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() nn = rpc(n.address) response = yield nn.run(function=dumps(lambda: 1)) assert response['status'] == 'OK' assert loads(response['result']) == 1
def test_update_state(loop): s = Scheduler() s.start(0) s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={ 'x': 1, 'y': (inc, 'x') }, keys=['y'], dependencies={ 'y': 'x', 'x': set() }, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int), compute_start=10, compute_stop=11) s.ensure_occupied(alice) assert set(s.processing[alice]) == {'y'} assert set(s.rprocessing['y']) == {alice} assert not s.ready assert s.who_wants == {'y': {'client'}} assert s.wants_what == {'client': {'y'}} s.update_graph(tasks={ 'a': 1, 'z': (add, 'y', 'a') }, keys=['z'], dependencies={'z': {'y', 'a'}}, client='client') assert s.tasks == {'x': 1, 'y': (inc, 'x'), 'a': 1, 'z': (add, 'y', 'a')} assert s.dependencies == { 'x': set(), 'a': set(), 'y': {'x'}, 'z': {'a', 'y'} } assert s.dependents == {'z': set(), 'y': {'z'}, 'a': {'z'}, 'x': {'y'}} assert s.waiting == {'z': {'a', 'y'}} assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'a': {'z'}, 'z': set()} assert s.who_wants == {'z': {'client'}, 'y': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} assert 'a' in s.ready or 'a' in s.processing[alice] s.stop()
def test_update_state_supports_recomputing_released_results(loop): s = Scheduler() s.start(0) s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={ 'x': 1, 'y': (inc, 'x'), 'z': (inc, 'x') }, keys=['z'], dependencies={ 'y': {'x'}, 'x': set(), 'z': {'y'} }, client='client') s.mark_task_finished('x', alice, nbytes=10, type=dumps(int)) s.mark_task_finished('y', alice, nbytes=10, type=dumps(int)) s.mark_task_finished('z', alice, nbytes=10, type=dumps(int)) assert not s.waiting assert not s.ready assert s.waiting_data == {'z': set()} assert s.who_has == {'z': {alice}} s.update_graph(tasks={ 'x': 1, 'y': (inc, 'x') }, keys=['y'], dependencies={'y': {'x'}}, client='client') assert s.waiting == {'y': {'x'}} assert s.waiting_data == {'x': {'y'}, 'y': set(), 'z': set()} assert s.who_wants == {'z': {'client'}, 'y': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} assert s.processing[alice] == {'x'} s.stop()
def test_update_state(s): s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={ 'x': 1, 'y': (inc, 'x') }, keys=['y'], dependencies={ 'y': 'x', 'x': set() }, client='client') s.ensure_occupied() r = s.transition('x', 'memory', nbytes=10, type=dumps(int), compute_start=10, compute_stop=11, worker=alice) s.transitions(r) s.ensure_occupied() assert set(s.processing[alice]) == {'y'} assert set(s.rprocessing['y']) == {alice} assert not s.ready assert s.who_wants == {'y': {'client'}} assert s.wants_what == {'client': {'y'}} s.update_graph(tasks={ 'a': 1, 'z': (add, 'y', 'a') }, keys=['z'], dependencies={'z': {'y', 'a'}}, client='client') assert s.tasks == {'x': 1, 'y': (inc, 'x'), 'a': 1, 'z': (add, 'y', 'a')} assert s.dependencies == { 'x': set(), 'a': set(), 'y': {'x'}, 'z': {'a', 'y'} } assert s.dependents == {'z': set(), 'y': {'z'}, 'a': {'z'}, 'x': {'y'}} assert s.waiting == {'z': {'a', 'y'}} assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'a': {'z'}, 'z': set()} assert s.who_wants == {'z': {'client'}, 'y': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} assert 'a' in s.ready or 'a' in s.processing[alice]
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 loads(response['result']) == 1 yield n._close()
def test_nanny_process_failure(): c = Center('127.0.0.1') c.listen(0) n = Nanny(c.ip, c.port, ncores=2, ip='127.0.0.1') yield n._start() nn = rpc(ip=n.ip, port=n.port) first_dir = n.worker_dir assert os.path.exists(first_dir) ww = rpc(ip=n.ip, port=n.worker_port) yield ww.update_data(data=valmap(dumps, {'x': 1, 'y': 2})) with ignoring(StreamClosedError): yield ww.compute(function=dumps(sys.exit), args=dumps((0,)), key='z') start = time() while n.process.is_alive(): # wait while process dies yield gen.sleep(0.01) assert time() - start < 2 start = time() while not n.process.is_alive(): # wait while process comes back yield gen.sleep(0.01) assert time() - start < 2 start = time() while n.worker_address not in c.ncores or n.worker_dir is None: yield gen.sleep(0.01) assert time() - start < 2 second_dir = n.worker_dir yield n._close() assert not os.path.exists(second_dir) assert not os.path.exists(first_dir) assert first_dir != n.worker_dir nn.close_streams() c.stop()
def f(c, a, b): aa = rpc(ip=a.ip, port=a.port) bb = rpc(ip=b.ip, port=b.port) result = yield aa.identity() assert not a.active response = yield aa.compute(key='x', function=dumps(add), args=dumps([1, 2]), who_has={}, close=True) assert not a.active assert response['status'] == 'OK' assert a.data['x'] == 3 assert c.who_has['x'] == {a.address} assert isinstance(response['compute-start'], float) assert isinstance(response['compute-stop'], float) assert isinstance(response['thread'], int) response = yield bb.compute(key='y', function=dumps(add), args=dumps(['x', 10]), who_has={'x': [a.address]}) assert response['status'] == 'OK' assert b.data['y'] == 13 assert c.who_has['y'] == {b.address} assert response['nbytes'] == sizeof(b.data['y']) assert isinstance(response['transfer-start'], float) assert isinstance(response['transfer-stop'], float) def bad_func(): 1 / 0 response = yield bb.compute(key='z', function=dumps(bad_func), args=dumps(()), close=True) assert not b.active assert response['status'] == 'error' 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) aa.close_streams() yield a._close() assert a.address not in c.ncores and b.address in c.ncores assert list(c.ncores.keys()) == [b.address] assert isinstance(b.address, str) assert b.ip in b.address assert str(b.port) in b.address bb.close_streams() yield b._close()
def test_nanny_process_failure(s): n = Nanny(s.ip, s.port, ncores=2, ip='127.0.0.1', loop=s.loop) yield n._start() nn = rpc(ip=n.ip, port=n.port) first_dir = n.worker_dir assert os.path.exists(first_dir) original_process = n.process ww = rpc(ip=n.ip, port=n.worker_port) yield ww.update_data(data=valmap(dumps, {'x': 1, 'y': 2})) with ignoring(StreamClosedError): yield ww.compute(function=dumps(sys.exit), args=dumps((0,)), key='z') start = time() while n.process is original_process: # wait while process dies yield gen.sleep(0.01) assert time() - start < 5 start = time() while not n.process.poll() is None: # wait while process comes back yield gen.sleep(0.01) assert time() - start < 5 start = time() while n.worker_address not in s.ncores or n.worker_dir is None: yield gen.sleep(0.01) assert time() - start < 5 second_dir = n.worker_dir yield n._close() assert not os.path.exists(second_dir) assert not os.path.exists(first_dir) assert first_dir != n.worker_dir nn.close_streams() s.stop()
def test_update_state_supports_recomputing_released_results(s): s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={'x': 1, 'y': dumps_task((inc, 'x')), 'z': dumps_task((inc, 'x'))}, keys=['z'], dependencies={'y': {'x'}, 'x': set(), 'z': {'y'}}, client='client') s.ensure_occupied() r = s.transition('x', 'memory', nbytes=10, type=dumps(int), compute_start=10, compute_stop=11, worker=alice) s.transitions(r) s.ensure_occupied() r = s.transition('y', 'memory', nbytes=10, type=dumps(int), compute_start=10, compute_stop=11, worker=alice) s.transitions(r) s.ensure_occupied() r = s.transition('z', 'memory', nbytes=10, type=dumps(int), compute_start=10, compute_stop=11, worker=alice) s.transitions(r) s.ensure_occupied() assert not s.waiting assert not s.ready assert s.waiting_data == {'z': set()} assert s.who_has == {'z': {alice}} s.update_graph(tasks={'x': 1, 'y': dumps_task((inc, 'x'))}, keys=['y'], dependencies={'y': {'x'}, 'x': set()}, client='client') assert s.waiting == {'y': {'x'}} assert s.waiting_data == {'x': {'y'}, 'y': set(), 'z': set()} assert s.who_wants == {'z': {'client'}, 'y': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} assert set(s.processing[alice]) == {'x'}
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 f(c, a, b): aa = rpc(ip=a.ip, port=a.port) bb = rpc(ip=b.ip, port=b.port) result = yield aa.identity() assert not a.active response = yield aa.compute(key='x', function=dumps(add), args=dumps([1, 2]), who_has={}, close=True) assert not a.active assert response['status'] == 'OK' assert a.data['x'] == 3 assert c.who_has['x'] == {a.address} assert isinstance(response['compute-start'], float) assert isinstance(response['compute-stop'], float) assert isinstance(response['thread'], int) response = yield bb.compute(key='y', function=dumps(add), args=dumps(['x', 10]), who_has={'x': [a.address]}) assert response['status'] == 'OK' assert b.data['y'] == 13 assert c.who_has['y'] == {b.address} assert response['nbytes'] == sizeof(b.data['y']) assert isinstance(response['transfer-start'], float) assert isinstance(response['transfer-stop'], float) def bad_func(): 1 / 0 response = yield bb.compute(key='z', function=dumps(bad_func), args=dumps(()), close=True) assert not b.active assert response['status'] == 'error' 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) aa.close_streams() yield a._close() assert a.address not in c.ncores and b.address in c.ncores assert list(c.ncores.keys()) == [b.address] assert isinstance(b.address, str) assert b.ip in b.address assert str(b.port) in b.address bb.close_streams() yield b._close()
def test_feed(s, a, b): def func(scheduler): return dumps((scheduler.processing, scheduler.stacks)) 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, s.stacks assert cloudpickle.loads(response) == expected stream.close()
def test_compute_stream(s, a, b): stream = yield connect(a.ip, a.port) yield write(stream, {'op': 'compute-stream'}) msgs = [{'op': 'compute-task', 'function': dumps(inc), 'args': dumps((i,)), 'key': 'x-%d' % i} for i in range(10)] bstream = BatchedStream(stream, 0) for msg in msgs[:5]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' for msg in msgs[5:]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' yield write(stream, {'op': 'close'})
def test_nanny_process_failure(s): n = Nanny(s.ip, s.port, ncores=2, ip='127.0.0.1', loop=s.loop) yield n._start() first_dir = n.worker_dir assert os.path.exists(first_dir) original_process = n.process ww = rpc(ip=n.ip, port=n.worker_port) yield ww.update_data(data=valmap(dumps, {'x': 1, 'y': 2})) with ignoring(StreamClosedError): yield ww.compute(function=dumps(sys.exit), args=dumps((0, )), key='z') start = time() while n.process is original_process: # wait while process dies yield gen.sleep(0.01) assert time() - start < 5 start = time() while not isalive(n.process): # wait while process comes back yield gen.sleep(0.01) assert time() - start < 5 start = time() while n.worker_address not in s.ncores or n.worker_dir is None: yield gen.sleep(0.01) assert time() - start < 5 second_dir = n.worker_dir yield n._close() assert not os.path.exists(second_dir) assert not os.path.exists(first_dir) assert first_dir != n.worker_dir ww.close_rpc() s.stop()
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_update_state(s): s.add_worker(address=alice, ncores=1, coerce_address=False) s.update_graph(tasks={'x': 1, 'y': (inc, 'x')}, keys=['y'], dependencies={'y': 'x', 'x': set()}, client='client') s.ensure_occupied() r = s.transition('x', 'memory', nbytes=10, type=dumps(int), compute_start=10, compute_stop=11, worker=alice) s.transitions(r) s.ensure_occupied() assert set(s.processing[alice]) == {'y'} assert set(s.rprocessing['y']) == {alice} assert not s.ready assert s.who_wants == {'y': {'client'}} assert s.wants_what == {'client': {'y'}} s.update_graph(tasks={'a': 1, 'z': (add, 'y', 'a')}, keys=['z'], dependencies={'z': {'y', 'a'}}, client='client') assert s.tasks == {'x': 1, 'y': (inc, 'x'), 'a': 1, 'z': (add, 'y', 'a')} assert s.dependencies == {'x': set(), 'a': set(), 'y': {'x'}, 'z': {'a', 'y'}} assert s.dependents == {'z': set(), 'y': {'z'}, 'a': {'z'}, 'x': {'y'}} assert s.waiting == {'z': {'a', 'y'}} assert s.waiting_data == {'x': {'y'}, 'y': {'z'}, 'a': {'z'}, 'z': set()} assert s.who_wants == {'z': {'client'}, 'y': {'client'}} assert s.wants_what == {'client': {'y', 'z'}} assert 'a' in s.ready or 'a' in s.processing[alice]
def func(scheduler): return dumps((scheduler.processing, scheduler.stacks))
def test_worker_task_data(s, a, b): aa = rpc(ip=a.ip, port=a.port) yield aa.compute(task=dumps(2), key='x', report=False) assert a.data['x'] == 2
def test_worker_task(s, a, b): with rpc(ip=a.ip, port=a.port) as aa: yield aa.compute(task=dumps((inc, 1)), key='x', report=False) assert a.data['x'] == 2
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_streams() bb.close_streams()