Make memoizer instance pluggable

This is the main breaking change of PR #3535 checkpoint plugins:
this PR reworks construction of the memoizer so that the user now supplies
an instance, plugin style, instead of the DFK constructing an instance.

Memoization and configuration options, which turn out to be specific
to the BasicMemoizer, are now supplied by the user directly to the
instance they construct, rather than being passed through the DFK.

Any checkpoint users will need to alter their Parsl configuration,
although in a fairly straightforward way.

If no memoizer is specified, the DFK constructs a BasicMemozier that
preserves the no-parameters-supplied default behaviour:

Memoization will occur, checkpointing will not occur.

Author: benclifford

docs/reference.rst

@@ -143,6 +143,16 @@ Batch jobs
     parsl.jobs.error_handlers.simple_error_handler
     parsl.jobs.error_handlers.windowed_error_handler
 
+Checkpointing and Memoization
+=============================
+
+.. autosummary::
+    :toctree: stubs
+    :nosignatures:
+
+    parsl.dataflow.memoization.Memoizer
+    parsl.dataflow.memoization.BasicMemoizer
+
 Exceptions
 ==========
 
@@ -207,8 +217,6 @@ Internal
     parsl.app.python.PythonApp
     parsl.dataflow.dflow.DataFlowKernel
     parsl.dataflow.memoization.id_for_memo
-    parsl.dataflow.memoization.Memoizer
-    parsl.dataflow.memoization.BasicMemoizer
     parsl.dataflow.states.FINAL_STATES
     parsl.dataflow.states.States
     parsl.dataflow.taskrecord.TaskRecord

docs/userguide/advanced/plugins.rst

@@ -56,6 +56,10 @@ will be used.
 Memoization/checkpointing
 -------------------------
 
+The default checkpoint/memoization system can be replaced by supplying an
+instance of the `Memoizer` class. This instance is queried before an app is
+launched, and is given details of completed apps.
+
 When parsl memoizes/checkpoints an app parameter, it does so by computing a
 hash of that parameter that should be the same if that parameter is the same
 on subsequent invocations. This isn't straightforward to do for arbitrary

docs/userguide/workflows/checkpoints.rst

@@ -8,10 +8,14 @@ reuse the result from the first invocation without executing the app again.
 
 This can save time and computational resources.
 
-This is done in two ways:
+The memoization and checkpointing system is pluggable, with basic behaviour
+provided by the `BasicMemoizer`. The rest of this chapter refers to the
+behaviour of the `BasicMemoizer`.
+
+Memoization and checkpointing is done in two ways:
 
 * Firstly, *app caching* will allow reuse of results and exceptions within
-  the same run.
+  the same run. This is also referred to as *memoization*.
 
 * Building on top of that, *checkpointing* will store results (but not
   exceptions) on the filesystem and reuse those results in later runs.
@@ -39,8 +43,8 @@ decorator to ``True`` (by default it is ``False``).
    def hello (msg, stdout=None):
        return 'echo {}'.format(msg)
 			
-App caching can be globally disabled by setting ``app_cache=False``
-in the :class:`~parsl.config.Config`.
+App caching can be globally disabled by supplying a new memoizer in
+:class:`~parsl.config.Config` defined as ``BasicMemoizer(memoize=False)``.
 
 App caching can be particularly useful when developing interactive programs such as when
 using a Jupyter notebook. In this case, cells containing apps are often re-executed
@@ -158,26 +162,25 @@ use the hash of the app and the invocation input parameters to identify previous
 results. If multiple checkpoints exist for an app (with the same hash)
 the most recent entry will be used.
 
-Parsl provides four checkpointing modes:
+Parsl provides four checkpointing modes, which can be specified using the ``checkpoint_mode``
+parameter to ``memoizer=BasicMemoizer(...)``
 
 1. ``task_exit``: a checkpoint is created each time an app completes or fails
    (after retries if enabled). This mode minimizes the risk of losing information
    from completed tasks.
 
    .. code-block:: python
 
-      from parsl.configs.local_threads import config
-      config.checkpoint_mode = 'task_exit'
+      BasicMemoizer(checkpoint_mode = 'task_exit')
 
 2. ``periodic``: a checkpoint is created periodically using a user-specified
    checkpointing interval. Results will be saved to the checkpoint file for
    all tasks that have completed during this period.
 
    .. code-block:: python
 
-      from parsl.configs.local_threads import config
-      config.checkpoint_mode = 'periodic'
-      config.checkpoint_period = "01:00:00"
+      BasicMemoizer(checkpoint_mode = 'periodic',
+                    checkpoint_period = '01:00:00')
 
 3. ``dfk_exit``: checkpoints are created when Parsl is
    about to exit. This reduces the risk of losing results due to
@@ -187,19 +190,14 @@ Parsl provides four checkpointing modes:
 
    .. code-block:: python
 
