type hint fixes for adaptive/runner.py

Author: basnijholt

adaptive/runner.py

@@ -9,36 +9,39 @@
 import traceback
 import warnings
 from _asyncio import Future, Task
-from concurrent.futures.process import ProcessPoolExecutor
 from contextlib import suppress
-from typing import Any, Callable, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
-from distributed.cfexecutor import ClientExecutor
-from distributed.client import Client
-from ipyparallel.client.asyncresult import AsyncResult
-from ipyparallel.client.view import ViewExecutor
-
-from adaptive.learner import BaseLearner
+from adaptive.learner.base_learner import BaseLearner
 from adaptive.notebook_integration import in_ipynb, live_info, live_plot
 
+_ThirdPartyClient = []
+_ThirdPartyExecutor = []
+
 try:
     import ipyparallel
+    from ipyparallel.client.asyncresult import AsyncResult
 
     with_ipyparallel = True
+    _ThirdPartyClient.append(ipyparallel.Client)
+    _ThirdPartyExecutor.append(ipyparallel.client.view.ViewExecutor)
 except ModuleNotFoundError:
     with_ipyparallel = False
 
 try:
     import distributed
 
     with_distributed = True
+    _ThirdPartyClient.append(distributed.client.Client)
+    _ThirdPartyExecutor.append(distributed.cfexecutor.ClientExecutor)
 except ModuleNotFoundError:
     with_distributed = False
 
 try:
     import mpi4py.futures
 
     with_mpi4py = True
+    _ThirdPartyExecutor.append(mpi4py.futures.MPIPoolExecutor)
 except ModuleNotFoundError:
     with_mpi4py = False
 
@@ -47,6 +50,8 @@
 
     asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
+ThirdPartyClient = Union[tuple(_ThirdPartyClient)]
+ThirdPartyExecutor = Union[tuple(_ThirdPartyExecutor)]
 
 if os.name == "nt":
     if with_distributed:
@@ -72,8 +77,80 @@ def _default_executor(*args, **kwargs):
     _default_executor_kwargs = {}
 
 
+# -- Internal executor-related, things
+
+
+class SequentialExecutor(concurrent.Executor):
+    """A trivial executor that runs functions synchronously.
+
+    This executor is mainly for testing.
+    """
+
+    def submit(self, fn: Callable, *args, **kwargs) -> Future:
+        fut = concurrent.Future()
+        try:
+            fut.set_result(fn(*args, **kwargs))
+        except Exception as e:
+            fut.set_exception(e)
+        return fut
+
+    def map(self, fn, *iterable, timeout=None, chunksize=1):
+        return map(fn, iterable)
+
+    def shutdown(self, wait=True):
+        pass
+
+
+def _ensure_executor(
+    executor: Optional[Union[ThirdPartyClient, concurrent.Executor]]
+) -> concurrent.Executor:
+    if executor is None:
+        executor = _default_executor(**_default_executor_kwargs)
+
+    if isinstance(executor, concurrent.Executor):
+        return executor
+    elif with_ipyparallel and isinstance(executor, ipyparallel.Client):
+        return executor.executor()
+    elif with_distributed and isinstance(executor, distributed.Client):
+        return executor.get_executor()
+    else:
+        raise TypeError(
+            "Only a concurrent.futures.Executor, distributed.Client,"
+            " or ipyparallel.Client can be used."
+        )
+
+
+def _get_ncores(
+    ex: Union[
+        ThirdPartyExecutor,
+        concurrent.ProcessPoolExecutor,
+        concurrent.ThreadPoolExecutor,
+        SequentialExecutor,
+    ]
+) -> int:
+    """Return the maximum  number of cores that an executor can use."""
+    if with_ipyparallel and isinstance(ex, ipyparallel.client.view.ViewExecutor):
+        return len(ex.view)
+    elif isinstance(
+        ex, (concurrent.ProcessPoolExecutor, concurrent.ThreadPoolExecutor)
+    ):
+        return ex._max_workers  # not public API!
+    elif isinstance(ex, SequentialExecutor):
+        return 1
+    elif with_distributed and isinstance(ex, distributed.cfexecutor.ClientExecutor):
+        return sum(n for n in ex._client.ncores().values())
+    elif with_mpi4py and isinstance(ex, mpi4py.futures.MPIPoolExecutor):
+        ex.bootup()  # wait until all workers are up and running
+        return ex._pool.size  # not public API!
+    else:
+        raise TypeError(f"Cannot get number of cores for {ex.__class__}")
+
+
+# -- Runner definitions
+
+
 class BaseRunner(metaclass=abc.ABCMeta):
