From 11710a7f68dce91de4fe8bafd97d09418692f042 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Mon, 16 Feb 2026 02:04:02 +0000 Subject: [PATCH 1/8] Modify quickrun to allow resuming --- environment.yml | 2 +- src/openfecli/commands/quickrun.py | 19 +++++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/environment.yml b/environment.yml index c45d8102a..f74301b03 100644 --- a/environment.yml +++ b/environment.yml @@ -53,7 +53,7 @@ dependencies: # Control blas/openmp threads - threadpoolctl - pip: - - git+https://github.com/OpenFreeEnergy/gufe@main + - git+https://github.com/OpenFreeEnergy/gufe@restart_execute - run_constrained: # drop this pin when handled upstream in espaloma-feedstock - smirnoff99frosst>=1.1.0.1 #https://github.com/openforcefield/smirnoff99Frosst/issues/109 diff --git a/src/openfecli/commands/quickrun.py b/src/openfecli/commands/quickrun.py index f34410d69..443e20a7d 100644 --- a/src/openfecli/commands/quickrun.py +++ b/src/openfecli/commands/quickrun.py @@ -49,9 +49,11 @@ def quickrun(transformation, work_dir, output): for each repeat of the sampling process (by default 3). """ import logging + from json import JSONDecodeError import os import sys + 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 +96,26 @@ 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 + if (work_dir / "protocol_dag.json").is_file(): + write("Attempting to recover edge simulations from file") + try: + dag = ProtocolDAG.from_json(work_dir / "protocol_dag.json") + except JSONDecodeError: + errmsg = "Recovery failed, please clean workdir before continuing" + raise click.ClickException(errmsg) + else: + # Create the DAG instead and then serialize for later resuming + write("Planning simulations for this edge...") + dag = trans.create() + dag.to_json(work_dir / "protocol_dag.json") + write("Starting the simulations for this edge...") dagresult = execute_DAG( dag, shared_basedir=work_dir, scratch_basedir=work_dir, + unitresults_basedir=work_dir, keep_shared=True, raise_error=False, n_retries=2, From 322bc23858bae0b187b1b38131fe526259723a87 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Mon, 16 Feb 2026 12:16:24 +0000 Subject: [PATCH 2/8] fix the gather tests --- .../protocols/openmm_abfe/test_abfe_protocol_results.py | 6 +++--- .../protocols/openmm_ahfe/test_ahfe_protocol_results.py | 6 +++--- .../tests/protocols/openmm_md/test_plain_md_protocol.py | 6 +++--- .../tests/protocols/openmm_rfe/test_hybrid_top_protocol.py | 6 +++--- .../tests/protocols/openmm_septop/test_septop_protocol.py | 6 +++--- 5 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/openfe/tests/protocols/openmm_abfe/test_abfe_protocol_results.py b/src/openfe/tests/protocols/openmm_abfe/test_abfe_protocol_results.py index 5d815c713..53574f972 100644 --- a/src/openfe/tests/protocols/openmm_abfe/test_abfe_protocol_results.py +++ b/src/openfe/tests/protocols/openmm_abfe/test_abfe_protocol_results.py @@ -79,12 +79,12 @@ def patcher(): yield -def test_gather(benzene_complex_dag, patcher, tmpdir): +def test_gather(benzene_complex_dag, patcher, tmp_path): # check that .gather behaves as expected dagres = gufe.protocols.execute_DAG( benzene_complex_dag, - shared_basedir=tmpdir, - scratch_basedir=tmpdir, + shared_basedir=tmp_path, + scratch_basedir=tmp_path, keep_shared=True, ) diff --git a/src/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol_results.py b/src/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol_results.py index 0cb2d2d25..619b199d1 100644 --- a/src/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol_results.py +++ b/src/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol_results.py @@ -99,12 +99,12 @@ def patcher(): yield -def test_gather(benzene_solvation_dag, patcher, tmpdir): +def test_gather(benzene_solvation_dag, patcher, tmp_path): # check that .gather behaves as expected dagres = gufe.protocols.execute_DAG( benzene_solvation_dag, - shared_basedir=tmpdir, - scratch_basedir=tmpdir, + shared_basedir=tmp_path, + scratch_basedir=tmp_path, keep_shared=True, ) diff --git a/src/openfe/tests/protocols/openmm_md/test_plain_md_protocol.py b/src/openfe/tests/protocols/openmm_md/test_plain_md_protocol.py index 60c7e8c47..b8e3153ff 100644 --- a/src/openfe/tests/protocols/openmm_md/test_plain_md_protocol.py +++ b/src/openfe/tests/protocols/openmm_md/test_plain_md_protocol.py @@ -508,7 +508,7 @@ def test_unit_tagging(solvent_protocol_dag, tmpdir): assert len(repeats) == 3 -def test_gather(solvent_protocol_dag, tmpdir): +def test_gather(solvent_protocol_dag, tmp_path): # check .gather behaves as expected with mock.patch( "openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run", @@ -519,8 +519,8 @@ def test_gather(solvent_protocol_dag, tmpdir): ): dagres = gufe.protocols.execute_DAG( solvent_protocol_dag, - shared_basedir=tmpdir, - scratch_basedir=tmpdir, + shared_basedir=tmp_path, + scratch_basedir=tmp_path, keep_shared=True, ) diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py index bd7a1f72f..c632505ea 100644 --- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py +++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py @@ -1225,7 +1225,7 @@ def test_unit_tagging(solvent_protocol_dag, tmpdir): assert len(setup_results) == len(sim_results) == len(analysis_results) == 3 -def test_gather(solvent_protocol_dag, tmpdir): +def test_gather(solvent_protocol_dag, tmp_path): # check .gather behaves as expected with ( mock.patch( @@ -1263,8 +1263,8 @@ def test_gather(solvent_protocol_dag, tmpdir): ): dagres = gufe.protocols.execute_DAG( solvent_protocol_dag, - shared_basedir=tmpdir, - scratch_basedir=tmpdir, + shared_basedir=tmp_path, + scratch_basedir=tmp_path, keep_shared=True, ) diff --git a/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py b/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py index e5e4a9f91..4a22aebcf 100644 --- a/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py +++ b/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py @@ -1295,7 +1295,7 @@ def test_unit_tagging(benzene_toluene_dag, tmpdir): assert len(complex_repeats) == len(solv_repeats) == 2 -def test_gather(benzene_toluene_dag, tmpdir): +def test_gather(benzene_toluene_dag, tmp_path): # check that .gather behaves as expected with ( mock.patch( @@ -1339,8 +1339,8 @@ def test_gather(benzene_toluene_dag, tmpdir): ): dagres = gufe.protocols.execute_DAG( benzene_toluene_dag, - shared_basedir=tmpdir, - scratch_basedir=tmpdir, + shared_basedir=tmp_path, + scratch_basedir=tmp_path, keep_shared=True, ) From a61598e79bfdabbdd12215b03fad2cdc1b3d8f72 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Feb 2026 13:02:17 +0000 Subject: [PATCH 3/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/openfecli/commands/quickrun.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/openfecli/commands/quickrun.py b/src/openfecli/commands/quickrun.py index 443e20a7d..cf809605c 100644 --- a/src/openfecli/commands/quickrun.py +++ b/src/openfecli/commands/quickrun.py @@ -49,9 +49,9 @@ def quickrun(transformation, work_dir, output): for each repeat of the sampling process (by default 3). """ import logging - from json import JSONDecodeError import os import sys + from json import JSONDecodeError from gufe import ProtocolDAG from gufe.protocols.protocoldag import execute_DAG From 0f43f6fcb1659c30b5e44f30d279e282902d6393 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 11 Mar 2026 14:04:29 -0700 Subject: [PATCH 4/8] add check for protocol_dag.json --- src/openfecli/tests/commands/test_quickrun.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/openfecli/tests/commands/test_quickrun.py b/src/openfecli/tests/commands/test_quickrun.py index 86fe00b26..2bab4c53b 100644 --- a/src/openfecli/tests/commands/test_quickrun.py +++ b/src/openfecli/tests/commands/test_quickrun.py @@ -34,6 +34,8 @@ def test_quickrun(extra_args, json_file): assert result.exit_code == 0 assert "Here is the result" in result.output + assert pathlib.Path(extra_args.get("-d", ""), "protocol_dag.json").exists() + if outfile := extra_args.get("-o"): assert pathlib.Path(outfile).exists() with open(outfile, mode="r") as outf: From 5e1a21c7bce95ff6e268113eede2722b8f3b321c Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 11 Mar 2026 14:26:44 -0700 Subject: [PATCH 5/8] add basic test --- src/openfecli/tests/commands/test_quickrun.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/openfecli/tests/commands/test_quickrun.py b/src/openfecli/tests/commands/test_quickrun.py index 2bab4c53b..64b487821 100644 --- a/src/openfecli/tests/commands/test_quickrun.py +++ b/src/openfecli/tests/commands/test_quickrun.py @@ -5,6 +5,7 @@ import click import pytest from click.testing import CliRunner +from gufe import Transformation from gufe.tokenization import JSON_HANDLER from openfecli.commands.quickrun import quickrun @@ -94,3 +95,16 @@ 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_resume(json_file): + trans = Transformation.from_json(json_file) + dag = trans.create() + + runner = CliRunner() + with runner.isolated_filesystem(): + dag.to_json("protocol_dag.json") + result = runner.invoke(quickrun, [json_file]) + + assert result.exit_code == 0 + assert "Attempting to recover edge simulations" in result.output From 182562fed57ed38f3946eabe3cc9edeb5ddbf915 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 11 Mar 2026 14:45:17 -0700 Subject: [PATCH 6/8] clearer language, hopefully --- src/openfecli/commands/quickrun.py | 7 ++++--- src/openfecli/tests/commands/test_quickrun.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/openfecli/commands/quickrun.py b/src/openfecli/commands/quickrun.py index cf809605c..3169c6d94 100644 --- a/src/openfecli/commands/quickrun.py +++ b/src/openfecli/commands/quickrun.py @@ -97,12 +97,13 @@ def quickrun(transformation, work_dir, output): output.parent.mkdir(exist_ok=True, parents=True) # Attempt to either deserialize or freshly create DAG - if (work_dir / "protocol_dag.json").is_file(): - write("Attempting to recover edge simulations from file") + dag_json = work_dir / "protocol_dag.json" + if dag_json.is_file(): + write(f"Attempting to resume execution using existing edges from '{dag_json}'") try: dag = ProtocolDAG.from_json(work_dir / "protocol_dag.json") except JSONDecodeError: - errmsg = "Recovery failed, please clean workdir before continuing" + errmsg = f"Recovery failed, please remove {dag_json} and any results from your working directory before continuing to create a new protocol." raise click.ClickException(errmsg) else: # Create the DAG instead and then serialize for later resuming diff --git a/src/openfecli/tests/commands/test_quickrun.py b/src/openfecli/tests/commands/test_quickrun.py index 64b487821..b7e89625e 100644 --- a/src/openfecli/tests/commands/test_quickrun.py +++ b/src/openfecli/tests/commands/test_quickrun.py @@ -107,4 +107,4 @@ def test_quickrun_resume(json_file): result = runner.invoke(quickrun, [json_file]) assert result.exit_code == 0 - assert "Attempting to recover edge simulations" in result.output + assert "Attempting to resume" in result.output From 3180b8c98848003625f6e2929a0da978f47c55e2 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Wed, 11 Mar 2026 15:30:08 -0700 Subject: [PATCH 7/8] store protocol dag using transformation key --- src/openfecli/commands/quickrun.py | 13 +++++++------ src/openfecli/tests/commands/test_quickrun.py | 10 +++++++--- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/openfecli/commands/quickrun.py b/src/openfecli/commands/quickrun.py index 3169c6d94..308d8f7b0 100644 --- a/src/openfecli/commands/quickrun.py +++ b/src/openfecli/commands/quickrun.py @@ -97,19 +97,20 @@ def quickrun(transformation, work_dir, output): output.parent.mkdir(exist_ok=True, parents=True) # Attempt to either deserialize or freshly create DAG - dag_json = work_dir / "protocol_dag.json" - if dag_json.is_file(): - write(f"Attempting to resume execution using existing edges from '{dag_json}'") + trans_DAG_json = work_dir / f"Transformation-{trans.key}-protocolDAG.json" + + if trans_DAG_json.is_file(): + write(f"Attempting to resume execution using existing edges from '{trans_DAG_json}'") try: - dag = ProtocolDAG.from_json(work_dir / "protocol_dag.json") + dag = ProtocolDAG.from_json(trans_DAG_json) except JSONDecodeError: - errmsg = f"Recovery failed, please remove {dag_json} and any results from your working directory before continuing to create a new protocol." + 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: # Create the DAG instead and then serialize for later resuming write("Planning simulations for this edge...") dag = trans.create() - dag.to_json(work_dir / "protocol_dag.json") + dag.to_json(trans_DAG_json) write("Starting the simulations for this edge...") dagresult = execute_DAG( diff --git a/src/openfecli/tests/commands/test_quickrun.py b/src/openfecli/tests/commands/test_quickrun.py index b7e89625e..6a290cbac 100644 --- a/src/openfecli/tests/commands/test_quickrun.py +++ b/src/openfecli/tests/commands/test_quickrun.py @@ -10,6 +10,8 @@ from openfecli.commands.quickrun import quickrun +# from ..utils import assert_click_success + @pytest.fixture def json_file(): @@ -34,8 +36,10 @@ def test_quickrun(extra_args, json_file): result = runner.invoke(quickrun, [json_file] + extras) assert result.exit_code == 0 assert "Here is the result" in result.output - - assert pathlib.Path(extra_args.get("-d", ""), "protocol_dag.json").exists() + trans = Transformation.from_json(json_file) + assert pathlib.Path( + extra_args.get("-d", ""), f"Transformation-{trans.key}-protocolDAG.json" + ).exists() if outfile := extra_args.get("-o"): assert pathlib.Path(outfile).exists() @@ -103,7 +107,7 @@ def test_quickrun_resume(json_file): runner = CliRunner() with runner.isolated_filesystem(): - dag.to_json("protocol_dag.json") + dag.to_json(f"Transformation-{trans.key}-protocolDAG.json") result = runner.invoke(quickrun, [json_file]) assert result.exit_code == 0 From 1c0fdf79c0aa066f4ea6aab97b3fe65ff943053c Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 12 Mar 2026 07:58:17 -0700 Subject: [PATCH 8/8] another tmpdir -> tmp_path fix --- .../tests/protocols/openmm_rfe/test_hybrid_top_protocol.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py index 7e257e865..83b795eb2 100644 --- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py +++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py @@ -1257,12 +1257,12 @@ def test_unit_tagging(solvent_protocol_dag, unit_mock_patcher, tmpdir): assert len(setup_results) == len(sim_results) == len(analysis_results) == 3 -def test_gather(solvent_protocol_dag, unit_mock_patcher, tmpdir): +def test_gather(solvent_protocol_dag, unit_mock_patcher, tmp_path): # check .gather behaves as expected dagres = gufe.protocols.execute_DAG( solvent_protocol_dag, - shared_basedir=tmpdir, - scratch_basedir=tmpdir, + shared_basedir=tmp_path, + scratch_basedir=tmp_path, keep_shared=True, )