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_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_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_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_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_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 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()
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()
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
