out-of-memory checkpointing

goal: results should not (never? in weak small cache?) be stored in an
in-memory memo table. so that memo table should be not present in this
implementation. instead all memo questions go to the sqlite3 database.

this drives some blurring between in-memory caching and disk-based
checkpointing: the previous disk based checkpointed model relied on
repopulating the in-memory memo table cache...

i hit some thread problems when using one sqlite3 connection across threads
and the docs are unclear about what I can/cannot do, so i made this open
the sqlite3 database on every access. that's probably got quite a performance
hit, but its probably enough for basically validating the idea.

the three usecases: sqlite3, in-memory and traditional on-disk checkpoint,
maybe needs some thoughts about the completion semantics and data structures
to go with it.

TODO: this breaks:
  pytest  parsl/tests/test_python_apps/test_memoize_exception.py  --config parsl/tests/configs/htex_local_alternate.py
because the behaviour asserted there is only behaviour for the default
memoizer. So it should probably become a --config local test with
known configuration(s).

Author: benclifford

parsl/benchmark/perf.py

@@ -48,7 +48,7 @@ def performance(*, resources: dict, target_t: float, args_extra_size: int, itera
 
     iteration = 1
 
-    args_extra_payload = "x" * args_extra_size
+    # args_extra_payload = "x" * args_extra_size
 
     iterate = True
 
@@ -59,7 +59,13 @@ def performance(*, resources: dict, target_t: float, args_extra_size: int, itera
 
         fs = []
         print("Submitting tasks / invoking apps")
-        for _ in range(n):
+        for index in range(n):
+            # this means there is a different argument for each iteration,
+            # which will make checkpointing/memo behave differently
+            # so this could be switchable in parsl-perf dev branch
+            # args_extra_payload = index   # always a new one (except for run repeats)
+
+            args_extra_payload = index % 10
             fs.append(app(args_extra_payload, parsl_resource_specification=resources))
 
         submitted_t = time.time()

parsl/dataflow/memosql.py

@@ -0,0 +1,145 @@
+import logging
+import pickle
+import sqlite3
+from concurrent.futures import Future
+from pathlib import Path
+from typing import Any, Optional
+
+from parsl.dataflow.memoization import Memoizer, make_hash
+from parsl.dataflow.taskrecord import TaskRecord
+
+logger = logging.getLogger(__name__)
+
+
+class SQLiteMemoizer(Memoizer):
+    """Memoize out of memory into an sqlite3 database.
+    """
+
+    def __init__(self, *, checkpoint_dir: str | None = None):
+        self.checkpoint_dir = checkpoint_dir
+
+    def start(self, *, run_dir: str) -> None:
+        """TODO: run_dir is the per-workflow run dir, but we need a broader checkpoint context... one level up
+        by default... get_all_checkpoints uses "runinfo/" as a relative path for that by default so replicating
+        that choice would do here. likewise I think for monitoring."""
+
+        self.run_dir = run_dir
+
+        dir = self.checkpoint_dir if self.checkpoint_dir is not None else self.run_dir
+
+        self.db_path = Path(dir) / "checkpoint.sqlite3"
+        logger.debug("starting with db_path %r", self.db_path)
+
+        connection = sqlite3.connect(self.db_path)
+        cursor = connection.cursor()
+
+        cursor.execute("CREATE TABLE IF NOT EXISTS checkpoints(key, result)")
+        # probably want some index on key because that's what we're doing all the access via.
+
+        connection.commit()
+        connection.close()
+        logger.debug("checkpoint table created")
+
+    def close(self):
+        """TODO: probably going to need some kind of shutdown now, to close the sqlite3 connection."""
+        pass
+
+    def checkpoint(self, *, task: TaskRecord | None = None) -> None:
+        """All the behaviour for this memoizer is in check_memo and update_memo.
+        """
+        logger.debug("Explicit checkpoint call is a no-op with this memoizer")
+
+    def check_memo(self, task: TaskRecord) -> Optional[Future]:
+        """TODO: document this: check_memo is required to set the task hashsum,
+        if that's how we're going to key checkpoints in update_memo. (that's not
+        a requirement though: other equalities are available."""
+        task_id = task['id']
+
+        if not task['memoize']:
+            task['hashsum'] = None
+            logger.debug("Task %s will not be memoized", task_id)
+            return None
+
+        hashsum = make_hash(task)
+        logger.debug("Task {} has memoization hash {}".format(task_id, hashsum))
+        task['hashsum'] = hashsum
+
+        connection = sqlite3.connect(self.db_path)
+        cursor = connection.cursor()
+        cursor.execute("SELECT result FROM checkpoints WHERE key = ?", (hashsum, ))
+        r = cursor.fetchone()
+
+        if r is None:
+            connection.close()
+            return None
+        else:
+            data = pickle.loads(r[0])
+            connection.close()
+
+            memo_fu: Future = Future()
+
+            if data['exception'] is None:
+                memo_fu.set_result(data['result'])
+            else:
+                assert data['result'] is None
+                memo_fu.set_exception(data['exception'])
+
+            return memo_fu
+
+    def update_memo_result(self, task: TaskRecord, result: Any) -> None:
+        logger.debug("updating memo")
+
+        if not task['memoize'] or 'hashsum' not in task:
+            logger.debug("preconditions for memo not satisfied")
+            return
+
+        if not isinstance(task['hashsum'], str):
+            logger.error(f"Attempting to update app cache entry but hashsum is not a string key: {task['hashsum']}")
+            return
+
+        hashsum = task['hashsum']
+
+        # this comes from the original concatenation-based checkpoint code:
+        # assert app_fu.done(), "assumption: update_memo is called after future has a result"
+        t = {'hash': hashsum, 'exception': None, 'result': result}
+        # else:
+        #    t = {'hash': hashsum, 'exception': app_fu.exception(), 'result': None}
+
+        value = pickle.dumps(t)
+
+        connection = sqlite3.connect(self.db_path)
+        cursor = connection.cursor()
+
+        cursor.execute("INSERT INTO checkpoints VALUES(?, ?)", (hashsum, value))
+
+        connection.commit()
+        connection.close()
+
+    def update_memo_exception(self, task: TaskRecord, exception: BaseException) -> None:
+        logger.debug("updating memo")
+
+        if not task['memoize'] or 'hashsum' not in task:
+            logger.debug("preconditions for memo not satisfied")
+            return
+
+        if not isinstance(task['hashsum'], str):
+            logger.error(f"Attempting to update app cache entry but hashsum is not a string key: {task['hashsum']}")
+            return
+
+        hashsum = task['hashsum']
+
+        # this comes from the original concatenation-based checkpoint code:
+        # assert app_fu.done(), "assumption: update_memo is called after future has a result"
+        # t = {'hash': hashsum, 'exception': None, 'result': app_fu.result()}
+        # else:
+        t = {'hash': hashsum, 'exception': exception, 'result': None}
+
+        value = pickle.dumps(t)
+
+        connection = sqlite3.connect(self.db_path)
+        cursor = connection.cursor()
+
+        cursor.execute("INSERT INTO checkpoints VALUES(?, ?)", (hashsum, value))
+
+        connection.commit()
+        connection.close()

parsl/tests/configs/htex_local_alternate.py

@@ -22,7 +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.dataflow.memosql import SQLiteMemoizer
 from parsl.executors import HighThroughputExecutor
 from parsl.launchers import SingleNodeLauncher
 
@@ -57,14 +57,14 @@ def fresh_config():
             )
         ],
         strategy='simple',
