def test_connection_pool_tls():
"""
Make sure connection args are supported.
"""
sec = tls_security()
connection_args = sec.get_connection_args('client')
listen_args = sec.get_listen_args('scheduler')
@gen.coroutine
def ping(comm, delay=0.01):
yield gen.sleep(delay)
raise gen.Return('pong')
servers = [Server({'ping': ping}) for i in range(10)]
for server in servers:
server.listen('tls://', listen_args=listen_args)
rpc = ConnectionPool(limit=5, connection_args=connection_args)
yield [rpc(s.address).ping() for s in servers[:5]]
yield [rpc(s.address).ping() for s in servers[::2]]
yield [rpc(s.address).ping() for s in servers]
assert rpc.active == 0
rpc.close()
def test_upload_file(c, 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(a.address)
bb = rpc(b.address)
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
future = c.submit(g, workers=a.address)
result = yield future
assert result == 123
yield a._close()
yield b._close()
aa.close_rpc()
bb.close_rpc()
assert not os.path.exists(os.path.join(a.local_dir, 'foobar.py'))
def test_connection_pool_remove():
@gen.coroutine
def ping(comm, delay=0.01):
yield gen.sleep(delay)
raise gen.Return('pong')
servers = [Server({'ping': ping}) for i in range(5)]
for server in servers:
server.listen(0)
rpc = ConnectionPool(limit=10)
serv = servers.pop()
yield [rpc(s.address).ping() for s in servers]
yield [rpc(serv.address).ping() for i in range(3)]
yield rpc.connect(serv.address)
assert sum(map(len, rpc.available.values())) == 6
assert sum(map(len, rpc.occupied.values())) == 1
assert rpc.active == 1
assert rpc.open == 7
rpc.remove(serv.address)
assert serv.address not in rpc.available
assert serv.address not in rpc.occupied
assert sum(map(len, rpc.available.values())) == 4
assert sum(map(len, rpc.occupied.values())) == 0
assert rpc.active == 0
assert rpc.open == 4
rpc.collect()
comm = yield rpc.connect(serv.address)
rpc.remove(serv.address)
rpc.reuse(serv.address, comm)
rpc.close()
def f(c, 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=g, key="x", args=(10,))
result = yield aa.get_data(keys=["x"])
assert result == {"x": 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_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 f(c, 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=b"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=g, key="x")
result = yield aa.get_data(keys=["x"])
assert result == {"x": 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_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 f(c, a, b):
aa = rpc(ip=a.ip, port=a.port)
bb = rpc(ip=b.ip, port=b.port)
response = yield aa.compute(key='x', function=add,
args=[1, 2], needed=[],
close=True)
assert response == b'OK'
assert a.data['x'] == 3
assert c.who_has['x'] == set([(a.ip, a.port)])
response = yield bb.compute(key='y', function=add,
args=['x', 10], needed=['x'])
assert response == b'OK'
assert b.data['y'] == 13
assert c.who_has['y'] == set([(b.ip, b.port)])
def bad_func():
1 / 0
response = yield bb.compute(key='z', function=bad_func,
args=(), needed=(), close=True)
assert response == b'error'
assert isinstance(b.data['z'], ZeroDivisionError)
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.ip, b.port)]
bb.close_streams()
yield b._close()
def test_upload_egg(c, 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)
future = c.submit(g, 10, workers=a.address)
result = yield future._result()
assert result == 10 + 1
yield a._close()
yield b._close()
aa.close_rpc()
bb.close_rpc()
assert not os.path.exists(os.path.join(a.local_dir, eggname))
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_rpc_inputs():
L = [rpc('127.0.0.1:8884'),
rpc(('127.0.0.1', 8884)),
rpc('tcp://127.0.0.1:8884'),
]
assert all(r.address == 'tcp://127.0.0.1:8884' for r in L), L
for r in L:
r.close_rpc()
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_upload_large_file(c, s, a, b):
pytest.importorskip('crick')
yield gen.sleep(0.05)
with rpc(a.address) as aa:
yield aa.upload_file(filename='myfile.dat', data=b'0' * 100000000)
yield gen.sleep(0.05)
assert a.digests['tick-duration'].components[0].max() < 0.050
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 f():
stream = yield TCPClient().connect('127.0.0.1', 8006)
cc = rpc(stream)
response = yield cc.register(address='alice', ncores=4)
assert 'alice' in c.has_what
assert c.ncores['alice'] == 4
response = yield cc.add_keys(address='alice', keys=['x', 'y'])
assert response == b'OK'
response = yield cc.register(address='bob', ncores=4)
response = yield cc.add_keys(address='bob', keys=['y', 'z'])
assert response == b'OK'
response = yield cc.who_has(keys=['x', 'y'])
assert response == {'x': set(['alice']), 'y': set(['alice', 'bob'])}
response = yield cc.remove_keys(address='bob', keys=['y'])
assert response == b'OK'
response = yield cc.has_what(keys=['alice', 'bob'])
assert response == {'alice': set(['x', 'y']), 'bob': set(['z'])}
response = yield cc.ncores()
assert response == {'alice': 4, 'bob': 4}
response = yield cc.ncores(addresses=['alice', 'charlie'])
assert response == {'alice': 4, 'charlie': None}
response = yield cc.unregister(address='alice', close=True)
assert response == b'OK'
assert 'alice' not in c.has_what
assert 'alice' not in c.ncores
c.stop()
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_nanny_worker_ports(loop):
try:
worker = Popen(['dworker', '127.0.0.1:8989', '--host', '127.0.0.1',
'--worker-port', '8788', '--nanny-port', '8789'],
stdout=PIPE, stderr=PIPE)
sched = Popen(['dscheduler', '--port', '8989'], stdout=PIPE, stderr=PIPE)
with Executor('127.0.0.1:8989', loop=loop) as e:
start = time()
while True:
d = sync(e.loop, e.scheduler.identity)
if d['workers']:
break
else:
assert time() - start < 5
sleep(0.1)
assert d['workers']['127.0.0.1:8788']['services']['nanny'] == 8789
finally:
with ignoring(Exception):
w = rpc('127.0.0.1:8789')
sync(loop, w.terminate)
with ignoring(Exception):
os.kill(sched.pid, signal.SIGINT)
with ignoring(Exception):
worker.kill()
def test_errors():
server = Server({'div': stream_div})
server.listen(0)
with rpc(('127.0.0.1', server.port)) as r:
with pytest.raises(ZeroDivisionError):
yield r.div(x=1, y=0)
def f(c, a, b):
aa = rpc(ip=a.ip, port=a.port)
bb = rpc(ip=b.ip, port=b.port)
assert not a.active
response, _ = yield aa.compute(key='x', function=add,
args=[1, 2], who_has={},
close=True)
assert not a.active
assert response == b'OK'
assert a.data['x'] == 3
assert c.who_has['x'] == set([(a.ip, a.port)])
response, info = yield bb.compute(key='y', function=add,
args=['x', 10], who_has={'x': {a.address}})
assert response == b'OK'
assert b.data['y'] == 13
assert c.who_has['y'] == set([(b.ip, b.port)])
assert info['nbytes'] == sizeof(b.data['y'])
def bad_func():
1 / 0
response, content = yield bb.compute(key='z',
function=bad_func, args=(), close=True)
assert not b.active
assert response == b'error'
assert isinstance(content['exception'], ZeroDivisionError)
if sys.version_info[0] >= 3:
assert any('1 / 0' in line
for line in pluck(3, traceback.extract_tb(content['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.ip, b.port)]
assert isinstance(b.address_string, str)
assert b.ip in b.address_string
assert str(b.port) in b.address_string
bb.close_streams()
yield b._close()
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_errors():
server = Server({'div': stream_div})
server.listen(0)
r = rpc(ip='127.0.0.1', port=server.port)
with pytest.raises(ZeroDivisionError):
yield r.div(x=1, y=0)
r.close_streams()
def test_gather(s, a, b):
b.data['x'] = 1
b.data['y'] = 2
with rpc(a.address) as aa:
resp = yield aa.gather(who_has={'x': [b.address], 'y': [b.address]})
assert resp['status'] == 'OK'
assert a.data['x'] == b.data['x']
assert a.data['y'] == b.data['y']
def f():
server = Server({})
server.listen(8887)
remote = rpc(ip='127.0.0.1', port=8887)
a = yield remote.identity()
b = yield remote.identity()
assert a['type'] == 'Server'
assert a['id'] == b['id']
def test_gather(s, a, b):
b.data['x'] = 1
b.data['y'] = 2
aa = rpc(ip=a.ip, port=a.port)
resp = yield aa.gather(who_has={'x': [b.address], 'y': [b.address]})
assert resp['status'] == 'OK'
assert a.data['x'] == b.data['x']
assert a.data['y'] == b.data['y']
def test_rpc_inputs():
L = [rpc('127.0.0.1:8887'),
rpc(b'127.0.0.1:8887'),
rpc(('127.0.0.1', 8887)),
rpc((b'127.0.0.1', 8887)),
rpc(ip='127.0.0.1', port=8887),
rpc(ip=b'127.0.0.1', port=8887),
rpc(addr='127.0.0.1:8887'),
rpc(addr=b'127.0.0.1:8887')]
assert all(r.ip == '127.0.0.1' and r.port == 8887 for r in L)
async def f():
server = Server({"echo": serialize})
await server.listen("tcp://")
with rpc(server.address) as r:
data = b"1" * 1000000
result = await r.echo(x=to_serialize(data))
assert result == {"result": data}
server.stop()
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 f():
server = Server({'echo': echo})
server.listen(8887)
data = b'0' * int(200e6) # slightly more than 100MB
conn = rpc(ip='127.0.0.1', port=8887)
result = yield conn.echo(x=data)
assert result == data
server.stop()
def test_cluster(loop):
with Cluster(scheduler_addr="127.0.0.1", scheduler_port=8786, worker_addrs=["127.0.0.1", "127.0.0.1"]) as c:
r = rpc(ip="127.0.0.1", port=8786)
result = []
while len(result) != 2:
result = loop.run_sync(r.ncores)
c.add_worker("127.0.0.1")
while len(result) != 3:
result = loop.run_sync(r.ncores)
def f():
server = Server({'echo': echo})
server.listen(8887)
data = b'0' * int(200e6) # slightly more than 100MB
conn = rpc(ip='127.0.0.1', port=8887)
result = yield conn.echo(x=data)
assert result == data
server.stop()
def check_identity(listen_arg):
server = Server({})
server.listen(listen_arg)
with rpc(server.address) as remote:
a = yield remote.identity()
b = yield remote.identity()
assert a['type'] == 'Server'
assert a['id'] == b['id']
server.stop()
def test_rpc_inputs():
L = [
rpc('127.0.0.1:8887'),
rpc(b'127.0.0.1:8887'),
rpc(('127.0.0.1', 8887)),
rpc((b'127.0.0.1', 8887)),
rpc(ip='127.0.0.1', port=8887),
rpc(ip=b'127.0.0.1', port=8887),
rpc(addr='127.0.0.1:8887'),
rpc(addr=b'127.0.0.1:8887')
]
assert all(r.ip == '127.0.0.1' and r.port == 8887 for r in L)
async def check_identity(listen_arg):
server = Server({})
await server.listen(listen_arg)
async with rpc(server.address) as remote:
a = await remote.identity()
b = await remote.identity()
assert a["type"] == "Server"
assert a["id"] == b["id"]
server.stop()
def f():
server = Server({})
server.listen(8887)
with rpc(ip='127.0.0.1', port=8887) as remote:
a = yield remote.identity()
b = yield remote.identity()
assert a['type'] == 'Server'
assert a['id'] == b['id']
server.stop()
def check_identity(listen_arg):
server = Server({})
server.listen(listen_arg)
with rpc(server.address) as remote:
a = yield remote.identity()
b = yield remote.identity()
assert a['type'] == 'Server'
assert a['id'] == b['id']
server.stop()
def check_identity(listen_arg):
server = Server({})
server.listen(listen_arg)
with rpc(server.address) as remote:
a = yield remote.identity()
b = yield remote.identity()
assert a["type"] == "Server"
assert a["id"] == b["id"]
server.stop()
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', 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 f(c, a, b):
aa = rpc(ip=a.ip, port=a.port)
bb = rpc(ip=b.ip, port=b.port)
response = yield aa.compute(key='x',
function=add,
args=[1, 2],
needed=[],
close=True)
assert response == b'OK'
assert a.data['x'] == 3
assert c.who_has['x'] == set([(a.ip, a.port)])
response = yield bb.compute(key='y',
function=add,
args=['x', 10],
needed=['x'])
assert response == b'OK'
assert b.data['y'] == 13
assert c.who_has['y'] == set([(b.ip, b.port)])
def bad_func():
1 / 0
response = yield bb.compute(key='z',
function=bad_func,
args=(),
needed=(),
close=True)
assert response == b'error'
assert isinstance(b.data['z'], ZeroDivisionError)
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.ip, b.port)]
bb.close_streams()
yield b._close()
def test_counters():
server = Server({'div': stream_div})
server.listen('tcp://')
with rpc(server.address) as r:
for i in range(2):
yield r.identity()
with pytest.raises(ZeroDivisionError):
yield r.div(x=1, y=0)
c = server.counters
assert c['op'].components[0] == {'identity': 2, 'div': 1}
async def test_counters():
server = Server({"div": stream_div})
await server.listen("tcp://")
async with rpc(server.address) as r:
for i in range(2):
await r.identity()
with pytest.raises(ZeroDivisionError):
await r.div(x=1, y=0)
c = server.counters
assert c["op"].components[0] == {"identity": 2, "div": 1}
def test_counters():
server = Server({'div': stream_div})
server.listen('tcp://')
with rpc(server.address) as r:
for i in range(2):
yield r.identity()
with pytest.raises(ZeroDivisionError):
yield r.div(x=1, y=0)
c = server.counters
assert c['op'].components[0] == {'identity': 2, 'div': 1}
def f():
server = Server({'ping': pingpong})
server.listen(8887)
remote = rpc(ip='127.0.0.1', port=8887)
response = yield remote.ping()
assert response == b'pong'
response = yield remote.ping(close=True)
assert response == b'pong'
server.stop()
def test_cluster():
with Cluster(center_addr='127.0.0.1',
center_port=8787,
worker_addrs=['127.0.0.1', '127.0.0.1']) as c:
c = rpc(ip='127.0.0.1', port=8787)
result = []
while len(result) != 2:
result = loop.run_sync(c.ncores)
c.add_worker('127.0.0.1')
while len(result) != 3:
result = loop.run_sync(c.ncores)
def test_cluster(loop):
with Cluster(scheduler_addr='127.0.0.1',
scheduler_port=8787,
worker_addrs=['127.0.0.1', '127.0.0.1']) as c:
r = rpc(ip='127.0.0.1', port=8787)
result = []
while len(result) != 2:
result = loop.run_sync(r.ncores)
c.add_worker('127.0.0.1')
while len(result) != 3:
result = loop.run_sync(r.ncores)
def f(c, a, b):
aa = rpc(ip=a.ip, port=a.port)
yield aa.update_data(data={'x': 1, 'y': 2})
assert a.data == {'x': 1, 'y': 2}
assert c.who_has == {'x': {(a.ip, a.port)}, 'y': {(a.ip, a.port)}}
assert c.has_what[(a.ip, a.port)] == {'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()
async def _start(self):
self.status: str
while self.status == "starting":
await asyncio.sleep(0.01)
if self.status == "running":
return
if self.status == "closed":
raise ValueError("Unexpected 'closed' status")
self.status = "starting"
self.scheduler_comm = rpc(
self.scheduler_address,
connection_args=self.security.get_connection_args("client"))
await super()._start()
async def test_connection_pool_tls():
"""
Make sure connection args are supported.
"""
sec = tls_security()
connection_args = sec.get_connection_args("client")
listen_args = sec.get_listen_args("scheduler")
async def ping(comm, delay=0.01):
await asyncio.sleep(delay)
return "pong"
servers = [Server({"ping": ping}) for i in range(10)]
for server in servers:
await server.listen("tls://", **listen_args)
rpc = await ConnectionPool(limit=5, connection_args=connection_args)
await asyncio.gather(*[rpc(s.address).ping() for s in servers[:5]])
await asyncio.gather(*[rpc(s.address).ping() for s in servers[::2]])
await asyncio.gather(*[rpc(s.address).ping() for s in servers])
assert rpc.active == 0
await rpc.close()
async def test_connection_pool_remove():
async def ping(comm, delay=0.01):
await asyncio.sleep(delay)
return "pong"
servers = [Server({"ping": ping}) for i in range(5)]
for server in servers:
await server.listen(0)
rpc = await ConnectionPool(limit=10)
serv = servers.pop()
await asyncio.gather(*[rpc(s.address).ping() for s in servers])
await asyncio.gather(*[rpc(serv.address).ping() for i in range(3)])
await rpc.connect(serv.address)
assert sum(map(len, rpc.available.values())) == 6
assert sum(map(len, rpc.occupied.values())) == 1
assert rpc.active == 1
assert rpc.open == 7
rpc.remove(serv.address)
assert serv.address not in rpc.available
assert serv.address not in rpc.occupied
assert sum(map(len, rpc.available.values())) == 4
assert sum(map(len, rpc.occupied.values())) == 0
assert rpc.active == 0
assert rpc.open == 4
rpc.collect()
# this pattern of calls (esp. `reuse` after `remove`)
# can happen in case of worker failures:
comm = await rpc.connect(serv.address)
rpc.remove(serv.address)
rpc.reuse(serv.address, comm)
await rpc.close()
async def check_large_packets(listen_arg):
""" tornado has a 100MB cap by default """
server = Server({"echo": echo})
await server.listen(listen_arg)
data = b"0" * int(200e6) # slightly more than 100MB
async with rpc(server.address) as conn:
result = await conn.echo(x=data)
assert result == data
d = {"x": data}
result = await conn.echo(x=d)
assert result == d
server.stop()
def f():
server = Server({'ping': pingpong})
server.listen(8887)
with rpc(ip='127.0.0.1', port=8887) as remote:
response = yield remote.ping()
assert response == b'pong'
response = yield remote.ping(close=True)
assert response == b'pong'
assert not remote.streams
assert remote.status == 'closed'
server.stop()
def from_name(cls, name, loop=None, asynchronous=False):
cluster = cls(asynchronous=asynchronous)
cluster.name = name
port = name.split("-")[-1]
cluster.scheduler_comm = rpc(f"tcp://localhost:{port}")
cluster._loop_runner = LoopRunner(loop=loop, asynchronous=asynchronous)
cluster.loop = cluster._loop_runner.loop
if not asynchronous:
cluster._loop_runner.start()
cluster.status = Status.starting
cluster.sync(cluster._start)
return cluster
def f():
server = Server({'echo': echo})
server.listen(8886)
data = b'0' * int(200e6) # slightly more than 100MB
conn = rpc('127.0.0.1:8886')
result = yield conn.echo(x=data)
assert result == data
d = {'x': data}
result = yield conn.echo(x=d)
assert result == d
conn.close_comms()
server.stop()
async def _init_internals(self):
self.scheduler_comm = rpc(
self.scheduler_address,
connection_args=self.security.get_connection_args("client"),
)
comm = None
try:
comm = await self.scheduler_comm.live_comm()
await comm.write({"op": "subscribe_worker_status"})
self.scheduler_info = await comm.read()
self._watch_worker_status_task = asyncio.ensure_future(
self._watch_worker_status(comm))
except Exception:
if comm is not None:
await comm.close()
def f(c, a, b):
aa = rpc(ip=a.ip, port=a.port)
response, content = yield aa.update_data(data={'x': 1, 'y': 2})
assert response == b'OK'
assert content['nbytes'] == {'x': sizeof(1), 'y': sizeof(2)}
assert a.data == {'x': 1, 'y': 2}
assert c.who_has == {'x': {(a.ip, a.port)}, 'y': {(a.ip, a.port)}}
assert c.has_what[(a.ip, a.port)] == {'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()
async def check_rpc_with_many_connections(listen_arg):
async def g():
for i in range(10):
await remote.ping()
server = Server({"ping": pingpong})
await server.listen(listen_arg)
async with rpc(server.address) as remote:
for i in range(10):
await g()
server.stop()
remote.close_comms()
assert all(comm.closed() for comm in remote.comms)
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 check_large_packets(listen_arg):
""" tornado has a 100MB cap by default """
server = Server({'echo': echo})
server.listen(listen_arg)
data = b'0' * int(200e6) # slightly more than 100MB
conn = rpc(server.address)
result = yield conn.echo(x=data)
assert result == data
d = {'x': data}
result = yield conn.echo(x=d)
assert result == d
conn.close_comms()
server.stop()
def check_rpc_with_many_connections(listen_arg):
@gen.coroutine
def g():
for i in range(10):
yield remote.ping()
server = Server({'ping': pingpong})
server.listen(listen_arg)
remote = rpc(server.address)
yield [g() for i in range(10)]
server.stop()
remote.close_comms()
assert all(comm.closed() for comm in remote.comms)
def dont_test_delete_data_with_missing_worker(c, a, b):
bad = '127.0.0.1:9001' # this worker doesn't exist
c.who_has['z'].add(bad)
c.who_has['z'].add(a.address)
c.has_what[bad].add('z')
c.has_what[a.address].add('z')
a.data['z'] = 5
cc = rpc(ip=c.ip, port=c.port)
yield cc.delete_data(keys=['z']) # TODO: this hangs for a while
assert 'z' not in a.data
assert not c.who_has['z']
assert not c.has_what[bad]
assert not c.has_what[a.address]
cc.close_rpc()
async def dont_test_delete_data_with_missing_worker(c, a, b):
bad = "127.0.0.1:9001" # this worker doesn't exist
c.who_has["z"].add(bad)
c.who_has["z"].add(a.address)
c.has_what[bad].add("z")
c.has_what[a.address].add("z")
a.data["z"] = 5
cc = rpc(ip=c.ip, port=c.port)
await cc.delete_data(keys=["z"]) # TODO: this hangs for a while
assert "z" not in a.data
assert not c.who_has["z"]
assert not c.has_what[bad]
assert not c.has_what[a.address]
await cc.close_rpc()
