def test_nested_schedulers():
class MyCallback(Callback):
def _start(self, dsk):
self.dsk = dsk
def _pretask(self, key, dsk, state):
assert key in self.dsk
inner_callback = MyCallback()
inner_dsk = {'x': (add, 1, 2),
'y': (add, 'x', 3)}
def nested_call(x):
assert not Callback.active
with inner_callback:
return get_threaded(inner_dsk, 'y') + x
outer_callback = MyCallback()
outer_dsk = {'a': (nested_call, 1),
'b': (add, 'a', 2)}
with outer_callback:
get_threaded(outer_dsk, 'b')
assert not Callback.active
assert outer_callback.dsk == outer_dsk
assert inner_callback.dsk == inner_dsk
assert not Callback.active
def test_progressbar(capsys):
with ProgressBar():
out = get_threaded(dsk, 'e')
assert out == 6
check_bar_completed(capsys)
with ProgressBar(width=20):
out = get_threaded(dsk, 'e')
check_bar_completed(capsys, 20)
def test_clean_exit(get):
dsk = {'a': (lambda: 1 / 0, )}
try:
with ProgressBar() as pbar:
get_threaded(dsk, 'a')
except ZeroDivisionError:
pass
assert not pbar._running
assert not pbar._timer.is_alive()
def test_with_alias(capsys):
dsk = {'a': 1,
'b': 2,
'c': (add, 'a', 'b'),
'd': (add, 1, 2),
'e': 'd',
'f': (mul, 'e', 'c')}
with ProgressBar():
get_threaded(dsk, 'f')
check_bar_completed(capsys)
def test_register(capsys):
try:
p = ProgressBar()
p.register()
assert Callback.active
get_threaded(dsk, 'e')
check_bar_completed(capsys)
p.unregister()
assert not Callback.active
finally:
Callback.active.clear()
def test_register(capsys):
try:
p = ProgressBar()
p.register()
assert _globals['callbacks']
get_threaded(dsk, 'e')
check_bar_completed(capsys)
p.unregister()
assert not _globals['callbacks']
finally:
_globals['callbacks'].clear()
def test_with_cache(capsys):
cachey = pytest.importorskip('cachey')
from dask.cache import Cache
c = cachey.Cache(10000)
cc = Cache(c)
with cc:
with ProgressBar():
assert get_threaded({'x': (mul, 1, 2)}, 'x') == 2
check_bar_completed(capsys)
assert c.data['x'] == 2
with cc:
with ProgressBar():
assert get_threaded({'x': (mul, 1, 2), 'y': (mul, 'x', 3)}, 'y') == 6
check_bar_completed(capsys)
def test_finish_always_called():
flag = [False]
class MyCallback(Callback):
def _finish(self, dsk, state, errored):
flag[0] = True
assert errored
dsk = {'x': (lambda: 1 / 0,)}
# `raise_on_exception=True`
try:
with MyCallback():
get_sync(dsk, 'x')
except Exception as e:
assert isinstance(e, ZeroDivisionError)
assert flag[0]
# `raise_on_exception=False`
flag[0] = False
try:
with MyCallback():
get_threaded(dsk, 'x')
except Exception as e:
assert isinstance(e, ZeroDivisionError)
assert flag[0]
# KeyboardInterrupt
def raise_keyboard():
raise KeyboardInterrupt()
dsk = {'x': (raise_keyboard,)}
flag[0] = False
try:
with MyCallback():
get_sync(dsk, 'x')
except BaseException as e:
assert isinstance(e, KeyboardInterrupt)
assert flag[0]
def test_finish_always_called():
flag = [False]
class MyCallback(Callback):
def _finish(self, dsk, state, errored):
flag[0] = True
assert errored
dsk = {"x": (lambda: 1 / 0, )}
# `raise_on_exception=True`
try:
with MyCallback():
get_sync(dsk, "x")
except Exception as e:
assert isinstance(e, ZeroDivisionError)
assert flag[0]
# `raise_on_exception=False`
flag[0] = False
try:
with MyCallback():
get_threaded(dsk, "x")
except Exception as e:
assert isinstance(e, ZeroDivisionError)
assert flag[0]
# KeyboardInterrupt
def raise_keyboard():
raise KeyboardInterrupt()
dsk = {"x": (raise_keyboard, )}
flag[0] = False
try:
with MyCallback():
get_sync(dsk, "x")
except BaseException as e:
assert isinstance(e, KeyboardInterrupt)
assert flag[0]
def test_minimum_time(capsys):
with ProgressBar(1.0):
out = get_threaded(dsk, "e")
out, err = capsys.readouterr()
assert out == "" and err == ""
def test_no_tasks(capsys):
with ProgressBar():
get_threaded({"x": 1}, "x")
check_bar_completed(capsys)
def nested_call(x):
assert not _globals['callbacks']
with inner_callback:
return get_threaded(inner_dsk, 'y') + x
def test_store_time():
p = ProgressBar()
with p:
get_threaded({"x": 1}, "x")
assert isinstance(p.last_duration, float)
def test_no_tasks(capsys):
with ProgressBar():
get_threaded({'x': 1}, 'x')
check_bar_completed(capsys)
def nested_call(x):
assert not Callback.active
with inner_callback:
return get_threaded(inner_dsk, 'y') + x
def nested_call(x):
assert not _globals['callbacks']
with inner_callback:
return get_threaded(inner_dsk, 'y') + x
def test_minimum_time(capsys):
with ProgressBar(1.0):
out = get_threaded(dsk, 'e')
out, err = capsys.readouterr()
assert out == '' and err == ''
def test_store_time():
p = ProgressBar()
with p:
get_threaded({'x': 1}, 'x')
assert isinstance(p.last_duration, float)
