make checkpoint move away from always using futures for result value

and happen in update memo

Author: benclifford

parsl/dataflow/dflow.py

@@ -560,19 +560,6 @@ def handle_app_update(self, task_record: TaskRecord, future: AppFuture) -> None:
         if not task_record['app_fu'] == future:
             logger.error("Internal consistency error: callback future is not the app_fu in task structure, for task {}".format(task_id))
 
-        # Cover all checkpointing cases here:
-        # Do we need to checkpoint now, or queue for later,
-        # or do nothing?
-        if self.checkpoint_mode == 'task_exit':
-            self.memoizer.checkpoint(task=task_record)
-        elif self.checkpoint_mode in ('manual', 'periodic', 'dfk_exit'):
-            with self._modify_checkpointable_tasks_lock:
-                self.memoizer.checkpointable_tasks.append(task_record)
-        elif self.checkpoint_mode is None:
-            pass
-        else:
-            raise InternalConsistencyError(f"Invalid checkpoint mode {self.checkpoint_mode}")
-
         self.wipe_task(task_id)
         return
 

parsl/dataflow/memoization.py

@@ -263,9 +263,29 @@ def check_memo(self, task: TaskRecord) -> Optional[Future[Any]]:
     def update_memo_result(self, task: TaskRecord, r: Any) -> None:
         self._update_memo(task)
 
+        if self.checkpoint_mode == 'task_exit':
+            self.checkpoint(task=task, result=r)
+        elif self.checkpoint_mode in ('manual', 'periodic', 'dfk_exit'):
+            # with self._modify_checkpointable_tasks_lock:  # TODO: sort out use of this lock
+            self.checkpointable_tasks.append(task)
+        elif self.checkpoint_mode is None:
+            pass
+        else:
+            assert False, "Invalid checkpoint mode {self.checkpoint_mode} - should have been validated at initialization"
+
     def update_memo_exception(self, task: TaskRecord, e: BaseException) -> None:
         self._update_memo(task)
 
+        if self.checkpoint_mode == 'task_exit':
+            self.checkpoint(task=task, exception=e)
+        elif self.checkpoint_mode in ('manual', 'periodic', 'dfk_exit'):
+            # with self._modify_checkpointable_tasks_lock:  # TODO: sort out use of this lock
+            self.checkpointable_tasks.append(task)
+        elif self.checkpoint_mode is None:
+            pass
+        else:
+            assert False, "Invalid checkpoint mode {self.checkpoint_mode} - should have been validated at initialization"
+
     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
@@ -357,7 +377,17 @@ def load_checkpoints(self, checkpointDirs: Optional[Sequence[str]]) -> Dict[str,
         else:
             return {}
 
-    def checkpoint(self, *, task: Optional[TaskRecord] = None) -> None:
+    # TODO: this call becomes even more multiplexed...
+    # called with no parameters, we write out the task
+    # called with a task record, we can now no longer expect to get the
+    # result from the task record future, because it will not be
+    # populated yet.
+    # so then either we can an exception, or if exception is None, then
+    # checkpoint result. it's possible that result can be None as a
+    # real result: in the case that exception is None.
+    # what a horrible API that needs refactoring...
+
+    def checkpoint(self, *, task: Optional[TaskRecord] = None, exception: Optional[BaseException] = None, result: Any = None) -> None:
         """Checkpoint the dfk incrementally to a checkpoint file.
 
         When called with no argument, all tasks registered in self.checkpointable_tasks
@@ -377,11 +407,6 @@ def checkpoint(self, *, task: Optional[TaskRecord] = None) -> None:
         """
         with self.checkpoint_lock:
 
-            if task:
-                checkpoint_queue = [task]
-            else:
-                checkpoint_queue = self.checkpointable_tasks
-
             checkpoint_dir = '{0}/checkpoint'.format(self.run_dir)
             checkpoint_tasks = checkpoint_dir + '/tasks.pkl'
 
@@ -391,22 +416,47 @@ def checkpoint(self, *, task: Optional[TaskRecord] = None) -> None:
             count = 0
 
             with open(checkpoint_tasks, 'ab') as f:
-                for task_record in checkpoint_queue:
-                    task_id = task_record['id']
 
-                    app_fu = task_record['app_fu']
+                if task:
+                    # TODO: refactor with below
 
-                    if app_fu.done() and app_fu.exception() is None:
-                        hashsum = task_record['hashsum']
+                    task_id = task['id']
+
+                    if exception is None:
+                        hashsum = task['hashsum']
                         if not hashsum:
-                            continue
-                        t = {'hash': hashsum, 'exception': None, 'result': app_fu.result()}
+                            pass  # TODO: log an error? see below discussion
+                        else:
+                            t = {'hash': hashsum, 'exception': None, 'result': result}
+
+                            # We are using pickle here since pickle dumps to a file in 'ab'
+                            # mode behave like a incremental log.
+                            pickle.dump(t, f)
+                            count += 1
 
-                        # We are using pickle here since pickle dumps to a file in 'ab'
-                        # mode behave like a incremental log.
-                        pickle.dump(t, f)
-                        count += 1
                         logger.debug("Task {} checkpointed".format(task_id))
+                else:
+                    checkpoint_queue = self.checkpointable_tasks
+
+                    for task_record in checkpoint_queue:
+                        task_id = task_record['id']
+
+                        app_fu = task_record['app_fu']
+
+                        assert app_fu.done(), "trying to checkpoint a task that is not done"
+
+                        if app_fu.done() and app_fu.exception() is None:
+                            hashsum = task_record['hashsum']
+                            if not hashsum:
+                                continue  # TODO: log an error? maybe some tasks don't have hashsums legitimately?
+                            t = {'hash': hashsum, 'exception': None, 'result': app_fu.result()}
+
+                            # We are using pickle here since pickle dumps to a file in 'ab'
+                            # mode behave like a incremental log.
+                            pickle.dump(t, f)
+                            count += 1
+
+                            logger.debug("Task {} checkpointed".format(task_id))
 
             self.checkpointed_tasks += count