You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
132 lines
3.1 KiB
132 lines
3.1 KiB
from operator import add, mul
|
|
|
|
import pytest
|
|
|
|
from dask.callbacks import Callback
|
|
from dask.local import get_sync
|
|
from dask.diagnostics import ProgressBar
|
|
from dask.diagnostics.progress import format_time
|
|
from dask.threaded import get as get_threaded
|
|
|
|
|
|
dsk = {'a': 1,
|
|
'b': 2,
|
|
'c': (add, 'a', 'b'),
|
|
'd': (mul, 'a', 'b'),
|
|
'e': (mul, 'c', 'd')}
|
|
|
|
|
|
def check_bar_completed(capsys, width=40):
|
|
out, err = capsys.readouterr()
|
|
assert out.count('100% Completed') == 1
|
|
bar, percent, time = [i.strip() for i in out.split('\r')[-1].split('|')]
|
|
assert bar == '[' + '#' * width + ']'
|
|
assert percent == '100% Completed'
|
|
|
|
|
|
def test_array_compute(capsys):
|
|
from dask.array import ones
|
|
data = ones((100,100),dtype='f4',chunks=(100,100))
|
|
with ProgressBar():
|
|
out = data.sum().compute()
|
|
assert out == 10000
|
|
check_bar_completed(capsys)
|
|
|
|
|
|
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_minimum_time(capsys):
|
|
with ProgressBar(1.0):
|
|
out = get_threaded(dsk, 'e')
|
|
out, err = capsys.readouterr()
|
|
assert out == '' and err == ''
|
|
|
|
|
|
@pytest.mark.parametrize('get', [get_threaded, get_sync])
|
|
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_format_time():
|
|
assert format_time(1.4) == ' 1.4s'
|
|
assert format_time(10.4) == '10.4s'
|
|
assert format_time(100.4) == ' 1min 40.4s'
|
|
assert format_time(1000.4) == '16min 40.4s'
|
|
assert format_time(10000.4) == ' 2hr 46min 40.4s'
|
|
|
|
|
|
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_no_tasks(capsys):
|
|
with ProgressBar():
|
|
get_threaded({'x': 1}, 'x')
|
|
check_bar_completed(capsys)
|
|
|
|
|
|
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_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_store_time():
|
|
p = ProgressBar()
|
|
with p:
|
|
get_threaded({'x': 1}, 'x')
|
|
|
|
assert isinstance(p.last_duration, float)
|