diff --git a/src/openfe/storage/warehouse.py b/src/openfe/storage/warehouse.py index c10baa12c..e4be520ca 100644 --- a/src/openfe/storage/warehouse.py +++ b/src/openfe/storage/warehouse.py @@ -166,6 +166,7 @@ def _get_store_for_key(self, key: GufeKey) -> ExternalStorage: ValueError If the key is not found in any store. """ + print(key) for name in self.stores: if key in self.stores[name]: return self.stores[name] @@ -313,7 +314,14 @@ class FileSystemWarehouse(WarehouseBaseClass): """ def __init__(self, root_dir: str = "warehouse"): + self.root_dir = root_dir setup_store = FileStorage(f"{root_dir}/setup") result_store = FileStorage(f"{root_dir}/result") stores = WarehouseStores(setup=setup_store, result=result_store) super().__init__(stores) + + def dump(self): + print(self.root_dir) + print(self.setup_store) + for item in self.setup_store: + print(item) diff --git a/src/openfecli/commands/plan_rbfe_network.py b/src/openfecli/commands/plan_rbfe_network.py index 5b73b5618..e74b06d1d 100644 --- a/src/openfecli/commands/plan_rbfe_network.py +++ b/src/openfecli/commands/plan_rbfe_network.py @@ -1,8 +1,8 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/openfe - import click +from openfe.storage.warehouse import FileSystemWarehouse from openfecli import OFECommandPlugin from openfecli.parameters import ( COFACTORS, @@ -12,6 +12,7 @@ OUTPUT_DIR, OVERWRITE, PROTEIN, + WAREHOUSE, YAML_OPTIONS, ) from openfecli.utils import print_duration, write @@ -130,6 +131,7 @@ def plan_rbfe_network_main( @N_PROTOCOL_REPEATS.parameter(multiple=False, required=False, default=3, help=N_PROTOCOL_REPEATS.kwargs["help"]) # fmt: skip @NCORES.parameter(help=NCORES.kwargs["help"], default=1) @OVERWRITE.parameter(help=OVERWRITE.kwargs["help"], default=OVERWRITE.kwargs["default"], is_flag=True) # fmt: skip +@WAREHOUSE.parameter(help=WAREHOUSE.kwargs["help"], is_flag=True) @print_duration def plan_rbfe_network( molecules: list[str], @@ -140,6 +142,7 @@ def plan_rbfe_network( n_protocol_repeats: int, n_cores: int, overwrite_charges: bool, + warehouse: bool, ): """ Plan a relative binding free energy network, saved as JSON files for use by @@ -243,10 +246,15 @@ def plan_rbfe_network( # OUTPUT write("Output:") write("\tSaving to: " + str(output_dir)) + warehouse_object = None + if warehouse: + warehouse_object = FileSystemWarehouse() + plan_alchemical_network_output( alchemical_network=alchemical_network, ligand_network=ligand_network, folder_path=OUTPUT_DIR.get(output_dir), + warehouse=warehouse_object, ) diff --git a/src/openfecli/commands/quickrun-warehouse.py b/src/openfecli/commands/quickrun-warehouse.py new file mode 100644 index 000000000..ac77e8af8 --- /dev/null +++ b/src/openfecli/commands/quickrun-warehouse.py @@ -0,0 +1,151 @@ +import json +import pathlib +from typing import Optional + +import click + +from openfe.storage.warehouse import FileSystemWarehouse, WarehouseBaseClass +from openfecli.commands.quickrun import _format_exception +from openfecli.plugins import OFECommandPlugin +from openfecli.utils import configure_logger, print_duration, write + + +@click.command("warehouse-quickrun", short_help="Run a given transformation, saved as a JSON file") +@click.option( + "--work-dir", "-d", default=None, + type=click.Path(dir_okay=True, file_okay=False, writable=True, path_type=pathlib.Path), + help=( + "Directory in which to store files in (defaults to current directory). " + "If the directory does not exist, it will be created at runtime." + ), +) # fmt: skip +@click.option( + "output", "-o", default=None, + 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.argument("transformation", type=click.STRING, required=True) +@print_duration +def quickrun(transformation, work_dir, output): + """Run the transformation (edge) in the given JSON file. + + Simulation JSON files can be created with the + :ref:`cli_plan-rbfe-network` + or from Python a :class:`.Transformation` can be saved using its to_json + method:: + + transformation.to_json("filename.json") + + That will save a JSON file suitable to be input for this command. + + Running this command will execute the simulation defined in the JSON file, + creating a directory for each individual task (``Unit``) in the workflow. + For example, when running the OpenMM HREX Protocol a directory will be created + for each repeat of the sampling process (by default 3). + """ + import logging + import os + import sys + + from gufe.protocols.protocoldag import execute_DAG + from gufe.tokenization import JSON_HANDLER + from gufe.transformations.transformation import Transformation + + from openfe.utils import logging_control + + # avoid problems with output not showing if queueing system kills a job + sys.stdout.reconfigure(line_buffering=True) + + stdout_handler = logging.StreamHandler(sys.stdout) + + configure_logger("gufekey", handler=stdout_handler) + configure_logger("gufe", handler=stdout_handler) + configure_logger("openfe", handler=stdout_handler) + + # silence the openmmtools.multistate API warning + logging_control._silence_message( + msg=[ + "The openmmtools.multistate API is experimental and may change in future releases", + ], + logger_names=[ + "openmmtools.multistate.multistatereporter", + "openmmtools.multistate.multistateanalyzer", + "openmmtools.multistate.multistatesampler", + ], + ) + # turn warnings into log message (don't show stack trace) + logging.captureWarnings(True) + + if work_dir is None: + work_dir = pathlib.Path(os.getcwd()) / "warehouse" + else: + work_dir.mkdir(exist_ok=True, parents=True) + + # Setup the warehouse + warehouse = FileSystemWarehouse() + trans: Transformation = warehouse.load_setup_tokenizable(transformation) + + write("Loading file...") + + output = work_dir / (str(trans.key) + "_results.json") + + write("Planning simulations for this edge...") + dag = trans.create() + write("Starting the simulations for this edge...") + dagresult = execute_DAG( + dag, + shared_basedir=work_dir, + scratch_basedir=work_dir, + keep_shared=True, + raise_error=False, + n_retries=2, + ) + # How this would change with new context + # dagresult = execute_DAG( + # dag, + # shared_storage=warehouse.shared, + # perm_storage=warehouse.perm + # scratch_basedir=work_dir, + # keep_shared=True, + # raise_error=False, + # n_retries=2, + # ) + warehouse.store_result_tokenizable(dagresult) + write("Done with all simulations! Analyzing the results....") + prot_result = trans.protocol.gather([dagresult]) + + if dagresult.ok(): + estimate = prot_result.get_estimate() + uncertainty = prot_result.get_uncertainty() + else: + estimate = uncertainty = None # for output file + + out_dict = { + "estimate": estimate, + "uncertainty": uncertainty, + "protocol_result": prot_result.to_dict(), + "unit_results": { + unit.key: unit.to_keyed_dict() for unit in dagresult.protocol_unit_results + }, + } + + with open(output, mode="w") as outf: + json.dump(out_dict, outf, cls=JSON_HANDLER.encoder) + + write(f"Here is the result:\n\tdG = {estimate} ± {uncertainty}\n") + write("") + + if not dagresult.ok(): + # there can be only one, MacCleod + failure = dagresult.protocol_unit_failures[-1] + raise click.ClickException( + f"The protocol unit '{failure.name}' failed with the error " + f"message:\n{_format_exception(failure.exception)}\n\n" + "Details provided in output." + ) + + +PLUGIN = OFECommandPlugin(command=quickrun, section="Quickrun Executor", requires_ofe=(0, 3)) + +if __name__ == "__main__": + quickrun() diff --git a/src/openfecli/commands/quickrun.py b/src/openfecli/commands/quickrun.py index f34410d69..fdc3377f7 100644 --- a/src/openfecli/commands/quickrun.py +++ b/src/openfecli/commands/quickrun.py @@ -7,6 +7,7 @@ import click from openfecli import OFECommandPlugin +from openfecli.parameters import WAREHOUSE from openfecli.utils import configure_logger, print_duration, write diff --git a/src/openfecli/parameters/__init__.py b/src/openfecli/parameters/__init__.py index fb8dcd0f9..81e96e18f 100644 --- a/src/openfecli/parameters/__init__.py +++ b/src/openfecli/parameters/__init__.py @@ -9,3 +9,4 @@ from .output_dir import OUTPUT_DIR from .plan_network_options import YAML_OPTIONS from .protein import PROTEIN +from .warehouse import WAREHOUSE diff --git a/src/openfecli/parameters/warehouse.py b/src/openfecli/parameters/warehouse.py new file mode 100644 index 000000000..5fb2f07f6 --- /dev/null +++ b/src/openfecli/parameters/warehouse.py @@ -0,0 +1,4 @@ +import click +from plugcli.params import Option + +WAREHOUSE = Option("--warehouse", type=click.BOOL, help="Use a warehouse", default=False) diff --git a/src/openfecli/plan_alchemical_networks_utils.py b/src/openfecli/plan_alchemical_networks_utils.py index 9636fb9b5..da0572603 100644 --- a/src/openfecli/plan_alchemical_networks_utils.py +++ b/src/openfecli/plan_alchemical_networks_utils.py @@ -4,8 +4,10 @@ import json import pathlib +from typing import Optional from openfe import AlchemicalNetwork, LigandNetwork +from openfe.storage.warehouse import WarehouseBaseClass from openfecli.utils import write @@ -13,26 +15,31 @@ def plan_alchemical_network_output( alchemical_network: AlchemicalNetwork, ligand_network: LigandNetwork, folder_path: pathlib.Path, + warehouse: Optional[WarehouseBaseClass], ): """Write the contents of an alchemical network into the structure""" - base_name = folder_path.name - folder_path.mkdir(parents=True, exist_ok=True) - - an_json = folder_path / f"{base_name}.json" - alchemical_network.to_json(an_json) - write("\t\t- " + base_name + ".json") - - ln_fname = "ligand_network.graphml" - with open(folder_path / ln_fname, mode="w") as f: - f.write(ligand_network.to_graphml()) - write(f"\t\t- {ln_fname}") - - transformations_dir = folder_path / "transformations" - transformations_dir.mkdir(parents=True, exist_ok=True) - - for transformation in alchemical_network.edges: - transformation_name = transformation.name or transformation.key - filename = f"{transformation_name}.json" - transformation.to_json(transformations_dir / filename) - write("\t\t\t\t- " + filename) + if warehouse: + warehouse.store_setup_tokenizable(alchemical_network) + warehouse.store_setup_tokenizable(ligand_network) + else: + base_name = folder_path.name + folder_path.mkdir(parents=True, exist_ok=True) + + an_json = folder_path / f"{base_name}.json" + alchemical_network.to_json(an_json) + write("\t\t- " + base_name + ".json") + + ln_fname = "ligand_network.graphml" + with open(folder_path / ln_fname, mode="w") as f: + f.write(ligand_network.to_graphml()) + write(f"\t\t- {ln_fname}") + + transformations_dir = folder_path / "transformations" + transformations_dir.mkdir(parents=True, exist_ok=True) + + for transformation in alchemical_network.edges: + transformation_name = transformation.name or transformation.key + filename = f"{transformation_name}.json" + transformation.to_json(transformations_dir / filename) + write("\t\t\t\t- " + filename)