def test_pickleable(client):
v = Variable('v')
def f(x):
v.set(x + 1)
client.submit(f, 10).result()
assert v.get() == 11
def test_progress_function(client, capsys):
f = client.submit(lambda: 1)
g = client.submit(lambda: 2)
progress([[f], [[g]]], notebook=False)
check_bar_completed(capsys)
progress(f)
check_bar_completed(capsys)
def test_picklability_sync(client):
q = Queue()
def f(x):
q.put(x + 1)
client.submit(f, 10).result()
assert q.get() == 11
def test_as_completed_with_results(client):
x = client.submit(throws, 1)
y = client.submit(inc, 5)
z = client.submit(inc, 1)
ac = as_completed([x, y, z], with_results=True)
y.cancel()
with pytest.raises(RuntimeError) as exc:
res = list(ac)
assert str(exc.value) == 'hello!'
def test_as_completed_error(client):
x = client.submit(throws, 1)
y = client.submit(inc, 1)
ac = as_completed([x, y])
result = set(ac)
assert result == {x, y}
assert x.status == 'error'
assert y.status == 'finished'
def test_as_completed_error(client):
x = client.submit(throws, 1)
y = client.submit(inc, 1)
ac = as_completed([x, y])
result = set(ac)
assert result == {x, y}
assert x.status == "error"
assert y.status == "finished"
def test_as_completed_with_results(client):
x = client.submit(throws, 1)
y = client.submit(inc, 5)
z = client.submit(inc, 1)
ac = as_completed([x, y, z], with_results=True)
y.cancel()
with pytest.raises(RuntimeError) as exc:
res = list(ac)
assert str(exc.value) == "hello!"
def test_as_completed(client):
x = client.submit(inc, 1)
y = client.submit(inc, 2)
z = client.submit(inc, 1)
seq = as_completed([x, y, z])
assert seq.count() == 3
assert isinstance(seq, Iterator)
assert set(seq) == {x, y, z}
assert seq.count() == 0
assert list(as_completed([])) == []
def test_as_completed(client):
x = client.submit(inc, 1)
y = client.submit(inc, 2)
z = client.submit(inc, 1)
seq = as_completed([x, y, z])
assert seq.count() == 3
assert isinstance(seq, Iterator)
assert set(seq) == {x, y, z}
assert seq.count() == 0
assert list(as_completed([])) == []
def test_values(client):
L = [client.submit(inc, i) for i in range(5)]
wait(L)
p = MultiProgressWidget(L)
client.sync(p.listen)
assert set(p.bars) == {'inc'}
assert p.status == 'finished'
assert p.comm.closed()
assert '5 / 5' in p.bar_texts['inc'].value
assert p.bars['inc'].value == 1.0
x = client.submit(throws, 1)
p = MultiProgressWidget([x])
client.sync(p.listen)
assert p.status == 'error'
def test_values(client):
L = [client.submit(inc, i) for i in range(5)]
wait(L)
p = MultiProgressWidget(L)
client.sync(p.listen)
assert set(p.bars) == {"inc"}
assert p.status == "finished"
assert p.comm.closed()
assert "5 / 5" in p.bar_texts["inc"].value
assert p.bars["inc"].value == 1.0
x = client.submit(throws, 1)
p = MultiProgressWidget([x])
client.sync(p.listen)
assert p.status == "error"
def test_values(client):
L = [client.submit(inc, i) for i in range(5)]
wait(L)
p = MultiProgressWidget(L)
client.sync(p.listen)
assert set(p.bars) == {'inc'}
assert p.status == 'finished'
assert p.comm.closed()
assert '5 / 5' in p.bar_texts['inc'].value
assert p.bars['inc'].value == 1.0
x = client.submit(throws, 1)
p = MultiProgressWidget([x])
client.sync(p.listen)
assert p.status == 'error'
def test_event_sync(client):
# Assert that we call the client.sync correctly
def wait_for_it_failing(x):
event = Event("x")
# Event is not set in another task so far
assert not event.wait(timeout=0.05)
assert not event.is_set()
def wait_for_it_ok(x):
event = Event("x")
# Event is set in another task
assert event.wait(timeout=0.5)
assert event.is_set()
def set_it():
event = Event("x")
event.set()
wait_futures = client.map(wait_for_it_failing, range(10))
client.gather(wait_futures)
set_future = client.submit(set_it)
client.gather(set_future)
wait_futures = client.map(wait_for_it_ok, range(10))
client.gather(wait_futures)
def test_as_completed_cancel_last(client):
w = client.submit(inc, 0.3)
x = client.submit(inc, 1)
y = client.submit(inc, 0.3)
async def _():
await asyncio.sleep(0.1)
await w.cancel(asynchronous=True)
await y.cancel(asynchronous=True)
client.loop.add_callback(_)
ac = as_completed([x, y])
result = set(ac)
assert result == {x, y}
def test_as_completed_with_actor(client):
class Counter:
n = 0
def __init__(self):
self.n = 0
def increment(self):
time.sleep(0.1)
self.n += 1
return self.n
future = client.submit(Counter, actor=True)
counter = future.result()
assert counter.n == 0
futures = []
for _ in range(3):
futures.append(counter.increment())
results = []
for future in as_completed(futures):
results.append(future.result())
assert counter.n == 3
assert sum(results) == 6 # i.e. future.result() returns 1, 2, 3
# ToDo: test other `as_completed` arguments, e.g.:
# for future, n in as_completed(futures, with_results=True):
# results.append(n)
def test_get_client_sync(client):
def f(x):
cc = get_client()
future = cc.submit(inc, x)
return future.result()
future = client.submit(f, 10)
assert future.result() == 11
def test_sync(client):
future = client.submit(lambda x: x + 1, 10)
x = Variable('x')
xx = Variable('x')
x.set(future)
future2 = xx.get()
assert future2.result() == 11
def test_as_completed_is_empty(client):
ac = as_completed()
assert ac.is_empty()
x = client.submit(inc, 1)
ac.add(x)
assert not ac.is_empty()
assert next(ac) is x
assert ac.is_empty()
def test_as_completed_cancel_last(client):
w = client.submit(inc, 0.3)
x = client.submit(inc, 1)
y = client.submit(inc, 0.3)
@gen.coroutine
def _():
yield gen.sleep(0.1)
yield w.cancel(asynchronous=True)
yield y.cancel(asynchronous=True)
client.loop.add_callback(_)
ac = as_completed([x, y])
result = set(ac)
assert result == {x, y}
def test_as_completed_cancel(client):
x = client.submit(inc, 1)
y = client.submit(inc, 1)
ac = as_completed([x, y])
x.cancel()
assert next(ac) is x or y
assert next(ac) is y or x
with pytest.raises(Empty):
ac.queue.get(timeout=0.1)
res = list(as_completed([x, y, x]))
assert len(res) == 3
assert set(res) == {x, y}
assert res.count(x) == 2
def test_as_completed_is_empty(client):
ac = as_completed()
assert ac.is_empty()
x = client.submit(inc, 1)
ac.add(x)
assert not ac.is_empty()
assert next(ac) is x
assert ac.is_empty()
def test_sync(client):
future = client.submit(lambda x: x + 1, 10)
x = Variable("x")
xx = Variable("x")
x.set(future)
future2 = xx.get()
assert future2.result() == 11
def test_timeout(client):
def func():
with worker_client(timeout=0) as wc:
print("hello")
future = client.submit(func)
with pytest.raises(EnvironmentError):
result = future.result()
def test_as_completed_cancel(client):
x = client.submit(inc, 1)
y = client.submit(inc, 1)
ac = as_completed([x, y])
x.cancel()
assert next(ac) is x or y
assert next(ac) is y or x
with pytest.raises(Empty):
ac.queue.get(timeout=0.1)
res = list(as_completed([x, y, x]))
assert len(res) == 3
assert set(res) == {x, y}
assert res.count(x) == 2
def test_as_completed_cancel_last(client):
w = client.submit(inc, 0.3)
x = client.submit(inc, 1)
y = client.submit(inc, 0.3)
@gen.coroutine
def _():
yield gen.sleep(0.1)
yield w.cancel(asynchronous=True)
yield y.cancel(asynchronous=True)
client.loop.add_callback(_)
ac = as_completed([x, y])
result = set(ac)
assert result == {x, y}
def test_get_client_sync(client):
def f(x):
cc = get_client()
future = cc.submit(inc, x)
return future.result()
future = client.submit(f, 10)
assert future.result() == 11
def test_timeout(client):
def func():
with worker_client(timeout=0) as wc:
print('hello')
future = client.submit(func)
with pytest.raises(EnvironmentError):
result = future.result()
def test_n_workers(self) -> None:
with LocalCluster(n_workers=2) as cluster:
with Client(cluster) as client:
workers = list(_get_client_workers(client).keys())
from sklearn.datasets import load_breast_cancer
X, y = load_breast_cancer(return_X_y=True)
dX = client.submit(da.from_array, X, workers=[workers[0]]).result()
dy = client.submit(da.from_array, y, workers=[workers[0]]).result()
train = xgb.dask.DaskDMatrix(client, dX, dy)
dX = dd.from_array(X)
dX = client.persist(dX, workers={dX: workers[1]})
dy = dd.from_array(y)
dy = client.persist(dy, workers={dy: workers[1]})
valid = xgb.dask.DaskDMatrix(client, dX, dy)
merged = xgb.dask._get_workers_from_data(train, evals=[(valid, 'Valid')])
assert len(merged) == 2
def test_progressbar_cancel(client):
import time
L = [client.submit(lambda: time.sleep(0.3), i) for i in range(5)]
p = ProgressWidget(L)
client.sync(p.listen)
L[-1].cancel()
wait(L[:-1])
assert p.status == 'error'
assert p.bar.value == 0 # no tasks finish before cancel is called
def test_progressbar_cancel(client):
import time
L = [client.submit(lambda: time.sleep(0.3), i) for i in range(5)]
p = ProgressWidget(L)
client.sync(p.listen)
L[-1].cancel()
wait(L[:-1])
assert p.status == 'error'
assert p.bar.value == 0 # no tasks finish before cancel is called
def test_as_completed_with_results_no_raise(client):
x = client.submit(throws, 1)
y = client.submit(inc, 5)
z = client.submit(inc, 1)
ac = as_completed([x, y, z], with_results=True, raise_errors=False)
y.cancel()
res = list(ac)
dd = {r[0]: r[1:] for r in res}
assert set(dd.keys()) == {y, x, z}
assert x.status == "error"
assert y.status == "cancelled"
assert z.status == "finished"
assert isinstance(dd[y][0], CancelledError) or dd[y][0] == 6
assert isinstance(dd[x][0][1], RuntimeError)
assert dd[z][0] == 2
def test_as_completed_with_results_no_raise(client):
x = client.submit(throws, 1)
y = client.submit(inc, 5)
z = client.submit(inc, 1)
ac = as_completed([x, y, z], with_results=True, raise_errors=False)
y.cancel()
res = list(ac)
dd = {r[0]: r[1:] for r in res}
assert set(dd.keys()) == {y, x, z}
assert x.status == 'error'
assert y.status == 'cancelled'
assert z.status == 'finished'
assert isinstance(dd[y][0], CancelledError)
assert isinstance(dd[x][0][1], RuntimeError)
assert dd[z][0] == 2
def test_sync(client):
future = client.submit(lambda x: x + 1, 10)
x = Queue('x')
xx = Queue('x')
x.put(future)
assert x.qsize() == 1
assert xx.qsize() == 1
future2 = xx.get()
assert future2.result() == 11
def test_sync(client):
future = client.submit(lambda x: x + 1, 10)
x = Queue("x")
xx = Queue("x")
x.put(future)
assert x.qsize() == 1
assert xx.qsize() == 1
future2 = xx.get()
assert future2.result() == 11
def test_future_erred_sync(client):
future = client.submit(div, 1, 0)
var = Variable()
var.set(future)
sleep(0.1)
future2 = var.get()
with pytest.raises(ZeroDivisionError):
future2.result()
def test_future_erred_sync(client):
future = client.submit(div, 1, 0)
var = Variable()
var.set(future)
sleep(0.1)
future2 = var.get()
with pytest.raises(ZeroDivisionError):
future2.result()
def test_as_completed_add(client):
total = 0
expected = sum(map(inc, range(10)))
futures = client.map(inc, range(10))
ac = as_completed(futures)
for future in ac:
result = future.result()
total += result
if random.random() < 0.5:
future = client.submit(add, future, 10)
ac.add(future)
expected += result + 10
assert total == expected
def test_as_completed_add(client):
total = 0
expected = sum(map(inc, range(10)))
futures = client.map(inc, range(10))
ac = as_completed(futures)
for future in ac:
result = future.result()
total += result
if random.random() < 0.5:
future = client.submit(add, future, 10)
ac.add(future)
expected += result + 10
assert total == expected
def test_as_completed_repeats(client):
ac = as_completed()
x = client.submit(inc, 1)
ac.add(x)
ac.add(x)
assert next(ac) is x
assert next(ac) is x
with pytest.raises(StopIteration):
next(ac)
ac.add(x)
assert next(ac) is x
def test_as_completed_repeats(client):
ac = as_completed()
x = client.submit(inc, 1)
ac.add(x)
ac.add(x)
assert next(ac) is x
assert next(ac) is x
with pytest.raises(StopIteration):
next(ac)
ac.add(x)
assert next(ac) is x
def test_progressbar_done(client):
L = [client.submit(inc, i) for i in range(5)]
wait(L)
p = ProgressWidget(L)
client.sync(p.listen)
assert p.status == 'finished'
assert p.bar.value == 1.0
assert p.bar.bar_style == 'success'
assert 'Finished' in p.elapsed_time.value
f = client.submit(throws, L)
wait([f])
p = ProgressWidget([f])
client.sync(p.listen)
assert p.status == 'error'
assert p.bar.value == 0.0
assert p.bar.bar_style == 'danger'
assert 'Exception' in p.elapsed_time.value
try:
throws(1)
except Exception as e:
assert repr(e) in p.elapsed_time.value
def test_progressbar_done(client):
L = [client.submit(inc, i) for i in range(5)]
wait(L)
p = ProgressWidget(L)
client.sync(p.listen)
assert p.status == "finished"
assert p.bar.value == 1.0
assert p.bar.bar_style == "success"
assert "Finished" in p.elapsed_time.value
f = client.submit(throws, L)
wait([f])
p = ProgressWidget([f])
client.sync(p.listen)
assert p.status == "error"
assert p.bar.value == 0.0
assert p.bar.bar_style == "danger"
assert "Exception" in p.elapsed_time.value
try:
throws(1)
except Exception as e:
assert repr(e) in p.elapsed_time.value
def test_sync(client):
counter = client.submit(Counter, actor=True)
counter = counter.result()
assert counter.n == 0
future = counter.increment()
n = future.result()
assert n == 1
assert counter.n == 1
assert future.result() == future.result()
assert "ActorFuture" in repr(future)
assert "distributed.actor" not in repr(future)
def test_sync(client):
counter = client.submit(Counter, actor=True)
counter = counter.result()
assert counter.n == 0
future = counter.increment()
n = future.result()
assert n == 1
assert counter.n == 1
assert future.result() == future.result()
assert 'ActorFuture' in repr(future)
assert 'distributed.actor' not in repr(future)
def test_data_initialization(self) -> None:
'''Assert each worker has the correct amount of data, and DMatrix initialization doesn't
generate unnecessary copies of data.
'''
with LocalCluster(n_workers=2) as cluster:
with Client(cluster) as client:
X, y, _ = generate_array()
n_partitions = X.npartitions
m = xgb.dask.DaskDMatrix(client, X, y)
workers = list(_get_client_workers(client).keys())
rabit_args = client.sync(xgb.dask._get_rabit_args,
len(workers), client)
n_workers = len(workers)
def worker_fn(worker_addr: str, data_ref: Dict) -> None:
with xgb.dask.RabitContext(rabit_args):
local_dtrain = xgb.dask._dmatrix_from_list_of_parts(
**data_ref)
total = np.array([local_dtrain.num_row()])
total = xgb.rabit.allreduce(total, xgb.rabit.Op.SUM)
assert total[0] == kRows
futures = []
for i in range(len(workers)):
futures.append(
client.submit(worker_fn,
workers[i],
m.create_fn_args(workers[i]),
pure=False,
workers=[workers[i]]))
client.gather(futures)
has_what = client.has_what()
cnt = 0
data = set()
for k, v in has_what.items():
for d in v:
cnt += 1
data.add(d)
assert len(data) == cnt
# Subtract the on disk resource from each worker
assert cnt - n_workers == n_partitions
async def test_pack_MaterializedLayer_handles_futures_in_graph_properly(
c, s, a, b):
fut = c.submit(inc, 1)
hlg = HighLevelGraph(
{
"l1": MaterializedLayer({
"x": fut,
"y": (inc, "x"),
"z": (inc, "y")
})
},
{"l1": set()},
)
# fill hlg.key_dependencies cache. This excludes known futures, so only
# includes a subset of all dependencies. Previously if the cache was present
# the future dependencies would be missing when packed.
hlg.get_all_dependencies()
packed = hlg.__dask_distributed_pack__(c, ["z"], {})
unpacked = HighLevelGraph.__dask_distributed_unpack__(packed)
assert unpacked["deps"] == {"x": {fut.key}, "y": {fut.key}, "z": {"y"}}
def f():
client = yield get_client()
future = client.submit(inc, 10)
result = yield future
raise gen.Return(result)
def f():
client = yield get_client()
future = client.submit(inc, 10)
result = yield future
raise gen.Return(result)
def test_progress_function_w_kwargs(client, capsys):
f = client.submit(lambda: 1)
g = client.submit(lambda: 2)
progress(f, interval="20ms")
check_bar_completed(capsys)
async def f():
client = await get_client()
future = client.submit(inc, 10)
result = await future
return result
