dev

Author: benclifford

parsl/dataflow/dflow.py

@@ -7,7 +7,6 @@
 import logging
 import os
 import pathlib
-import pickle
 import random
 import sys
 import threading
@@ -99,8 +98,6 @@ def __init__(self, config: Config) -> None:
 
         logger.info("Parsl version: {}".format(get_version()))
 
-        self.checkpoint_lock = threading.Lock()
-
         self.usage_tracker = UsageTracker(self)
         self.usage_tracker.send_start_message()
 
@@ -177,7 +174,8 @@ def __init__(self, config: Config) -> None:
             checkpoint_files = []
 
         self.memoizer = Memoizer(self, memoize=config.app_cache, checkpoint_files=checkpoint_files)
-        self.checkpointed_tasks = 0
+        self.memoizer.run_dir = self.run_dir
+
         self._checkpoint_timer = None
         self.checkpoint_mode = config.checkpoint_mode
 
@@ -575,7 +573,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.checkpoint(tasks=[task_record])
+            self.memoizer.checkpoint(tasks=[task_record])
         elif self.checkpoint_mode in ('manual', 'periodic', 'dfk_exit'):
             with self._modify_checkpointable_tasks_lock:
                 self.checkpointable_tasks.append(task_record)
@@ -1259,7 +1257,7 @@ def cleanup(self) -> None:
 
             # TODO: accesses to self.checkpointable_tasks should happen
             # under a lock?
-            self.checkpoint(self.checkpointable_tasks)
+            self.memoizer.checkpoint(self.checkpointable_tasks)
 
             if self._checkpoint_timer:
                 logger.info("Stopping checkpoint timer")
@@ -1334,76 +1332,10 @@ def cleanup(self) -> None:
 
     def invoke_checkpoint(self):
         with self._modify_checkpointable_tasks_lock:
-            r = self.checkpoint(self.checkpointable_tasks)
+            r = self.memoizer.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
-        checkpointed is checkpointed to a file.
-
-        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.
-
-        .. 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, dfk.pkl}
-        """
-        with self.checkpoint_lock:
-            checkpoint_queue = tasks
-
-            checkpoint_dir = '{0}/checkpoint'.format(self.run_dir)
-            checkpoint_dfk = checkpoint_dir + '/dfk.pkl'
-            checkpoint_tasks = checkpoint_dir + '/tasks.pkl'
-
-            if not os.path.exists(checkpoint_dir):
-                os.makedirs(checkpoint_dir, exist_ok=True)
-
-            with open(checkpoint_dfk, 'wb') as f:
-                state = {'rundir': self.run_dir,
-                         'task_count': self.task_count
-                         }
-                pickle.dump(state, f)
-
-            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 app_fu.done() and app_fu.exception() is None:
-                        hashsum = task_record['hashsum']
-                        if not hashsum:
-                            continue
-                        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
-
-            if count == 0:
-                if self.checkpointed_tasks == 0:
-                    logger.warning("No tasks checkpointed so far in this run. Please ensure caching is enabled")
-                else:
-                    logger.debug("No tasks checkpointed in this pass.")
-            else:
-                logger.info("Done checkpointing {} tasks".format(count))
-
-            return checkpoint_dir
-
     @staticmethod
     def _log_std_streams(task_record: TaskRecord) -> None:
         tid = task_record['id']

parsl/dataflow/memoization.py

@@ -4,6 +4,7 @@
 import logging
 import os
 import pickle
+import threading
 from functools import lru_cache, singledispatch
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence
 
@@ -150,6 +151,8 @@ class Memoizer:
 
     """
 
+    run_dir: str
+
     def __init__(self, dfk: DataFlowKernel, *, memoize: bool = True, checkpoint_files: Sequence[str]):
         """Initialize the memoizer.
 
@@ -163,6 +166,10 @@ def __init__(self, dfk: DataFlowKernel, *, memoize: bool = True, checkpoint_file
         self.dfk = dfk
         self.memoize = memoize
 
+        self.checkpointed_tasks = 0
+
+        self.checkpoint_lock = threading.Lock()
+
         # TODO: we always load checkpoints even if we then discard them...
         # this is more obvious here, less obvious in previous Parsl...
         checkpoint = self.load_checkpoints(checkpoint_files)
@@ -350,3 +357,69 @@ def load_checkpoints(self, checkpointDirs: Optional[Sequence[str]]) -> Dict[str,
             return self._load_checkpoints(checkpointDirs)
         else:
             return {}
+
+    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
+        checkpointed is checkpointed to a file.
+
+        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.
+
+        .. 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, dfk.pkl}
+        """
+        with self.checkpoint_lock:
+            checkpoint_queue = tasks
+
+            checkpoint_dir = '{0}/checkpoint'.format(self.run_dir)
+            checkpoint_dfk = checkpoint_dir + '/dfk.pkl'
+            checkpoint_tasks = checkpoint_dir + '/tasks.pkl'
+
+            if not os.path.exists(checkpoint_dir):
+                os.makedirs(checkpoint_dir, exist_ok=True)
+
+            with open(checkpoint_dfk, 'wb') as f:
+                state = {'rundir': self.run_dir,
+                         # TODO: this isn't relevant to checkpointing? 'task_count': self.task_count
+                         }
+                pickle.dump(state, f)
+
+            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 app_fu.done() and app_fu.exception() is None:
+                        hashsum = task_record['hashsum']
+                        if not hashsum:
+                            continue
+                        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
+
+            if count == 0:
+                if self.checkpointed_tasks == 0:
+                    logger.warning("No tasks checkpointed so far in this run. Please ensure caching is enabled")
+                else:
+                    logger.debug("No tasks checkpointed in this pass.")
+            else:
+                logger.info("Done checkpointing {} tasks".format(count))
+
+            return checkpoint_dir