-        memoizer=BasicMemoizer(memoize=True, checkpoint_mode='task_exit'),
         retries=2,
         monitoring=MonitoringHub(
                         monitoring_debug=False,
                         resource_monitoring_interval=1,
         ),
         usage_tracking=3,
-        project_name="parsl htex_local_alternate test configuration"
+        project_name="parsl htex_local_alternate test configuration",
+        memoizer=SQLiteMemoizer()
     )
 
 

parsl/tests/test_checkpointing/test_python_checkpoint_2_sqlite.py

@@ -0,0 +1,35 @@
+import contextlib
+import os
+
+import pytest
+
+import parsl
+from parsl import python_app
+from parsl.config import Config
+from parsl.dataflow.memosql import SQLiteMemoizer
+
+
+def parsl_configured(run_dir, memoizer):
+    return parsl.load(Config(
+        run_dir=str(run_dir),
+        memoizer=memoizer
+    ))
+
+
+@python_app(cache=True)
+def uuid_app():
+    import uuid
+    return uuid.uuid4()
+
+
+@pytest.mark.local
+def test_loading_checkpoint(tmpd_cwd):
+    """Load memoization table from previous checkpoint
+    """
+    with parsl_configured(tmpd_cwd, SQLiteMemoizer(checkpoint_dir=tmpd_cwd)):
+        result = uuid_app().result()
+
+    with parsl_configured(tmpd_cwd, SQLiteMemoizer(checkpoint_dir=tmpd_cwd)):
+        relaunched = uuid_app().result()
+
+    assert result == relaunched, "Expected following call to uuid_app to return cached uuid"