remove DFK's use of checkpoint mode coniguration, into the memoizer code

Author: benclifford

parsl/dataflow/dflow.py

@@ -49,7 +49,7 @@
 from parsl.monitoring.remote import monitor_wrapper
 from parsl.process_loggers import wrap_with_logs
 from parsl.usage_tracking.usage import UsageTracker
-from parsl.utils import Timer, get_std_fname_mode, get_version
+from parsl.utils import get_std_fname_mode, get_version
 
 logger = logging.getLogger(__name__)
 
@@ -167,13 +167,13 @@ def __init__(self, config: Config) -> None:
 
         # TODO: the parameters that remain here should be parameters that are going to be configured by
         # the user as part of checkpoint/memo configuration object.
-        self.memoizer = Memoizer(memoize=config.app_cache, checkpoint_mode=config.checkpoint_mode, checkpoint_files=config.checkpoint_files)
+        self.memoizer = Memoizer(memoize=config.app_cache,
+                                 checkpoint_mode=config.checkpoint_mode,
+                                 checkpoint_files=config.checkpoint_files,
+                                 checkpoint_period=config.checkpoint_period)
         self.memoizer.run_dir = self.run_dir
         self.memoizer.start()
 
-        # TODO: this block needs to move too
-        self._checkpoint_timer = None
-        self.checkpoint_mode = config.checkpoint_mode
         self._modify_checkpointable_tasks_lock = threading.Lock()
 
         # this must be set before executors are added since add_executors calls
@@ -190,17 +190,6 @@ def __init__(self, config: Config) -> None:
         self.add_executors(config.executors)
         self.add_executors([parsl_internal_executor])
 
-        if self.checkpoint_mode == "periodic":
-            if config.checkpoint_period is None:
-                raise ConfigurationError("Checkpoint period must be specified with periodic checkpoint mode")
-            else:
-                try:
-                    h, m, s = map(int, config.checkpoint_period.split(':'))
-                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.task_count = 0
         self.tasks: Dict[int, TaskRecord] = {}
         self.submitter_lock = threading.Lock()
@@ -1187,16 +1176,7 @@ def cleanup(self) -> None:
 
         self.log_task_states()
 
-        # checkpoint if any valid checkpoint method is specified
-        if self.checkpoint_mode is not None:
-
-            # TODO: accesses to self.checkpointable_tasks should happen
-            # under a lock?
-            self.memoizer.checkpoint()
-
-            if self._checkpoint_timer:
-                logger.info("Stopping checkpoint timer")
-                self._checkpoint_timer.close()
+        self.memoizer.close()
 
         # Send final stats
         self.usage_tracker.send_end_message()

parsl/dataflow/memoization.py

@@ -14,7 +14,8 @@
 
 from parsl.dataflow.errors import BadCheckpoint
 from parsl.dataflow.taskrecord import TaskRecord
-from parsl.utils import get_all_checkpoints
+from parsl.errors import ConfigurationError
+from parsl.utils import Timer, get_all_checkpoints
 
 logger = logging.getLogger(__name__)
 
@@ -153,6 +154,7 @@ class Memoizer:
     def __init__(self, *,
                  memoize: bool = True,
                  checkpoint_files: Sequence[str] | None,
+                 checkpoint_period: Optional[str],
                  checkpoint_mode: Literal['task_exit', 'periodic', 'dfk_exit', 'manual'] | None):
         """Initialize the memoizer.
 
@@ -168,9 +170,12 @@ def __init__(self, *,
 
         self.checkpoint_files = checkpoint_files
         self.checkpoint_mode = checkpoint_mode
+        self.checkpoint_period = checkpoint_period
 
         self.checkpointable_tasks: List[TaskRecord] = []
 
+        self._checkpoint_timer: Timer | None = None
+
     def start(self) -> None:
         if self.checkpoint_files is not None:
             checkpoint_files = self.checkpoint_files
@@ -190,6 +195,29 @@ def start(self) -> None:
             logger.info("App caching disabled for all apps")
             self.memo_lookup_table = {}
 
+        if self.checkpoint_mode == "periodic":
+            if self.checkpoint_period is None:
+                raise ConfigurationError("Checkpoint period must be specified with periodic checkpoint mode")
+            else:
+                try:
+                    h, m, s = map(int, self.checkpoint_period.split(':'))
+                except Exception:
+                    raise ConfigurationError("invalid checkpoint_period provided: {0} expected HH:MM:SS".format(self.checkpoint_period))
+                checkpoint_period = (h * 3600) + (m * 60) + s
+                self._checkpoint_timer = Timer(self.checkpoint, interval=checkpoint_period, name="Checkpoint")
+
+    def close(self) -> None:
+        # checkpoint if any valid checkpoint method is specified
+        if self.checkpoint_mode is not None:
+
+            # TODO: accesses to self.checkpointable_tasks should happen
+            # under a lock?
+            self.checkpoint()
+
+            if self._checkpoint_timer:
+                logger.info("Stopping checkpoint timer")
+                self._checkpoint_timer.close()
+
     def make_hash(self, task: TaskRecord) -> str:
         """Create a hash of the task inputs.