Modify quickrun to allow resuming

news/quickrun_resume.rst

@@ -0,0 +1,24 @@
+**Added:**
+
+* Added ``--resume`` flag to ``openfe quickrun``.
+  Quickrun now temporarily caches ``protocolDAG`` information and when used with the ``--resume`` flag, quickrun will attempt resume execution of an incomplete transformation.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/openfecli/commands/quickrun.py

@@ -3,6 +3,7 @@
 
 import json
 import pathlib
+import warnings
 
 import click
 
@@ -30,8 +31,14 @@ def _format_exception(exception) -> str:
     type=click.Path(dir_okay=False, file_okay=False, path_type=pathlib.Path),
     help="Filepath at which to create and write the JSON-formatted results.",
 )  # fmt: skip
+@click.option(
+    "--resume",
+    is_flag=True,
+    default=False,
+    help=("Attempt to resume this transformation's execution using the cache."),
+)
 @print_duration
-def quickrun(transformation, work_dir, output):
+def quickrun(transformation, work_dir, output, resume):
     """Run the transformation (edge) in the given JSON file.
 
     Simulation JSON files can be created with the
@@ -51,7 +58,9 @@ def quickrun(transformation, work_dir, output):
     import logging
     import os
     import sys
+    from json import JSONDecodeError
 
+    from gufe import ProtocolDAG
     from gufe.protocols.protocoldag import execute_DAG
     from gufe.tokenization import JSON_HANDLER
     from gufe.transformations.transformation import Transformation
@@ -94,13 +103,40 @@ def quickrun(transformation, work_dir, output):
     else:
         output.parent.mkdir(exist_ok=True, parents=True)
 
-    write("Planning simulations for this edge...")
-    dag = trans.create()
+    # Attempt to either deserialize or freshly create DAG
+    cache_basedir = work_dir / "quickrun_cache"
+    trans_DAG_json = cache_basedir / f"{trans.key}-ProtocolDAG.json"
+
+    if trans_DAG_json.is_file():
+        if resume:
+            write(f"Attempting to resume execution using existing edges from '{trans_DAG_json}'")
+            try:
+                dag = ProtocolDAG.from_json(trans_DAG_json)
+            except JSONDecodeError:
+                errmsg = f"Recovery failed, please remove {trans_DAG_json} and any results from your working directory before continuing to create a new protocol."
+                raise click.ClickException(errmsg)
+        else:
+            errmsg = f"Transformation has been started but is incomplete. Please remove {trans_DAG_json} and rerun, or resume execution using the ``--resume`` flag."
+            raise click.ClickException(errmsg)
+
+    else:
+        if resume:
+            click.echo(
+                f"openfe quickrun was run with --resume, but no checkpoint found at {trans_DAG_json}. Starting new execution."
+            )
+
+        # Create the DAG instead and then serialize for later resuming
+        write("Planning simulations for this edge...")
+        dag = trans.create()
+        cache_basedir.mkdir(exist_ok=True)
+        dag.to_json(trans_DAG_json)
+
     write("Starting the simulations for this edge...")
     dagresult = execute_DAG(
         dag,
         shared_basedir=work_dir,
         scratch_basedir=work_dir,
+        cache_basedir=cache_basedir,
         keep_shared=True,
         raise_error=False,
         n_retries=2,
@@ -126,6 +162,9 @@ def quickrun(transformation, work_dir, output):
     with open(output, mode="w") as outf:
         json.dump(out_dict, outf, cls=JSON_HANDLER.encoder)
 
+    # remove the checkpoint since the job has completed
+    os.remove(trans_DAG_json)
+
     write(f"Here is the result:\n\tdG = {estimate} ± {uncertainty}\n")
     write("")
 

src/openfecli/tests/commands/test_quickrun.py

@@ -1,6 +1,8 @@
 import json
+import os
 import pathlib
 from importlib import resources
+from unittest import mock
 
 import pytest
 from click.testing import CliRunner
@@ -9,6 +11,8 @@
 
 from openfecli.commands.quickrun import quickrun
 
+from ..utils import assert_click_success
+
 
 @pytest.fixture
 def json_file():
@@ -18,13 +22,7 @@ def json_file():
     return json_file
 
 
-@pytest.mark.parametrize(
-    "extra_args",
-    [
-        {},
-        {"-d": "foo_dir", "-o": "foo.json"},
-    ],
-)
+@pytest.mark.parametrize("extra_args", [{}, {"-d": "foo_dir", "-o": "foo.json"}])
 def test_quickrun(extra_args, json_file):
     extras = sum([list(kv) for kv in extra_args.items()], [])
 
@@ -37,8 +35,14 @@ def test_quickrun(extra_args, json_file):
         assert not pathlib.Path(outfile).exists()
 
         result = runner.invoke(quickrun, [json_file] + extras)
-        assert result.exit_code == 0
+
+        assert_click_success(result)
         assert "Here is the result" in result.output
+        trans = Transformation.from_json(json_file)
+        # checkpoint should be deleted when job is complete
+        assert not pathlib.Path(
+            extra_args.get("-d", ""), "quickrun_cache", f"{trans.key}-ProtocolDAG.json"
+        ).exists()
 
         # output json should exist and have results after execution
         assert pathlib.Path(outfile).exists()
@@ -55,6 +59,23 @@ def test_quickrun(extra_args, json_file):
         #     assert len(list(dirpath.iterdir())) > 0
 
 
+@pytest.mark.parametrize("extra_args", [{}, {"-d": "foo_dir", "-o": "foo.json"}])
+def test_quickrun_interrupted(extra_args, json_file):
+    """If a quickrun is unable to complete, the ProtocolDAG.json checkpoint should exist."""
+    extras = sum([list(kv) for kv in extra_args.items()], [])
+
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        with mock.patch("gufe.protocols.protocoldag.execute_DAG", side_effect=RuntimeError):
+            result = runner.invoke(quickrun, [json_file] + extras)
+
+        assert "Here is the result" not in result.output
+        trans = Transformation.from_json(json_file)
+        assert pathlib.Path(
+            extra_args.get("-d", ""), "quickrun_cache", f"{trans.key}-ProtocolDAG.json"
+        ).exists()
+
+
 def test_quickrun_output_file_exists(json_file):
     """Fail if the output file already exists."""
     runner = CliRunner()
@@ -71,7 +92,7 @@ def test_quickrun_output_file_in_nonexistent_directory(json_file):
     with runner.isolated_filesystem():
         outfile = pathlib.Path("not_dir/foo.json")
         result = runner.invoke(quickrun, [json_file, "-o", outfile])
-        assert result.exit_code == 0
+        assert_click_success(result)
         assert outfile.parent.is_dir()
 
 
@@ -82,7 +103,7 @@ def test_quickrun_dir_created_at_runtime(json_file):
         outdir = "not_dir"
         outfile = outdir + "foo.json"
         result = runner.invoke(quickrun, [json_file, "-d", outdir, "-o", outfile])
-        assert result.exit_code == 0
+        assert_click_success(result)
 
 
 def test_quickrun_unit_error():
@@ -98,3 +119,56 @@ def test_quickrun_unit_error():
         # to be stored in JSON
         # not sure whether that means we should always be storing all
         # protocol dag results maybe?
+
+
+def test_quickrun_existing_checkpoint(json_file):
+    """In the default case where resume=False, if the checkpoint exists, quickrun should error out and not attempt to execute."""
+    trans = Transformation.from_json(json_file)
+    dag = trans.create()
+
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        pathlib.Path("quickrun_cache").mkdir()
+        dag.to_json(pathlib.Path("quickrun_cache", f"{trans.key}-ProtocolDAG.json"))
+        result = runner.invoke(quickrun, [json_file])
+        assert result.exit_code == 1
+        assert "Attempting to resume" not in result.output
+        assert "Transformation has been started but is incomplete." in result.stderr
+
+
+def test_quickrun_resume_from_checkpoint(json_file):
+    trans = Transformation.from_json(json_file)
+    dag = trans.create()
+
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        pathlib.Path("quickrun_cache").mkdir()
+        dag.to_json(pathlib.Path("quickrun_cache", f"{trans.key}-ProtocolDAG.json"))
+        result = runner.invoke(quickrun, [json_file, "--resume"])
+
+        assert_click_success(result)
+        assert "Attempting to resume" in result.output
+
+
+def test_quickrun_resume_invalid_checkpoint(json_file):
+    """Fail if the output file doesn't load properly."""
+    trans = Transformation.from_json(json_file)
+
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        pathlib.Path("quickrun_cache").mkdir()
+        pathlib.Path("quickrun_cache", f"{trans.key}-ProtocolDAG.json").touch()
+        result = runner.invoke(quickrun, [json_file, "--resume"])
+
+        assert result.exit_code == 1
+        assert "Attempting to resume" in result.output
+        assert "Recovery failed" in result.stderr
+
+
+def test_quickrun_resume_missing_checkpoint(json_file):
+    """If --resume is passed but there's not checkpoint, just echo a message and keep going."""
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        result = runner.invoke(quickrun, [json_file, "--resume"])
+        assert_click_success(result)
+        assert "openfe quickrun was run with --resume, but no checkpoint found at" in result.output