Pass result or exception into update_memo (#4017)

This follows the Future API, which has two completion variants, one for
success and one for exception, and #4006 which introduces that
distinction in DFK task completion helpers.

Right now, the two new update_memo variants do the same thing. However,
an upcoming PR will move checkpoint code out of handle_app_update (which
is race-prone - #3762) and into update_memo which runs before the
result/exception is exposed to the user. That checkpoint code will then
need to see the result/exception explicitly.

# Changed Behaviour

none

## Type of change

- Code maintenance/cleanup

Author: benclifford

parsl/dataflow/dflow.py

@@ -539,7 +539,7 @@ def _complete_task_result(self, task_record: TaskRecord, new_state: States, resu
         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)
 
         self._send_task_log_info(task_record)
 
@@ -558,7 +558,7 @@ def _complete_task_exception(self, task_record: TaskRecord, new_state: States, e
         logger.info(f"Task {task_record['id']} failed ({old_state.name} -> {new_state.name})")
         task_record['time_returned'] = datetime.datetime.now()
 
-        self.memoizer.update_memo(task_record)
+        self.memoizer.update_memo_exception(task_record, exception)
 
         self._send_task_log_info(task_record)
 

parsl/dataflow/memoization.py

@@ -283,7 +283,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.