def cluster(**kwargs):
loop = asyncio.new_event_loop()
c = Center('127.0.0.1', 8100, loop=loop)
a = Worker('127.0.0.1', 8101, c.ip, c.port, loop=loop, **kwargs)
b = Worker('127.0.0.1', 8102, c.ip, c.port, loop=loop, **kwargs)
kill_q = Queue()
@asyncio.coroutine
def stop():
while kill_q.empty():
yield from asyncio.sleep(0.01, loop=loop)
kill_q.get()
cor = asyncio.gather(c.go(), a.go(), b.go(), loop=loop)
cor2 = asyncio.wait([stop(), cor],
loop=loop,
return_when=asyncio.FIRST_COMPLETED)
thread, loop = spawn_loop(cor, loop)
try:
yield c, a, b
finally:
if a.status != 'closed':
a.close()
if b.status != 'closed':
b.close()
c.close()
kill_q.put(b'')
def test_metadata():
c = Center('127.0.0.1', 8006, loop=loop)
@asyncio.coroutine
def f():
reader, writer = yield from connect('127.0.0.1', 8006, loop=loop)
cc = rpc(reader, writer)
response = yield from cc.register(address='alice', ncores=4)
assert 'alice' in c.has_what
assert c.ncores['alice'] == 4
response = yield from cc.add_keys(address='alice', keys=['x', 'y'])
assert response == b'OK'
response = yield from cc.register(address='bob', ncores=4)
response = yield from cc.add_keys(address='bob', keys=['y', 'z'])
assert response == b'OK'
response = yield from cc.who_has(keys=['x', 'y'])
assert response == {'x': set(['alice']), 'y': set(['alice', 'bob'])}
response = yield from cc.remove_keys(address='bob', keys=['y'])
assert response == b'OK'
response = yield from cc.has_what(keys=['alice', 'bob'])
assert response == {'alice': set(['x', 'y']), 'bob': set(['z'])}
response = yield from cc.ncores()
assert response == {'alice': 4, 'bob': 4}
response = yield from cc.unregister(address='alice', close=True)
assert response == b'OK'
assert 'alice' not in c.has_what
assert 'alice' not in c.ncores
yield from c._close()
loop.run_until_complete(asyncio.gather(c.go(), f()))
def test_map_locality():
c = Center('127.0.0.1', 8017, loop=loop)
a = Worker('127.0.0.1', 8018, c.ip, c.port, loop=loop, ncores=4)
b = Worker('127.0.0.1', 8019, c.ip, c.port, loop=loop, ncores=4)
p = Pool(c.ip, c.port, loop=loop, start=False)
@asyncio.coroutine
def f():
while len(c.ncores) < 2:
yield from asyncio.sleep(0.01, loop=loop)
yield from p._sync_center()
results = yield from p._map(lambda x: x * 1000, list(range(20)))
assert p.has_what[(a.ip, a.port)].issuperset(a.data)
assert p.has_what[(b.ip, b.port)].issuperset(b.data)
s = {(a.ip, a.port), (b.ip, b.port)}
assert all(p.who_has[result.key].issubset(s) for result in results)
results2 = yield from p._map(lambda x: -x, results)
aval = set(a.data.values())
bval = set(b.data.values())
try:
assert sum(-v in aval for v in aval) > 0.8 * len(aval)
assert sum(-v in bval for v in bval) > 0.8 * len(bval)
finally:
yield from p._close_connections()
yield from a._close()
yield from b._close()
yield from c._close()
loop.run_until_complete(asyncio.gather(c.go(), a.go(), b.go(), f()))
def test_delete_data():
from distributed import Worker
c = Center('127.0.0.1', 8037, loop=loop)
a = Worker('127.0.0.1', 8038, c.ip, c.port, loop=loop)
@asyncio.coroutine
def f():
while len(c.ncores) < 1:
yield from asyncio.sleep(0.01, loop=loop)
yield from rpc(a.ip, a.port).update_data(data={'x': 1, 'y': 2})
assert a.data == {'x': 1, 'y': 2}
yield from rpc(c.ip, c.port).add_keys(address=(a.ip, a.port),
keys=['x', 'y'])
yield from rpc(c.ip, c.port).delete_data(keys=['x'])
assert a.data == {'y': 2}
assert not c.who_has['x']
assert list(c.has_what[(a.ip, a.port)]) == ['y']
yield from a._close()
yield from c._close()
loop.run_until_complete(asyncio.gather(c.go(), a.go(), f()))
def test_pool():
c = Center('127.0.0.1', 8017, loop=loop)
a = Worker('127.0.0.1', 8018, c.ip, c.port, loop=loop, ncores=1)
b = Worker('127.0.0.1', 8019, c.ip, c.port, loop=loop, ncores=1)
p = Pool(c.ip, c.port, loop=loop, start=False)
@asyncio.coroutine
def f():
yield from p._sync_center()
computation = yield from p._apply_async(add, [1, 2])
assert computation.status == b'running'
assert set(p.available_cores.values()) == set([0, 1])
x = yield from computation._get()
assert computation.status == x.status == b'success'
assert list(p.available_cores.values()) == [1, 1]
result = yield from x._get()
assert result == 3
computation = yield from p._apply_async(add, [x, 10])
y = yield from computation._get()
result = yield from y._get()
assert result == 13
assert set((len(a.data), len(b.data))) == set((0, 2))
x = yield from p._apply_async(add, [1, 2])
y = yield from p._apply_async(add, [1, 2])
assert list(p.available_cores.values()) == [0, 0]
xx = yield from x._get()
yield from xx._get()
assert set(p.available_cores.values()) == set([0, 1])
yy = yield from y._get()
yield from yy._get()
assert list(p.available_cores.values()) == [1, 1]
seq = yield from p._map(lambda x: x * 100, [1, 2, 3])
result = yield from seq[0]._get(False)
assert result == 100
result = yield from seq[1]._get(False)
assert result == 200
result = yield from seq[2]._get(True)
assert result == 300
# Handle errors gracefully
results = yield from p._map(lambda x: 3 / x, [0, 1, 2, 3])
assert all(isinstance(result, RemoteData) for result in results)
try:
yield from results[0]._get()
assert False
except ZeroDivisionError:
pass
yield from p._close_connections()
yield from a._close()
yield from b._close()
yield from c._close()
loop.run_until_complete(asyncio.gather(c.go(), a.go(), b.go(), f()))