-      from parsl.configs.local_threads import config
-      config.checkpoint_mode = 'dfk_exit'
+      BasicMemoizer(checkpoint_mode = 'dfk_exit')
 
 4. ``manual``: in addition to these automated checkpointing modes, it is also possible
    to manually initiate a checkpoint by calling ``DataFlowKernel.checkpoint()`` in the
    Parsl program code.
 
    .. code-block:: python
 
-      import parsl
-      from parsl.configs.local_threads import config
-      dfk = parsl.load(config)
-      ....
       dfk.checkpoint()
 
 In all cases the checkpoint file is written out to the ``runinfo/RUN_ID/checkpoint/`` directory.
@@ -286,11 +284,11 @@ from the checkpoint file are are immediately returned.
     from parsl.tests.configs.local_threads import config
     from parsl.utils import get_all_checkpoints
 
-    config.checkpoint_files = get_all_checkpoints()
+    config.memoizer = BasicMemoizer(checkpoint_files = get_all_checkpoints())
 
     parsl.load(config)
 		
-		# Rerun the same workflow
+    # Rerun the same workflow
     d = []
     for i in range(5):
         d.append(slow_double(i))

parsl/config.py

@@ -5,6 +5,7 @@
 from typing_extensions import Literal
 
 from parsl.dataflow.dependency_resolvers import DependencyResolver
+from parsl.dataflow.memoization import Memoizer
 from parsl.dataflow.taskrecord import TaskRecord
 from parsl.errors import ConfigurationError
 from parsl.executors.base import ParslExecutor
@@ -27,17 +28,6 @@ class Config(RepresentationMixin, UsageInformation):
     executors : sequence of ParslExecutor, optional
         List (or other iterable) of `ParslExecutor` instances to use for executing tasks.
         Default is (:class:`~parsl.executors.threads.ThreadPoolExecutor()`,).
-    app_cache : bool, optional
-        Enable app caching. Default is True.
-    checkpoint_files : sequence of str, optional
-        List of paths to checkpoint files. See :func:`parsl.utils.get_all_checkpoints` and
-        :func:`parsl.utils.get_last_checkpoint` for helpers. Default is None.
-    checkpoint_mode : str, optional
-        Checkpoint mode to use, can be ``'dfk_exit'``, ``'task_exit'``, ``'periodic'`` or ``'manual'``.
-        If set to `None`, checkpointing will be disabled. Default is None.
-    checkpoint_period : str, optional
-        Time interval (in "HH:MM:SS") at which to checkpoint completed tasks. Only has an effect if
-        ``checkpoint_mode='periodic'``.
     dependency_resolver: plugin point for custom dependency resolvers. Default: only resolve Futures,
         using the `SHALLOW_DEPENDENCY_RESOLVER`.
     exit_mode: str, optional
@@ -100,14 +90,7 @@ class Config(RepresentationMixin, UsageInformation):
     @typeguard.typechecked
     def __init__(self,
                  executors: Optional[Iterable[ParslExecutor]] = None,
-                 app_cache: bool = True,
-                 checkpoint_files: Optional[Sequence[str]] = None,
-                 checkpoint_mode: Union[None,
-                                        Literal['task_exit'],
-                                        Literal['periodic'],
-                                        Literal['dfk_exit'],
-                                        Literal['manual']] = None,
-                 checkpoint_period: Optional[str] = None,
+                 memoizer: Optional[Memoizer] = None,
                  dependency_resolver: Optional[DependencyResolver] = None,
                  exit_mode: Literal['cleanup', 'skip', 'wait'] = 'cleanup',
                  garbage_collect: bool = True,
@@ -131,21 +114,7 @@ def __init__(self,
         self._executors: Sequence[ParslExecutor] = executors
         self._validate_executors()
 
-        self.app_cache = app_cache
-        self.checkpoint_files = checkpoint_files
-        self.checkpoint_mode = checkpoint_mode
-        if checkpoint_period is not None:
-            if checkpoint_mode is None:
-                logger.debug('The requested `checkpoint_period={}` will have no effect because `checkpoint_mode=None`'.format(
-                    checkpoint_period)
-                )
-            elif checkpoint_mode != 'periodic':
-                logger.debug("Requested checkpoint period of {} only has an effect with checkpoint_mode='periodic'".format(
-                    checkpoint_period)
-                )
-        if checkpoint_mode == 'periodic' and checkpoint_period is None:
-            checkpoint_period = "00:30:00"
-        self.checkpoint_period = checkpoint_period
+        self.memoizer = memoizer
         self.dependency_resolver = dependency_resolver
         self.exit_mode = exit_mode
         self.garbage_collect = garbage_collect

parsl/configs/ASPIRE1.py

@@ -1,5 +1,6 @@
 from parsl.addresses import address_by_interface
 from parsl.config import Config
+from parsl.dataflow.memoization import BasicMemoizer
 from parsl.executors import HighThroughputExecutor
 from parsl.launchers import MpiRunLauncher
 from parsl.monitoring.monitoring import MonitoringHub
@@ -38,7 +39,6 @@
         ),
         strategy='simple',
         retries=3,
-        app_cache=True,
-        checkpoint_mode='task_exit',
+        memoizer=BasicMemoizer(checkpoint_mode='task_exit'),
         usage_tracking=LEVEL_1,
 )

