def test_publish_bag(s, a, b):
db = pytest.importorskip('dask.bag')
c = Client((s.ip, s.port), start=False)
yield c._start()
f = Client((s.ip, s.port), start=False)
yield f._start()
bag = db.from_sequence([0, 1, 2])
bagp = c.persist(bag)
assert len(futures_of(bagp)) == 3
keys = {f.key for f in futures_of(bagp)}
assert keys == set(bag.dask)
yield c._publish_dataset(data=bagp)
# check that serialization didn't affect original bag's dask
assert len(futures_of(bagp)) == 3
result = yield f._get_dataset('data')
assert set(result.dask.keys()) == set(bagp.dask.keys())
assert {f.key for f in result.dask.values()} == {f.key for f in bagp.dask.values()}
out = yield f.compute(result)._result()
assert out == [0, 1, 2]
yield c._shutdown()
yield f._shutdown()
def test_publish_bag(s, a, b):
db = pytest.importorskip('dask.bag')
c = Client((s.ip, s.port), start=False)
yield c._start()
f = Client((s.ip, s.port), start=False)
yield f._start()
bag = db.from_sequence([0, 1, 2])
bagp = c.persist(bag)
assert len(futures_of(bagp)) == 3
keys = {f.key for f in futures_of(bagp)}
assert keys == set(bag.dask)
yield c._publish_dataset(data=bagp)
# check that serialization didn't affect original bag's dask
assert len(futures_of(bagp)) == 3
result = yield f._get_dataset('data')
assert set(result.dask.keys()) == set(bagp.dask.keys())
assert {f.key
for f in result.dask.values()
} == {f.key
for f in bagp.dask.values()}
out = yield f.compute(result)._result()
assert out == [0, 1, 2]
yield c._shutdown()
yield f._shutdown()
def test_scale_up_and_down():
loop = IOLoop.current()
cluster = LocalCluster(0,
scheduler_port=0,
nanny=False,
silence_logs=False,
diagnostics_port=None,
loop=loop,
start=False)
c = Client(cluster, start=False, loop=loop)
yield c._start()
assert not cluster.workers
yield cluster.scale_up(2)
assert len(cluster.workers) == 2
assert len(cluster.scheduler.ncores) == 2
addr = cluster.workers[0].address
yield cluster.scale_down([addr])
assert len(cluster.workers) == 1
assert addr not in cluster.scheduler.ncores
yield c._shutdown()
yield cluster._close()
def test_multiple_maxlen(c, s, a, b):
c2 = Client((s.ip, s.port), start=False)
yield c2._start()
x = c.channel('x', maxlen=10)
assert x.data.maxlen == 10
x2 = c2.channel('x', maxlen=20)
assert x2.data.maxlen == 20
for i in range(10):
x.append(c.submit(inc, i))
while len(s.wants_what[c2.id]) < 10:
yield gen.sleep(0.01)
for i in range(10, 20):
x.append(c.submit(inc, i))
while len(x2) < 20:
yield gen.sleep(0.01)
yield gen.sleep(0.1)
assert len(x2) == 20 # They stay this long after a delay
assert len(s.task_state) == 20
yield c2._shutdown()
def test_adaptive_local_cluster_multi_workers():
loop = IOLoop.current()
cluster = LocalCluster(0, scheduler_port=0, silence_logs=False, nanny=False,
diagnostics_port=None, loop=loop, start=False)
cluster.scheduler.allowed_failures = 1000
alc = Adaptive(cluster.scheduler, cluster, interval=100)
c = Client(cluster, start=False, loop=loop)
yield c._start()
futures = c.map(slowinc, range(100), delay=0.01)
start = time()
while not cluster.scheduler.worker_info:
yield gen.sleep(0.01)
assert time() < start + 15
yield c._gather(futures)
del futures
start = time()
while cluster.workers:
yield gen.sleep(0.01)
assert time() < start + 5
assert not cluster.workers
yield gen.sleep(0.2)
assert not cluster.workers
futures = c.map(slowinc, range(100), delay=0.01)
yield c._gather(futures)
yield c._shutdown()
yield cluster._close()
def test_multiple_maxlen(c, s, a, b):
c2 = Client((s.ip, s.port), start=False)
yield c2._start()
x = c.channel('x', maxlen=10)
assert x.futures.maxlen == 10
x2 = c2.channel('x', maxlen=20)
assert x2.futures.maxlen == 20
for i in range(10):
x.append(c.submit(inc, i))
while len(s.wants_what[c2.id]) < 10:
yield gen.sleep(0.01)
for i in range(10, 20):
x.append(c.submit(inc, i))
while len(x2) < 20:
yield gen.sleep(0.01)
yield gen.sleep(0.1)
assert len(x2) == 20 # They stay this long after a delay
assert len(s.task_state) == 20
yield c2._shutdown()
def test_multiple_clients_restart(s, a, b):
e1 = Client((s.ip, s.port), start=False)
yield e1._start()
e2 = Client((s.ip, s.port), start=False)
yield e2._start()
x = e1.submit(inc, 1)
y = e2.submit(inc, 2)
xx = yield x._result()
yy = yield y._result()
assert xx == 2
assert yy == 3
yield e1._restart()
assert x.cancelled()
assert y.cancelled()
yield e1._shutdown(fast=True)
yield e2._shutdown(fast=True)
def test_scheduler_file():
with tmpfile() as fn:
s = Scheduler(scheduler_file=fn)
s.start(0)
with open(fn) as f:
data = json.load(f)
assert data['address'] == s.address
c = Client(scheduler_file=fn, loop=s.loop, start=False)
yield c._start()
yield s.close()
def test_multiple_clients_restart(s, a, b):
e1 = Client((s.ip, s.port), start=False)
yield e1._start()
e2 = Client((s.ip, s.port), start=False)
yield e2._start()
x = e1.submit(inc, 1)
y = e2.submit(inc, 2)
xx = yield x._result()
yy = yield y._result()
assert xx == 2
assert yy == 3
yield e1._restart()
assert x.cancelled()
assert y.cancelled()
yield e1._shutdown(fast=True)
yield e2._shutdown(fast=True)
def test_publish_simple(s, a, b):
c = Client((s.ip, s.port), start=False)
yield c._start()
f = Client((s.ip, s.port), start=False)
yield f._start()
data = yield c._scatter(range(3))
out = yield c._publish_dataset(data=data)
assert 'data' in s.extensions['publish'].datasets
with pytest.raises(KeyError) as exc_info:
out = yield c._publish_dataset(data=data)
assert "exists" in str(exc_info.value)
assert "data" in str(exc_info.value)
result = yield c.scheduler.publish_list()
assert result == ['data']
result = yield f.scheduler.publish_list()
assert result == ['data']
yield c._shutdown()
yield f._shutdown()
def test_publish_roundtrip(s, a, b):
c = Client((s.ip, s.port), start=False)
yield c._start()
f = Client((s.ip, s.port), start=False)
yield f._start()
data = yield c._scatter([0, 1, 2])
yield c._publish_dataset(data=data)
assert 'published-data' in s.who_wants[data[0].key]
result = yield f._get_dataset(name='data')
assert len(result) == len(data)
out = yield f._gather(result)
assert out == [0, 1, 2]
with pytest.raises(KeyError) as exc_info:
result = yield f._get_dataset(name='nonexistent')
assert "not found" in str(exc_info.value)
assert "nonexistent" in str(exc_info.value)
yield c._shutdown()
yield f._shutdown()
def test_publish_simple(s, a, b):
c = Client((s.ip, s.port), start=False)
yield c._start()
f = Client((s.ip, s.port), start=False)
yield f._start()
data = yield c._scatter(range(3))
out = yield c._publish_dataset(data=data)
assert 'data' in s.extensions['publish'].datasets
with pytest.raises(KeyError) as exc_info:
out = yield c._publish_dataset(data=data)
assert "exists" in str(exc_info.value)
assert "data" in str(exc_info.value)
result = yield c.scheduler.publish_list()
assert result == ['data']
result = yield f.scheduler.publish_list()
assert result == ['data']
yield c._shutdown()
yield f._shutdown()
def test_publish_roundtrip(s, a, b):
c = Client((s.ip, s.port), start=False)
yield c._start()
f = Client((s.ip, s.port), start=False)
yield f._start()
data = yield c._scatter([0, 1, 2])
yield c._publish_dataset(data=data)
assert 'published-data' in s.who_wants[data[0].key]
result = yield f._get_dataset(name='data')
assert len(result) == len(data)
out = yield f._gather(result)
assert out == [0, 1, 2]
with pytest.raises(KeyError) as exc_info:
result = yield f._get_dataset(name='nonexistent')
assert "not found" in str(exc_info.value)
assert "nonexistent" in str(exc_info.value)
yield c._shutdown()
yield f._shutdown()
def test_adaptive_local_cluster_multi_workers():
loop = IOLoop.current()
cluster = LocalCluster(0, scheduler_port=0, silence_logs=False, nanny=False,
diagnostic_port=None, loop=loop, start=False)
alc = Adaptive(cluster.scheduler, cluster, interval=100)
c = Client(cluster, start=False, loop=loop)
yield c._start()
for i in range(20):
futures = c.map(slowinc, range(100), delay=0.01)
yield c._gather(futures)
del futures
yield gen.sleep(0.1)
yield c._shutdown()
yield cluster._close()
def test_values(c, s, a, b):
c2 = Client((s.ip, s.port), start=False)
yield c2._start()
x = c.channel('x')
x2 = c2.channel('x')
data = [123, 'Hello!', {'x': [1, 2, 3]}]
for item in data:
x.append(item)
while len(x2.data) < 3:
yield gen.sleep(0.01)
assert list(x2.data) == data
yield c2._shutdown()
def test_scale_up_and_down():
loop = IOLoop.current()
cluster = LocalCluster(0, scheduler_port=0, nanny=False, silence_logs=False,
diagnostics_port=None, loop=loop, start=False)
c = Client(cluster, start=False, loop=loop)
yield c._start()
assert not cluster.workers
yield cluster.scale_up(2)
assert len(cluster.workers) == 2
assert len(cluster.scheduler.ncores) == 2
addr = cluster.workers[0].address
yield cluster.scale_down([addr])
assert len(cluster.workers) == 1
assert addr not in cluster.scheduler.ncores
yield c._shutdown()
yield cluster._close()
def test_adaptive_local_cluster_multi_workers():
loop = IOLoop.current()
cluster = LocalCluster(0,
scheduler_port=0,
silence_logs=False,
nanny=False,
diagnostic_port=None,
loop=loop,
start=False)
alc = Adaptive(cluster.scheduler, cluster, interval=100)
c = Client(cluster, start=False, loop=loop)
yield c._start()
for i in range(20):
futures = c.map(slowinc, range(100), delay=0.01)
yield c._gather(futures)
del futures
yield gen.sleep(0.1)
yield c._shutdown()
yield cluster._close()
def test_adaptive_local_cluster_multi_workers():
loop = IOLoop.current()
cluster = LocalCluster(0,
scheduler_port=0,
silence_logs=False,
processes=False,
diagnostics_port=None,
loop=loop,
start=False)
cluster.scheduler.allowed_failures = 1000
alc = Adaptive(cluster.scheduler, cluster, interval=100)
c = Client(cluster, start=False, loop=loop)
yield c._start()
futures = c.map(slowinc, range(100), delay=0.01)
start = time()
while not cluster.scheduler.worker_info:
yield gen.sleep(0.01)
assert time() < start + 15
yield c._gather(futures)
del futures
start = time()
while cluster.workers:
yield gen.sleep(0.01)
assert time() < start + 5
assert not cluster.workers
assert not cluster.scheduler.workers
yield gen.sleep(0.2)
assert not cluster.workers
assert not cluster.scheduler.workers
futures = c.map(slowinc, range(100), delay=0.01)
yield c._gather(futures)
yield c._shutdown()
yield cluster._close()
