diff --git a/docs/userguide/workflows/checkpoints.rst b/docs/userguide/workflows/checkpoints.rst index 8867107b7a..d4185c73a7 100644 --- a/docs/userguide/workflows/checkpoints.rst +++ b/docs/userguide/workflows/checkpoints.rst @@ -10,10 +10,11 @@ This can save time and computational resources. This is done in two ways: -* Firstly, *app caching* will allow reuse of results within the same run. +* Firstly, *app caching* will allow reuse of results and exceptions within + the same run. -* Building on top of that, *checkpointing* will store results on the filesystem - and reuse those results in later runs. +* Building on top of that, *checkpointing* will store results (but not + exceptions) on the filesystem and reuse those results in later runs. .. _label-appcaching: @@ -264,8 +265,7 @@ of the ``slow_double`` app. # Wait for the results [i.result() for i in d] - cpt_dir = dfk.checkpoint() - print(cpt_dir) # Prints the checkpoint dir + dfk.checkpoint() Resuming from a checkpoint diff --git a/parsl/config.py b/parsl/config.py index 1358e99d28..8489d109ec 100644 --- a/parsl/config.py +++ b/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 @@ -101,6 +102,7 @@ class Config(RepresentationMixin, UsageInformation): def __init__(self, executors: Optional[Iterable[ParslExecutor]] = None, app_cache: bool = True, + memoizer: Optional[Memoizer] = None, checkpoint_files: Optional[Sequence[str]] = None, checkpoint_mode: Union[None, Literal['task_exit'], @@ -131,6 +133,7 @@ def __init__(self, self._executors: Sequence[ParslExecutor] = executors self._validate_executors() + self.memoizer = memoizer self.app_cache = app_cache self.checkpoint_files = checkpoint_files self.checkpoint_mode = checkpoint_mode diff --git a/parsl/dataflow/dflow.py b/parsl/dataflow/dflow.py index 3b3c0ca6e4..0d1e23c20d 100644 --- a/parsl/dataflow/dflow.py +++ b/parsl/dataflow/dflow.py @@ -6,7 +6,6 @@ import inspect import logging import os -import pickle import random import sys import threading @@ -30,7 +29,7 @@ from parsl.dataflow.dependency_resolvers import SHALLOW_DEPENDENCY_RESOLVER from parsl.dataflow.errors import DependencyError, JoinError from parsl.dataflow.futures import AppFuture -from parsl.dataflow.memoization import Memoizer +from parsl.dataflow.memoization import BasicMemoizer, Memoizer from parsl.dataflow.rundirs import make_rundir from parsl.dataflow.states import FINAL_FAILURE_STATES, FINAL_STATES, States from parsl.dataflow.taskrecord import TaskRecord @@ -96,8 +95,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() @@ -160,6 +157,8 @@ def __init__(self, config: Config) -> None: self.monitoring.send((MessageType.WORKFLOW_INFO, workflow_info)) + # TODO: this configuration should become part of the particular memoizer code + # - this is a checkpoint-implementation-specific parameter if config.checkpoint_files is not None: checkpoint_files = config.checkpoint_files elif config.checkpoint_files is None and config.checkpoint_mode is not None: @@ -167,10 +166,20 @@ def __init__(self, config: Config) -> None: else: checkpoint_files = [] - self.memoizer = Memoizer(self, memoize=config.app_cache, checkpoint_files=checkpoint_files) - self.checkpointed_tasks = 0 + # self.memoizer: Memoizer = BasicMemoizer(self, memoize=config.app_cache, checkpoint_files=checkpoint_files) + # the memoize flag might turn into the user choosing different instances + # of the Memoizer interface + self.memoizer: Memoizer + if config.memoizer is not None: + self.memoizer = config.memoizer + else: + self.memoizer = BasicMemoizer() + + self.memoizer.start(dfk=self, memoize=config.app_cache, checkpoint_files=checkpoint_files, run_dir=self.run_dir) self._checkpoint_timer = None self.checkpoint_mode = config.checkpoint_mode + + self._modify_checkpointable_tasks_lock = threading.Lock() self.checkpointable_tasks: List[TaskRecord] = [] # this must be set before executors are added since add_executors calls @@ -186,6 +195,10 @@ def __init__(self, config: Config) -> None: self.add_executors(config.executors) self.add_executors([parsl_internal_executor]) + # TODO: these checkpoint modes should move into the memoizer implementation + # they're (probably?) checkpointer specific: for example the sqlite3-pure-memoizer + # doesn't have a notion of building up an in-memory checkpoint table that needs to be + # flushed on a separate policy if self.checkpoint_mode == "periodic": if config.checkpoint_period is None: raise ConfigurationError("Checkpoint period must be specified with periodic checkpoint mode") @@ -195,7 +208,7 @@ def __init__(self, config: Config) -> None: 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._checkpoint_timer = Timer(self.invoke_checkpoint, interval=checkpoint_period, name="Checkpoint") self.task_count = 0 self.tasks: Dict[int, TaskRecord] = {} @@ -558,9 +571,9 @@ 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.checkpoint_lock: + with self._modify_checkpointable_tasks_lock: self.checkpointable_tasks.append(task_record) elif self.checkpoint_mode is None: pass @@ -1190,15 +1203,23 @@ def cleanup(self) -> None: self.log_task_states() + # TODO: do this in the basic memoizer # Checkpointing takes priority over the rest of the tasks # checkpoint if any valid checkpoint method is specified if self.checkpoint_mode is not None: - self.checkpoint() + + # TODO: accesses to self.checkpointable_tasks should happen + # under a lock? + self.memoizer.checkpoint(self.checkpointable_tasks) if self._checkpoint_timer: logger.info("Stopping checkpoint timer") self._checkpoint_timer.close() + logger.info("Closing memoizer") + self.memoizer.close() + logger.info("Closed memoizer") + # Send final stats self.usage_tracker.send_end_message() self.usage_tracker.close() @@ -1247,68 +1268,10 @@ def cleanup(self) -> None: logger.info("DFK cleanup complete") - def checkpoint(self, tasks: Optional[Sequence[TaskRecord]] = None) -> 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 - """ - with self.checkpoint_lock: - if tasks: - checkpoint_queue = tasks - else: - checkpoint_queue = self.checkpointable_tasks - self.checkpointable_tasks = [] - - checkpoint_dir = '{0}/checkpoint'.format(self.run_dir) - checkpoint_tasks = checkpoint_dir + '/tasks.pkl' - - if not os.path.exists(checkpoint_dir): - os.makedirs(checkpoint_dir, exist_ok=True) - - 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 + def invoke_checkpoint(self) -> None: + with self._modify_checkpointable_tasks_lock: + self.memoizer.checkpoint(self.checkpointable_tasks) + self.checkpointable_tasks = [] @staticmethod def _log_std_streams(task_record: TaskRecord) -> None: diff --git a/parsl/dataflow/memoization.py b/parsl/dataflow/memoization.py index 14ff9d90cb..c9ddb4a71a 100644 --- a/parsl/dataflow/memoization.py +++ b/parsl/dataflow/memoization.py @@ -4,12 +4,14 @@ import logging import os import pickle +import threading from functools import lru_cache, singledispatch -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence import typeguard from parsl.dataflow.errors import BadCheckpoint +from parsl.dataflow.futures import AppFuture from parsl.dataflow.taskrecord import TaskRecord if TYPE_CHECKING: @@ -119,7 +121,60 @@ def id_for_memo_function(f: types.FunctionType, output_ref: bool = False) -> byt return pickle.dumps(["types.FunctionType", f.__name__, f.__module__]) +def make_hash(task: TaskRecord) -> str: + """Create a hash of the task inputs. + + Args: + - task (dict) : Task dictionary from dfk.tasks + + Returns: + - hash (str) : A unique hash string + """ + + t: List[bytes] = [] + + # if kwargs contains an outputs parameter, that parameter is removed + # and normalised differently - with output_ref set to True. + # kwargs listed in ignore_for_cache will also be removed + + filtered_kw = task['kwargs'].copy() + + ignore_list = task['ignore_for_cache'] + + logger.debug("Ignoring these kwargs for checkpointing: %s", ignore_list) + for k in ignore_list: + logger.debug("Ignoring kwarg %s", k) + del filtered_kw[k] + + if 'outputs' in task['kwargs']: + outputs = task['kwargs']['outputs'] + del filtered_kw['outputs'] + t.append(id_for_memo(outputs, output_ref=True)) + + t.extend(map(id_for_memo, (filtered_kw, task['func'], task['args']))) + + x = b''.join(t) + return hashlib.md5(x).hexdigest() + + class Memoizer: + def start(self, *, dfk: DataFlowKernel, memoize: bool = True, checkpoint_files: Sequence[str], run_dir: str) -> None: + raise NotImplementedError + + def update_memo(self, task: TaskRecord, r: Future) -> None: + raise NotImplementedError + + def checkpoint(self, tasks: Sequence[TaskRecord]) -> None: + raise NotImplementedError + + def check_memo(self, task: TaskRecord) -> Optional[Future]: + raise NotImplementedError + + def close(self) -> None: + raise NotImplementedError + + +class BasicMemoizer(Memoizer): """Memoizer is responsible for ensuring that identical work is not repeated. When a task is repeated, i.e., the same function is called with the same exact arguments, the @@ -150,7 +205,12 @@ class Memoizer: """ - def __init__(self, dfk: DataFlowKernel, *, memoize: bool = True, checkpoint_files: Sequence[str]): + run_dir: str + + def __init__(self) -> None: + pass + + def start(self, *, dfk: DataFlowKernel, memoize: bool = True, checkpoint_files: Sequence[str], run_dir: str) -> None: """Initialize the memoizer. Args: @@ -162,7 +222,14 @@ def __init__(self, dfk: DataFlowKernel, *, memoize: bool = True, checkpoint_file """ self.dfk = dfk self.memoize = memoize + self.run_dir = run_dir + + 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) if self.memoize: @@ -172,42 +239,10 @@ def __init__(self, dfk: DataFlowKernel, *, memoize: bool = True, checkpoint_file logger.info("App caching disabled for all apps") self.memo_lookup_table = {} - def make_hash(self, task: TaskRecord) -> str: - """Create a hash of the task inputs. - - Args: - - task (dict) : Task dictionary from dfk.tasks - - Returns: - - hash (str) : A unique hash string - """ - - t: List[bytes] = [] + def close(self) -> None: + pass # nothing to close but more should move here - # if kwargs contains an outputs parameter, that parameter is removed - # and normalised differently - with output_ref set to True. - # kwargs listed in ignore_for_cache will also be removed - - filtered_kw = task['kwargs'].copy() - - ignore_list = task['ignore_for_cache'] - - logger.debug("Ignoring these kwargs for checkpointing: %s", ignore_list) - for k in ignore_list: - logger.debug("Ignoring kwarg %s", k) - del filtered_kw[k] - - if 'outputs' in task['kwargs']: - outputs = task['kwargs']['outputs'] - del filtered_kw['outputs'] - t.append(id_for_memo(outputs, output_ref=True)) - - t.extend(map(id_for_memo, (filtered_kw, task['func'], task['args']))) - - x = b''.join(t) - return hashlib.md5(x).hexdigest() - - def check_memo(self, task: TaskRecord) -> Optional[Future[Any]]: + def check_memo(self, task: TaskRecord) -> Optional[Future]: """Create a hash of the task and its inputs and check the lookup table for this hash. If present, the results are returned. @@ -228,7 +263,7 @@ def check_memo(self, task: TaskRecord) -> Optional[Future[Any]]: logger.debug("Task {} will not be memoized".format(task_id)) return None - hashsum = self.make_hash(task) + hashsum = make_hash(task) logger.debug("Task {} has memoization hash {}".format(task_id, hashsum)) result = None if hashsum in self.memo_lookup_table: @@ -242,7 +277,7 @@ def check_memo(self, task: TaskRecord) -> Optional[Future[Any]]: assert isinstance(result, Future) or result is None return result - def hash_lookup(self, hashsum: str) -> Future[Any]: + def hash_lookup(self, hashsum: str) -> Future: """Lookup a hash in the memoization table. Args: @@ -256,7 +291,7 @@ def hash_lookup(self, hashsum: str) -> Future[Any]: """ return self.memo_lookup_table[hashsum] - def update_memo(self, task: TaskRecord, r: Future[Any]) -> None: + def update_memo(self, task: TaskRecord, r: Future) -> None: """Updates the memoization lookup table with the result from a task. Args: @@ -281,7 +316,7 @@ def update_memo(self, task: TaskRecord, r: Future[Any]) -> None: logger.debug(f"Storing app cache entry {task['hashsum']} with result from task {task_id}") self.memo_lookup_table[task['hashsum']] = r - def _load_checkpoints(self, checkpointDirs: Sequence[str]) -> Dict[str, Future[Any]]: + def _load_checkpoints(self, checkpointDirs: Sequence[str]) -> Dict[str, Future]: """Load a checkpoint file into a lookup table. The data being loaded from the pickle file mostly contains input @@ -309,8 +344,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: @@ -348,3 +387,69 @@ def load_checkpoints(self, checkpointDirs: Optional[Sequence[str]]) -> Dict[str, return self._load_checkpoints(checkpointDirs) else: return {} + + def checkpoint(self, tasks: Sequence[TaskRecord]) -> None: + """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 + """ + with self.checkpoint_lock: + checkpoint_queue = tasks + + checkpoint_dir = '{0}/checkpoint'.format(self.run_dir) + checkpoint_tasks = checkpoint_dir + '/tasks.pkl' + + if not os.path.exists(checkpoint_dir): + os.makedirs(checkpoint_dir, exist_ok=True) + + 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 self.filter_for_checkpoint(app_fu): + + hashsum = task_record['hashsum'] + if not hashsum: + continue + + if app_fu.exception() is None: + t = {'hash': hashsum, 'exception': None, 'result': app_fu.result()} + else: + t = {'hash': hashsum, 'exception': app_fu.exception(), 'result': None} + + # 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 as result".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)) + + def filter_for_checkpoint(self, app_fu: AppFuture) -> bool: + """Overridable method to decide if an entry should be checkpointed""" + return app_fu.exception() is None diff --git a/parsl/dataflow/memosql.py b/parsl/dataflow/memosql.py new file mode 100644 index 0000000000..42bddd505a --- /dev/null +++ b/parsl/dataflow/memosql.py @@ -0,0 +1,118 @@ +import logging +import pickle +import sqlite3 +from concurrent.futures import Future +from pathlib import Path +from typing import Optional, Sequence + +from parsl.dataflow.dflow import DataFlowKernel +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. + + TODO: probably going to need some kind of shutdown now, to close + the sqlite3 connection. + which might also be useful for driving final checkpoints in the + original impl? + """ + + def start(self, *, dfk: DataFlowKernel, memoize: bool = True, checkpoint_files: Sequence[str], 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.db_path = Path(dfk.config.run_dir) / "checkpoint.sqlite3" + logger.debug("starting with db_path %r", self.db_path) + + # TODO: api wart... turning memoization on or off should not be part of the plugin API + self.memoize = memoize + + 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): + pass + + def checkpoint(self, tasks: Sequence[TaskRecord]) -> 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 self.memoize or 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(self, task: TaskRecord, r: Future) -> None: + logger.debug("updating memo") + + if not self.memoize or 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 + + app_fu = task['app_fu'] + hashsum = task['hashsum'] + + # this comes from the original concatenation-based checkpoint code: + if app_fu.exception() is None: + t = {'hash': hashsum, 'exception': None, 'result': app_fu.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() diff --git a/parsl/tests/configs/htex_local_alternate.py b/parsl/tests/configs/htex_local_alternate.py index cc69d56186..638dc818bb 100644 --- a/parsl/tests/configs/htex_local_alternate.py +++ b/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.memosql import SQLiteMemoizer from parsl.executors import HighThroughputExecutor from parsl.launchers import SingleNodeLauncher @@ -64,7 +65,8 @@ def fresh_config(): 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() ) diff --git a/parsl/tests/test_checkpointing/test_python_checkpoint_1.py b/parsl/tests/test_checkpointing/test_python_checkpoint_1.py index 7b8f1f0697..9042c39315 100644 --- a/parsl/tests/test_checkpointing/test_python_checkpoint_1.py +++ b/parsl/tests/test_checkpointing/test_python_checkpoint_1.py @@ -1,4 +1,5 @@ import os +from pathlib import Path import pytest @@ -20,12 +21,14 @@ def uuid_app(): @pytest.mark.local -def test_initial_checkpoint_write(): +def test_initial_checkpoint_write() -> None: """1. Launch a few apps and write the checkpoint once a few have completed """ uuid_app().result() - cpt_dir = parsl.dfk().checkpoint() + parsl.dfk().invoke_checkpoint() - cptpath = cpt_dir + '/tasks.pkl' + cpt_dir = Path(parsl.dfk().run_dir) / 'checkpoint' + + cptpath = cpt_dir / 'tasks.pkl' assert os.path.exists(cptpath), f"Tasks checkpoint missing: {cptpath}" diff --git a/parsl/tests/test_checkpointing/test_python_checkpoint_2_sqlite.py b/parsl/tests/test_checkpointing/test_python_checkpoint_2_sqlite.py new file mode 100644 index 0000000000..756dcad113 --- /dev/null +++ b/parsl/tests/test_checkpointing/test_python_checkpoint_2_sqlite.py @@ -0,0 +1,44 @@ +import contextlib +import os + +import pytest + +import parsl +from parsl import python_app +from parsl.dataflow.memosql import SQLiteMemoizer +from parsl.tests.configs.local_threads_checkpoint import fresh_config + + +@contextlib.contextmanager +def parsl_configured(run_dir, **kw): + c = fresh_config() + c.memoizer = SQLiteMemoizer() + c.run_dir = run_dir + for config_attr, config_val in kw.items(): + setattr(c, config_attr, config_val) + dfk = parsl.load(c) + for ex in dfk.executors.values(): + ex.working_dir = run_dir + yield dfk + + parsl.dfk().cleanup() + + +@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, checkpoint_mode="task_exit"): + checkpoint_files = [os.path.join(parsl.dfk().run_dir, "checkpoint")] + result = uuid_app().result() + + with parsl_configured(tmpd_cwd, checkpoint_files=checkpoint_files): + relaunched = uuid_app().result() + + assert result == relaunched, "Expected following call to uuid_app to return cached uuid" diff --git a/parsl/tests/test_checkpointing/test_python_checkpoint_exceptions.py b/parsl/tests/test_checkpointing/test_python_checkpoint_exceptions.py new file mode 100644 index 0000000000..1eca421562 --- /dev/null +++ b/parsl/tests/test_checkpointing/test_python_checkpoint_exceptions.py @@ -0,0 +1,66 @@ +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_for_checkpoint(self, app_fu): + # checkpoint everything, rather than selecting only futures with + # results, not exceptions. + + # task record is available from app_fu.task_record + assert app_fu.task_record is not None + + return True + + +def fresh_config(): + return Config( + memoizer=CheckpointExceptionsMemoizer(), + executors=[ + ThreadPoolExecutor( + label='local_threads_checkpoint', + ) + ] + ) + + +@contextlib.contextmanager +def parsl_configured(run_dir, **kw): + c = fresh_config() + c.run_dir = run_dir + for config_attr, config_val in kw.items(): + setattr(c, config_attr, config_val) + dfk = parsl.load(c) + for ex in dfk.executors.values(): + ex.working_dir = run_dir + yield dfk + + parsl.dfk().cleanup() + + +@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_configured(tmpd_cwd, checkpoint_mode="task_exit"): + checkpoint_files = [os.path.join(parsl.dfk().run_dir, "checkpoint")] + result = uuid_app().exception() + + with parsl_configured(tmpd_cwd, 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" diff --git a/parsl/tests/test_python_apps/test_memoize_exception.py b/parsl/tests/test_python_apps/test_memoize_exception.py new file mode 100644 index 0000000000..e152bcfba5 --- /dev/null +++ b/parsl/tests/test_python_apps/test_memoize_exception.py @@ -0,0 +1,22 @@ +import argparse + +import parsl +from parsl.app.app import python_app + + +@python_app(cache=True) +def raise_exception(x, cache=True): + raise RuntimeError("exception from raise_exception") + + +def test_python_memoization(n=2): + """Testing python memoization with exceptions.""" + x = raise_exception(0) + + # wait for x to be done + x.exception() + + for i in range(0, n): + foo = raise_exception(0) + print(foo.exception()) + assert foo.exception() == x.exception(), "Memoized exceptions were not used" diff --git a/parsl/tests/test_python_apps/test_memoize_plugin.py b/parsl/tests/test_python_apps/test_memoize_plugin.py new file mode 100644 index 0000000000..724facf165 --- /dev/null +++ b/parsl/tests/test_python_apps/test_memoize_plugin.py @@ -0,0 +1,53 @@ +import argparse + +import pytest + +import parsl +from parsl.app.app import python_app +from parsl.config import Config +from parsl.dataflow.memoization import BasicMemoizer +from parsl.dataflow.taskrecord import TaskRecord + + +class DontReuseSevenMemoizer(BasicMemoizer): + def check_memo(self, task_record: TaskRecord): + if task_record['args'][0] == 7: + return None # we didn't find a suitable memo record... + else: + return super().check_memo(task_record) + + +def local_config(): + return Config(memoizer=DontReuseSevenMemoizer()) + + +@python_app(cache=True) +def random_uuid(x, cache=True): + import uuid + return str(uuid.uuid4()) + + +@pytest.mark.local +def test_python_memoization(n=10): + """Testing python memoization disable + """ + + # TODO: this .result() needs to be here, not in the loop + # because otherwise we race to complete... and then + # we might sometimes get a memoization before the loop + # and sometimes not... + x = random_uuid(0).result() + + for i in range(0, n): + foo = random_uuid(0) + print(i) + print(foo.result()) + assert foo.result() == x, "Memoized results were incorrectly not used" + + y = random_uuid(7).result() + + for i in range(0, n): + foo = random_uuid(7) + print(i) + print(foo.result()) + assert foo.result() != y, "Memoized results were incorrectly used"