-    r"""Base class for runners that use `concurrent.futures.Executors`.
+    r"""Base class for runners that use `concurrent.futures.Executor`\'s.
 
     Parameters
     ----------
@@ -133,12 +210,17 @@ def __init__(
         learner: BaseLearner,
         goal: Callable,
         *,
-        executor=None,
-        ntasks=None,
-        log=False,
-        shutdown_executor=False,
-        retries=0,
-        raise_if_retries_exceeded=True,
+        executor: Union[
+            ThirdPartyExecutor,
+            concurrent.ProcessPoolExecutor,
+            concurrent.ThreadPoolExecutor,
+            SequentialExecutor,
+        ] = None,
+        ntasks: int = None,
+        log: bool = False,
+        shutdown_executor: bool = False,
+        retries: int = 0,
+        raise_if_retries_exceeded: bool = True,
     ) -> None:
 
         self.executor = _ensure_executor(executor)
@@ -153,7 +235,7 @@ def __init__(
         self.shutdown_executor = shutdown_executor or (executor is None)
 
         self.learner = learner
-        self.log = [] if log else None
+        self.log: Optional[list] = [] if log else None
 
         # Timing
         self.start_time = time.time()
@@ -216,7 +298,12 @@ def overhead(self):
         return (1 - t_function / t_total) * 100
 
     def _process_futures(
-        self, done_futs: Union[Set[Future], Set[Future], Set[AsyncResult], Set[Task]]
+        self,
+        done_futs: Union[
+            Set[Future],
+            Set[AsyncResult],  # XXX: AsyncResult might not be imported
+            Set[Task],
+        ],
     ) -> None:
         for fut in done_futs:
             x = self.pending_points.pop(fut)
@@ -240,7 +327,11 @@ def _process_futures(
 
     def _get_futures(
         self,
-    ) -> Union[List[Task], List[Future], List[Future], List[AsyncResult]]:
+    ) -> Union[
+        List[Task],
+        List[Future],
+        List[AsyncResult],  # XXX: AsyncResult might not be imported
+    ]:
         # Launch tasks to replace the ones that completed
         # on the last iteration, making sure to fill workers
         # that have started since the last iteration.
@@ -363,8 +454,13 @@ def __init__(
         learner: BaseLearner,
         goal: Callable,
         *,
-        executor=None,
-        ntasks=None,
+        executor: Union[
+            ThirdPartyExecutor,
+            concurrent.ProcessPoolExecutor,
+            concurrent.ThreadPoolExecutor,
+            SequentialExecutor,
+        ] = None,
+        ntasks: Optional[int] = None,
         log=False,
         shutdown_executor=False,
         retries=0,
@@ -386,9 +482,7 @@ def __init__(
         )
         self._run()
 
-    def _submit(
-        self, x: Union[Tuple[int, int], int, Tuple[float, float], float]
-    ) -> Union[Future, AsyncResult]:
+    def _submit(self, x: Union[Tuple[float, ...], float, int]) -> Future:
         return self.executor.submit(self.learner.function, x)
 
     def _run(self) -> None:
@@ -494,13 +588,18 @@ def __init__(
         learner: BaseLearner,
         goal: Optional[Callable] = None,
         *,
-        executor=None,
-        ntasks=None,
-        log=False,
-        shutdown_executor=False,
+        executor: Union[
+            ThirdPartyExecutor,
+            concurrent.ProcessPoolExecutor,
+            concurrent.ThreadPoolExecutor,
+            SequentialExecutor,
+        ] = None,
+        ntasks: Optional[int] = None,
+        log: bool = False,
+        shutdown_executor: bool = False,
         ioloop=None,
-        retries=0,
-        raise_if_retries_exceeded=True,
+        retries: int = 0,
+        raise_if_retries_exceeded: bool = True,
     ) -> None:
 
         if goal is None:
@@ -640,7 +739,7 @@ async def _run(self) -> None:
                 await asyncio.wait(remaining)
             self._cleanup()
 
-    def elapsed_time(self):
+    def elapsed_time(self) -> float:
         """Return the total time elapsed since the runner
         was started."""
         if self.task.done():
@@ -653,7 +752,7 @@ def elapsed_time(self):
             end_time = time.time()
         return end_time - self.start_time
 
-    def start_periodic_saving(self, save_kwargs, interval):
+    def start_periodic_saving(self, save_kwargs: Dict[str, Any], interval: int):
         """Periodically save the learner's data.
 
         Parameters
@@ -711,16 +810,7 @@ def simple(learner: BaseLearner, goal: Callable) -> None:
             learner.tell(x, y)
 
 
-def replay_log(
-    learner: BaseLearner,
-    log: List[
-        Union[
-            Tuple[str, int],
-            Tuple[str, Tuple[int, int, int], float],
-            Tuple[str, Tuple[float, float, float], float],
-        ]
-    ],
-) -> None:
+def replay_log(learner: BaseLearner, log) -> None:
     """Apply a sequence of method calls to a learner.
 
     This is useful for debugging runners.
@@ -771,67 +861,3 @@ def stop_after(*, seconds=0, minutes=0, hours=0) -> Callable:
     """
     stop_time = time.time() + seconds + 60 * minutes + 3600 * hours
     return lambda _: time.time() > stop_time
-
-
-# -- Internal executor-related, things
-
-
-class SequentialExecutor(concurrent.Executor):
-    """A trivial executor that runs functions synchronously.
-
-    This executor is mainly for testing.
-    """
-
-    def submit(self, fn: Callable, *args, **kwargs) -> Future:
-        fut = concurrent.Future()
-        try:
-            fut.set_result(fn(*args, **kwargs))
-        except Exception as e:
-            fut.set_exception(e)
-        return fut
-
-    def map(self, fn, *iterable, timeout=None, chunksize=1):
-        return map(fn, iterable)
-
-    def shutdown(self, wait=True):
-        pass
-
-
-def _ensure_executor(
-    executor: Optional[Union[Client, Client, ProcessPoolExecutor, SequentialExecutor]]
-) -> Union[SequentialExecutor, ProcessPoolExecutor, ViewExecutor, ClientExecutor]:
-    if executor is None:
-        executor = _default_executor(**_default_executor_kwargs)
-
-    if isinstance(executor, concurrent.Executor):
-        return executor
-    elif with_ipyparallel and isinstance(executor, ipyparallel.Client):
-        return executor.executor()
-    elif with_distributed and isinstance(executor, distributed.Client):
-        return executor.get_executor()
-    else:
-        raise TypeError(
-            "Only a concurrent.futures.Executor, distributed.Client,"
-            " or ipyparallel.Client can be used."
-        )
-
-
-def _get_ncores(
-    ex: Union[SequentialExecutor, ProcessPoolExecutor, ViewExecutor, ClientExecutor]
-) -> int:
-    """Return the maximum  number of cores that an executor can use."""
-    if with_ipyparallel and isinstance(ex, ipyparallel.client.view.ViewExecutor):
-        return len(ex.view)
-    elif isinstance(
-        ex, (concurrent.ProcessPoolExecutor, concurrent.ThreadPoolExecutor)
-    ):
-        return ex._max_workers  # not public API!
-    elif isinstance(ex, SequentialExecutor):
-        return 1
-    elif with_distributed and isinstance(ex, distributed.cfexecutor.ClientExecutor):
-        return sum(n for n in ex._client.ncores().values())
-    elif with_mpi4py and isinstance(ex, mpi4py.futures.MPIPoolExecutor):
-        ex.bootup()  # wait until all workers are up and running
-        return ex._pool.size  # not public API!
-    else:
-        raise TypeError(f"Cannot get number of cores for {ex.__class__}")