async def test_asyncio_run():
async with AioClient(processes=False) as c:
results = await c.run(inc, 1)
assert len(results) > 0
assert [value == 2 for value in results.values()]
results = await c.run(inc, 1, workers=[])
assert results == {}
async def test_asyncio_as_completed():
async with AioClient(processes=False) as c:
futures = c.map(inc, range(10))
results = []
async for future in as_completed(futures):
results.append(await future)
assert set(results) == set(range(1, 11))
async def test_asyncio_restart():
c = AioClient(processes=False)
await c.start()
assert c.status == 'running'
x = c.submit(inc, 1)
assert x.key in c.refcount
await c.restart()
assert x.key not in c.refcount
key = x.key
del x
import gc
gc.collect()
assert key not in c.refcount
await c.shutdown()
async def test_asyncio_run_on_scheduler():
def f(dask_scheduler=None):
return dask_scheduler.address
async with AioClient(processes=False) as c:
address = await c.run_on_scheduler(f)
assert address == c.cluster.scheduler.address
with pytest.raises(ZeroDivisionError):
await c.run_on_scheduler(div, 1, 0)
async def test_asyncio_exceptions():
async with AioClient(processes=False) as c:
result = await c.submit(div, 1, 2)
assert result == 1 / 2
with pytest.raises(ZeroDivisionError):
result = await c.submit(div, 1, 0)
result = await c.submit(div, 10, 2) # continues to operate
assert result == 10 / 2
async def test_asyncio_exception_on_exception():
async with AioClient(processes=False) as c:
x = c.submit(lambda: 1 / 0)
y = c.submit(inc, x)
with pytest.raises(ZeroDivisionError):
await y
z = c.submit(inc, y)
with pytest.raises(ZeroDivisionError):
await z
async def test_asyncio_gather():
async with AioClient(processes=False) as c:
x = c.submit(inc, 10)
y = c.submit(inc, x)
result = await c.gather(x)
assert result == 11
result = await c.gather([x])
assert result == [11]
result = await c.gather({'x': x, 'y': [y]})
assert result == {'x': 11, 'y': [12]}
async def test_asyncio_start_close():
async with AioClient(processes=False, dashboard_address=False) as c:
assert c.status == "running"
# AioClient has installed its AioLoop shim.
assert isinstance(IOLoop.current(instance=False), BaseAsyncIOLoop)
result = await c.submit(inc, 10)
assert result == 11
await c.close()
assert c.status == "closed"
async def test_asyncio_start_close():
async with AioClient(processes=False, diagnostics_port=False) as c:
assert c.status == 'running'
# AioClient has installed its AioLoop shim.
assert isinstance(IOLoop.current(instance=False), BaseAsyncIOLoop)
result = await c.submit(inc, 10)
assert result == 11
await c.close()
assert c.status == 'closed'
async def test_asyncio_wait():
async with AioClient(processes=False) as c:
x = c.submit(inc, 1)
y = c.submit(inc, 2)
z = c.submit(inc, 3)
await wait(x)
assert x.done() is True
await wait([y, z])
assert y.done() is True
assert z.done() is True
async def test_asyncio_start_shutdown():
c = AioClient(processes=False)
assert c.status is None
await c.start()
assert c.status == 'running'
result = await c.submit(inc, 10)
assert result == 11
await c.shutdown()
assert c.status == 'closed'
async def test_asyncio_get():
async with AioClient(processes=False) as c:
result = await c.get({'x': (inc, 1)}, 'x')
assert result == 2
result = await c.get({'x': (inc, 1)}, ['x'])
assert result == [2]
result = await c.get({}, [])
assert result == []
result = await c.get({('x', 1): (inc, 1), ('x', 2): (inc, ('x', 1))},
('x', 2))
assert result == 3
async def test_asyncio_restart():
async with AioClient(processes=False) as c:
assert c.status == "running"
x = c.submit(inc, 1)
assert x.key in c.refcount
await c.restart()
assert x.key not in c.refcount
key = x.key
del x
import gc
gc.collect()
assert key not in c.refcount
async def test_asyncio_get():
async with AioClient(processes=False) as c:
result = await c.get({"x": (inc, 1)}, "x")
assert result == 2
result = await c.get({"x": (inc, 1)}, ["x"])
assert result == [2]
result = await c.get({}, [])
assert result == []
result = await c.get({
("x", 1): (inc, 1),
("x", 2): (inc, ("x", 1))
}, ("x", 2))
assert result == 3
async def test_asyncio_future_await():
async with AioClient(processes=False) as c:
x = c.submit(inc, 10)
assert not x.done()
assert isinstance(x, AioFuture)
assert x.client is c
result = await x
assert result == 11
assert x.done()
y = c.submit(inc, 20)
z = c.submit(add, x, y)
result = await z
assert result == 11 + 21
async def test_asyncio_run_coroutine():
async def aioinc(x, delay=0.02):
await asyncio.sleep(delay)
return x + 1
async def aiothrows(x, delay=0.02):
await asyncio.sleep(delay)
raise RuntimeError('hello')
async with AioClient(processes=False) as c:
results = await c.run_coroutine(aioinc, 1, delay=0.05)
assert len(results) > 0
assert [value == 2 for value in results.values()]
results = await c.run_coroutine(aioinc, 1, workers=[])
assert results == {}
with pytest.raises(RuntimeError) as exc_info:
await c.run_coroutine(aiothrows, 1)
exc_info.match("hello")
async def test_asyncio_map():
async with AioClient(processes=False) as c:
L1 = c.map(inc, range(5))
assert len(L1) == 5
assert isdistinct(x.key for x in L1)
assert all(isinstance(x, AioFuture) for x in L1)
result = await L1[0]
assert result == inc(0)
L2 = c.map(inc, L1)
result = await L2[1]
assert result == inc(inc(1))
total = c.submit(sum, L2)
result = await total
assert result == sum(map(inc, map(inc, range(5))))
L3 = c.map(add, L1, L2)
result = await L3[1]
assert result == inc(1) + inc(inc(1))
L4 = c.map(add, range(3), range(4))
results = await c.gather(L4)
assert results == list(map(add, range(3), range(4)))
def f(x, y=10):
return x + y
L5 = c.map(f, range(5), y=5)
results = await c.gather(L5)
assert results == list(range(5, 10))
y = c.submit(f, 10)
L6 = c.map(f, range(5), y=y)
results = await c.gather(L6)
assert results == list(range(20, 25))
async def test_asyncio_channels():
async with AioClient(processes=False) as c:
x = c.channel('x')
y = c.channel('y')
assert len(x) == 0
while set(c.extensions['channels'].channels) != {'x', 'y'}:
await asyncio.sleep(0.01)
xx = c.channel('x')
yy = c.channel('y')
assert len(x) == 0
await asyncio.sleep(0.1)
assert set(c.extensions['channels'].channels) == {'x', 'y'}
future = c.submit(inc, 1)
x.append(future)
while not x.data:
await asyncio.sleep(0.01)
assert len(x) == 1
assert xx.data[0].key == future.key
xxx = c.channel('x')
while not xxx.data:
await asyncio.sleep(0.01)
assert xxx.data[0].key == future.key
assert 'x' in repr(x)
assert '1' in repr(x)
async def test_asyncio_cancel():
async with AioClient(processes=False) as c:
s = c.cluster.scheduler
x = c.submit(slowinc, 1)
y = c.submit(slowinc, x)
while y.key not in s.tasks:
await asyncio.sleep(0.01)
await c.cancel([x])
assert x.cancelled()
assert 'cancel' in str(x)
s.validate_state()
start = time()
while not y.cancelled():
await asyncio.sleep(0.01)
assert time() < start + 5
assert not s.tasks
assert not s.who_has
s.validate_state()
async def test_asyncio_nanny_workers():
async with AioClient(n_workers=2) as c:
assert await c.submit(inc, 1) == 2