split update_memo into two notify-exception and notify-result forms, that run before the future is set and which

means checkpointing can run in that place rather than racing with the user workflow.

this is needed to address an existing race condition - that task_exit tasks aren't checkpointed by the time
the workflow observes they have finished. it's debatable whether that is actually the intended functionality
or not, but i think it should be. otherwise what's the behaviour? a little bit later?

this manifests as awkwardness in implementation in later checkpoint methods too.

and moving more checkpoint code into update_memo makes the checkpoint/memo API unification cleaner.

Author: benclifford

parsl/dataflow/dflow.py

@@ -370,7 +370,7 @@ def handle_exec_update(self, task_record: TaskRecord, future: Future) -> None:
                 logger.info("Task {} failed due to dependency failure so skipping retries".format(task_id))
                 task_record['time_returned'] = datetime.datetime.now()
                 self._send_task_log_info(task_record)
-                self.memoizer.update_memo(task_record)
+                self.memoizer.update_memo_exception(task_record, e)
                 with task_record['app_fu']._update_lock:
                     task_record['app_fu'].set_exception(e)
 
@@ -396,7 +396,7 @@ def handle_exec_update(self, task_record: TaskRecord, future: Future) -> None:
                 self.update_task_state(task_record, States.failed)
                 task_record['time_returned'] = datetime.datetime.now()
                 self._send_task_log_info(task_record)
-                self.memoizer.update_memo(task_record)
+                self.memoizer.update_memo_exception(task_record, e)
                 with task_record['app_fu']._update_lock:
                     task_record['app_fu'].set_exception(e)
 
@@ -443,10 +443,10 @@ def handle_exec_update(self, task_record: TaskRecord, future: Future) -> None:
                         self.update_task_state(task_record, States.failed)
                         task_record['time_returned'] = datetime.datetime.now()
                         self._send_task_log_info(task_record)
-                        self.memoizer.update_memo(task_record)
+                        ex = TypeError(f"join_app body must return a Future or list of Futures, got {joinable} of type {type(joinable)}")
+                        self.memoizer.update_memo_exception(task_record, ex)
                         with task_record['app_fu']._update_lock:
-                            task_record['app_fu'].set_exception(
-                                TypeError(f"join_app body must return a Future or list of Futures, got {joinable} of type {type(joinable)}"))
+                            task_record['app_fu'].set_exception(ex)
 
         self._log_std_streams(task_record)
 
@@ -519,7 +519,7 @@ def handle_join_update(self, task_record: TaskRecord, inner_app_future: Optional
 
                 self.update_task_state(task_record, States.failed)
                 task_record['time_returned'] = datetime.datetime.now()
-                self.memoizer.update_memo(task_record)
+                self.memoizer.update_memo_exception(task_record, e)
                 with task_record['app_fu']._update_lock:
                     task_record['app_fu'].set_exception(e)
 
@@ -588,7 +588,7 @@ def _complete_task(self, task_record: TaskRecord, new_state: States, result: Any
         logger.info(f"Task {task_record['id']} completed ({old_state.name} -> {new_state.name})")
         task_record['time_returned'] = datetime.datetime.now()
 
-        self.memoizer.update_memo(task_record)
+        self.memoizer.update_memo_result(task_record, result)
         with task_record['app_fu']._update_lock:
             task_record['app_fu'].set_result(result)
 

parsl/dataflow/memoization.py

@@ -260,7 +260,13 @@ def check_memo(self, task: TaskRecord) -> Optional[Future[Any]]:
         assert isinstance(result, Future) or result is None
         return result
 
-    def update_memo(self, task: TaskRecord) -> None:
+    def update_memo_result(self, task: TaskRecord, r: Any) -> None:
+        self._update_memo(task)
+
+    def update_memo_exception(self, task: TaskRecord, e: BaseException) -> None:
+        self._update_memo(task)
+
+    def _update_memo(self, task: TaskRecord) -> None:
         """Updates the memoization lookup table with the result from a task.
         This doesn't move any values around but associates the memoization
         hashsum with the completed (by success or failure) AppFuture.