def test_submit(self):
e = LocalDaskExecutor()
with e.start():
assert e.submit(lambda: 1).compute(scheduler="sync") == 1
assert e.submit(lambda x: x, 1).compute(scheduler="sync") == 1
assert e.submit(lambda x: x, x=1).compute(scheduler="sync") == 1
assert e.submit(lambda: prefect).compute(scheduler="sync") is prefect
def test_interrupt_stops_running_tasks_quickly(self, scheduler):
# Windows implements `queue.get` using polling,
# which means we can set an exception to interrupt the call to `get`.
# Python 3 on other platforms requires sending SIGINT to the main thread.
if os.name == "nt":
from _thread import interrupt_main
else:
main_thread = threading.get_ident()
def interrupt_main():
import signal
signal.pthread_kill(main_thread, signal.SIGINT)
def long_task():
for i in range(50):
time.sleep(0.1)
e = LocalDaskExecutor(scheduler)
try:
interrupter = threading.Timer(0.5, interrupt_main)
interrupter.start()
start = time.time()
with e.start():
e.wait(e.submit(long_task))
except KeyboardInterrupt:
pass
except Exception:
assert False, "Failed to interrupt"
stop = time.time()
if stop - start > 4:
assert False, "Failed to interrupt"
def test_wait(self):
e = LocalDaskExecutor()
with e.start():
assert e.wait(1) == 1
assert e.wait(prefect) is prefect
assert e.wait(e.submit(lambda: 1)) == 1
assert e.wait(e.submit(lambda x: x, 1)) == 1
assert e.wait(e.submit(lambda x: x, x=1)) == 1
assert e.wait(e.submit(lambda: prefect)) is prefect
def test_configurable_scheduler(self):
e = LocalDaskExecutor(scheduler="synchronous")
assert e.scheduler == "synchronous"
def check_scheduler(val):
assert dask.config.get("scheduler") == val
with dask.config.set(scheduler="threads"):
check_scheduler("threads")
with e.start():
e.submit(check_scheduler, "synchronous")
def test_submit_sets_task_name(self):
e = LocalDaskExecutor()
with e.start():
f = e.submit(lambda x: x + 1, 1, extra_context={"task_name": "inc"})
(res,) = e.wait([f])
assert f.key.startswith("inc-")
assert res == 2
f = e.submit(
lambda x: x + 1, 1, extra_context={"task_name": "inc", "task_index": 1}
)
(res,) = e.wait([f])
assert f.key.startswith("inc-1-")
assert res == 2
def test_temporary_pool_created_of_proper_size_and_kind(
self, scheduler, num_workers):
from dask.system import CPU_COUNT
from multiprocessing.pool import Pool, ThreadPool
e = LocalDaskExecutor(scheduler, num_workers=num_workers)
with e.start():
if scheduler == "synchronous":
assert e._pool is None
else:
sol = num_workers or CPU_COUNT
kind = ThreadPool if scheduler == "threads" else Pool
assert isinstance(e._pool, kind)
assert e._pool._processes == sol
assert e._pool is None
def test_temporary_pool_created_of_proper_size_and_kind(
self, scheduler, num_workers):
from dask.system import CPU_COUNT
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
e = LocalDaskExecutor(scheduler, num_workers=num_workers)
with e.start():
if scheduler == "synchronous":
assert e._pool is None
else:
sol = num_workers or CPU_COUNT
kind = (ThreadPoolExecutor
if scheduler == "threads" else ProcessPoolExecutor)
assert e._pool._max_workers == sol
assert isinstance(e._pool, kind)
assert e._pool is None
def test_only_compute_once(self, scheduler, tmpdir):
e = LocalDaskExecutor(scheduler)
def inc(x, path):
if os.path.exists(path):
raise ValueError("Should only run once!")
with open(path, "wb"):
pass
return x + 1
with e.start():
f1 = e.submit(inc, 0, str(tmpdir.join("f1")))
f2 = e.submit(inc, f1, str(tmpdir.join("f2")))
f3 = e.submit(inc, f2, str(tmpdir.join("f3")))
assert e.wait([f1]) == [1]
assert e.wait([f2]) == [2]
assert e.wait([f3]) == [3]
assert e.wait([f1, f2, f3]) == [1, 2, 3]
def test_interrupt_stops_running_tasks_quickly(self, scheduler):
# TODO: remove this skip
if scheduler == "processes":
pytest.skip(
"This test periodically hangs for some reason on circleci, but passes "
"locally. We should debug this later, but squashing it for now"
)
main_thread = threading.get_ident()
def interrupt():
if sys.platform == "win32":
# pthread_kill is Windows only
from _thread import interrupt_main
interrupt_main()
else:
import signal
signal.pthread_kill(main_thread, signal.SIGINT)
def long_task():
for i in range(50):
time.sleep(0.1)
e = LocalDaskExecutor(scheduler)
try:
interrupter = threading.Timer(0.5, interrupt)
interrupter.start()
start = time.time()
with e.start():
e.wait(e.submit(long_task))
except KeyboardInterrupt:
pass # Don't exit test on the interrupt
stop = time.time()
# Defining "quickly" here as 4 seconds generally and 6 seconds on
# Windows which tends to be a little slower
assert (stop - start) < (6 if sys.platform == "win32" else 4)
def test_interrupt_stops_running_tasks_quickly(self, scheduler):
# TODO: remove this skip
if scheduler == "processes":
pytest.skip(
"This test periodically hangs for some reason on circleci, but passes "
"locally. We should debug this later, but squashing it for now"
)
# Windows implements `queue.get` using polling,
# which means we can set an exception to interrupt the call to `get`.
# Python 3 on other platforms requires sending SIGINT to the main thread.
if os.name == "nt":
from _thread import interrupt_main
else:
main_thread = threading.get_ident()
def interrupt_main():
import signal
signal.pthread_kill(main_thread, signal.SIGINT)
def long_task():
for i in range(50):
time.sleep(0.1)
e = LocalDaskExecutor(scheduler)
try:
interrupter = threading.Timer(0.5, interrupt_main)
interrupter.start()
start = time.time()
with e.start():
e.wait(e.submit(long_task))
except KeyboardInterrupt:
pass
except Exception:
assert False, "Failed to interrupt"
stop = time.time()
if stop - start > 4:
assert False, "Failed to interrupt"
def test_is_pickleable_after_start(self):
e = LocalDaskExecutor()
with e.start():
post = cloudpickle.loads(cloudpickle.dumps(e))
assert isinstance(post, LocalDaskExecutor)
assert post._pool is None
