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_submit(self): e = LocalDaskExecutor() with e.start(): assert e.submit(lambda: 1).compute() == 1 assert e.submit(lambda x: x, 1).compute() == 1 assert e.submit(lambda x: x, x=1).compute() == 1 assert e.submit(lambda: prefect).compute() is prefect
def test_map_doesnt_do_anything_for_empty_list_input(self): def map_fn(*args): raise ValueError("map_fn was called") e = LocalDaskExecutor() with e.start(): res = e.wait(e.map(map_fn)) assert res == []
def test_map_iterates_over_multiple_args(self): def map_fn(x, y): return x + y e = LocalDaskExecutor() with e.start(): res = e.wait(e.map(map_fn, [1, 2], [1, 3])) assert res == [2, 5]
def test_map_fails_with_processes_executor(self): def map_fn(x, y): return x + y e = LocalDaskExecutor(scheduler="processes") with pytest.raises(RuntimeError): with e.start(): res = e.wait(e.map(map_fn, [1, 2], [1, 3]))
def test_map_with_threads_scheduler(self): def map_fn(x, y): return x + y e = LocalDaskExecutor(scheduler="threads") with e.start(): res = e.wait(e.map(map_fn, [1, 2], [1, 3])) assert res == [2, 5]
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_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
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_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_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_only_compute_once(self): e = LocalDaskExecutor() count = 0 def inc(x): nonlocal count count += 1 return x + 1 with e.start(): f1 = e.submit(inc, 0) f2 = e.submit(inc, f1) f3 = e.submit(inc, f2) assert e.wait([f1]) == [1] assert count == 1 assert e.wait([f2]) == [2] assert count == 2 assert e.wait([f3]) == [3] assert count == 3 assert e.wait([f1, f2, f3]) == [1, 2, 3] assert count == 3
def test_is_pickleable_after_start(self): e = LocalDaskExecutor() with e.start(): post = cloudpickle.loads(cloudpickle.dumps(e)) assert isinstance(post, LocalDaskExecutor)