move checkpointable_tasks list a bit more

Author: benclifford

parsl/dataflow/dflow.py

@@ -564,7 +564,7 @@ def handle_app_update(self, task_record: TaskRecord, future: AppFuture) -> None:
         # Do we need to checkpoint now, or queue for later,
         # or do nothing?
         if self.checkpoint_mode == 'task_exit':
-            self.memoizer.checkpoint(tasks=[task_record])
+            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)
@@ -1205,7 +1205,7 @@ def cleanup(self) -> None:
 
             # TODO: accesses to self.checkpointable_tasks should happen
             # under a lock?
-            self.memoizer.checkpoint(self.memoizer.checkpointable_tasks)
+            self.memoizer.checkpoint()
 
             if self._checkpoint_timer:
                 logger.info("Stopping checkpoint timer")
@@ -1269,7 +1269,7 @@ def cleanup(self) -> None:
 
     def checkpoint(self) -> None:
         with self._modify_checkpointable_tasks_lock:
-            self.memoizer.checkpoint(self.memoizer.checkpointable_tasks)
+            self.memoizer.checkpoint()
             self.memoizer.checkpointable_tasks = []
 
     @staticmethod

parsl/dataflow/memoization.py

@@ -351,26 +351,30 @@ def load_checkpoints(self, checkpointDirs: Optional[Sequence[str]]) -> Dict[str,
         else:
             return {}
 
-    def checkpoint(self, tasks: Sequence[TaskRecord]) -> None:
+    def checkpoint(self, *, task: Optional[TaskRecord] = None) -> None:
         """Checkpoint the dfk incrementally to a checkpoint file.
 
-        When called, every task that has been completed yet not
-        checkpointed is checkpointed to a file.
+        When called with no argument, all tasks registered in self.checkpointable_tasks
+        will be checkpointed. When called with a single TaskRecord argument, that task will be
+        checkpointed.
+
+        By default the checkpoints are written to the RUNDIR of the current
+        run under RUNDIR/checkpoints/tasks.pkl
 
         Kwargs:
-            - tasks (List of task records) : List of task ids to checkpoint. Default=None
-                                         if set to None, we iterate over all tasks held by the DFK.
+            - task (Optional task records) : A task to checkpoint. Default=None, meaning all
+              tasks registered for checkpointing.
 
         .. note::
             Checkpointing only works if memoization is enabled
 
-        Returns:
-            Checkpoint dir if checkpoints were written successfully.
-            By default the checkpoints are written to the RUNDIR of the current
-            run under RUNDIR/checkpoints/tasks.pkl
         """
         with self.checkpoint_lock:
-            checkpoint_queue = tasks
+
+            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'