checkpoint exceptions

Author: benclifford

parsl/dataflow/memoization.py

@@ -388,8 +388,12 @@ def _load_checkpoints(self, checkpointDirs: Sequence[str]) -> Dict[str, Future[A
                             data = pickle.load(f)
                             # Copy and hash only the input attributes
                             memo_fu: Future = Future()
-                            assert data['exception'] is None
-                            memo_fu.set_result(data['result'])
+
+                            if data['exception'] is None:
+                                memo_fu.set_result(data['result'])
+                            else:
+                                assert data['result'] is None
+                                memo_fu.set_exception(data['exception'])
                             memo_lookup_table[data['hash']] = memo_fu
 
                         except EOFError:
@@ -472,20 +476,22 @@ def checkpoint(self, *, task: Optional[TaskRecord] = None, exception: Optional[B
                     # TODO: refactor with below
 
                     task_id = task['id']
-
-                    if exception is None:
-                        hashsum = task['hashsum']
-                        if not hashsum:
-                            pass  # TODO: log an error? see below discussion
-                        else:
+                    hashsum = task['hashsum']
+                    if not hashsum:
+                        pass  # TODO: log an error? see below discussion
+                    else:
+                        if exception is None and self.filter_result_for_checkpoint(result):
                             t = {'hash': hashsum, 'exception': None, 'result': 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))
+                            logger.debug("Task {} checkpointed result".format(task_id))
+                        elif exception is not None and self.filter_exception_for_checkpoint(exception):
+                            t = {'hash': hashsum, 'exception': exception, 'result': None}
+                            pickle.dump(t, f)
+                            count += 1
+                            logger.debug("Task {} checkpointed exception".format(task_id))
+                        else:
+                            pass  # no checkpoint - maybe debug log? TODO
                 else:
                     checkpoint_queue = self.checkpointable_tasks
 
@@ -496,18 +502,22 @@ def checkpoint(self, *, task: Optional[TaskRecord] = None, exception: Optional[B
 
                         assert app_fu.done(), "trying to checkpoint a task that is not done"
 
-                        if app_fu.done() and app_fu.exception() is None:
-                            hashsum = task_record['hashsum']
-                            if not hashsum:
-                                continue  # TODO: log an error? maybe some tasks don't have hashsums legitimately?
-                            t = {'hash': hashsum, 'exception': None, 'result': app_fu.result()}
+                        hashsum = task_record['hashsum']
+                        if not hashsum:
+                            continue  # TODO: log an error? maybe some tasks don't have hashsums legitimately?
 
-                            # We are using pickle here since pickle dumps to a file in 'ab'
-                            # mode behave like a incremental log.
+                        if app_fu.exception() is None and self.filter_result_for_checkpoint(app_fu.result()):
+                            t = {'hash': hashsum, 'exception': None, 'result': app_fu.result()}
                             pickle.dump(t, f)
                             count += 1
-
-                            logger.debug("Task {} checkpointed".format(task_id))
+                            logger.debug("Task {} checkpointed result".format(task_id))
+                        elif (e := app_fu.exception()) is not None and self.filter_exception_for_checkpoint(e):
+                            t = {'hash': hashsum, 'exception': app_fu.exception(), 'result': None}
+                            pickle.dump(t, f)
+                            count += 1
+                            logger.debug("Task {} checkpointed exception".format(task_id))
+                        else:
+                            pass  # TODO: maybe log at debug level
 
             self.checkpointed_tasks += count
 
@@ -525,3 +535,11 @@ def checkpoint(self, *, task: Optional[TaskRecord] = None, exception: Optional[B
             # Or maybe a failure of iteration if the list is appended to while checkpointing is happening?
             if not task:
                 self.checkpointable_tasks = []
+
+    def filter_result_for_checkpoint(self, result: Any) -> bool:
+        """Overridable method to decide if an task that ended with a successful result should be checkpointed"""
+        return True
+
+    def filter_exception_for_checkpoint(self, exception: BaseException) -> bool:
+        """Overridable method to decide if an entry that ended with an exception should be checkpointed"""
+        return False

parsl/tests/test_checkpointing/test_python_checkpoint_exceptions.py

@@ -0,0 +1,49 @@
+import contextlib
+import os
+
+import pytest
+
+import parsl
+from parsl import python_app
+from parsl.config import Config
+from parsl.dataflow.memoization import BasicMemoizer
+from parsl.executors.threads import ThreadPoolExecutor
+
+
+class CheckpointExceptionsMemoizer(BasicMemoizer):
+    def filter_exception_for_checkpoint(self, ex):
+        # TODO: this used to be the case, but in moving to results-only mode,
+        # the task record is lost. Maybe it's useful to pass it in? What
+        # are the use cases for this deciding function?
+        # task record is available from app_fu.task_record
+        # assert app_fu.task_record is not None
+
+        # override the default always-False, to be always-True
+        return True
+
+
+def fresh_config(run_dir, memoizer):
+    return Config(
+        memoizer=memoizer,
+        run_dir=str(run_dir)
+    )
+
+
+@python_app(cache=True)
+def uuid_app():
+    import uuid
+    raise RuntimeError(str(uuid.uuid4()))
+
+
+@pytest.mark.local
+def test_loading_checkpoint(tmpd_cwd):
+    """Load memoization table from previous checkpoint
+    """
+    with parsl.load(fresh_config(tmpd_cwd, CheckpointExceptionsMemoizer(checkpoint_mode="task_exit"))):
+        checkpoint_files = [os.path.join(parsl.dfk().run_dir, "checkpoint")]
+        result = uuid_app().exception()
+
+    with parsl.load(fresh_config(tmpd_cwd, CheckpointExceptionsMemoizer(checkpoint_files=checkpoint_files))):
+        relaunched = uuid_app().exception()
+
+    assert result.args == relaunched.args, "Expected following call to uuid_app to return cached uuid in exception"