parsl/dataflow/dflow.py

@@ -165,13 +165,8 @@ def __init__(self, config: Config) -> None:
             self.monitoring_radio.send((MessageType.WORKFLOW_INFO,
                                        workflow_info))
 
-        self.memoizer: Memoizer = BasicMemoizer(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()
+        self.memoizer: Memoizer = config.memoizer if config.memoizer is not None else BasicMemoizer()
+        self.memoizer.start(run_dir=self.run_dir)
 
         # this must be set before executors are added since add_executors calls
         # job_status_poller.add_executors.

parsl/dataflow/memoization.py

@@ -179,6 +179,17 @@ def update_memo_result(self, task: TaskRecord, r: Any) -> None:
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def start(self, *, run_dir: str) -> None:
+        """Called by the DFK when it starts up.
+
+        This is an opportunity for the memoization/checkpoint system to
+        initialize itself.
+
+        The path to the base run directory is passed as a parameter.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def checkpoint_queue(self) -> None:
         """Called by the DFK when the user calls dfk.checkpoint(). This
@@ -238,18 +249,26 @@ class BasicMemoizer(Memoizer):
     run_dir: str
 
     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):
+                 checkpoint_files: Sequence[str] | None = None,
+                 checkpoint_period: Optional[str] = None,
+                 checkpoint_mode: Literal['task_exit', 'periodic', 'dfk_exit', 'manual'] | None = None,
+                 memoize: bool = True):
         """Initialize the memoizer.
 
         KWargs:
-            - memoize (Bool): enable memoization or not.
-            - checkpoint (Dict): A checkpoint loaded as a dict.
-        """
-        self.memoize = memoize
 
+            - checkpoint_files : sequence of str, optional
+                  List of paths to checkpoint files. See :func:`parsl.utils.get_all_checkpoints` and
+                  :func:`parsl.utils.get_last_checkpoint` for helpers. Default is None.
+            - checkpoint_period : str, optional
+                  Time interval (in "HH:MM:SS") at which to checkpoint completed tasks. Only has an effect if
+                  ``checkpoint_mode='periodic'``.
+            - checkpoint_mode : str, optional
+                  Checkpoint mode to use, can be ``'dfk_exit'``, ``'task_exit'``, ``'periodic'`` or ``'manual'``.
+                  If set to `None`, checkpointing will be disabled. Default is None.
+            - memoize : str, enable memoization or not.
+
+        """
         self.checkpointed_tasks = 0
 
         # this lock must be held when:
@@ -264,8 +283,12 @@ def __init__(self, *,
         self.checkpointable_tasks: List[CheckpointCommand] = []
 
         self._checkpoint_timer: Timer | None = None
+        self.memoize = memoize
+
+    def start(self, *, run_dir: str) -> None:
+
+        self.run_dir = run_dir
 
-    def start(self) -> None:
         if self.checkpoint_files is not None:
             checkpoint_files = self.checkpoint_files
         elif self.checkpoint_files is None and self.checkpoint_mode is not None:

parsl/tests/configs/htex_local_alternate.py

@@ -22,6 +22,7 @@
 from parsl.data_provider.ftp import FTPInTaskStaging
 from parsl.data_provider.http import HTTPInTaskStaging
 from parsl.data_provider.zip import ZipFileStaging
+from parsl.dataflow.memoization import BasicMemoizer
 from parsl.executors import HighThroughputExecutor
 from parsl.launchers import SingleNodeLauncher
 
@@ -56,7 +57,7 @@ def fresh_config():
             )
         ],
         strategy='simple',
-        app_cache=True, checkpoint_mode='task_exit',
+        memoizer=BasicMemoizer(memoize=True, checkpoint_mode='task_exit'),
         retries=2,
         monitoring=MonitoringHub(
                         monitoring_debug=False,

parsl/tests/configs/local_threads_checkpoint.py

@@ -1,4 +1,5 @@
 from parsl.config import Config
+from parsl.dataflow.memoization import BasicMemoizer
 from parsl.executors.threads import ThreadPoolExecutor
 
 config = Config(
@@ -7,5 +8,5 @@
             label='local_threads_checkpoint_dfk_exit',
         )
     ],
-    checkpoint_mode='dfk_exit'
+    memoizer=BasicMemoizer(checkpoint_mode='dfk_exit')
 )