Add lots of simple type annotations throughout parsl/dataflow (#2414)

This comes from the benc-mypy branch

The benc-mypy branch contains some more complicated situations,
which are not merged by this PR.

Author: benclifford

parsl/dataflow/dflow.py

@@ -244,7 +244,7 @@ def _create_task_log_info(self, task_record):
             task_log_info['task_joins'] = None
         return task_log_info
 
-    def _count_deps(self, depends):
+    def _count_deps(self, depends: Sequence[Future]) -> int:
         """Count the number of unresolved futures in the list depends.
         """
         count = 0
@@ -256,7 +256,7 @@ def _count_deps(self, depends):
         return count
 
     @property
-    def config(self):
+    def config(self) -> Config:
         """Returns the fully initialized config that the DFK is actively using.
 
         Returns:
@@ -483,14 +483,14 @@ def _unwrap_remote_exception_wrapper(future: Future) -> Any:
             result.reraise()
         return result
 
-    def wipe_task(self, task_id):
+    def wipe_task(self, task_id: int) -> None:
         """ Remove task with task_id from the internal tasks table
         """
         if self.config.garbage_collect:
             del self.tasks[task_id]
 
     @staticmethod
-    def check_staging_inhibited(kwargs):
+    def check_staging_inhibited(kwargs: Dict[str, Any]) -> bool:
         return kwargs.get('_parsl_staging_inhibit', False)
 
     def launch_if_ready(self, task_record: TaskRecord) -> None:
@@ -642,7 +642,7 @@ def launch_task(self, task_record: TaskRecord, executable: Callable, *args: Tupl
 
         return exec_fu
 
-    def _add_input_deps(self, executor, args, kwargs, func):
+    def _add_input_deps(self, executor: str, args: Sequence[Any], kwargs: Dict[str, Any], func: Callable) -> Tuple[Sequence[Any], Dict[str, Any], Callable]:
         """Look for inputs of the app that are files. Give the data manager
         the opportunity to replace a file with a data future for that file,
         for example wrapping the result of a staging action.
@@ -672,7 +672,7 @@ def _add_input_deps(self, executor, args, kwargs, func):
 
         return tuple(newargs), kwargs, func
 
-    def _add_output_deps(self, executor, args, kwargs, app_fut, func):
+    def _add_output_deps(self, executor: str, args: Sequence[Any], kwargs: Dict[str, Any], app_fut: AppFuture, func: Callable) -> Callable:
         logger.debug("Adding output dependencies")
         outputs = kwargs.get('outputs', [])
         app_fut._outputs = []
@@ -720,7 +720,7 @@ def _gather_all_deps(self, args: Sequence[Any], kwargs: Dict[str, Any]) -> List[
         """
         depends: List[Future] = []
 
-        def check_dep(d):
+        def check_dep(d: Any) -> None:
             if isinstance(d, Future):
                 depends.extend([d])
 
@@ -953,7 +953,7 @@ def submit(self, func, app_args, executors='all', cache=False, ignore_for_cache=
 
         for d in depends:
 
-            def callback_adapter(dep_fut):
+            def callback_adapter(dep_fut: Future) -> None:
                 self.launch_if_ready(task_def)
 
             try:
@@ -972,7 +972,7 @@ def callback_adapter(dep_fut):
     # and a drain function might look like this.
     # If tasks have their states changed, this won't work properly
     # but we can validate that...
-    def log_task_states(self):
+    def log_task_states(self) -> None:
         logger.info("Summary of tasks in DFK:")
 
         with self.task_state_counts_lock:
@@ -1221,7 +1221,7 @@ def checkpoint(self, tasks: Optional[Sequence[TaskRecord]] = None) -> str:
 
             return checkpoint_dir
 
-    def _load_checkpoints(self, checkpointDirs):
+    def _load_checkpoints(self, checkpointDirs: Sequence[str]) -> Dict[str, Future[Any]]:
         """Load a checkpoint file into a lookup table.
 
         The data being loaded from the pickle file mostly contains input
@@ -1248,7 +1248,7 @@ def _load_checkpoints(self, checkpointDirs):
                         try:
                             data = pickle.load(f)
                             # Copy and hash only the input attributes
-                            memo_fu = Future()
+                            memo_fu: Future = Future()
                             assert data['exception'] is None
                             memo_fu.set_result(data['result'])
                             memo_lookup_table[data['hash']] = memo_fu

parsl/dataflow/flow_control.py

@@ -2,6 +2,9 @@
 import threading
 import time
 
+from typing import Sequence
+
+from parsl.executors.base import ParslExecutor
 from parsl.dataflow.task_status_poller import TaskStatusPoller
 
 logger = logging.getLogger(__name__)
@@ -67,7 +70,7 @@ def __init__(self, dfk, *args, threshold=20, interval=5):
         self._thread.daemon = True
         self._thread.start()
 
-    def _wake_up_timer(self, kill_event):
+    def _wake_up_timer(self, kill_event: threading.Event) -> None:
         """Internal. This is the function that the thread will execute.
         waits on an event so that the thread can make a quick exit when close() is called
 
@@ -95,7 +98,7 @@ def notify(self, event_id):
             logger.debug("Eventcount >= threshold")
             self.make_callback()
 
-    def make_callback(self):
+    def make_callback(self) -> None:
         """Makes the callback and resets the timer.
         """
         self._wake_up_time = time.time() + self.interval
@@ -105,10 +108,10 @@ def make_callback(self):
             logger.error("Flow control callback threw an exception - logging and proceeding anyway", exc_info=True)
         self._event_buffer = []
 
-    def add_executors(self, executors):
+    def add_executors(self, executors: Sequence[ParslExecutor]) -> None:
         self.task_status_poller.add_executors(executors)
 
-    def close(self):
+    def close(self) -> None:
         """Merge the threads and terminate."""
         self._kill_event.set()
         self._thread.join()
@@ -161,7 +164,7 @@ def __init__(self, callback, *args, interval=5, name=None):
         self._thread.daemon = True
         self._thread.start()
 
-    def _wake_up_timer(self, kill_event):
+    def _wake_up_timer(self, kill_event: threading.Event) -> None:
         """Internal. This is the function that the thread will execute.
         waits on an event so that the thread can make a quick exit when close() is called
 
@@ -185,13 +188,13 @@ def _wake_up_timer(self, kill_event):
             else:
                 print("Sleeping a bit more")
 
-    def make_callback(self):
+    def make_callback(self) -> None:
         """Makes the callback and resets the timer.
         """
         self._wake_up_time = time.time() + self.interval
         self.callback(*self.cb_args)
 
-    def close(self):
+    def close(self) -> None:
         """Merge the threads and terminate.
         """
         self._kill_event.set()

parsl/dataflow/strategy.py

@@ -130,18 +130,18 @@ def add_executors(self, executors):
         for executor in executors:
             self.executors[executor.label] = {'idle_since': None, 'config': executor.label}
 
-    def _strategy_noop(self, status: List[ExecutorStatus], tasks):
+    def _strategy_noop(self, status: List[ExecutorStatus], tasks: List[int]) -> None:
         """Do nothing.
 
         Args:
             - tasks (task_ids): Not used here.
         """
         logger.debug("strategy_noop: doing nothing")
 
-    def _strategy_simple(self, status_list, tasks):
+    def _strategy_simple(self, status_list, tasks: List[int]) -> None:
         self._general_strategy(status_list, tasks, strategy_type='simple')
 
-    def _strategy_htex_auto_scale(self, status_list, tasks):
+    def _strategy_htex_auto_scale(self, status_list, tasks: List[int]) -> None:
         """HTEX specific auto scaling strategy
 
         This strategy works only for HTEX. This strategy will scale out by

parsl/dataflow/task_status_poller.py

@@ -36,10 +36,10 @@ def __init__(self, executor: ParslExecutor, dfk: "parsl.dataflow.dflow.DataFlowK
             self.hub_channel.connect("tcp://{}:{}".format(hub_address, hub_port))
             logger.info("Monitoring enabled on task status poller")
 
-    def _should_poll(self, now: float):
+    def _should_poll(self, now: float) -> bool:
         return now >= self._last_poll_time + self._interval
 
-    def poll(self, now: float):
+    def poll(self, now: float) -> None:
         if self._should_poll(now):
             previous_status = self._status
             self._status = self._executor.status()
@@ -53,7 +53,7 @@ def poll(self, now: float):
             if delta_status:
                 self.send_monitoring_info(delta_status)
 
-    def send_monitoring_info(self, status=None):
+    def send_monitoring_info(self, status: Dict):
         # Send monitoring info for HTEX when monitoring enabled
         if self.monitoring_enabled:
             msg = self._executor.create_monitoring_info(status)
@@ -95,7 +95,7 @@ def scale_out(self, n):
             self._status.update(new_status)
         return block_ids
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return self._status.__repr__()
 
 
@@ -111,12 +111,12 @@ def poll(self, tasks=None):
         self._error_handler.run(self._poll_items)
         self._strategy.strategize(self._poll_items, tasks)
 
-    def _update_state(self):
+    def _update_state(self) -> None:
         now = time.time()
         for item in self._poll_items:
             item.poll(now)
 
-    def add_executors(self, executors: Sequence[ParslExecutor]):
+    def add_executors(self, executors: Sequence[ParslExecutor]) -> None:
         for executor in executors:
             if executor.status_polling_interval > 0:
                 logger.debug("Adding executor {}".format(executor.label))