make checkpoint call not use dfk state, in prep for moving into memoizer

Author: benclifford

parsl/dataflow/dflow.py

@@ -176,6 +176,8 @@ def __init__(self, config: Config) -> None:
         self.checkpointed_tasks = 0
         self._checkpoint_timer = None
         self.checkpoint_mode = config.checkpoint_mode
+
+        self._modify_checkpointable_tasks_lock = threading.Lock()
         self.checkpointable_tasks: List[TaskRecord] = []
 
         # this must be set before executors are added since add_executors calls
@@ -200,7 +202,7 @@ def __init__(self, config: Config) -> None:
                 except Exception:
                     raise ConfigurationError("invalid checkpoint_period provided: {0} expected HH:MM:SS".format(config.checkpoint_period))
                 checkpoint_period = (h * 3600) + (m * 60) + s
-                self._checkpoint_timer = Timer(self.checkpoint, interval=checkpoint_period, name="Checkpoint")
+                self._checkpoint_timer = Timer(self.invoke_checkpoint, interval=checkpoint_period, name="Checkpoint")
 
         self.task_count = 0
         self.tasks: Dict[int, TaskRecord] = {}
@@ -571,7 +573,7 @@ def handle_app_update(self, task_record: TaskRecord, future: AppFuture) -> None:
         if self.checkpoint_mode == 'task_exit':
             self.checkpoint(tasks=[task_record])
         elif self.checkpoint_mode in ('manual', 'periodic', 'dfk_exit'):
-            with self.checkpoint_lock:
+            with self._modify_checkpointable_tasks_lock:
                 self.checkpointable_tasks.append(task_record)
         elif self.checkpoint_mode is None:
             pass
@@ -1250,7 +1252,10 @@ def cleanup(self) -> None:
         # Checkpointing takes priority over the rest of the tasks
         # checkpoint if any valid checkpoint method is specified
         if self.checkpoint_mode is not None:
-            self.checkpoint()
+
+            # TODO: accesses to self.checkpointable_tasks should happen
+            # under a lock?
+            self.checkpoint(self.checkpointable_tasks)
 
             if self._checkpoint_timer:
                 logger.info("Stopping checkpoint timer")
@@ -1323,7 +1328,13 @@ def cleanup(self) -> None:
 
         logger.info("DFK cleanup complete")
 
-    def checkpoint(self, tasks: Optional[Sequence[TaskRecord]] = None) -> str:
+    def invoke_checkpoint(self):
+        with self._modify_checkpointable_tasks_lock:
+            r = self.checkpoint(self.checkpointable_tasks)
+            self.checkpointable_tasks = []
+        return r
+
+    def checkpoint(self, tasks: Sequence[TaskRecord]) -> str:
         """Checkpoint the dfk incrementally to a checkpoint file.
 
         When called, every task that has been completed yet not
@@ -1342,11 +1353,7 @@ def checkpoint(self, tasks: Optional[Sequence[TaskRecord]] = None) -> str:
             run under RUNDIR/checkpoints/{tasks.pkl, dfk.pkl}
         """
         with self.checkpoint_lock:
-            if tasks:
-                checkpoint_queue = tasks
-            else:
-                checkpoint_queue = self.checkpointable_tasks
-                self.checkpointable_tasks = []
+            checkpoint_queue = tasks
 
             checkpoint_dir = '{0}/checkpoint'.format(self.run_dir)
             checkpoint_dfk = checkpoint_dir + '/dfk.pkl'

parsl/tests/test_checkpointing/test_python_checkpoint_1.py

@@ -25,7 +25,7 @@ def test_initial_checkpoint_write():
     """
     uuid_app().result()
 
-    cpt_dir = parsl.dfk().checkpoint()
+    cpt_dir = parsl.dfk().invoke_checkpoint()
 
     cptpath = cpt_dir + '/dfk.pkl'
     assert os.path.exists(cptpath), f"DFK checkpoint missing: {cptpath}"

parsl/tests/test_checkpointing/test_python_checkpoint_3.py

@@ -36,7 +36,7 @@ def test_checkpointing():
     # Wait for the results
     [i.result() for i in d]
 
-    checkpoint_dir = dfk.checkpoint()
+    checkpoint_dir = dfk.invoke_checkpoint()
     print(checkpoint_dir)
 
     assert os.path.exists(checkpoint_dir), "Checkpoint dir does not exist"