import functools import os import random import sys import threading import time import unittest import warnings from concurrent.futures._base import ( CANCELLED, CANCELLED_AND_NOTIFIED, FINISHED, PENDING, RUNNING, ) from pathlib import Path sys.path.append( os.fspath(Path(__file__).resolve().parent.parent.parent / "test") ) import mpitestutil as testutil from mpi4py import MPI, futures SHARED_POOL = futures._core.SharedPool is not None WORLD_SIZE = MPI.COMM_WORLD.Get_size() def create_future(state=PENDING, exception=None, result=None): f = futures.Future() f._state = state f._exception = exception f._result = result return f PENDING_FUTURE = create_future(state=PENDING) RUNNING_FUTURE = create_future(state=RUNNING) CANCELLED_FUTURE = create_future(state=CANCELLED) CANCELLED_AND_NOTIFIED_FUTURE = create_future(state=CANCELLED_AND_NOTIFIED) EXCEPTION_FUTURE = create_future(state=FINISHED, exception=OSError()) SUCCESSFUL_FUTURE = create_future(state=FINISHED, result=42) def mul(x, y): # noqa: FURB118 return x * y def sleep_and_raise(t): time.sleep(t) raise RuntimeError("this is an exception") def check_global_var(x): return global_var == x # noqa: F821 def check_run_name(name): return __name__ == name def check_comm_workers(): comm = futures.get_comm_workers() return comm.Get_size() def sys_flags_get(name): return getattr(sys.flags, name) class ExecutorMixin: worker_count = 2 def setUp(self): self.t1 = time.time() try: self.executor = self.executor_type(max_workers=self.worker_count) except NotImplementedError as exc: self.skipTest(str(exc)) self._prime_executor() def tearDown(self): self.executor.shutdown(wait=True) dt = time.time() - self.t1 self.assertLess(dt, 60, "synchronization issue: test lasted too long") def _prime_executor(self): # Make sure that the executor is ready to do work before running the # tests. This should reduce the probability of timeouts in the tests. futures = [ self.executor.submit(time.sleep, 0) for _ in range(self.worker_count) ] for f in futures: f.result() class ProcessPoolMixin(ExecutorMixin): executor_type = futures.MPIPoolExecutor if "coverage" in sys.modules: executor_type = staticmethod( functools.partial( executor_type, python_args="-m coverage run".split(), ) ) @unittest.skipIf(not SHARED_POOL, "not-shared-pool") class SharedPoolInitTest(unittest.TestCase): executor_type = futures.MPIPoolExecutor def test_initializer_0(self): executor = self.executor_type( initializer=time.sleep, initargs=(0,), ) executor.bootup() executor.submit(time.sleep, 0).result() executor.shutdown() def test_initializer_1(self): for _ in range(2): executor = self.executor_type( initializer=sleep_and_raise, initargs=(0.2,), ) executor.submit(time.sleep, 0).cancel() future = executor.submit(time.sleep, 0) with self.assertRaises(futures.BrokenExecutor): executor.submit(time.sleep, 0).result() with self.assertRaises(futures.BrokenExecutor): future.result() with self.assertRaises(futures.BrokenExecutor): executor.submit(time.sleep, 0) def test_initializer_2(self): executor = self.executor_type( initializer=time.sleep, initargs=(0,), ) executor.bootup() executor.submit(time.sleep, 0).result() executor.shutdown() def test_initializer_3(self): executor = self.executor_type() executor.submit(time.sleep, 0).result() executor.shutdown() def test_initializer_4(self): rnd = random.random def test(tid): with self.executor_type( initializer=time.sleep, initargs=(rnd() / 100,), ) as executor: futures.as_completed([ executor.submit(time.sleep, rnd() / 100) for _ in range(executor.num_workers + tid) ]) ts = [threading.Thread(target=test, args=(i,)) for i in range(5)] for t in ts: t.start() for t in ts: t.join() class ProcessPoolInitTest(ProcessPoolMixin, unittest.TestCase): def setUp(self): pass def tearDown(self): pass def _prime_executor(self): pass def test_init(self): self.executor_type() def test_init_args(self): self.executor_type(1) def test_init_kwargs(self): executor = self.executor_type( python_exe=sys.executable, max_workers=None, mpi_info={"soft": "0:1"}, globals=None, main=False, path=[], wdir=Path.cwd(), env={}, use_pkl5=None, backoff=0.001, ) futures = [ executor.submit(time.sleep, 0) for _ in range(self.worker_count) ] for f in futures: f.result() executor.shutdown() def test_init_pyargs(self): executor_type = futures.MPIPoolExecutor executor = executor_type(python_args=["-B", "-Wi"]) executor.submit(time.sleep, 0).result() executor.shutdown() @unittest.skipIf(SHARED_POOL, "shared-pool") def test_init_sys_flags(self): executor_type = futures.MPIPoolExecutor sys_flags = [ ("debug", "-d", True), ("optimize", "-O", 1), ("optimize", "-OO", 2), ("dont_write_bytecode", "-B", True), ("dev_mode", "-Xdev", True), ("utf8_mode", "-Xutf8", True), ] if sys.version_info >= (3, 11): sys_flags.extend([ ("safe_path", "-P", True), ]) for name, flag, value in sys_flags: if not isinstance(value, bool): if isinstance(value, int): value += getattr(sys.flags, name) with executor_type(python_args=[flag]) as executor: result = executor.submit(sys_flags_get, name).result() if isinstance(value, bool): result = bool(result) self.assertEqual(value, result, f"sys.flags.{name}") @unittest.skipIf(SHARED_POOL, "shared-pool") def test_init_globals(self): executor = self.executor_type(globals={"global_var": 42}) future1 = executor.submit(check_global_var, 42) future2 = executor.submit(check_global_var, 24) self.assertTrue(future1.result()) self.assertFalse(future2.result()) executor.shutdown() @unittest.skipIf(SHARED_POOL and WORLD_SIZE == 1, "shared-pool") def test_run_name(self): executor = self.executor_type() run_name = futures._core.MAIN_RUN_NAME future = executor.submit(check_run_name, run_name) self.assertTrue(future.result()) @unittest.skipIf(SHARED_POOL and WORLD_SIZE > 2, "shared-pool") def test_max_workers(self): executor = self.executor_type(max_workers=1) self.assertEqual(executor.num_workers, 1) self.assertEqual(executor.num_workers, executor._max_workers) executor.shutdown() self.assertEqual(executor.num_workers, 0) self.assertEqual(executor.num_workers, executor._max_workers) @unittest.skipIf(SHARED_POOL and WORLD_SIZE > 2, "shared-pool") def test_max_workers_environ(self): save = os.environ.get("MPI4PY_FUTURES_MAX_WORKERS") os.environ["MPI4PY_FUTURES_MAX_WORKERS"] = "1" try: executor = self.executor_type() executor.submit(time.sleep, 0).result() executor.shutdown() executor = self.executor_type() self.assertEqual(executor.num_workers, 1) executor.shutdown() self.assertEqual(executor.num_workers, 0) finally: del os.environ["MPI4PY_FUTURES_MAX_WORKERS"] if save is not None: os.environ["MPI4PY_FUTURES_MAX_WORKERS"] = save def test_max_workers_negative(self): for number in (0, -1): with self.assertRaises(ValueError): self.executor_type(max_workers=number) def test_get_comm_workers(self): executor = self.executor_type() num_workers = executor.submit(check_comm_workers).result() self.assertEqual(executor.num_workers, num_workers) self.assertRaises(RuntimeError, check_comm_workers) @unittest.skipIf(SHARED_POOL, "shared-pool") def test_use_pkl5_kwarg(self): executor = self.executor_type(use_pkl5=True) executor.submit(time.sleep, 0).result() executor.shutdown() @unittest.skipIf(SHARED_POOL, "shared-pool") def test_use_pkl5_environ(self): save = os.environ.get("MPI4PY_FUTURES_USE_PKL5") try: for value in ("false", "true"): os.environ["MPI4PY_FUTURES_USE_PKL5"] = value executor = self.executor_type() executor.submit(time.sleep, 0).result() executor.shutdown() with warnings.catch_warnings(record=True) as wlist: warnings.simplefilter("always") os.environ["MPI4PY_FUTURES_USE_PKL5"] = "foobar" executor = self.executor_type() executor.submit(time.sleep, 0).result() executor.shutdown() self.assertTrue(wlist) msg = wlist[0].message self.assertIsInstance(msg, RuntimeWarning) self.assertIn("foobar", msg.args[0]) finally: del os.environ["MPI4PY_FUTURES_USE_PKL5"] if save is not None: os.environ["MPI4PY_FUTURES_USE_PKL5"] = save def test_initializer(self): executor = self.executor_type( initializer=time.sleep, initargs=(0,), ) executor.submit(time.sleep, 0).result() def test_initializer_bad(self): with self.assertRaises(TypeError): self.executor_type(initializer=123) def test_initializer_error(self): executor = self.executor_type( initializer=sleep_and_raise, initargs=(0.2,), ) executor.submit(time.sleep, 0).cancel() future = executor.submit(time.sleep, 0) with self.assertRaises(futures.BrokenExecutor): executor.submit(time.sleep, 0).result() with self.assertRaises(futures.BrokenExecutor): future.result() with self.assertRaises(futures.BrokenExecutor): executor.submit(time.sleep, 0) self.assertEqual(executor.num_workers, 0) def test_initializer_error_del(self): executor = self.executor_type( initializer=sleep_and_raise, initargs=(0.2,), ) executor.bootup() del executor def test_initializer_error_del_nowait(self): executor = self.executor_type( initializer=sleep_and_raise, initargs=(1.2,), ) executor.bootup(wait=False) executor.shutdown(wait=False) del executor class ProcessPoolBootupTest(ProcessPoolMixin, unittest.TestCase): def _prime_executor(self): pass def test_bootup(self): executor = self.executor_type(1) executor.bootup() executor.bootup() executor.shutdown() with self.assertRaises(RuntimeError): executor.bootup() def test_bootup_wait(self): executor = self.executor_type(1) executor.bootup(wait=True) executor.bootup(wait=True) executor.shutdown(wait=True) with self.assertRaises(RuntimeError): executor.bootup(wait=True) def test_bootup_nowait(self): executor = self.executor_type(1) executor.bootup(wait=False) executor.bootup(wait=False) executor.shutdown(wait=False) with self.assertRaises(RuntimeError): executor.bootup(wait=False) executor.shutdown(wait=True) def test_bootup_nowait_wait(self): executor = self.executor_type(1) executor.bootup(wait=False) executor.bootup(wait=True) executor.shutdown() with self.assertRaises(RuntimeError): executor.bootup() def test_bootup_shutdown_nowait(self): executor = self.executor_type(1) executor.bootup(wait=False) executor.shutdown(wait=False) worker = executor._pool del executor worker.join() class ExecutorShutdownTestMixin: def test_run_after_shutdown(self): self.executor.shutdown() with self.assertRaises(RuntimeError): self.executor.submit(pow, 2, 5) def test_hang_issue12364(self): fs = [self.executor.submit(time.sleep, 0.01) for _ in range(50)] self.executor.shutdown() for f in fs: f.result() class ProcessPoolShutdownTest( ProcessPoolMixin, ExecutorShutdownTestMixin, unittest.TestCase ): def _prime_executor(self): pass def test_shutdown(self): executor = self.executor_type(max_workers=1) self.assertIsNone(executor._pool) self.assertFalse(executor._shutdown) executor.submit(mul, 21, 2) executor.submit(mul, 6, 7) executor.submit(mul, 3, 14) self.assertIsNotNone(executor._pool.thread) self.assertFalse(executor._shutdown) executor.shutdown(wait=False) self.assertIsNotNone(executor._pool.thread) self.assertTrue(executor._shutdown) executor.shutdown(wait=True) self.assertIsNone(executor._pool) self.assertTrue(executor._shutdown) def test_submit_shutdown_cancel(self): executor = self.executor_type(max_workers=1) executor.bootup() num_workers = executor.num_workers for _ in range(num_workers * 100): executor.submit(time.sleep, 0.1) fut = executor.submit(time.sleep, 0) executor.shutdown(wait=False, cancel_futures=False) self.assertFalse(fut.cancelled()) executor.shutdown(wait=True, cancel_futures=True) self.assertTrue(fut.cancelled()) def test_submit_shutdown_cancel_wait(self): executor = self.executor_type(max_workers=1) executor.bootup() num_workers = executor.num_workers fut1 = executor.submit(time.sleep, 0.1) for _ in range(num_workers * 100): executor.submit(time.sleep, 0.1) fut2 = executor.submit(time.sleep, 0) fut3 = executor.submit(time.sleep, 0) time.sleep(0.2) executor.shutdown(wait=False, cancel_futures=True) _done, not_done = futures.wait({fut1, fut2, fut3}) self.assertEqual(len(not_done), 0) self.assertFalse(fut1.cancelled()) self.assertTrue(fut2.cancelled()) self.assertTrue(fut3.cancelled()) executor.shutdown(wait=True, cancel_futures=True) def test_shutdown_cancel(self): executor = self.executor_type(max_workers=1) executor.bootup() executor._pool.cancel() executor.shutdown(wait=False, cancel_futures=False) executor.shutdown(wait=False, cancel_futures=False) executor.shutdown(wait=False, cancel_futures=True) executor.shutdown(wait=False, cancel_futures=True) executor.shutdown(wait=True, cancel_futures=True) executor.shutdown(wait=True, cancel_futures=True) def test_init_bootup_shutdown(self): executor = self.executor_type(max_workers=1) self.assertIsNone(executor._pool) self.assertFalse(executor._shutdown) executor.bootup() self.assertTrue(executor._pool.event.is_set()) self.assertFalse(executor._shutdown) executor.shutdown() self.assertIsNone(executor._pool) self.assertTrue(executor._shutdown) def test_context_manager_shutdown(self): with self.executor_type(max_workers=1) as e: self.assertEqual( list(e.map(abs, range(-5, 5))), [5, 4, 3, 2, 1, 0, 1, 2, 3, 4] ) threads = [e._pool.thread] queues = [e._pool.queue] events = [e._pool.event] for t in threads: t.join() for q in queues: with self.assertRaises(LookupError): q.pop() for e in events: self.assertTrue(e.is_set()) def test_del_shutdown(self): executor = self.executor_type(max_workers=1) list(executor.map(abs, range(-5, 5))) threads = [executor._pool.thread] queues = [executor._pool.queue] events = [executor._pool.event] if hasattr(sys, "pypy_version_info"): executor.shutdown(False) else: del executor for t in threads: t.join() for q in queues: with self.assertRaises(LookupError): q.pop() for e in events: self.assertTrue(e.is_set()) class WaitTestMixin: def test_first_completed(self): future1 = self.executor.submit(mul, 21, 2) future2 = self.executor.submit(time.sleep, 0.25) done, not_done = futures.wait( [CANCELLED_FUTURE, future1, future2], return_when=futures.FIRST_COMPLETED, ) self.assertEqual({future1}, done) self.assertEqual({CANCELLED_FUTURE, future2}, not_done) def test_first_completed_some_already_completed(self): future1 = self.executor.submit(time.sleep, 0.5) finished, pending = futures.wait( [CANCELLED_AND_NOTIFIED_FUTURE, SUCCESSFUL_FUTURE, future1], return_when=futures.FIRST_COMPLETED, ) self.assertEqual( {CANCELLED_AND_NOTIFIED_FUTURE, SUCCESSFUL_FUTURE}, finished ) self.assertEqual({future1}, pending) def test_first_exception(self): future1 = self.executor.submit(mul, 2, 21) future2 = self.executor.submit(sleep_and_raise, 0.25) future3 = self.executor.submit(time.sleep, 0.5) finished, pending = futures.wait( [future1, future2, future3], return_when=futures.FIRST_EXCEPTION ) self.assertEqual({future1, future2}, finished) self.assertEqual({future3}, pending) def test_first_exception_some_already_complete(self): future1 = self.executor.submit(divmod, 21, 0) future2 = self.executor.submit(time.sleep, 0.5) finished, pending = futures.wait( [ SUCCESSFUL_FUTURE, CANCELLED_FUTURE, CANCELLED_AND_NOTIFIED_FUTURE, future1, future2, ], return_when=futures.FIRST_EXCEPTION, ) self.assertEqual( {SUCCESSFUL_FUTURE, CANCELLED_AND_NOTIFIED_FUTURE, future1}, finished, ) self.assertEqual({CANCELLED_FUTURE, future2}, pending) def test_first_exception_one_already_failed(self): future1 = self.executor.submit(time.sleep, 0.25) finished, pending = futures.wait( [EXCEPTION_FUTURE, future1], return_when=futures.FIRST_EXCEPTION ) self.assertEqual({EXCEPTION_FUTURE}, finished) self.assertEqual({future1}, pending) def test_all_completed(self): future1 = self.executor.submit(divmod, 2, 0) future2 = self.executor.submit(mul, 2, 21) finished, pending = futures.wait( [ SUCCESSFUL_FUTURE, CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, future1, future2, ], return_when=futures.ALL_COMPLETED, ) self.assertEqual( { SUCCESSFUL_FUTURE, CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, future1, future2, }, finished, ) self.assertEqual(set(), pending) def test_timeout(self): future1 = self.executor.submit(mul, 6, 7) future2 = self.executor.submit(time.sleep, 0.75) finished, pending = futures.wait( [ CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, future1, future2, ], timeout=0.5, return_when=futures.ALL_COMPLETED, ) self.assertEqual( { CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, future1, }, finished, ) self.assertEqual({future2}, pending) class ProcessPoolWaitTest(ProcessPoolMixin, WaitTestMixin, unittest.TestCase): pass class AsCompletedTestMixin: def test_no_timeout(self): future1 = self.executor.submit(mul, 2, 21) future2 = self.executor.submit(mul, 7, 6) completed = set( futures.as_completed([ CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, future1, future2, ]) ) self.assertEqual( { CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, future1, future2, }, completed, ) def test_zero_timeout(self): future1 = self.executor.submit(time.sleep, 0.5) completed_futures = set() try: completed_futures.update( futures.as_completed( [ CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, future1, ], timeout=0, ) ) except futures.TimeoutError: pass self.assertEqual( { CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, }, completed_futures, ) def test_nonzero_timeout(self): future1 = self.executor.submit(time.sleep, 0.0) self.executor.submit(time.sleep, 0.5) completed_futures = set() try: completed_futures.update( futures.as_completed( [ CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, future1, ], timeout=0.2, ) ) except futures.TimeoutError: pass self.assertEqual( { CANCELLED_AND_NOTIFIED_FUTURE, EXCEPTION_FUTURE, SUCCESSFUL_FUTURE, future1, }, completed_futures, ) def test_duplicate_futures(self): # Issue 20367. Duplicate futures should not raise exceptions or give # duplicate responses. future1 = self.executor.submit(time.sleep, 0.1) completed = list(futures.as_completed([future1, future1])) self.assertEqual(len(completed), 1) class ProcessPoolAsCompletedTest( ProcessPoolMixin, AsCompletedTestMixin, unittest.TestCase ): pass class ExecutorTestMixin: def test_submit(self): future = self.executor.submit(pow, 2, 8) self.assertEqual(256, future.result()) def test_submit_keyword(self): future = self.executor.submit(mul, 2, y=8) self.assertEqual(16, future.result()) future = self.executor.submit(mul, x=2, y=8) self.assertEqual(16, future.result()) def test_submit_cancel(self): fs = [] num_workers = self.executor.num_workers for _ in range(num_workers * 100): f = self.executor.submit(time.sleep, 0.1) fs.append(f) future = self.executor.submit(time.sleep, 0) future.cancel() self.assertTrue(future.cancelled()) for f in fs: f.cancel() def test_map(self): self.assertEqual( list(self.executor.map(pow, range(10), range(10))), list(map(pow, range(10), range(10))), ) def test_starmap(self): sequence = [(a, a) for a in range(10)] self.assertEqual( list(self.executor.starmap(pow, sequence)), list(map(pow, range(10), range(10))), ) self.assertEqual( list(self.executor.starmap(pow, iter(sequence))), list(map(pow, range(10), range(10))), ) def test_map_exception(self): i = self.executor.map(divmod, [1, 1, 1, 1], [2, 3, 0, 5]) self.assertEqual(next(i), (0, 1)) self.assertEqual(next(i), (0, 1)) with self.assertRaises(ZeroDivisionError): next(i) def test_map_timeout(self): results = [] try: for i in self.executor.map(time.sleep, [0, 0, 1], timeout=0.25): results.append(i) except futures.TimeoutError: pass else: self.fail("expected TimeoutError") self.assertEqual([None, None], results) def test_map_timeout_one(self): results = [] for i in self.executor.map(time.sleep, [0, 0, 0], timeout=1): results.append(i) self.assertEqual([None, None, None], results) class ProcessPoolExecutorTest( ProcessPoolMixin, ExecutorTestMixin, unittest.TestCase ): def test_map_chunksize(self): ref = list(map(pow, range(40), range(40))) self.assertEqual( list(self.executor.map(pow, range(40), range(40), chunksize=6)), ref, ) self.assertEqual( list(self.executor.map(pow, range(40), range(40), chunksize=50)), ref, ) self.assertEqual( list(self.executor.map(pow, range(40), range(40), chunksize=40)), ref, ) with self.assertRaises(ValueError): list(self.executor.map(pow, range(40), range(40), chunksize=-1)) def test_starmap_chunksize(self): ref = list(map(pow, range(40), range(40))) sequence = [(a, a) for a in range(40)] self.assertEqual( list(self.executor.starmap(pow, sequence, chunksize=6)), ref ) self.assertEqual( list(self.executor.starmap(pow, sequence, chunksize=50)), ref ) self.assertEqual( list(self.executor.starmap(pow, sequence, chunksize=40)), ref ) self.assertEqual( list(self.executor.starmap(pow, iter(sequence), chunksize=6)), ref ) self.assertEqual( list(self.executor.starmap(pow, iter(sequence), chunksize=50)), ref ) self.assertEqual( list(self.executor.starmap(pow, iter(sequence), chunksize=40)), ref ) with self.assertRaises(ValueError): list(self.executor.starmap(pow, sequence, chunksize=-1)) def test_map_buffersize(self): ref = list(map(pow, range(10), range(10))) for bs in (1, 2, 9, 10, 11): for cs in (1, 2, 9, 10, 11): self.assertEqual( list( self.executor.map( pow, range(10), range(10), chunksize=cs, buffersize=bs, ) ), ref, ) with self.assertRaises(TypeError): self.executor.map( pow, range(10), range(10), buffersize=str(1), ) with self.assertRaises(ValueError): self.executor.map( pow, range(10), range(10), buffersize=-1, ) with self.assertRaises(ValueError): self.executor.map( pow, range(10), range(10), buffersize=1, unordered=True, ) def test_map_unordered(self): map_unordered = functools.partial(self.executor.map, unordered=True) self.assertEqual( set(map_unordered(pow, range(10), range(10))), set(map(pow, range(10), range(10))), ) def test_map_unordered_timeout(self): map_unordered = functools.partial(self.executor.map, unordered=True) num_workers = self.executor.num_workers results = [] try: args = [1] + [0] * (num_workers - 1) for i in map_unordered(time.sleep, args, timeout=0.25): results.append(i) except futures.TimeoutError: pass else: self.fail("expected TimeoutError") self.assertEqual([None] * (num_workers - 1), results) def test_map_unordered_timeout_one(self): map_unordered = functools.partial(self.executor.map, unordered=True) results = [] for i in map_unordered(time.sleep, [0, 0, 0], timeout=1): results.append(i) self.assertEqual([None, None, None], results) def test_map_unordered_exception(self): map_unordered = functools.partial(self.executor.map, unordered=True) i = map_unordered(divmod, [1, 1, 1, 1], [2, 3, 0, 5]) try: self.assertEqual(next(i), (0, 1)) except ZeroDivisionError: return def test_map_unordered_chunksize(self): map_unordered = functools.partial(self.executor.map, unordered=True) ref = set(map(pow, range(40), range(40))) self.assertEqual( set(map_unordered(pow, range(40), range(40), chunksize=6)), ref ) self.assertEqual( set(map_unordered(pow, range(40), range(40), chunksize=50)), ref ) self.assertEqual( set(map_unordered(pow, range(40), range(40), chunksize=40)), ref ) with self.assertRaises(ValueError): set(map_unordered(pow, range(40), range(40), chunksize=-1)) class ProcessPoolSubmitTest(unittest.TestCase): @unittest.skipIf(MPI.get_vendor()[0] == "Microsoft MPI", "msmpi") def test_multiple_executors(self): executor1 = futures.MPIPoolExecutor(1).bootup(wait=True) executor2 = futures.MPIPoolExecutor(1).bootup(wait=True) executor3 = futures.MPIPoolExecutor(1).bootup(wait=True) fs1 = [executor1.submit(abs, i) for i in range(100, 200)] fs2 = [executor2.submit(abs, i) for i in range(200, 300)] fs3 = [executor3.submit(abs, i) for i in range(300, 400)] futures.wait(fs3 + fs2 + fs1) for i, f in enumerate(fs1): self.assertEqual(f.result(), i + 100) for i, f in enumerate(fs2): self.assertEqual(f.result(), i + 200) for i, f in enumerate(fs3): self.assertEqual(f.result(), i + 300) executor1 = executor2 = executor3 = None def test_mpi_serialized_support(self): futures._core.setup_mpi_threads() serialized = futures._core.serialized lock_save = serialized.lock try: if lock_save is None: serialized.lock = threading.Lock() executor = futures.MPIPoolExecutor(1).bootup() executor.submit(abs, 0).result() executor.shutdown() serialized.lock = lock_save else: serialized.lock = None with lock_save: executor = futures.MPIPoolExecutor(1).bootup() executor.submit(abs, 0).result() executor.shutdown() serialized.lock = lock_save finally: serialized.lock = lock_save def test_shared_executors(self): if not SHARED_POOL: return executors = [futures.MPIPoolExecutor() for _ in range(16)] fs = [] for i in range(128): fs.extend( e.submit(abs, i * 16 + j) for j, e in enumerate(executors) ) self.assertEqual(sorted(f.result() for f in fs), list(range(16 * 128))) world_size = MPI.COMM_WORLD.Get_size() num_workers = max(1, world_size - 1) for e in executors: self.assertEqual(e.num_workers, num_workers) del e, executors def inout(arg): return arg class GoodPickle: def __init__(self, value=0): self.value = value self.pickled = False self.unpickled = False def __getstate__(self): self.pickled = True return (self.value,) def __setstate__(self, state): self.unpickled = True self.value = state[0] class BadPickle: def __init__(self): self.pickled = False def __getstate__(self): self.pickled = True _ = 1 / 0 def __setstate__(self, state): pass class BadUnpickle: def __init__(self): self.pickled = False def __getstate__(self): self.pickled = True return (None,) def __setstate__(self, state): if state[0] is not None: raise ValueError _ = 1 / 0 @unittest.skipIf(SHARED_POOL and WORLD_SIZE == 1, "shared-pool") class ProcessPoolPickleTest(unittest.TestCase): def setUp(self): self.executor = futures.MPIPoolExecutor(1) def tearDown(self): self.executor.shutdown() def test_good_pickle(self): o = GoodPickle(42) r = self.executor.submit(inout, o).result() self.assertEqual(o.value, r.value) self.assertTrue(o.pickled) self.assertTrue(r.unpickled) r = self.executor.submit(GoodPickle, 77).result() self.assertEqual(r.value, 77) self.assertTrue(r.unpickled) def test_bad_pickle(self): o = BadPickle() self.assertFalse(o.pickled) f = self.executor.submit(inout, o) with self.assertRaises(ZeroDivisionError): f.result() self.assertTrue(o.pickled) f = self.executor.submit(BadPickle) with self.assertRaises(ZeroDivisionError): f.result() f = self.executor.submit(abs, 42) self.assertEqual(f.result(), 42) def test_bad_unpickle(self): o = BadUnpickle() self.assertFalse(o.pickled) f = self.executor.submit(inout, o) with self.assertRaises(ZeroDivisionError): f.result() self.assertTrue(o.pickled) f = self.executor.submit(BadUnpickle) with self.assertRaises(ZeroDivisionError): f.result() f = self.executor.submit(abs, 42) self.assertEqual(f.result(), 42) def test_exc_pickle(self): o = BadPickle() f = self.executor.submit(inout, o) exc = f.exception() self.assertIsInstance(exc, ZeroDivisionError) cause = exc.__cause__ self.assertIsNone(cause) def test_exc_unpickle(self): o = BadUnpickle() f = self.executor.submit(inout, o) exc = f.exception() self.assertIsInstance(exc, ZeroDivisionError) cause = exc.__cause__ self.assertIsInstance(cause, futures._core.RemoteTraceback) class MPICommExecutorTest(unittest.TestCase): MPICommExecutor = futures.MPICommExecutor def test_default(self): with self.MPICommExecutor() as executor: if executor is not None: executor.bootup() future1 = executor.submit(time.sleep, 0) future2 = executor.submit(time.sleep, 0) executor.shutdown() self.assertIsNone(future1.result()) self.assertIsNone(future2.result()) def test_self(self): with self.MPICommExecutor(MPI.COMM_SELF) as executor: future = executor.submit(time.sleep, 0) self.assertIsNone(future.result()) self.assertIsNone(future.exception()) future = executor.submit(sleep_and_raise, 0) with self.assertRaises(RuntimeError): future.result() self.assertEqual(RuntimeError, type(future.exception())) list(executor.map(time.sleep, [0, 0])) list(executor.map(time.sleep, [0, 0], timeout=1)) iterator = executor.map(time.sleep, [0.2, 0], timeout=0) with self.assertRaises(futures.TimeoutError): list(iterator) def test_args(self): with self.MPICommExecutor(MPI.COMM_SELF) as executor: self.assertIsNotNone(executor) with self.MPICommExecutor(MPI.COMM_SELF, 0) as executor: self.assertIsNotNone(executor) def test_kwargs(self): with self.MPICommExecutor(comm=MPI.COMM_SELF) as executor: self.assertIsNotNone(executor) with self.MPICommExecutor(comm=MPI.COMM_SELF, root=0) as executor: self.assertIsNotNone(executor) @unittest.skipIf(SHARED_POOL, "shared-pool") def test_arg_root(self): comm = MPI.COMM_WORLD rank = comm.Get_rank() for root in range(comm.Get_size()): with self.MPICommExecutor(comm, root) as executor: if rank == root: self.assertIsNotNone(executor) else: self.assertIsNone(executor) with self.MPICommExecutor(root=root) as executor: if rank == root: self.assertIsNotNone(executor) else: self.assertIsNone(executor) @unittest.skipIf(SHARED_POOL, "shared-pool") def test_arg_bad_root(self): size = MPI.COMM_WORLD.Get_size() with self.assertRaises(ValueError): self.MPICommExecutor(root=-size) with self.assertRaises(ValueError): self.MPICommExecutor(root=-1) with self.assertRaises(ValueError): self.MPICommExecutor(root=+size) @unittest.skipIf(SHARED_POOL, "shared-pool") def test_arg_bad_comm(self): if MPI.COMM_WORLD.Get_size() == 1: return intercomm, intracomm = futures._core.comm_split(MPI.COMM_WORLD, 0) try: with self.assertRaises(ValueError): self.MPICommExecutor(intercomm) finally: intercomm.Free() if intracomm: intracomm.Free() def test_with_bad(self): mpicommexecutor = self.MPICommExecutor(MPI.COMM_SELF) with mpicommexecutor: try: with mpicommexecutor: pass except RuntimeError: pass else: self.fail("expected RuntimeError") def test_initializer(self): mpicommexecutor = self.MPICommExecutor( initializer=time.sleep, initargs=(0,), ) with mpicommexecutor as executor: if executor is not None: executor.bootup() del executor with mpicommexecutor as executor: if executor is not None: executor.submit(time.sleep, 0).result() def test_initializer_error(self): mpicommexecutor = self.MPICommExecutor( initializer=sleep_and_raise, initargs=(0.2,), ) with mpicommexecutor as executor: if executor is not None: executor.submit(time.sleep, 0).cancel() future = executor.submit(time.sleep, 0) with self.assertRaises(futures.BrokenExecutor): executor.submit(time.sleep, 0).result() with self.assertRaises(futures.BrokenExecutor): future.result() def test_initializer_error_del(self): mpicommexecutor = self.MPICommExecutor( initializer=sleep_and_raise, initargs=(0.2,), ) with mpicommexecutor as executor: if executor is not None: executor.bootup() del executor def test_initializer_error_del_nowait(self): mpicommexecutor = self.MPICommExecutor( initializer=sleep_and_raise, initargs=(0.2,), ) with mpicommexecutor as executor: if executor is not None: executor.bootup(wait=False) executor.shutdown(wait=False) del executor def test_get_comm_workers(self): for _comm in (MPI.COMM_SELF, MPI.COMM_WORLD): with self.MPICommExecutor(MPI.COMM_SELF) as executor: num_workers = executor.submit(check_comm_workers).result() self.assertEqual(executor.num_workers, num_workers) self.assertRaises(RuntimeError, check_comm_workers) class ThreadPoolMixin(ExecutorMixin): executor_type = futures.ThreadPoolExecutor class ThreadPoolTest( ThreadPoolMixin, ExecutorTestMixin, ExecutorShutdownTestMixin, unittest.TestCase, ): pass from mpi4py.futures.aplus import ThenableFuture # noqa: E402 class ThenTest(unittest.TestCase): assert_ = unittest.TestCase.assertTrue def test_cancel_base(self): base_f = ThenableFuture() new_f = base_f.then() self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.cancel() self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(base_f.cancelled()) self.assertTrue(new_f.cancelled()) def test_cancel_new(self): base_f = ThenableFuture() new_f = base_f.then() self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) new_f.cancel() self.assertTrue(not base_f.done()) self.assertTrue(new_f.done()) base_f.set_result(1) self.assertTrue(base_f.done()) self.assertTrue(new_f.cancelled()) def test_then_multiple(self): base_f = ThenableFuture() new_f1 = base_f.then() new_f2 = base_f.then() new_f3 = base_f.then() self.assertIsNot(base_f, new_f1) self.assertIsNot(base_f, new_f2) self.assertIsNot(base_f, new_f3) self.assertTrue(not base_f.done()) self.assertTrue(not new_f1.done()) self.assertTrue(not new_f2.done()) self.assertTrue(not new_f3.done()) base_f.set_result("done") self.assertTrue(base_f.done()) self.assertTrue(new_f1.done()) self.assertTrue(new_f2.done()) self.assertTrue(new_f3.done()) self.assertTrue(not new_f1.exception()) self.assertTrue(not new_f2.exception()) self.assertTrue(not new_f3.exception()) self.assertEqual(new_f1.result(), "done") self.assertEqual(new_f2.result(), "done") self.assertEqual(new_f3.result(), "done") def test_no_callbacks_and_success(self): base_f = ThenableFuture() new_f = base_f.then() self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_result("done") self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(not new_f.exception()) self.assertEqual(new_f.result(), "done") def test_no_callbacks_and_failure(self): class MyException(Exception): pass base_f = ThenableFuture() new_f = base_f.then() self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_exception(MyException("sad")) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(new_f.exception()) with self.assertRaises(MyException) as catcher: new_f.result() self.assertEqual(catcher.exception.args[0], "sad") def test_success_callback_and_success(self): base_f = ThenableFuture() new_f = base_f.then(lambda result: result + " manipulated") self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_result("done") self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(not new_f.exception()) self.assertEqual(new_f.result(), "done manipulated") def test_err_callback_and_failure_repackage(self): class MyException(Exception): pass class MyRepackagedException(Exception): pass class NotMatched(Exception): pass def on_failure(ex): if isinstance(ex, MyException): return MyRepackagedException(ex.args[0] + " repackaged") else: return NotMatched("?") base_f = ThenableFuture() new_f = base_f.then(None, on_failure) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_exception(MyException("sad")) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(new_f.exception()) with self.assertRaises(MyRepackagedException) as catcher: new_f.result() self.assertEqual(catcher.exception.args[0], "sad repackaged") def test_err_callback_and_failure_raised(self): class MyException(Exception): pass class MyRepackagedException(Exception): pass def raise_something_else(ex): raise MyRepackagedException(ex.args[0] + " repackaged") base_f = ThenableFuture() new_f = base_f.then(None, raise_something_else) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_exception(MyException("sad")) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(new_f.exception()) with self.assertRaises(MyRepackagedException) as catcher: new_f.result() self.assertEqual(catcher.exception.args[0], "sad repackaged") def test_err_callback_convert_to_success(self): class MyException(Exception): pass class NotMatched(Exception): pass def on_failure(ex): if isinstance(ex, MyException): return ex.args[0] + " repackaged" else: return NotMatched("?") base_f = ThenableFuture() new_f = base_f.catch(on_failure) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_exception(MyException("sad")) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(not new_f.exception()) self.assertEqual(new_f.result(), "sad repackaged") def test_err_catch_ignore(self): base_f = ThenableFuture() new_f = base_f.catch() self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_exception(Exception("sad")) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertIsNone(new_f.exception()) self.assertIsNone(new_f.result()) def test_success_callback_and_failure_raised(self): class MyException(Exception): pass def raise_something_else(value): raise MyException(value + " repackaged") base_f = ThenableFuture() new_f = base_f.then(raise_something_else) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_result("sad") self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(new_f.exception()) with self.assertRaises(MyException) as catcher: new_f.result() assert catcher.exception.args[0] == "sad repackaged" def test_chained_success_callback_and_success(self): def transform(value): f = ThenableFuture() if value < 5: f.set_result(transform(value + 1)) else: f.set_result(value) return f base_f = ThenableFuture() new_f = base_f.then(transform) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_result(1) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(not new_f.exception()) self.assertEqual(new_f.result(), 5) def test_chained_failure_callback_and_success(self): def transform(exc): self.assertIsInstance(exc, RuntimeError) f = ThenableFuture() f.set_result(5) return f base_f = ThenableFuture() new_f = base_f.catch(transform) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_exception(RuntimeError()) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(not new_f.exception()) self.assertEqual(new_f.result(), 5) def test_detect_cycle_chain(self): f1 = ThenableFuture() f2 = ThenableFuture() chain = [f1, f2, f1] def transform(a): try: f = chain.pop(0) r = transform(a) f.__init__() f.set_result(r) except IndexError: return 42 else: return f base_f = ThenableFuture() new_f = base_f.then(transform) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_result(1) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(new_f.exception()) with self.assertRaises(RuntimeError) as catcher: new_f.result() self.assertIn( "chain cycle detected", catcher.exception.args[0], ) def test_detect_self_chain(self): base_f = ThenableFuture() new_f = base_f.then(lambda _: new_f) self.assertIsNot(base_f, new_f) self.assertTrue(not base_f.done()) self.assertTrue(not new_f.done()) base_f.set_result(1) self.assertTrue(base_f.done()) self.assertTrue(new_f.done()) self.assertTrue(new_f.exception()) with self.assertRaises(RuntimeError) as catcher: new_f.result() self.assertIn( "chain cycle detected", catcher.exception.args[0], ) class CollectTest(unittest.TestCase): def test_empty(self): future = futures.collect([]) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertTrue(future.done()) self.assertEqual(future.result(), []) def test_item_success(self): fs = [futures.Future() for _ in range(5)] future = futures.collect(fs) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertFalse(future.done()) for i in range(5): fs[i].set_result(i) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertTrue(future.done()) self.assertEqual(future.result(), list(range(5))) def test_item_failure(self): fs = [futures.Future() for _ in range(5)] future = futures.collect(fs) for i in range(2, 4): fs[i].set_result(i) fs[-1].set_exception(RuntimeError()) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertTrue(future.done()) self.assertIsInstance(future.exception(), RuntimeError) for i in range(2): self.assertTrue(fs[i].cancelled()) for i in range(2, 4): self.assertFalse(fs[i].cancelled()) self.assertFalse(fs[-1].cancelled()) def test_item_done(self): fs = [futures.Future() for _ in range(5)] for i in range(5): fs[i].set_result(i) future = futures.collect(fs) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertTrue(future.done()) self.assertEqual(future.result(), list(range(5))) def test_item_cancel(self): fs = [futures.Future() for _ in range(5)] future = futures.collect(fs) for i in range(2, 4): fs[i].set_result(i) fs[-1].cancel() self.assertTrue(future.cancelled()) self.assertFalse(future.running()) self.assertTrue(future.done()) for i in range(2): self.assertTrue(fs[i].cancelled()) for i in range(2, 4): self.assertFalse(fs[i].cancelled()) self.assertTrue(fs[-1].cancelled()) def test_cancel(self): fs = [futures.Future() for _ in range(5)] future = futures.collect(fs) future.cancel() for f in fs: self.assertTrue(f.cancelled()) def test_cancel_pending(self): class MyFuture(futures.Future): def cancel(self): pass fs = [MyFuture() for _ in range(5)] future = futures.collect(fs) self.assertIs(type(future), MyFuture) super(MyFuture, future).cancel() for f in fs: self.assertFalse(f.cancelled()) f.set_result(None) class ComposeTest(unittest.TestCase): def test_result(self): base = futures.Future() future = futures.compose(base) self.assertIs(type(future), type(base)) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertFalse(future.done()) base.set_result(42) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertTrue(future.done()) self.assertEqual(future.result(), 42) def test_except(self): base = futures.Future() future = futures.compose(base) self.assertIs(type(future), type(base)) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertFalse(future.done()) base.set_exception(RuntimeError(42)) self.assertFalse(future.cancelled()) self.assertFalse(future.running()) self.assertTrue(future.done()) self.assertIs(type(future.exception()), RuntimeError) self.assertEqual(future.exception().args, (42,)) def test_cancel_new(self): base = futures.Future() future = futures.compose(base) base.cancel() self.assertTrue(future.cancelled()) def test_cancel_old(self): base = futures.Future() future = futures.compose(base) future.cancel() self.assertTrue(base.cancelled()) def test_result_hook(self): base = futures.Future() future = futures.compose(base, int) base.set_result("42") self.assertEqual(future.result(), 42) def test_result_hook_failure(self): base = futures.Future() future = futures.compose(base, resulthook=lambda _: 1 / 0) base.set_result(42) self.assertIs(type(future.exception()), ZeroDivisionError) def test_except_hook(self): base = futures.Future() future = futures.compose(base, excepthook=lambda exc: exc.args[0]) base.set_exception(RuntimeError(42)) self.assertEqual(future.result(), 42) def test_except_hook_except(self): base = futures.Future() future = futures.compose( base, excepthook=lambda exc: RuntimeError(exc.args[0]) ) base.set_exception(ValueError(42)) self.assertIs(type(future.exception()), RuntimeError) self.assertEqual(future.exception().args, (42,)) def test_except_hook_failure(self): base = futures.Future() future = futures.compose(base, excepthook=lambda _: 1 / 0) base.set_exception(ValueError(42)) self.assertIs(type(future.exception()), ZeroDivisionError) SKIP_POOL_TEST = testutil.disable_mpi_spawn() if SHARED_POOL: del ProcessPoolInitTest.test_init_sys_flags del ProcessPoolInitTest.test_init_globals del ProcessPoolInitTest.test_use_pkl5_kwarg del ProcessPoolInitTest.test_use_pkl5_environ if WORLD_SIZE == 1: del ProcessPoolInitTest.test_run_name if WORLD_SIZE > 2: del ProcessPoolInitTest.test_max_workers del ProcessPoolInitTest.test_max_workers_environ if WORLD_SIZE == 1: del ProcessPoolPickleTest del MPICommExecutorTest.test_arg_root del MPICommExecutorTest.test_arg_bad_root del MPICommExecutorTest.test_arg_bad_comm elif WORLD_SIZE > 1 or SKIP_POOL_TEST: del ProcessPoolInitTest del ProcessPoolBootupTest del ProcessPoolShutdownTest del ProcessPoolWaitTest del ProcessPoolAsCompletedTest del ProcessPoolExecutorTest del ProcessPoolSubmitTest del ProcessPoolPickleTest if not SHARED_POOL: del SharedPoolInitTest if __name__ == "__main__": unittest.main()