diff --git a/docs/environment.yaml b/docs/environment.yaml index 6c221e121..80bc95e54 100644 --- a/docs/environment.yaml +++ b/docs/environment.yaml @@ -33,6 +33,8 @@ dependencies: - git+https://github.com/OpenFreeEnergy/kartograf@main - git+https://github.com/OpenFreeEnergy/konnektor@main - git+https://github.com/OpenFreeEnergy/lomap@main + - git+https://github.com/OpenFreeEnergy/exorcist@main + # These are added automatically by RTD, so we include them here # for a consistent environment. diff --git a/docs/guide/execution/index.rst b/docs/guide/execution/index.rst index 0f5f2cf33..e17b39a46 100644 --- a/docs/guide/execution/index.rst +++ b/docs/guide/execution/index.rst @@ -12,3 +12,4 @@ then :ref:`reading on the available Python functions`. .. toctree:: quickrun_execution execution_theory + warehouse diff --git a/docs/guide/execution/warehouse.rst b/docs/guide/execution/warehouse.rst new file mode 100644 index 000000000..359c26a98 --- /dev/null +++ b/docs/guide/execution/warehouse.rst @@ -0,0 +1,73 @@ +Data Handling with Warehouse +============================== + +**openfe**'s ``Warehouse`` defines the interface for an execution engine to store and access data during execution. + +A Warehouse is any instance of a derived class of the abstract :class:`.WarehouseBaseClass`. In other words, *where* the data is stored is decided by the derived class, but *how* the data is accessed is defined by ``WarehouseBaseClass``. + +You can think of the ``WarehouseBaseClass`` as a set of specifications that must be met by a Warehouse implementation (subclass), such that any openfe Protocol can then interact appropriately with its data. + +For example, **openfe** Protocols require several types of storage - scratch, setup, and result. + +Naively, we could require that all three of these storage types be filesystem directories that can be accessed locally by the Protocol, but this significantly limits the ways in which the Protocol can be executed, e.g. a Protocol could not store **result** data on a remote machine or cloud storage. + +Where a Warehouse stores its data is defined by a :class:`.WarehouseStores` object, which is a small `TypedDict `_ containing ``'setup'`` and ``'result'`` keys (note that ``scratch`` is not a key, *must* be locally-accessible file storage) that correspond to gufe ``ExternalStorage`` objects. +The type of :class:`.ExternalStorage` objects used by a Warehouse implementation are where the the code author has flexibility to choose *where* data is stored. + + +The below example implementation, :class:`.FileSystemWarehouse`, is a derived class that inherits from ``WarehouseBaseClass``. +This is a simple example of how to construct a Warehouse given a root directory, which is uses to create ``WarehouseStores``. + +.. TODO: reference FileStorage gufe docs + +.. code-block:: + + class FileSystemWarehouse(WarehouseBaseClass): + """Warehouse implementation using local filesystem storage. + + Provides a file-based storage backend for GufeTokenizable objects + organized in a directory structure. + + Parameters + ---------- + root_dir : str, optional + Root directory for the warehouse storage, by default "warehouse/". + + Notes + ----- + Creates a "setup/" subdirectory within the root directory for storing + setup-related objects. Future versions may include additional stores + for results and other data types. + """ + + def __init__(self, root_dir: str = "warehouse"): + 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) + + +Using this new ``FileSystemWarehouse``, we can now access the data cleanly and reproducibly in a way that abstracts away the use of a filesystem. + +without Warehouse, using only the filesystem: + +.. code-block:: + + root_dir = os.Path("my_warehouse") + setup = root / "setup" + result = root / "result" + + +with Warehouse: + +.. code-block:: + + from openfe.storage import FileSystemWarehouse + my_warehouse = FileSystemWarehouse(root_dir="my_warehouse") + + ... + + my_warehouse.store_result_tokenizable(result) +.. TODO: add example of dropping in non-filesystem storage once gufe supports it + +.. For more information about what types of storage are available, see . \ No newline at end of file diff --git a/docs/reference/api/defining_and_executing_simulations.rst b/docs/reference/api/defining_and_executing_simulations.rst index 182aabdd4..b09fdccd0 100644 --- a/docs/reference/api/defining_and_executing_simulations.rst +++ b/docs/reference/api/defining_and_executing_simulations.rst @@ -12,10 +12,11 @@ Executing Simulations :noindex: .. autosummary:: - :nosignatures: :toctree: generated/ + :recursive: execute_DAG + storage General classes --------------- @@ -32,10 +33,10 @@ General classes ProtocolUnitFailure ProtocolDAGResult -Specialised classes +Specialized classes ------------------- -These classes are abstract classes that are specialised (subclassed) for an individual Protocol. +These classes are abstract classes that are specialized (subclassed) for an individual Protocol. .. module:: openfe :noindex: diff --git a/environment.yml b/environment.yml index e788eab3c..e6736a22b 100644 --- a/environment.yml +++ b/environment.yml @@ -56,6 +56,7 @@ dependencies: - git+https://github.com/OpenFreeEnergy/kartograf@main # needs gufe which pins to pydantic < 2.13. update this once gufe v1.13 is released on conda-forge - git+https://github.com/OpenFreeEnergy/konnektor@main - git+https://github.com/OpenFreeEnergy/gufe@main + - git+https://github.com/OpenFreeEnergy/exorcist@main - 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/news/warehouse.rst b/news/warehouse.rst new file mode 100644 index 000000000..7cb85dc38 --- /dev/null +++ b/news/warehouse.rst @@ -0,0 +1,23 @@ +**Added:** + +* Added ``openfe.Warehouse``, an interface for storing and accessing data during simulation execution (`PR #1864 `_). + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* Removed the unused methods ``metadatastore``, ``resultclient``, and ``resultserver`` from ``openfe.storage``. + +**Fixed:** + +* + +**Security:** + +* diff --git a/src/openfe/orchestration/__init__.py b/src/openfe/orchestration/__init__.py index e69de29bb..14999e2f2 100644 --- a/src/openfe/orchestration/__init__.py +++ b/src/openfe/orchestration/__init__.py @@ -0,0 +1,244 @@ +"""Task orchestration utilities backed by Exorcist and a warehouse.""" + +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path + +from exorcist.taskdb import TaskStatusDB +from gufe.protocols.protocoldag import _pu_to_pur +from gufe.protocols.protocolunit import ( + Context, + ProtocolUnit, + ProtocolUnitResult, +) +from gufe.storage.externalresource.base import ExternalStorage +from gufe.storage.externalresource.filestorage import FileStorage +from gufe.tokenization import GufeKey + +from openfe.storage.warehouse import FileSystemWarehouse + +from .exorcist_utils import ( + alchemical_network_to_task_graph, + build_task_db_from_alchemical_network, +) + + +@dataclass +class Worker: + """Execute protocol units from an Exorcist task database. + + Parameters + ---------- + warehouse : FileSystemWarehouse + Warehouse used to load queued tasks and store execution results. + task_db_path : pathlib.Path, default=Path("./warehouse/tasks.db") + Path to the Exorcist SQLite task database. + """ + + warehouse: FileSystemWarehouse + task_db_path: Path = Path("./warehouse/tasks.db") + + _RESULT_INDEX_PREFIX = "protocol_unit_results" + _TASK_WORKDIR_PREFIX = "task_workdirs" + + @staticmethod + def _collect_protocol_unit_keys(value: object) -> set[GufeKey]: + """Collect `ProtocolUnit` keys from nested unit inputs.""" + + if isinstance(value, ProtocolUnit): + return {value.key} + + found: set[GufeKey] = set() + items: Iterable # TODO: update this to dict_values | list after python 3.13 min? + if isinstance(value, dict): + items = value.values() + elif isinstance(value, list): + items = value + else: + return found + + for item in items: + found.update(Worker._collect_protocol_unit_keys(item)) + return found + + @classmethod + def _result_index_location(cls, source_key: GufeKey) -> str: + return f"{cls._RESULT_INDEX_PREFIX}/{source_key}" + + @classmethod + def _task_workdir_name(cls, taskid: str) -> str: + return taskid.replace(":", "__") + + def _task_workspace_paths( + self, taskid: str, scratch_root: Path, shared_root: Path + ) -> tuple[Path, Path]: + workdir_name = self._task_workdir_name(taskid) + task_scratch = scratch_root / self._TASK_WORKDIR_PREFIX / workdir_name + task_shared = shared_root / self._TASK_WORKDIR_PREFIX / workdir_name + return task_scratch, task_shared + + def _store_result_index(self, result: ProtocolUnitResult) -> None: + shared_store: ExternalStorage = self.warehouse.stores["shared"] + location = self._result_index_location(result.source_key) + shared_store.store_bytes(location, str(result.key).encode("utf-8")) + + def _load_result_from_index(self, source_key: GufeKey) -> ProtocolUnitResult | None: + shared_store: ExternalStorage = self.warehouse.stores["shared"] + location = self._result_index_location(source_key) + + if not shared_store.exists(location): + return None + + with shared_store.load_stream(location) as stream: + result_key = stream.read().decode("utf-8").strip() + + loaded = self.warehouse.load_result_tokenizable(GufeKey(result_key)) + if isinstance(loaded, ProtocolUnitResult): + return loaded + + return None + + def _scan_result_store_for_sources( + self, source_keys: set[GufeKey] + ) -> dict[GufeKey, ProtocolUnitResult]: + found: dict[GufeKey, ProtocolUnitResult] = {} + + for location in self.warehouse.result_store.iter_contents(): + if len(found) == len(source_keys): + break + + loaded = self.warehouse.load_result_tokenizable(GufeKey(location)) + if not isinstance(loaded, ProtocolUnitResult): + continue + + source_key = loaded.source_key + if source_key in source_keys and source_key not in found: + found[source_key] = loaded + + return found + + def _build_input_result_mapping(self, unit: ProtocolUnit) -> dict[GufeKey, ProtocolUnitResult]: + required_keys = self._collect_protocol_unit_keys(unit.inputs) + if not required_keys: + return {} + + results: dict[GufeKey, ProtocolUnitResult] = {} + unresolved = set(required_keys) + + for source_key in required_keys: + loaded = self._load_result_from_index(source_key) + if loaded is not None: + results[source_key] = loaded + unresolved.discard(source_key) + + if unresolved: + scanned = self._scan_result_store_for_sources(unresolved) + for source_key, loaded in scanned.items(): + results[source_key] = loaded + self._store_result_index(loaded) + unresolved.discard(source_key) + + if unresolved: + missing_keys = ", ".join(sorted(str(k) for k in unresolved)) + raise RuntimeError( + "Missing ProtocolUnitResult(s) for dependency key(s): " + f"{missing_keys}. Ensure upstream tasks completed successfully." + ) + + return results + + def _checkout_task(self) -> tuple[TaskStatusDB, str, ProtocolUnit] | None: + """Check out one available task and load its protocol unit. + + Returns + ------- + tuple[TaskStatusDB, str, ProtocolUnit] or None + The open database connection, checked-out task ID, and corresponding + protocol unit, or ``None`` if no task is currently available. + The caller is responsible for calling ``mark_task_completed`` on the + returned database using the returned task ID. + """ + + db: TaskStatusDB = TaskStatusDB.from_filename(self.task_db_path) + # The format for the taskid is "ProtocolUnit-" + taskid = db.check_out_task() + if taskid is None: + return None + + protocol_unit_key = taskid + unit = self.warehouse.load_task(GufeKey(protocol_unit_key)) + return db, taskid, unit + + def _get_task(self) -> tuple[str, ProtocolUnit]: + """Return the next available task ID and protocol unit. + + Returns + ------- + tuple[str, ProtocolUnit] + The checked-out task ID and corresponding protocol unit. + + Raises + ------ + RuntimeError + Raised when no task is available in the task database. + """ + + task = self._checkout_task() + if task is None: + raise RuntimeError("No AVAILABLE tasks found in the task database.") + db, taskid, unit = task + return taskid, unit + + def execute_unit(self, scratch: Path) -> tuple[str, ProtocolUnitResult] | None: + """Execute one checked-out protocol unit and persist its result. + + Parameters + ---------- + scratch : pathlib.Path + Scratch directory passed to the protocol execution context. + + Returns + ------- + tuple[str, ProtocolUnitResult] or None + The task ID and execution result for the processed task, or + ``None`` if no task is currently available. + + Raises + ------ + Exception + Re-raises any exception thrown during protocol unit execution after + marking the task as failed. + """ + + # 1. Get task/unit + task = self._checkout_task() + if task is None: + return None + db, taskid, unit = task + # 2. Construct the context + # NOTE: On changes to context (gufe PR #753), this can easily be replaced with external storage objects + # However, to satisfy the current work, we will use this implementation where we + # force the use of a FileSystemWarehouse and in turn can assert that an object is FileStorage. + shared_store = self.warehouse.stores["shared"] + if not isinstance(shared_store, FileStorage): + raise TypeError("Expected a FileStorage backend for the shared store") + shared_root_dir = shared_store.root_dir + task_scratch, task_shared = self._task_workspace_paths(taskid, scratch, shared_root_dir) + task_scratch.mkdir(parents=True, exist_ok=True) + task_shared.mkdir(parents=True, exist_ok=True) + ctx = Context(task_scratch, shared=task_shared) + # 3. Execute unit + try: + results = self._build_input_result_mapping(unit) + inputs = _pu_to_pur(unit.inputs, results) + result = unit.execute(context=ctx, **inputs) + except Exception: + db.mark_task_completed(taskid, success=False) + raise + + db.mark_task_completed(taskid, success=result.ok()) + # 4. output result to warehouse + # TODO: we may need to end up handling namespacing on the warehouse side for tokenizables + self.warehouse.store_result_tokenizable(result) + self._store_result_index(result) + return taskid, result diff --git a/src/openfe/orchestration/exorcist_utils.py b/src/openfe/orchestration/exorcist_utils.py new file mode 100644 index 000000000..459877b9c --- /dev/null +++ b/src/openfe/orchestration/exorcist_utils.py @@ -0,0 +1,98 @@ +"""Utilities for building Exorcist task graphs and task databases. + +This module translates an :class:`gufe.AlchemicalNetwork` into Exorcist task +structures and can initialize an Exorcist task database from that graph. +""" + +from pathlib import Path + +import exorcist +import networkx as nx +import pandas as pd +from gufe import AlchemicalNetwork + +from openfe.storage.warehouse import WarehouseBaseClass + + +def alchemical_network_to_task_graph( + alchemical_network: AlchemicalNetwork, warehouse: WarehouseBaseClass +) -> nx.DiGraph: + """Build a global task DAG from an alchemical network. + + Parameters + ---------- + alchemical_network : AlchemicalNetwork + Network containing transformations to execute. + warehouse : WarehouseBaseClass + Warehouse used to persist protocol units as tasks while the graph is + constructed. + + Returns + ------- + nx.DiGraph + A directed acyclic graph where each node is a task ID in the form + ``":"`` and edges encode + protocol-unit dependencies. + + Raises + ------ + ValueError + Raised if the assembled task graph is not acyclic. + """ + + global_dag = nx.DiGraph() + for transformation in alchemical_network.edges: + dag = transformation.create() + for unit in dag.protocol_units: + node_id = str(unit.key) + global_dag.add_node(node_id) + warehouse.store_task(unit) + # store the protocol_dag as a shallow dict, since all its units are + # already written to disk + warehouse.store_protocol_dag(dag) + for dependent_unit, dependency_unit in dag.graph.edges: + upstream_id = str(dependency_unit.key) + downstream_id = str(dependent_unit.key) + global_dag.add_edge(upstream_id, downstream_id) + + if not nx.is_directed_acyclic_graph(global_dag): + raise ValueError("AlchemicalNetwork produced a task graph that is not a DAG.") + + return global_dag + + +# TODO: do we test adding a multiple alchemical networks to the same task graph? +def build_task_db_from_alchemical_network( + alchemical_network: AlchemicalNetwork, + warehouse: WarehouseBaseClass, + db_path: Path | None = None, + max_tries: int = 1, +) -> exorcist.TaskStatusDB: + """Create and populate a task database from an alchemical network. + + Parameters + ---------- + alchemical_network : AlchemicalNetwork + Network containing transformations to convert into task records. + warehouse : WarehouseBaseClass + Warehouse used to persist protocol units while building the task DAG. + db_path : pathlib.Path or None, optional + Location of the SQLite-backed Exorcist database. If ``None``, defaults + to {warehouse.name}.db in the current working directory. + max_tries : int, default=1 + Maximum number of retries for each task before Exorcist marks it as + ``TOO_MANY_RETRIES``. + + Returns + ------- + exorcist.TaskStatusDB + Initialized task database populated with graph nodes and dependency + edges derived from ``alchemical_network``. + """ + if db_path is None: + db_path = Path(f"{warehouse.name}.db") + + global_dag: nx.DiGraph = alchemical_network_to_task_graph(alchemical_network, warehouse) + db = exorcist.TaskStatusDB.from_filename(db_path) + db.add_task_network(global_dag, max_tries) + return db diff --git a/src/openfe/storage/metadatastore.py b/src/openfe/storage/metadatastore.py deleted file mode 100644 index 6c2f29e7e..000000000 --- a/src/openfe/storage/metadatastore.py +++ /dev/null @@ -1,100 +0,0 @@ -# This code is part of OpenFE and is licensed under the MIT license. -# For details, see https://github.com/OpenFreeEnergy/gufe -import abc -import collections -import json -from typing import Dict, Tuple - -from gufe.storage.errors import ChangedExternalResourceError, MissingExternalResourceError -from gufe.storage.externalresource.base import Metadata - - -class MetadataStore(collections.abc.Mapping): - def __init__(self, external_store): - self.external_store = external_store - self._metadata_cache = self.load_all_metadata() - - @abc.abstractmethod - def store_metadata(self, location: str, metadata: Metadata): - raise NotImplementedError() - - @abc.abstractmethod - def load_all_metadata(self) -> Dict[str, Metadata]: - raise NotImplementedError() - - @abc.abstractmethod - def __delitem__(self, location): - raise NotImplementedError() - - def __getitem__(self, location): - return self._metadata_cache[location] - - def __iter__(self): - return iter(self._metadata_cache) - - def __len__(self): - return len(self._metadata_cache) - - -class JSONMetadataStore(MetadataStore): - # Using JSON for now because it is easy to write this class and doesn't - # require any external dependencies. It is NOT the right way to go in - # the long term. API will probably stay the same, though. - def _dump_file(self): - metadata_dict = {key: val.to_dict() for key, val in self._metadata_cache.items()} - metadata_bytes = json.dumps(metadata_dict).encode("utf-8") - self.external_store.store_bytes("metadata.json", metadata_bytes) - - def store_metadata(self, location: str, metadata: Metadata): - self._metadata_cache[location] = metadata - self._dump_file() - - def load_all_metadata(self): - if not self.external_store.exists("metadata.json"): - return {} - - with self.external_store.load_stream("metadata.json") as json_f: - all_metadata_dict = json.loads(json_f.read().decode("utf-8")) - - all_metadata = {key: Metadata(**val) for key, val in all_metadata_dict.items()} - - return all_metadata - - def __delitem__(self, location): - del self._metadata_cache[location] - self._dump_file() - - -class PerFileJSONMetadataStore(MetadataStore): - _metadata_prefix = "metadata/" - - def _metadata_path(self, location): - return self._metadata_prefix + location + ".json" - - def store_metadata(self, location: str, metadata: Metadata): - self._metadata_cache[location] = metadata - path = self._metadata_path(location) - dct = { - "path": location, - "metadata": metadata.to_dict(), - } - metadata_bytes = json.dumps(dct).encode("utf-8") - self.external_store.store_bytes(path, metadata_bytes) - - def load_all_metadata(self): - metadata_cache = {} - prefix = self._metadata_prefix - for location in self.external_store.iter_contents(prefix=prefix): - if location.endswith(".json"): - with self.external_store.load_stream(location) as f: - dct = json.loads(f.read().decode("utf-8")) - - if set(dct) != {"path", "metadata"}: - raise ChangedExternalResourceError(f"Bad metadata file: '{location}'") - metadata_cache[dct["path"]] = Metadata(**dct["metadata"]) - - return metadata_cache - - def __delitem__(self, location): - del self._metadata_cache[location] - self.external_store.delete(self._metadata_path(location)) diff --git a/src/openfe/storage/resultclient.py b/src/openfe/storage/resultclient.py deleted file mode 100644 index ff3912997..000000000 --- a/src/openfe/storage/resultclient.py +++ /dev/null @@ -1,287 +0,0 @@ -# This code is part of OpenFE and is licensed under the MIT license. -# For details, see https://github.com/OpenFreeEnergy/gufe -import abc -import json -import re -from typing import Any - -from gufe.tokenization import ( - JSON_HANDLER, - from_dict, - get_all_gufe_objs, - key_decode_dependencies, -) - -from .metadatastore import JSONMetadataStore -from .resultserver import ResultServer - -GUFEKEY_JSON_REGEX = re.compile('":gufe-key:": "(?P[A-Za-z0-9_]+-[0-9a-f]+)"') - - -class _ResultContainer(abc.ABC): - """ - Abstract class, represents all data under some level of the hierarchy. - """ - - def __init__(self, parent, path_component): - self.parent = parent - self._path_component = self._to_path_component(path_component) - self._cache = {} - - def __eq__(self, other): - return isinstance(other, self.__class__) and self.path == other.path - - @staticmethod - def _to_path_component(item: Any) -> str: - """Convert input (object or string) to path string""" - if isinstance(item, str): - return item - - # TODO: instead of str(hash(...)), this should return the digest - # that is being introduced in another PR; Python hash is not stable - # across sessions - return str(hash(item)) - - def __getitem__(self, item): - # code for the case this is a file - if item in self.result_server: - return self.result_server.load_stream(item) - - # code for the case this is a "directory" - hash_item = self._to_path_component(item) - - if hash_item not in self._cache: - self._cache[hash_item] = self._load_next_level(item) - - return self._cache[hash_item] - - def __truediv__(self, item): - return self[item] - - @abc.abstractmethod - def _load_next_level(self, item): - raise NotImplementedError() - - def __iter__(self): - for loc in self.result_server: - if loc.startswith(self.path): - yield loc - - def load_stream(self, location, *, allow_changed=False): - return self.result_server.load_stream(location, allow_changed) - - def load_bytes(self, location, *, allow_changed=False): - with self.load_stream(location, allow_changed=allow_changed) as f: - byte_data = f.read() - - return byte_data - - @property - def path(self): - return self.parent.path + "/" + self._path_component - - @property - def result_server(self): - return self.parent.result_server - - def __repr__(self): - # probably should include repr of external store, too - return f"{self.__class__.__name__}({self.path})" - - -class ResultClient(_ResultContainer): - def __init__(self, external_store): - # default client is using JSONMetadataStore with the given external - # result store; users could easily write a subblass that behaves - # differently - metadata_store = JSONMetadataStore(external_store) - self._result_server = ResultServer(external_store, metadata_store) - super().__init__(parent=self, path_component=None) - - def delete(self, location): - self._result_server.delete(location) - - @staticmethod - def _gufe_key_to_storage_key(prefix: str, key: str): - """Create the storage key from the gufe key. - - Parameters - ---------- - prefix : str - the prefix defining which section of storage should be used for - this (e.g., ``setup``, ...) - key : str - the GufeKey for a GufeTokenizable (technically, is likely to be - passed as a :class:`.GufeKey`, which is a subclass of ``str``) - - Returns - ------- - str : - storage key (string identifier used by storage to locate this - object) - """ - pref = prefix.split("/") # remove this if we switch to tuples - cls, token = key.split("-") - tup = tuple(list(pref) + [cls, f"{token}.json"]) - # right now we're using strings, but we've talked about switching - # that to tuples - return "/".join(tup) - - def _store_gufe_tokenizable(self, prefix, obj): - """generic function for deduplicating/storing a GufeTokenizable""" - for o in get_all_gufe_objs(obj): - key = self._gufe_key_to_storage_key(prefix, o.key) - - # we trust that if we get the same key, it's the same object, so - # we only store on keys that we don't already know - if key not in self.result_server: - data = json.dumps( - o.to_keyed_dict(), cls=JSON_HANDLER.encoder, sort_keys=True - ).encode("utf-8") - self.result_server.store_bytes(key, data) - - def store_transformation(self, transformation): - """Store a :class:`.Transformation`. - - Parameters - --------- - transformation: :class:`.Transformation` - the transformation to store - """ - self._store_gufe_tokenizable("setup", transformation) - - def store_network(self, network): - """Store a :class:`.AlchemicalNetwork`. - - Parameters - --------- - network: :class:`.AlchemicalNetwork` - the network to store - """ - self._store_gufe_tokenizable("setup", network) - - def _load_gufe_tokenizable(self, prefix, gufe_key): - """generic function to load deduplicated object from a key""" - registry = {} - - def recursive_build_object_cache(gufe_key): - """DFS to rebuild object hierarchy""" - # This implementation is a bit fragile, because ensuring that we - # don't duplicate objects in memory depends on the fact that - # `key_decode_dependencies` gets keyencoded objects from a cache - # (they are cached on creation). - storage_key = self._gufe_key_to_storage_key(prefix, gufe_key) - with self.load_stream(storage_key) as f: - keyencoded_json = f.read().decode("utf-8") - - dct = json.loads(keyencoded_json, cls=JSON_HANDLER.decoder) - # this implementation may seem strange, but it will be a - # faster than traversing the dict - key_encoded = set(GUFEKEY_JSON_REGEX.findall(keyencoded_json)) - - # this approach takes the dct instead of the json str - # found = [] - # modify_dependencies(dct, found.append, is_gufe_key_dict) - # key_encoded = {d[":gufe-key:"] for d in found} - - for key in key_encoded: - # we're actually only doing this for the side effect of - # generating the objects and adding them to the registry - recursive_build_object_cache(key) - - if len(key_encoded) == 0: - # fast path for objects that don't contain other gufe - # objects (these tend to be larger dicts; avoid walking - # them) - obj = from_dict(dct) - else: - # objects that contain other gufe objects need be walked to - # replace everything - obj = key_decode_dependencies(dct, registry) - - registry[obj.key] = obj - return obj - - return recursive_build_object_cache(gufe_key) - - def load_transformation(self, key: str): - """Load a :class:`.Transformation` from its GufeKey - - Parameters - ---------- - key: str - the gufe key for this object - - Returns - ------- - :class:`.Transformation` - the desired transformation - """ - return self._load_gufe_tokenizable("setup", key) - - def load_network(self, key: str): - """Load a :class:`.AlchemicalNetwork` from its GufeKey - - Parameters - ---------- - key: str - the gufe key for this object - - Returns - ------- - :class:`.AlchemicalNetwork` - the desired network - """ - return self._load_gufe_tokenizable("setup", key) - - def _load_next_level(self, transformation): - return TransformationResult(self, transformation) - - # override these two inherited properties since this is always the end of - # the recursive chain - @property - def path(self): - return "transformations" - - @property - def result_server(self): - return self._result_server - - -class TransformationResult(_ResultContainer): - def __init__(self, parent, transformation): - super().__init__(parent, transformation) - self.transformation = transformation - - def _load_next_level(self, clone): - return CloneResult(self, clone) - - -class CloneResult(_ResultContainer): - def __init__(self, parent, clone): - super().__init__(parent, clone) - self.clone = clone - - @staticmethod - def _to_path_component(item): - return str(item) - - def _load_next_level(self, extension): - return ExtensionResult(self, extension) - - -class ExtensionResult(_ResultContainer): - def __init__(self, parent, extension): - super().__init__(parent, str(extension)) - self.extension = extension - - @staticmethod - def _to_path_component(item): - return str(item) - - def __getitem__(self, filename): - # different here -- we don't cache the actual file objects - return self._load_next_level(filename) - - def _load_next_level(self, filename): - return self.result_server.load_stream(self.path + "/" + filename) diff --git a/src/openfe/storage/resultserver.py b/src/openfe/storage/resultserver.py deleted file mode 100644 index 73c0f6a6c..000000000 --- a/src/openfe/storage/resultserver.py +++ /dev/null @@ -1,60 +0,0 @@ -# This code is part of OpenFE and is licensed under the MIT license. -# For details, see https://github.com/OpenFreeEnergy/gufe -import warnings -from typing import ClassVar - -from gufe.storage.errors import ChangedExternalResourceError, MissingExternalResourceError - - -class ResultServer: - """Class to manage communication between metadata and data storage. - - At this level, we provide an abstraction where client code no longer - needs to be aware of the nature of the metadata, or even that it exists. - """ - - def __init__(self, external_store, metadata_store): - self.external_store = external_store - self.metadata_store = metadata_store - - def _store_metadata(self, location): - metadata = self.external_store.get_metadata(location) - self.metadata_store.store_metadata(location, metadata) - - def store_bytes(self, location, byte_data): - self.external_store.store_bytes(location, byte_data) - self._store_metadata(location) - - def store_path(self, location, path): - self.external_store.store_path(location, path) - self._store_metadata(location) - - def delete(self, location): - del self.metadata_store[location] - self.external_store.delete(location) - - def validate(self, location, allow_changed=False): - try: - metadata = self.metadata_store[location] - except KeyError: - raise MissingExternalResourceError(f"Metadata for '{location}' not found") - - if not self.external_store.get_metadata(location) == metadata: - msg = f"Metadata mismatch for {location}: this object may have changed." - if not allow_changed: - raise ChangedExternalResourceError( - msg + " To allow this, set ExternalStorage.allow_changed = True" - ) - else: - warnings.warn(msg) - - def __iter__(self): - return iter(self.metadata_store) - - def find_missing_files(self): - """Identify files listed in metadata but unavailable in storage""" - return [f for f in self if not self.external_store.exists(f)] - - def load_stream(self, location, allow_changed=False): - self.validate(location, allow_changed) - return self.external_store.load_stream(location) diff --git a/src/openfe/storage/warehouse.py b/src/openfe/storage/warehouse.py new file mode 100644 index 000000000..9844e0e0d --- /dev/null +++ b/src/openfe/storage/warehouse.py @@ -0,0 +1,427 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/gufe +import json +import pathlib +import re +from typing import Generator, Literal, TypedDict + +from gufe.protocols.protocoldag import ProtocolDAG +from gufe.protocols.protocolunit import ProtocolUnit +from gufe.storage.externalresource import ExternalStorage, FileStorage +from gufe.tokenization import ( + JSON_HANDLER, + GufeKey, + GufeTokenizable, + from_dict, + get_all_gufe_objs, + key_decode_dependencies, +) + +GUFEKEY_JSON_REGEX = re.compile('":gufe-key:": "(?P[A-Za-z0-9_]+-[0-9a-f]+)"') + + +class WarehouseStores(TypedDict): + """Typed dictionary for accessing warehouse storage locations. + + Parameters + ---------- + setup : Required[ExternalStorage] + Storage location for setup-related objects and configurations. + result : Required[ExternalStorage] + Storage location for result-related object. + shared : ExternalStorage + Storage location for non-permanent shared data. + tasks: ExternalStorage + Storage location for execution tasks. + protocol_dags: ExternalStorage + Storage location for ProtocolDAGs that correspond to the ProtocolUnits stored in 'tasks'. + + Notes + ----- + Additional stores for results and tasks may be added in future versions. + """ + + setup: ExternalStorage + result: ExternalStorage + shared: ExternalStorage + tasks: ExternalStorage + protocol_dags: ExternalStorage + + +class WarehouseBaseClass: + """Base class for warehouse storage management. + + Provides functionality to store, load, and manage GufeTokenizable objects + across different storage backends. + + Parameters + ---------- + stores : WarehouseStores + Typed dictionary containing the storage locations for different + types of objects. + + Attributes + ---------- + stores : WarehouseStores + The storage locations managed by this warehouse instance. + """ + + def __init__(self, stores: WarehouseStores, name: str): + self.stores = stores + if not isinstance(name, str) or len(name) == 0: + raise ValueError("Warehouse name must be a string.") + self.name = name + + def __eq__(self, other): + return isinstance(other, self.__class__) and self.stores == other.stores + + def __repr__(self): + # probably should include repr of external store, too + return f"{self.__class__.__name__}({self.stores})" + + def delete(self, store_name: Literal["setup", "result"], location: str): + """Delete an object from a specific store. + + Parameters + ---------- + store_name : Literal["setup"] + Name of the store to delete from. + location : str + Location/path of the object to delete. + + Raises + ------- + MissingExternalResourceError + If the object cannot be deleted from the store. + """ + store: ExternalStorage = self.stores[store_name] + store.delete(location) + + def store_task(self, obj: ProtocolUnit): + self._store_gufe_tokenizable("tasks", obj) + + def load_task(self, obj: GufeKey) -> ProtocolUnit: + unit = self._load_gufe_tokenizable(obj) + if not isinstance(unit, ProtocolUnit): + raise ValueError("Unable to load ProtocolUnit") + return unit + + def store_setup_tokenizable(self, obj: GufeTokenizable): + """Store a GufeTokenizable object in the setup store. + + Parameters + ---------- + obj : GufeTokenizable + The object to store. + """ + self._store_gufe_tokenizable("setup", obj) + + def load_setup_tokenizable(self, obj: GufeKey) -> GufeTokenizable: + # TODO: this doesn't actually look specifically in the setup store, which is misleading + """Load a GufeTokenizable object from the setup store. + + Parameters + ---------- + obj : GufeKey + The key of the object to load. + + Returns + ------- + GufeTokenizable + The loaded object. + """ + return self._load_gufe_tokenizable(gufe_key=obj) + + def store_result_tokenizable(self, obj: GufeTokenizable): + """Store a GufeTokenizable object from the result store. + + Parameters + ---------- + obj : GufeKey + The key of the object to store. + """ + return self._store_gufe_tokenizable("result", obj) + + def load_result_tokenizable(self, obj: GufeKey) -> GufeTokenizable: + # TODO: this doesn't actually look specifically in the result store, which is misleading + """Load a GufeTokenizable object from the result store. + + Parameters + ---------- + obj : GufeKey + The key of the object to load. + + Returns + ------- + GufeTokenizable + The loaded object. + """ + return self._load_gufe_tokenizable(gufe_key=obj) + + def store_protocol_dag(self, dag: ProtocolDAG): + """Store a ProtocolDAG in the "protocol_dags" store of this warehouse. + Parameters + ---------- + dag : ProtocolDAG + The ProtocolDAG object to store. + + Raises + ------ + ValueError + If `dag` is not a ProtocolDAG instance. + """ + if not isinstance(dag, ProtocolDAG): + raise ValueError("Only ProtocolDAGs may be written to the 'protocol_dags' store.") + self._store_gufe_tokenizable("protocol_dags", dag) + + def load_protocol_dag(self, gufe_key=GufeKey) -> GufeTokenizable: + """Load a GufeTokenizable object from the protocol_dag store. + + Parameters + ---------- + obj : GufeKey + The key of the protocoldag to load. + + Returns + ------- + GufeTokenizable + The loaded object. + """ + # TODO: type check that it is a protocol dag before returning? + return self._load_gufe_tokenizable(gufe_key=gufe_key) + + def exists(self, key: GufeKey) -> bool: + """Check if an object with the given key exists in any store that holds tokenizables. + + Parameters + ---------- + key : GufeKey + The key to check for existence. + + Returns + ------- + bool + True if the object exists, False otherwise. + """ + # TODO: resolve type checking + return any(key in store for store in self.stores.values()) # type: ignore + + def _get_store_for_key(self, key: GufeKey) -> ExternalStorage: + """Function to find the store in which a gufe key is stored in. + + Parameters + ---------- + key : GufeKey + The key to locate. + + Returns + ------- + ExternalStorage + The store containing the key. + + Raises + ------ + ValueError + If the key is not found in any store. + """ + # TODO: resolve mypy Literal/str conflict here + # https://mypy.readthedocs.io/en/stable/literal_types.html + for name in self.stores: + if key in self.stores[name]: # type: ignore + return self.stores[name] # type: ignore + raise ValueError(f"GufeKey {key} is not stored") + + def _store_gufe_tokenizable( + self, + store_name: Literal["setup", "result", "tasks", "protocol_dags"], + obj: GufeTokenizable, + name: str | None = None, + ): + """Store a GufeTokenizable object with deduplication. + + Parameters + ---------- + store_name : Literal["setup"] + Name of the store to store the object in. + obj : GufeTokenizable + The object to store. + + Notes + ----- + This function performs deduplication by checking if the object + already exists in any store before storing. + """ + # Try and get the key for the given store + target: ExternalStorage = self.stores[store_name] + # Get all of the sub-objects + chain = obj.to_keyed_chain() + for item in chain: + gufe_key = GufeKey(item[0]) + keyed_dict = item[1] + if not self.exists(gufe_key): + data = json.dumps(keyed_dict, cls=JSON_HANDLER.encoder, sort_keys=True).encode( + "utf-8" + ) + if name: + target.store_bytes(name, data) + else: + target.store_bytes(gufe_key, data) + + def _load_gufe_tokenizable(self, gufe_key: GufeKey) -> GufeTokenizable: + """Load a deduplicated object from a GufeKey. + + Parameters + ---------- + gufe_key : GufeKey + The key of the object to load. + + Returns + ------- + GufeTokenizable + The loaded object with all dependencies resolved. + + Notes + ----- + Uses depth-first search to rebuild object hierarchy and ensure + proper deduplication in memory. + """ + registry: dict[GufeKey, GufeTokenizable] = {} + + def recursive_build_object_cache(key: GufeKey) -> GufeTokenizable: + """DFS to rebuild object hierarchy. + + Parameters + ---------- + key : GufeKey + The key of the object to build. + + Returns + ------- + GufeTokenizable + The reconstructed object. + """ + # This implementation is a bit fragile, because ensuring that we + # don't duplicate objects in memory depends on the fact that + # `key_decode_dependencies` gets keyencoded objects from a cache + # (they are cached on creation). + store = self._get_store_for_key(key=key) + + with store.load_stream(key) as f: + keyencoded_json = f.read().decode("utf-8") + + dct = json.loads(keyencoded_json, cls=JSON_HANDLER.decoder) + # this implementation may seem strange, but it will be a + # faster than traversing the dict + key_encoded = set(GUFEKEY_JSON_REGEX.findall(keyencoded_json)) + + # this approach takes the dct instead of the json str + # found = [] + # modify_dependencies(dct, found.append, is_gufe_key_dict) + # key_encoded = {d[":gufe-key:"] for d in found} + + for key in key_encoded: + # obj = GufeTokenizable.from_dict(dct) + recursive_build_object_cache(key) + # obj = GufeTokenizable.from_json(content=keyencoded_json) + + if len(key_encoded) == 0: + # fast path for objects that don't contain other gufe + # objects (these tend to be larger dicts; avoid walking + # them) + obj = GufeTokenizable.from_dict(dct) + # objects that contain other gufe objects need be walked to + # replace everything + else: + obj = key_decode_dependencies(dct, registry) + # + registry[obj.key] = obj + return obj + + return recursive_build_object_cache(gufe_key) + + def get_protocol_dags(self) -> Generator[ProtocolDAG, None, None]: + """Yield the protocol dags present in the Warehouse's 'protocol_dags' store. + + Note that this requires the name of the item to start with 'ProtocolDAG'. + + Yields + ------ + Generator[ProtocolDAG] + The ProtocolDAGs found in this Warehouse's 'protocol_dags' store. + """ + # NOTE: this can be made more robust (but slower) by using isinstance(obj, openfe.ProtocolDAG) + # _after_ loading each item, rather than filtering by name + for item in self.stores["protocol_dags"]: + if item.startswith("ProtocolDAG"): + dag = self.load_protocol_dag(item) + yield dag + + @property + def setup_store(self) -> ExternalStorage: + """Get the setup store. + + Returns + ------- + ExternalStorage + The setup storage location. + """ + return self.stores["setup"] + + @property + def result_store(self) -> ExternalStorage: + """Get the result store. + + Returns + ------- + ExternalStorage + The result storage location. + """ + return self.stores["result"] + + @property + def shared_store(self): + """Get the shared store. + + Returns + ------- + ExternalStorage + The shared storage location + """ + return self.stores["shared"] + + +class FileSystemWarehouse(WarehouseBaseClass): + """Warehouse implementation using local filesystem storage. + + Provides a file-based storage backend for GufeTokenizable objects + organized in a directory structure. + + Parameters + ---------- + root_dir : str, optional + Root directory for the warehouse storage, by default "warehouse/". + + Notes + ----- + Creates a "setup/" subdirectory within the root directory for storing + setup-related objects. Future versions may include additional stores + for results and other data types. + """ + + def __init__(self, name): + # TODO: should name and location be different? + self.root_dir = pathlib.Path(f"{name}") + setup_store = FileStorage(f"{self.root_dir}/setup") + result_store = FileStorage(f"{self.root_dir}/result") + shared_store = FileStorage(f"{self.root_dir}/shared") + tasks_store = FileStorage(f"{self.root_dir}/tasks") + # TODO: we can store dags in setup if we have a performant way of accessing them + protocol_dag_store = FileStorage(f"{self.root_dir}/protocol_dags") + stores = WarehouseStores( + setup=setup_store, + result=result_store, + shared=shared_store, + tasks=tasks_store, + protocol_dags=protocol_dag_store, + ) + super().__init__(stores, name) diff --git a/src/openfe/tests/orchestration/__init__.py b/src/openfe/tests/orchestration/__init__.py new file mode 100644 index 000000000..efae32ddb --- /dev/null +++ b/src/openfe/tests/orchestration/__init__.py @@ -0,0 +1,2 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe diff --git a/src/openfe/tests/orchestration/conftest.py b/src/openfe/tests/orchestration/conftest.py new file mode 100644 index 000000000..1851b7c05 --- /dev/null +++ b/src/openfe/tests/orchestration/conftest.py @@ -0,0 +1,118 @@ +import gufe +import pytest +from gufe import ChemicalSystem, SolventComponent +from gufe.tests.test_protocol import DummyProtocol +from openff.units import unit + + +@pytest.fixture +def solv_comp(): + yield SolventComponent(positive_ion="K", negative_ion="Cl", ion_concentration=0.0 * unit.molar) + + +@pytest.fixture +def solvated_complex(T4_protein_component, benzene_transforms, solv_comp): + return ChemicalSystem( + { + "ligand": benzene_transforms["toluene"], + "protein": T4_protein_component, + "solvent": solv_comp, + } + ) + + +@pytest.fixture +def solvated_ligand(benzene_transforms, solv_comp): + return ChemicalSystem( + { + "ligand": benzene_transforms["toluene"], + "solvent": solv_comp, + } + ) + + +@pytest.fixture +def absolute_transformation(solvated_ligand, solvated_complex): + return gufe.Transformation( + solvated_ligand, + solvated_complex, + protocol=DummyProtocol(settings=DummyProtocol.default_settings()), + mapping=None, + ) + + +@pytest.fixture +def complex_equilibrium(solvated_complex): + return gufe.NonTransformation( + solvated_complex, + protocol=DummyProtocol(settings=DummyProtocol.default_settings()), + ) + + +@pytest.fixture +def benzene_variants_star_map(benzene_transforms, solv_comp, T4_protein_component): + variants = ["toluene", "phenol", "benzonitrile", "anisole", "benzaldehyde", "styrene"] + + # define the solvent chemical systems and transformations between + # benzene and the others + solvated_ligands = {} + solvated_ligand_transformations = {} + + solvated_ligands["benzene"] = ChemicalSystem( + { + "solvent": solv_comp, + "ligand": benzene_transforms["benzene"], + }, + name="benzene-solvent", + ) + + for ligand in variants: + solvated_ligands[ligand] = ChemicalSystem( + { + "solvent": solv_comp, + "ligand": benzene_transforms[ligand], + }, + name=f"{ligand}-solvent", + ) + + solvated_ligand_transformations[("benzene", ligand)] = gufe.Transformation( + solvated_ligands["benzene"], + solvated_ligands[ligand], + protocol=DummyProtocol(settings=DummyProtocol.default_settings()), + mapping=None, + ) + + # define the complex chemical systems and transformations between + # benzene and the others + solvated_complexes = {} + solvated_complex_transformations = {} + + solvated_complexes["benzene"] = gufe.ChemicalSystem( + { + "protein": T4_protein_component, + "solvent": solv_comp, + "ligand": benzene_transforms["benzene"], + }, + name="benzene-complex", + ) + + for ligand in variants: + solvated_complexes[ligand] = gufe.ChemicalSystem( + { + "protein": T4_protein_component, + "solvent": solv_comp, + "ligand": benzene_transforms[ligand], + }, + name=f"{ligand}-complex", + ) + solvated_complex_transformations[("benzene", ligand)] = gufe.Transformation( + solvated_complexes["benzene"], + solvated_complexes[ligand], + protocol=DummyProtocol(settings=DummyProtocol.default_settings()), + mapping=None, + ) + + return gufe.AlchemicalNetwork( + list(solvated_ligand_transformations.values()) + + list(solvated_complex_transformations.values()) + ) diff --git a/src/openfe/tests/orchestration/test_exorcist_utils.py b/src/openfe/tests/orchestration/test_exorcist_utils.py new file mode 100644 index 000000000..aea4df87f --- /dev/null +++ b/src/openfe/tests/orchestration/test_exorcist_utils.py @@ -0,0 +1,231 @@ +from pathlib import Path +from typing import cast +from unittest import mock + +import exorcist +import networkx as nx +import pytest +import sqlalchemy as sqla +from gufe.tokenization import GufeKey + +from openfe.orchestration.exorcist_utils import ( + alchemical_network_to_task_graph, + build_task_db_from_alchemical_network, +) +from openfe.storage.warehouse import FileSystemWarehouse, WarehouseBaseClass + + +class _RecordingWarehouse: + def __init__(self): + self.stored_tasks = [] + + def store_task(self, task): + self.stored_tasks.append(task) + + def store_setup_tokenizable(self, obj): + # TODO: add tests for tokenizable storage? + pass + + def store_protocol_dag(self, dag): + pass + + +def _network_units(benzene_variants_star_map): + units = [] + for transformation in benzene_variants_star_map.edges: + units.extend(transformation.create().protocol_units) + return units + + +@pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) +def test_alchemical_network_to_task_graph_stores_all_units(request, fixture): + warehouse = _RecordingWarehouse() + network = request.getfixturevalue(fixture) + expected_units = _network_units(network) + alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) + + stored_unit_names = [str(unit.name) for unit in warehouse.stored_tasks] + expected_unit_names = [str(unit.name) for unit in expected_units] + + assert len(stored_unit_names) == len(expected_unit_names) + assert sorted(stored_unit_names) == sorted(expected_unit_names) + + +@pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) +def test_alchemical_network_to_task_graph_uses_canonical_task_ids(request, fixture): + warehouse = _RecordingWarehouse() + network = request.getfixturevalue(fixture) + + graph = alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) + + expected_protocol_unit_keys = sorted(str(unit.key) for unit in warehouse.stored_tasks) + observed_protocol_unit_keys = [] + + for node in graph.nodes: + protocol_unit_key = node + observed_protocol_unit_keys.append(protocol_unit_key) + + assert sorted(observed_protocol_unit_keys) == expected_protocol_unit_keys + + +@pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) +def test_alchemical_network_to_task_graph_edges_reference_existing_nodes(request, fixture): + warehouse = _RecordingWarehouse() + network = request.getfixturevalue(fixture) + + graph = alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) + + assert len(graph.edges) > 0 + for u, v in graph.edges: + assert u in graph.nodes + assert v in graph.nodes + + +@pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) +def test_alchemical_network_to_task_graph_edge_direction_matches_dependencies(request, fixture): + warehouse = _RecordingWarehouse() + network = request.getfixturevalue(fixture) + + graph = alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) + units_by_key = {str(unit.key): unit for unit in warehouse.stored_tasks} + + for upstream_id, downstream_id in graph.edges: + # as of now this is true, but 'node' may contain more info + upstream_key = upstream_id + downstream_key = downstream_id + upstream_unit = units_by_key[upstream_key] + downstream_unit = units_by_key[downstream_key] + assert upstream_unit in downstream_unit.dependencies + + +def test_alchemical_network_to_task_graph_raises_for_cycle(): + class _Unit: + def __init__(self, name: str, key: str): + self.name = name + self.key = key + + class _Transformation: + name = "cyclic" + key = "Transformation-cycle" + + def create(self): + unit_a = _Unit("unit-a", "ProtocolUnit-a") + unit_b = _Unit("unit-b", "ProtocolUnit-b") + dag = mock.Mock() + dag.protocol_units = [unit_a, unit_b] + dag.graph = nx.DiGraph() + dag.graph.add_nodes_from([unit_a, unit_b]) + dag.graph.add_edges_from([(unit_a, unit_b), (unit_b, unit_a)]) + return dag + + network = mock.Mock() + network.edges = [_Transformation()] + warehouse = mock.Mock() + + with pytest.raises(ValueError, match="not a DAG"): + alchemical_network_to_task_graph(network, warehouse) + + +@pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) +def test_build_task_db_checkout_order_is_dependency_safe(tmp_path, request, fixture): + network = request.getfixturevalue(fixture) + warehouse = FileSystemWarehouse(str(tmp_path / "warehouse")) + # Build the real sqlite task DB from a real alchemical network fixture. + db = build_task_db_from_alchemical_network( + network, + warehouse, + db_path=tmp_path / "tasks.db", + ) + + # Read task IDs and dependency edges from the persisted DB state. + initial_task_rows = list(db.get_all_tasks()) + graph_taskids = {row.taskid for row in initial_task_rows} + with db.engine.connect() as conn: + dep_rows = conn.execute(sqla.select(db.dependencies_table)).all() + graph_edges = {(row._mapping["from"], row._mapping["to"]) for row in dep_rows} + + checkout_order = [] + # Hard upper bound prevents infinite checkout loops. + max_checkouts = len(graph_taskids) + for _ in range(max_checkouts): + taskid = db.check_out_task() + if taskid is None: + break + + checkout_order.append(taskid) + protocol_unit_key = taskid + loaded_unit = warehouse.load_task(GufeKey(protocol_unit_key)) + assert str(loaded_unit.key) == protocol_unit_key + db.mark_task_completed(taskid, success=True) + + # Coverage/completion: every task is checked out exactly once. + observed_taskids = set(checkout_order) + assert observed_taskids == graph_taskids + assert len(checkout_order) == len(graph_taskids) + + # Dependency safety: upstream tasks must appear before downstream tasks. + checkout_index = {taskid: idx for idx, taskid in enumerate(checkout_order)} + for upstream, downstream in graph_edges: + assert checkout_index[upstream] < checkout_index[downstream] + + # Final DB state: all tasks are completed. + task_rows = list(db.get_all_tasks()) + assert len(task_rows) == len(graph_taskids) + assert {row.taskid for row in task_rows} == graph_taskids + assert {row.status for row in task_rows} == {exorcist.TaskStatus.COMPLETED.value} + + +@pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) +def test_build_task_db_default_path(request, fixture): + network = request.getfixturevalue(fixture) + warehouse = mock.Mock() + fake_graph = nx.DiGraph() + fake_db = mock.Mock() + + with ( + mock.patch( + "openfe.orchestration.exorcist_utils.alchemical_network_to_task_graph", + return_value=fake_graph, + ) as task_graph_mock, + mock.patch( + "openfe.orchestration.exorcist_utils.exorcist.TaskStatusDB.from_filename", + return_value=fake_db, + ) as db_ctor, + ): + result = build_task_db_from_alchemical_network(network, warehouse) + + task_graph_mock.assert_called_once_with(network, warehouse) + db_ctor.assert_called_once_with(Path(f"{warehouse.name}.db")) + fake_db.add_task_network.assert_called_once_with(fake_graph, 1) + assert result is fake_db + + +@pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) +def test_build_task_db_forwards_graph_and_max_tries(request, tmp_path, fixture): + network = request.getfixturevalue(fixture) + warehouse = mock.Mock() + fake_graph = nx.DiGraph() + fake_db = mock.Mock() + db_path = tmp_path / "custom_tasks.db" + + with ( + mock.patch( + "openfe.orchestration.exorcist_utils.alchemical_network_to_task_graph", + return_value=fake_graph, + ) as task_graph_mock, + mock.patch( + "openfe.orchestration.exorcist_utils.exorcist.TaskStatusDB.from_filename", + return_value=fake_db, + ) as db_ctor, + ): + result = build_task_db_from_alchemical_network( + network, + warehouse, + db_path=db_path, + max_tries=7, + ) + + task_graph_mock.assert_called_once_with(network, warehouse) + db_ctor.assert_called_once_with(db_path) + fake_db.add_task_network.assert_called_once_with(fake_graph, 7) + assert result is fake_db diff --git a/src/openfe/tests/orchestration/test_worker.py b/src/openfe/tests/orchestration/test_worker.py new file mode 100644 index 000000000..221e8de54 --- /dev/null +++ b/src/openfe/tests/orchestration/test_worker.py @@ -0,0 +1,273 @@ +from pathlib import Path +from unittest import mock + +import exorcist +import gufe +import networkx as nx +import pytest +from gufe.protocols.protocolunit import ProtocolUnit + +from openfe.orchestration import Worker +from openfe.orchestration.exorcist_utils import build_task_db_from_alchemical_network +from openfe.storage.warehouse import FileSystemWarehouse + + +def _result_store_files(warehouse: FileSystemWarehouse) -> set[str]: + result_root = Path(warehouse.result_store.root_dir) + return {str(path.relative_to(result_root)) for path in result_root.rglob("*") if path.is_file()} + + +def _contains_protocol_unit(value) -> bool: + if isinstance(value, ProtocolUnit): + return True + if isinstance(value, dict): + return any(_contains_protocol_unit(item) for item in value.values()) + if isinstance(value, list): + return any(_contains_protocol_unit(item) for item in value) + return False + + +class _ToyProtocolUnit(ProtocolUnit): + @staticmethod + def _execute(ctx, **inputs) -> dict[str, int]: + increment = inputs["increment"] + upstream = inputs.get("upstream") + base = 0 if upstream is None else upstream.outputs["value"] + return {"value": base + increment} + + +class _FileWritingUnit(ProtocolUnit): + @staticmethod + def _execute(ctx, **inputs) -> dict[str, str]: + shared_file = ctx.shared / "simulation.nc" + shared_file.parent.mkdir(parents=True, exist_ok=True) + shared_file.write_text("unit output", encoding="utf-8") + return {"shared_file": str(shared_file)} + + +def _get_dependency_free_unit(absolute_transformation): + for unit in absolute_transformation.create().protocol_units: + if not _contains_protocol_unit(unit.inputs): + return unit + raise ValueError("No dependency-free protocol unit found for execution test setup.") + + +@pytest.fixture +def worker_with_real_db(tmp_path, absolute_transformation): + warehouse_root = tmp_path / "warehouse" + db_path = warehouse_root / "tasks.db" + warehouse = FileSystemWarehouse(str(warehouse_root)) + network = gufe.AlchemicalNetwork([absolute_transformation]) + db = build_task_db_from_alchemical_network(network, warehouse, db_path=db_path) + worker = Worker(warehouse=warehouse, task_db_path=db_path) + return worker, warehouse, db + + +@pytest.fixture +def worker_with_executable_task_db(tmp_path, absolute_transformation): + warehouse_root = tmp_path / "warehouse" + db_path = warehouse_root / "tasks.db" + warehouse = FileSystemWarehouse(str(warehouse_root)) + unit = _get_dependency_free_unit(absolute_transformation) + warehouse.store_task(unit) + + taskid = unit.key + task_graph = nx.DiGraph() + task_graph.add_node(taskid) + + db = exorcist.TaskStatusDB.from_filename(db_path) + db.add_task_network(task_graph, 1) + + worker = Worker(warehouse=warehouse, task_db_path=db_path) + return worker, warehouse, db, unit + + +def test_get_task_uses_default_db_path_without_patching( + tmp_path, monkeypatch, absolute_transformation +): + monkeypatch.chdir(tmp_path) + warehouse = FileSystemWarehouse("warehouse") + db_path = Path("warehouse/tasks.db") + network = gufe.AlchemicalNetwork([absolute_transformation]) + db = build_task_db_from_alchemical_network(network, warehouse, db_path=db_path) + + worker = Worker(warehouse=warehouse) + taskid, loaded = worker._get_task() + + expected_keys = {task_row.taskid for task_row in db.get_all_tasks()} + assert worker.task_db_path == Path("./warehouse/tasks.db") + assert str(loaded.key) in expected_keys + assert taskid == loaded.key + + +def test_get_task_returns_task_with_canonical_protocol_unit_suffix(worker_with_real_db): + worker, warehouse, db = worker_with_real_db + + task_ids = [row.taskid for row in db.get_all_tasks()] + expected_protocol_unit_keys = {task_id for task_id in task_ids} + + taskid, loaded = worker._get_task() + reloaded = warehouse.load_task(loaded.key) + + assert str(loaded.key) in expected_protocol_unit_keys + assert loaded == reloaded + assert taskid == loaded.key + + +def test_execute_unit_stores_real_result(worker_with_executable_task_db, tmp_path): + worker, warehouse, db, _ = worker_with_executable_task_db + before = _result_store_files(warehouse) + + execution = worker.execute_unit(scratch=tmp_path / "scratch") + assert execution is not None + taskid, _ = execution + + after = _result_store_files(warehouse) + assert len(after) > len(before) + rows = list(db.get_all_tasks()) + status_by_taskid = {row.taskid: row.status for row in rows} + assert status_by_taskid[taskid] == exorcist.TaskStatus.COMPLETED.value + + +def test_execute_unit_propagates_execute_error_without_store( + worker_with_executable_task_db, tmp_path +): + worker, warehouse, db, unit = worker_with_executable_task_db + before = _result_store_files(warehouse) + taskid = list(db.get_all_tasks())[0].taskid + + with mock.patch.object( + type(unit), + "execute", + autospec=True, + side_effect=RuntimeError("unit execution failed"), + ): + with pytest.raises(RuntimeError, match="unit execution failed"): + worker.execute_unit(scratch=tmp_path / "scratch") + + after = _result_store_files(warehouse) + assert after == before + rows = list(db.get_all_tasks()) + status_by_taskid = {row.taskid: row.status for row in rows} + assert status_by_taskid[taskid] == exorcist.TaskStatus.TOO_MANY_RETRIES.value + + +def test_checkout_task_returns_none_when_no_available_tasks(tmp_path): + warehouse_root = tmp_path / "warehouse" + db_path = warehouse_root / "tasks.db" + warehouse_root.mkdir(parents=True, exist_ok=True) + warehouse = FileSystemWarehouse(str(warehouse_root)) + exorcist.TaskStatusDB.from_filename(db_path) + worker = Worker(warehouse=warehouse, task_db_path=db_path) + + assert worker._checkout_task() is None + + +def test_execute_unit_returns_none_when_no_available_tasks(tmp_path): + warehouse_root = tmp_path / "warehouse" + db_path = warehouse_root / "tasks.db" + warehouse_root.mkdir(parents=True, exist_ok=True) + warehouse = FileSystemWarehouse(str(warehouse_root)) + exorcist.TaskStatusDB.from_filename(db_path) + worker = Worker(warehouse=warehouse, task_db_path=db_path) + + assert worker.execute_unit(scratch=tmp_path / "scratch") is None + + +def test_execute_unit_resolves_dependency_results(tmp_path): + warehouse_root = tmp_path / "warehouse" + db_path = warehouse_root / "tasks.db" + warehouse = FileSystemWarehouse(str(warehouse_root)) + + first_unit = _ToyProtocolUnit(name="first", increment=1) + second_unit = _ToyProtocolUnit(name="second", upstream=first_unit, increment=2) + + warehouse.store_task(first_unit) + warehouse.store_task(second_unit) + + first_taskid = first_unit.key + second_taskid = second_unit.key + + task_graph = nx.DiGraph() + task_graph.add_edge(first_taskid, second_taskid) + + db = exorcist.TaskStatusDB.from_filename(db_path) + db.add_task_network(task_graph, max_tries=1) + + worker = Worker(warehouse=warehouse, task_db_path=db_path) + + first_execution = worker.execute_unit(scratch=tmp_path / "scratch") + second_execution = worker.execute_unit(scratch=tmp_path / "scratch") + + assert first_execution is not None + assert first_execution[0] == first_taskid + assert second_execution is not None + assert second_execution[0] == second_taskid + assert second_execution[1].outputs["value"] == 3 + + status_by_taskid = {row.taskid: row.status for row in db.get_all_tasks()} + assert status_by_taskid[first_taskid] == exorcist.TaskStatus.COMPLETED.value + assert status_by_taskid[second_taskid] == exorcist.TaskStatus.COMPLETED.value + + +def test_execute_unit_marks_missing_dependency_as_failed(tmp_path): + warehouse_root = tmp_path / "warehouse" + db_path = warehouse_root / "tasks.db" + warehouse = FileSystemWarehouse(str(warehouse_root)) + + missing_upstream = _ToyProtocolUnit(name="missing", increment=1) + dependent_unit = _ToyProtocolUnit(name="dependent", upstream=missing_upstream, increment=2) + warehouse.store_task(dependent_unit) + + taskid = dependent_unit.key + task_graph = nx.DiGraph() + task_graph.add_node(taskid) + + db = exorcist.TaskStatusDB.from_filename(db_path) + db.add_task_network(task_graph, max_tries=1) + + worker = Worker(warehouse=warehouse, task_db_path=db_path) + + with pytest.raises(RuntimeError, match="Missing ProtocolUnitResult"): + worker.execute_unit(scratch=tmp_path / "scratch") + + status_by_taskid = {row.taskid: row.status for row in db.get_all_tasks()} + assert status_by_taskid[taskid] == exorcist.TaskStatus.TOO_MANY_RETRIES.value + + +def test_execute_unit_uses_isolated_shared_workspace_per_task(tmp_path): + warehouse_root = tmp_path / "warehouse" + db_path = warehouse_root / "tasks.db" + warehouse = FileSystemWarehouse(str(warehouse_root)) + + first_unit = _FileWritingUnit(name="first") + second_unit = _FileWritingUnit(name="second") + + warehouse.store_task(first_unit) + warehouse.store_task(second_unit) + + first_taskid = first_unit.key + second_taskid = second_unit.key + + task_graph = nx.DiGraph() + task_graph.add_node(first_taskid) + task_graph.add_node(second_taskid) + + db = exorcist.TaskStatusDB.from_filename(db_path) + db.add_task_network(task_graph, max_tries=1) + + worker = Worker(warehouse=warehouse, task_db_path=db_path) + + first_execution = worker.execute_unit(scratch=tmp_path / "scratch") + second_execution = worker.execute_unit(scratch=tmp_path / "scratch") + + assert first_execution is not None + assert second_execution is not None + + first_path = Path(first_execution[1].outputs["shared_file"]) + second_path = Path(second_execution[1].outputs["shared_file"]) + + assert first_path != second_path + assert first_path.name == "simulation.nc" + assert second_path.name == "simulation.nc" + assert first_path.parent != second_path.parent diff --git a/src/openfe/tests/storage/test_metadatastore.py b/src/openfe/tests/storage/test_metadatastore.py deleted file mode 100644 index 9f3dcd167..000000000 --- a/src/openfe/tests/storage/test_metadatastore.py +++ /dev/null @@ -1,147 +0,0 @@ -import json -import pathlib - -import pytest -from gufe.storage.errors import ChangedExternalResourceError, MissingExternalResourceError -from gufe.storage.externalresource import FileStorage -from gufe.storage.externalresource.base import Metadata - -from openfe.storage.metadatastore import JSONMetadataStore, PerFileJSONMetadataStore - - -@pytest.fixture -def json_metadata(tmp_path): - metadata_dict = {"path/to/foo.txt": {"md5": "bar"}} - external_store = FileStorage(tmp_path) - with open(tmp_path / "metadata.json", mode="wb") as f: - f.write(json.dumps(metadata_dict).encode("utf-8")) - json_metadata = JSONMetadataStore(external_store) - return json_metadata - - -@pytest.fixture -def per_file_metadata(tmp_path): - metadata_dict = {"path": "path/to/foo.txt", "metadata": {"md5": "bar"}} - metadata_loc = "metadata/path/to/foo.txt.json" - external_store = FileStorage(tmp_path) - metadata_path = tmp_path / pathlib.Path(metadata_loc) - metadata_path.parent.mkdir(parents=True, exist_ok=True) - with open(metadata_path, mode="wb") as f: - f.write(json.dumps(metadata_dict).encode("utf-8")) - - per_file_metadata = PerFileJSONMetadataStore(external_store) - return per_file_metadata - - -class MetadataTests: - """Mixin with a few tests for any subclass of MetadataStore""" - - def test_store_metadata(self, metadata): - raise NotImplementedError() - - def test_load_all_metadata(self): - raise NotImplementedError("This should call self._test_load_all_metadata") - - def test_delete(self): - raise NotImplementedError("This should call self._test_delete") - - def _test_load_all_metadata(self, metadata): - expected = {"path/to/foo.txt": Metadata(md5="bar")} - metadata._metadata_cache = {} - loaded = metadata.load_all_metadata() - assert loaded == expected - - def _test_delete(self, metadata): - assert "path/to/foo.txt" in metadata - assert len(metadata) == 1 - del metadata["path/to/foo.txt"] - assert "path/to/foo.txt" not in metadata - assert len(metadata) == 0 - - def _test_iter(self, metadata): - assert list(metadata) == ["path/to/foo.txt"] - - def _test_len(self, metadata): - assert len(metadata) == 1 - - def _test_getitem(self, metadata): - assert metadata["path/to/foo.txt"] == Metadata(md5="bar") - - -class TestJSONMetadataStore(MetadataTests): - def test_store_metadata(self, json_metadata): - meta = Metadata(md5="other") - json_metadata.store_metadata("path/to/other.txt", meta) - base_path = json_metadata.external_store.root_dir - metadata_json = base_path / "metadata.json" - assert metadata_json.exists() - with open(metadata_json, mode="r") as f: - metadata_dict = json.load(f) - - metadata = {key: Metadata(**val) for key, val in metadata_dict.items()} - - assert metadata == json_metadata._metadata_cache - assert json_metadata["path/to/other.txt"] == meta - assert len(metadata) == 2 - - def test_load_all_metadata(self, json_metadata): - self._test_load_all_metadata(json_metadata) - - def test_load_all_metadata_nofile(self, tmp_path): - json_metadata = JSONMetadataStore(FileStorage(tmp_path)) - # implicitly called on init anyway - assert json_metadata._metadata_cache == {} - # but we also call explicitly - assert json_metadata.load_all_metadata() == {} - - def test_delete(self, json_metadata): - self._test_delete(json_metadata) - - def test_iter(self, json_metadata): - self._test_iter(json_metadata) - - def test_len(self, json_metadata): - self._test_len(json_metadata) - - def test_getitem(self, json_metadata): - self._test_getitem(json_metadata) - - -class TestPerFileJSONMetadataStore(MetadataTests): - def test_store_metadata(self, per_file_metadata): - expected_loc = "metadata/path/to/other.txt.json" - root = per_file_metadata.external_store.root_dir - expected_path = root / expected_loc - assert not expected_path.exists() - meta = Metadata(md5="other") - per_file_metadata.store_metadata("path/to/other.txt", meta) - assert expected_path.exists() - expected = {"path": "path/to/other.txt", "metadata": {"md5": "other"}} - with open(expected_path, mode="r") as f: - assert json.load(f) == expected - - def test_load_all_metadata(self, per_file_metadata): - self._test_load_all_metadata(per_file_metadata) - - def test_delete(self, per_file_metadata): - self._test_delete(per_file_metadata) - # TODO: add additional test that the file is gone - - def test_iter(self, per_file_metadata): - self._test_iter(per_file_metadata) - - def test_len(self, per_file_metadata): - self._test_len(per_file_metadata) - - def test_getitem(self, per_file_metadata): - self._test_getitem(per_file_metadata) - - def test_bad_metadata_contents(self, tmp_path): - loc = tmp_path / "metadata/foo.txt.json" - loc.parent.mkdir(parents=True, exist_ok=True) - bad_dict = {"foo": "bar"} - with open(loc, mode="wb") as f: - f.write(json.dumps(bad_dict).encode("utf-8")) - - with pytest.raises(ChangedExternalResourceError, match="Bad metadata"): - PerFileJSONMetadataStore(FileStorage(tmp_path)) diff --git a/src/openfe/tests/storage/test_resultclient.py b/src/openfe/tests/storage/test_resultclient.py deleted file mode 100644 index 52695958f..000000000 --- a/src/openfe/tests/storage/test_resultclient.py +++ /dev/null @@ -1,304 +0,0 @@ -import os -from unittest import mock - -import pytest -from gufe.storage.externalresource import MemoryStorage -from gufe.tokenization import TOKENIZABLE_REGISTRY - -from openfe.storage.resultclient import ( - CloneResult, - ExtensionResult, - ResultClient, - TransformationResult, -) - - -@pytest.fixture -def result_client(): - external = MemoryStorage() - result_client = ResultClient(external) - - # store one file with contents "foo" - result_client.result_server.store_bytes( - "transformations/MAIN_TRANS/0/0/file.txt", - "foo".encode("utf-8"), - ) - - # create some empty files as well - empty_files = [ - "transformations/MAIN_TRANS/0/0/other.txt", - "transformations/MAIN_TRANS/0/1/file.txt", - "transformations/MAIN_TRANS/1/0/file.txt", - "transformations/OTHER_TRANS/0/0/file.txt", - "other_dir/file.txt", - ] - - for file in empty_files: - result_client.result_server.store_bytes(file, b"") # empty - - return result_client - - -def _make_mock_transformation(hash_str): - return mock.Mock( - # TODO: fill this in so that it mocks out the digest we use - ) - - -def test_load_file(result_client): - file_handler = result_client / "MAIN_TRANS" / "0" / 0 / "file.txt" - with file_handler as f: - assert f.read().decode("utf-8") == "foo" - - -class _ResultContainerTest: - @staticmethod - def get_container(result_client): - raise NotImplementedError() - - def _getitem_object(self, container): - raise NotImplementedError() - - def test_iter(self, result_client): - container = self.get_container(result_client) - assert set(container) == set(self.expected_files) - - def _get_key(self, as_object, container): - # TODO: this isn't working yet -- need an interface that allows me - # to patch the hex digest that we'll be using - if as_object: - pytest.skip("Waiting on hex digest patching") - obj = self._getitem_object(container) - # next line uses some internal implementation - key = obj if as_object else obj._path_component - return key, obj - - @pytest.mark.parametrize("as_object", [True, False]) - def test_getitem(self, as_object, result_client): - container = self.get_container(result_client) - key, obj = self._get_key(as_object, container) - assert container[key] == obj - - @pytest.mark.parametrize("as_object", [True, False]) - def test_div(self, as_object, result_client): - container = self.get_container(result_client) - key, obj = self._get_key(as_object, container) - assert container / key == obj - - @pytest.mark.parametrize("load_with", ["div", "getitem"]) - def test_caching(self, result_client, load_with): - # used to test caching regardless of how first loaded was loaded - container = self.get_container(result_client) - key, obj = self._get_key(False, container) - - if load_with == "div": - loaded = container / key - elif load_with == "getitem": - loaded = container[key] - else: # -no-cov- - raise RuntimeError(f"Bad input: can't load with '{load_with}'") - - assert loaded == obj - assert loaded is not obj - reloaded_div = container / key - reloaded_getitem = container[key] - - assert loaded is reloaded_div - assert reloaded_div is reloaded_getitem - - def test_load_stream(self, result_client): - container = self.get_container(result_client) - loc = "transformations/MAIN_TRANS/0/0/file.txt" - with container.load_stream(loc) as f: - assert f.read().decode("utf-8") == "foo" - - def test_load_bytes(self, result_client): - container = self.get_container(result_client) - loc = "transformations/MAIN_TRANS/0/0/file.txt" - assert container.load_bytes(loc).decode("utf-8") == "foo" - - def test_path(self, result_client): - container = self.get_container(result_client) - assert container.path == self.expected_path - - def test_result_server(self, result_client): - container = self.get_container(result_client) - assert container.result_server == result_client.result_server - - -class TestResultClient(_ResultContainerTest): - expected_files = [ - "transformations/MAIN_TRANS/0/0/file.txt", - "transformations/MAIN_TRANS/0/0/other.txt", - "transformations/MAIN_TRANS/0/1/file.txt", - "transformations/MAIN_TRANS/1/0/file.txt", - "transformations/OTHER_TRANS/0/0/file.txt", - ] - expected_path = "transformations" - - @staticmethod - def get_container(result_client): - return result_client - - def _getitem_object(self, container): - return TransformationResult( - parent=container, - transformation=_make_mock_transformation("MAIN_TRANS"), - ) - - def test_store_protocol_dag_result(self): - pytest.skip("Not implemented yet") - - @staticmethod - def _test_store_load_same_process(obj, store_func_name, load_func_name): - store = MemoryStorage() - client = ResultClient(store) - store_func = getattr(client, store_func_name) - load_func = getattr(client, load_func_name) - assert store._data == {} - store_func(obj) - assert store._data != {} - reloaded = load_func(obj.key) - assert reloaded is obj - - @staticmethod - def _test_store_load_different_process(obj, store_func_name, load_func_name): - store = MemoryStorage() - client = ResultClient(store) - store_func = getattr(client, store_func_name) - load_func = getattr(client, load_func_name) - assert store._data == {} - store_func(obj) - assert store._data != {} - # make it look like we have an empty cache, as if this was a - # different process - key = obj.key - registry_dict = "gufe.tokenization.TOKENIZABLE_REGISTRY" - with mock.patch.dict(registry_dict, {}, clear=True): - reload = load_func(key) - assert reload == obj - assert reload is not obj - - @pytest.mark.parametrize( - "fixture", - ["absolute_transformation", "complex_equilibrium"], - ) - def test_store_load_transformation_same_process(self, request, fixture): - transformation = request.getfixturevalue(fixture) - self._test_store_load_same_process( - transformation, - "store_transformation", - "load_transformation", - ) - - @pytest.mark.parametrize( - "fixture", - ["absolute_transformation", "complex_equilibrium"], - ) - def test_store_load_transformation_different_process(self, request, fixture): - transformation = request.getfixturevalue(fixture) - self._test_store_load_different_process( - transformation, - "store_transformation", - "load_transformation", - ) - - @pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) - def test_store_load_network_same_process(self, request, fixture): - network = request.getfixturevalue(fixture) - self._test_store_load_same_process(network, "store_network", "load_network") - - @pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) - def test_store_load_network_different_process(self, request, fixture): - network = request.getfixturevalue(fixture) - self._test_store_load_different_process(network, "store_network", "load_network") - - def test_delete(self, result_client): - file_to_delete = self.expected_files[0] - storage = result_client.result_server.external_store - assert storage.exists(file_to_delete) - result_client.delete(file_to_delete) - assert not storage.exists(file_to_delete) - - -class TestTransformationResults(_ResultContainerTest): - expected_files = [ - "transformations/MAIN_TRANS/0/0/file.txt", - "transformations/MAIN_TRANS/0/0/other.txt", - "transformations/MAIN_TRANS/0/1/file.txt", - "transformations/MAIN_TRANS/1/0/file.txt", - ] - expected_path = "transformations/MAIN_TRANS" - - @staticmethod - def get_container(result_client): - container = TransformationResult( - parent=TestResultClient.get_container(result_client), - transformation=_make_mock_transformation("MAIN_TRANS"), - ) - container._path_component = "MAIN_TRANS" - return container - - def _getitem_object(self, container): - return CloneResult(parent=container, clone=0) - - -class TestCloneResults(_ResultContainerTest): - expected_files = [ - "transformations/MAIN_TRANS/0/0/file.txt", - "transformations/MAIN_TRANS/0/0/other.txt", - "transformations/MAIN_TRANS/0/1/file.txt", - ] - expected_path = "transformations/MAIN_TRANS/0" - - @staticmethod - def get_container(result_client): - return CloneResult( - parent=TestTransformationResults.get_container(result_client), - clone=0, - ) - - def _getitem_object(self, container): - return ExtensionResult(parent=container, extension=0) - - -class TestExtensionResults(_ResultContainerTest): - expected_files = [ - "transformations/MAIN_TRANS/0/0/file.txt", - "transformations/MAIN_TRANS/0/0/other.txt", - ] - expected_path = "transformations/MAIN_TRANS/0/0" - - @staticmethod - def get_container(result_client): - return ExtensionResult( - parent=TestCloneResults.get_container(result_client), - extension=0, - ) - - def _get_key(self, as_object, container): - if self.as_object: # -no-cov- - raise RuntimeError("TestExtensionResults does not support as_object=True") - path = "transformations/MAIN_TRANS/0/0/" - fname = "file.txt" - return fname, container.result_server.load_stream(path + fname) - - # things involving div and getitem need custom treatment - def test_div(self, result_client): - container = self.get_container(result_client) - with container / "file.txt" as f: - assert f.read().decode("utf-8") == "foo" - - def test_getitem(self, result_client): - container = self.get_container(result_client) - with container["file.txt"] as f: - assert f.read().decode("utf-8") == "foo" - - def test_caching(self, result_client): - # this one does not cache results; the cache should remain empty - container = self.get_container(result_client) - assert container._cache == {} - from_div = container / "file.txt" - assert container._cache == {} - from_getitem = container["file.txt"] - assert container._cache == {} diff --git a/src/openfe/tests/storage/test_resultserver.py b/src/openfe/tests/storage/test_resultserver.py deleted file mode 100644 index 919205191..000000000 --- a/src/openfe/tests/storage/test_resultserver.py +++ /dev/null @@ -1,117 +0,0 @@ -import pathlib -from unittest import mock - -import pytest -from gufe.storage.errors import ChangedExternalResourceError, MissingExternalResourceError -from gufe.storage.externalresource import FileStorage -from gufe.storage.externalresource.base import Metadata - -from openfe.storage.metadatastore import JSONMetadataStore -from openfe.storage.resultserver import ResultServer - - -@pytest.fixture -def result_server(tmp_path): - external = FileStorage(tmp_path) - metadata = JSONMetadataStore(external) - result_server = ResultServer(external, metadata) - result_server.store_bytes("path/to/foo.txt", "foo".encode("utf-8")) - return result_server - - -class TestResultServer: - def test_store_bytes(self, result_server): - # first check the thing stored during the fixture - metadata_store = result_server.metadata_store - foo_loc = "path/to/foo.txt" - assert len(metadata_store) == 1 - assert foo_loc in metadata_store - assert result_server.external_store.exists(foo_loc) - - # also explicitly test storing here - mock_hash = mock.Mock( - return_value=mock.Mock( - hexdigest=mock.Mock(return_value="deadbeef"), - ) - ) - bar_loc = "path/to/bar.txt" - with mock.patch("hashlib.md5", mock_hash): - result_server.store_bytes(bar_loc, "bar".encode("utf-8")) - - assert len(metadata_store) == 2 - assert bar_loc in metadata_store - assert result_server.external_store.exists(bar_loc) - assert metadata_store[bar_loc].to_dict() == {"md5": "deadbeef"} - external = result_server.external_store - with external.load_stream(bar_loc) as f: - assert f.read().decode("utf-8") == "bar" - - def test_store_path(self, result_server, tmp_path): - orig_file = tmp_path / ".hidden" / "bar.txt" - orig_file.parent.mkdir(parents=True, exist_ok=True) - with open(orig_file, mode="wb") as f: - f.write("bar".encode("utf-8")) - - mock_hash = mock.Mock( - return_value=mock.Mock( - hexdigest=mock.Mock(return_value="deadc0de"), - ) - ) - bar_loc = "path/to/bar.txt" - - assert len(result_server.metadata_store) == 1 - assert bar_loc not in result_server.metadata_store - - with mock.patch("hashlib.md5", mock_hash): - result_server.store_path(bar_loc, orig_file) - - assert len(result_server.metadata_store) == 2 - assert bar_loc in result_server.metadata_store - metadata_dict = result_server.metadata_store[bar_loc].to_dict() - assert metadata_dict == {"md5": "deadc0de"} - external = result_server.external_store - with external.load_stream(bar_loc) as f: - assert f.read().decode("utf-8") == "bar" - - def test_iter(self, result_server): - assert list(result_server) == ["path/to/foo.txt"] - - def test_find_missing_files(self, result_server): - meta = Metadata(md5="1badc0de") - result_server.metadata_store.store_metadata("fake/file.txt", meta) - - assert result_server.find_missing_files() == ["fake/file.txt"] - - def test_load_stream(self, result_server): - with result_server.load_stream("path/to/foo.txt") as f: - contents = f.read() - - assert contents.decode("utf-8") == "foo" - - def test_delete(self, result_server, tmp_path): - location = "path/to/foo.txt" - path = tmp_path / pathlib.Path(location) - assert path.exists() - assert location in result_server.metadata_store - result_server.delete(location) - assert not path.exists() - assert location not in result_server.metadata_store - - def test_load_stream_missing(self, result_server): - with pytest.raises(MissingExternalResourceError, match="not found"): - result_server.load_stream("path/does/not/exist.txt") - - def test_load_stream_error_bad_hash(self, result_server): - meta = Metadata(md5="1badc0de") - result_server.metadata_store.store_metadata("path/to/foo.txt", meta) - with pytest.raises(ChangedExternalResourceError): - result_server.load_stream("path/to/foo.txt") - - def test_load_stream_allow_bad_hash(self, result_server): - meta = Metadata(md5="1badc0de") - result_server.metadata_store.store_metadata("path/to/foo.txt", meta) - with pytest.warns(UserWarning, match="Metadata mismatch"): - file = result_server.load_stream("path/to/foo.txt", allow_changed=True) - - with file as f: - assert f.read().decode("utf-8") == "foo" diff --git a/src/openfe/tests/storage/test_warehouse.py b/src/openfe/tests/storage/test_warehouse.py new file mode 100644 index 000000000..bacec0b10 --- /dev/null +++ b/src/openfe/tests/storage/test_warehouse.py @@ -0,0 +1,251 @@ +import os +import tempfile +from pathlib import Path +from typing import Literal +from unittest import mock + +import pytest +from gufe.storage.externalresource import MemoryStorage +from gufe.tokenization import GufeTokenizable + +from openfe.storage.warehouse import ( + FileSystemWarehouse, + WarehouseBaseClass, + WarehouseStores, +) + + +class TestWarehouseBaseClass: + def test_store_protocol_dag_result(self): + pytest.skip("Not implemented yet") + + @staticmethod + def _build_stores() -> WarehouseStores: + return WarehouseStores( + setup=MemoryStorage(), + result=MemoryStorage(), + shared=MemoryStorage(), + tasks=MemoryStorage(), + protocol_dags=MemoryStorage(), + ) + + @staticmethod + def _get_protocol_unit(transformation): + dag = transformation.create() + return next(iter(dag.protocol_units)) + + @staticmethod + def _test_store_load_same_process( + obj, store_func_name, load_func_name, store_name: Literal["setup", "result"] + ) -> tuple[GufeTokenizable, WarehouseBaseClass]: + setup_store = MemoryStorage() + result_store = MemoryStorage() + stores = WarehouseStores(setup=setup_store, result=result_store) + client = WarehouseBaseClass(stores) + store_func = getattr(client, store_func_name) + load_func = getattr(client, load_func_name) + assert stores["setup"]._data == {} + assert stores["result"]._data == {} + assert stores["shared"]._data == {} + assert stores["tasks"]._data == {} + store_func(obj) + store_under_test: MemoryStorage = stores[store_name] + assert store_under_test._data != {} + reloaded: GufeTokenizable = load_func(obj.key) + assert reloaded is obj + return reloaded, client + + @staticmethod + def _test_store_load_different_process( + obj: GufeTokenizable, + store_func_name, + load_func_name, + store_name: Literal["setup", "result"], + ) -> None: + setup_store = MemoryStorage() + result_store = MemoryStorage() + stores = WarehouseStores(setup=setup_store, result=result_store) + client = WarehouseBaseClass(stores) + store_func = getattr(client, store_func_name) + load_func = getattr(client, load_func_name) + assert stores["setup"]._data == {} + assert stores["result"]._data == {} + assert stores["shared"]._data == {} + assert stores["tasks"]._data == {} + store_func(obj) + store_under_test: MemoryStorage = stores[store_name] + assert store_under_test._data != {} + # make it look like we have an empty cache, as if this was a + # different process + key = obj.key + registry_dict = "gufe.tokenization.TOKENIZABLE_REGISTRY" + with mock.patch.dict(registry_dict, {}, clear=True): + reload = load_func(key) + assert reload == obj + assert reload is not obj + + def test_store_load_task_same_process(self, absolute_transformation): + unit = self._get_protocol_unit(absolute_transformation) + self._test_store_load_same_process(unit, "store_task", "load_task", "tasks") + + def test_store_load_task_different_process(self, absolute_transformation): + unit = self._get_protocol_unit(absolute_transformation) + self._test_store_load_different_process(unit, "store_task", "load_task", "tasks") + + def test_store_task_writes_to_tasks_store(self, absolute_transformation): + unit = self._get_protocol_unit(absolute_transformation) + stores = self._build_stores() + client = WarehouseBaseClass(stores, name="test_warehouse") + client.store_task(unit) + + assert stores["tasks"]._data != {} + assert stores["setup"]._data == {} + assert stores["result"]._data == {} + assert stores["shared"]._data == {} + + def test_exists_finds_task_key(self, absolute_transformation): + unit = self._get_protocol_unit(absolute_transformation) + stores = self._build_stores() + client = WarehouseBaseClass(stores, "test_warehouse") + + client.store_task(unit) + + assert client.exists(unit.key) + + def test_load_task_returns_object(self, absolute_transformation): + unit = self._get_protocol_unit(absolute_transformation) + stores = self._build_stores() + client = WarehouseBaseClass(stores, name="test_warehouse") + + client.store_task(unit) + loaded = client.load_task(unit.key) + + assert loaded is not None + assert isinstance(loaded, GufeTokenizable) + + @pytest.mark.parametrize( + "fixture", + ["absolute_transformation", "complex_equilibrium"], + ) + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_store_load_transformation_same_process(self, request, fixture, store): + transformation = request.getfixturevalue(fixture) + store_func_name = f"store_{store}_tokenizable" + load_func_name = f"load_{store}_tokenizable" + self._test_store_load_same_process(transformation, store_func_name, load_func_name, store) + + @pytest.mark.parametrize( + "fixture", + ["absolute_transformation", "complex_equilibrium"], + ) + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_store_load_transformation_different_process(self, request, fixture, store): + transformation = request.getfixturevalue(fixture) + store_func_name = f"store_{store}_tokenizable" + load_func_name = f"load_{store}_tokenizable" + self._test_store_load_different_process( + transformation, store_func_name, load_func_name, store + ) + + # + @pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_store_load_network_same_process(self, request, fixture, store): + network = request.getfixturevalue(fixture) + assert isinstance(network, GufeTokenizable) + store_func_name = f"store_{store}_tokenizable" + load_func_name = f"load_{store}_tokenizable" + self._test_store_load_same_process(network, store_func_name, load_func_name, store) + + @pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_store_load_network_different_process(self, request, fixture, store): + network = request.getfixturevalue(fixture) + assert isinstance(network, GufeTokenizable) + store_func_name = f"store_{store}_tokenizable" + load_func_name = f"load_{store}_tokenizable" + self._test_store_load_different_process(network, store_func_name, load_func_name, store) + + @pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_delete(self, request, fixture, store): + network = request.getfixturevalue(fixture) + store_func_name = f"store_{store}_tokenizable" + load_func_name = f"load_{store}_tokenizable" + obj, client = self._test_store_load_same_process( + network, store_func_name, load_func_name, store + ) + client.delete(store, obj.key) + assert not client.exists(obj.key) + + +class TestFileSystemWarehouse: + @staticmethod + def _test_store_load_same_process(obj, store_func_name, load_func_name): + with tempfile.TemporaryDirectory() as tmpdir: + client = FileSystemWarehouse(tmpdir) + store_func = getattr(client, store_func_name) + load_func = getattr(client, load_func_name) + assert not any(Path(f"{tmpdir}").iterdir()) + store_func(obj) + assert any(Path(f"{tmpdir}").iterdir()) + reloaded = load_func(obj.key) + assert reloaded is obj + + @staticmethod + def _test_store_load_different_process(obj: GufeTokenizable, store_func_name, load_func_name): + with tempfile.TemporaryDirectory() as tmpdir: + client = FileSystemWarehouse(tmpdir) + store_func = getattr(client, store_func_name) + load_func = getattr(client, load_func_name) + assert not any(Path(f"{tmpdir}").iterdir()) + store_func(obj) + assert any(Path(f"{tmpdir}").iterdir()) + # make it look like we have an empty cache, as if this was a + # different process + key = obj.key + registry_dict = "gufe.tokenization.TOKENIZABLE_REGISTRY" + with mock.patch.dict(registry_dict, {}, clear=True): + reload = load_func(key) + assert reload == obj + assert reload is not obj + + @pytest.mark.parametrize( + "fixture", + ["absolute_transformation", "complex_equilibrium"], + ) + def test_store_load_transformation_same_process(self, request, fixture): + transformation = request.getfixturevalue(fixture) + self._test_store_load_same_process( + transformation, + "store_setup_tokenizable", + "load_setup_tokenizable", + ) + + def test_filesystemwarehouse_has_shared_and_tasks_stores(self, absolute_transformation): + unit = TestWarehouseBaseClass._get_protocol_unit(absolute_transformation) + + with tempfile.TemporaryDirectory() as tmpdir: + client = FileSystemWarehouse(tmpdir) + + assert "shared" in client.stores + assert "tasks" in client.stores + + client.stores["shared"].store_bytes("sentinel", b"shared-data") + with client.stores["shared"].load_stream("sentinel") as f: + assert f.read() == b"shared-data" + + client.store_task(unit) + assert client.exists(unit.key) + + @pytest.mark.parametrize( + "fixture", + ["absolute_transformation", "complex_equilibrium"], + ) + def test_store_load_transformation_different_process(self, request, fixture): + transformation = request.getfixturevalue(fixture) + self._test_store_load_different_process( + transformation, + "store_setup_tokenizable", + "load_setup_tokenizable", + ) diff --git a/src/openfecli/commands/plan_rbfe_network.py b/src/openfecli/commands/plan_rbfe_network.py index 25c454e87..4c764250f 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, @@ -13,6 +13,7 @@ OVERWRITE, PROTEIN, PROTEIN_MEMBRANE, + WAREHOUSE, YAML_OPTIONS, ) from openfecli.utils import print_duration, write @@ -134,6 +135,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], @@ -145,6 +147,7 @@ def plan_rbfe_network( n_protocol_repeats: int, n_cores: int, overwrite_charges: bool, + warehouse: bool, ): """ Plan a relative binding free energy AlchemicalNetwork, saved as JSON files for use by the quickrun command. @@ -247,10 +250,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/worker.py b/src/openfecli/commands/worker.py new file mode 100644 index 000000000..c312e3ff0 --- /dev/null +++ b/src/openfecli/commands/worker.py @@ -0,0 +1,133 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +import pathlib + +import click + +from openfecli import OFECommandPlugin +from openfecli.utils import configure_logger, print_duration, write + + +def _build_worker(warehouse_path: pathlib.Path, db_path: pathlib.Path): + from openfe.orchestration import Worker + from openfe.storage.warehouse import FileSystemWarehouse + + warehouse = FileSystemWarehouse(str(warehouse_path)) + return Worker(warehouse=warehouse, task_db_path=db_path) + + +def _write_failure_result_details(taskid: str, result) -> None: + source_key = getattr(result, "source_key", None) + exception = getattr(result, "exception", None) + traceback_text = getattr(result, "traceback", None) + + write(f"Task '{taskid}' returned a failure result.") + if source_key is not None: + write(f"Failed unit source key: {source_key}") + + if isinstance(exception, tuple) and len(exception) == 2: + exc_type, exc_args = exception + write(f"Protocol unit exception: {exc_type}: {exc_args}") + + if isinstance(traceback_text, str) and traceback_text: + write("Protocol unit traceback:") + write(traceback_text) + + +def worker_main(warehouse_path: pathlib.Path, scratch: pathlib.Path | None): + import logging + import os + import sys + import traceback + + 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) + db_path = warehouse_path / "tasks.db" + if not db_path.is_file(): + raise click.ClickException(f"Task database not found at: {db_path}") + + if scratch is None: + scratch = pathlib.Path.cwd() + + scratch.mkdir(parents=True, exist_ok=True) + + worker = _build_worker(warehouse_path, db_path) + + try: + write("Executing unit...") + execution = worker.execute_unit(scratch=scratch) + except Exception as exc: + write(traceback.format_exc()) + raise click.ClickException(f"Task execution failed: {exc}") from exc + + if execution is None: + write("No available task in task graph.") + return None + + taskid, result = execution + if not result.ok(): + _write_failure_result_details(taskid, result) + raise click.ClickException(f"Task '{taskid}' returned a failure result.") + + write(f"Completed task: {taskid}") + return result + + +@click.command("worker", short_help="Execute one available task from a filesystem warehouse") +@click.argument( + "warehouse_path", + type=click.Path( + exists=True, + readable=True, + file_okay=False, + dir_okay=True, + path_type=pathlib.Path, + ), +) +@click.option( + "--scratch", + "-s", + default=None, + type=click.Path( + writable=True, + file_okay=False, + dir_okay=True, + path_type=pathlib.Path, + ), + help="Directory for scratch files. Defaults to current working directory.", +) +@print_duration +def worker(warehouse_path: pathlib.Path, scratch: pathlib.Path | None): + """ + Execute one available task from a warehouse task graph. + + The warehouse directory must contain a ``tasks.db`` task database and task + payloads under ``tasks/`` created via OpenFE orchestration setup. + """ + worker_main(warehouse_path=warehouse_path, scratch=scratch) + + +PLUGIN = OFECommandPlugin(command=worker, section="Quickrun Executor", requires_ofe=(0, 3)) diff --git a/src/openfecli/parameters/__init__.py b/src/openfecli/parameters/__init__.py index c25ff4fcb..db3814c33 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, PROTEIN_MEMBRANE +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 0b50da135..eea016f94 100644 --- a/src/openfecli/plan_alchemical_networks_utils.py +++ b/src/openfecli/plan_alchemical_networks_utils.py @@ -3,8 +3,12 @@ from __future__ import annotations import pathlib +from pathlib import Path +from typing import Optional from openfe import AlchemicalNetwork, LigandNetwork +from openfe.orchestration.exorcist_utils import build_task_db_from_alchemical_network +from openfe.storage.warehouse import FileSystemWarehouse from openfecli.utils import write @@ -12,26 +16,32 @@ def plan_alchemical_network_output( alchemical_network: AlchemicalNetwork, ligand_network: LigandNetwork, folder_path: pathlib.Path, + warehouse: Optional[FileSystemWarehouse] = None, ): """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) + db_path = Path(warehouse.root_dir) / "tasks.db" + _ = build_task_db_from_alchemical_network(alchemical_network, warehouse, db_path) + 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) diff --git a/src/openfecli/tests/commands/test_worker.py b/src/openfecli/tests/commands/test_worker.py new file mode 100644 index 000000000..b060b57e2 --- /dev/null +++ b/src/openfecli/tests/commands/test_worker.py @@ -0,0 +1,143 @@ +from pathlib import Path +from unittest import mock + +from click.testing import CliRunner + +from openfecli.commands.worker import worker + + +class _SuccessfulResult: + def ok(self): + return True + + +class _FailedResult: + def ok(self): + return False + + +class _FailedResultWithDetails: + source_key = "HybridTopologyMultiStateSimulationUnit-deadbeef" + exception = ("RuntimeError", ("simulation blew up",)) + traceback = 'Traceback (most recent call last):\n File "sim.py", line 1\nRuntimeError: simulation blew up' + + def ok(self): + return False + + +def test_worker_requires_task_database(): + runner = CliRunner() + with runner.isolated_filesystem(): + Path("warehouse").mkdir() + result = runner.invoke(worker, ["warehouse"]) + assert result.exit_code == 1 + assert "Task database not found at" in result.output + + +def test_worker_no_available_task_exits_zero(): + runner = CliRunner() + with runner.isolated_filesystem(): + warehouse_path = Path("warehouse") + warehouse_path.mkdir() + (warehouse_path / "tasks.db").touch() + + mock_worker = mock.Mock() + mock_worker.execute_unit.return_value = None + + with mock.patch( + "openfecli.commands.worker._build_worker", return_value=mock_worker + ) as build_worker: + result = runner.invoke(worker, ["warehouse"]) + + assert result.exit_code == 0 + assert "No available task in task graph." in result.output + build_worker.assert_called_once_with(warehouse_path, warehouse_path / "tasks.db") + kwargs = mock_worker.execute_unit.call_args.kwargs + assert kwargs["scratch"] == Path.cwd() + + +def test_worker_executes_one_task_and_reports_completion(): + runner = CliRunner() + with runner.isolated_filesystem(): + warehouse_path = Path("warehouse") + warehouse_path.mkdir() + (warehouse_path / "tasks.db").touch() + + mock_worker = mock.Mock() + mock_worker.execute_unit.return_value = ( + "Transformation-abc:ProtocolUnit-def", + _SuccessfulResult(), + ) + + with mock.patch("openfecli.commands.worker._build_worker", return_value=mock_worker): + result = runner.invoke(worker, ["warehouse", "--scratch", "scratch"]) + + assert result.exit_code == 0 + assert "Completed task: Transformation-abc:ProtocolUnit-def" in result.output + assert Path("scratch").is_dir() + kwargs = mock_worker.execute_unit.call_args.kwargs + assert kwargs["scratch"] == Path("scratch") + + +def test_worker_raises_when_result_is_failure(): + runner = CliRunner() + with runner.isolated_filesystem(): + warehouse_path = Path("warehouse") + warehouse_path.mkdir() + (warehouse_path / "tasks.db").touch() + + mock_worker = mock.Mock() + mock_worker.execute_unit.return_value = ( + "Transformation-abc:ProtocolUnit-def", + _FailedResult(), + ) + + with mock.patch("openfecli.commands.worker._build_worker", return_value=mock_worker): + result = runner.invoke(worker, ["warehouse"]) + + assert result.exit_code == 1 + assert "returned a failure result" in result.output + + +def test_worker_prints_failure_result_details_when_available(): + runner = CliRunner() + with runner.isolated_filesystem(): + warehouse_path = Path("warehouse") + warehouse_path.mkdir() + (warehouse_path / "tasks.db").touch() + + mock_worker = mock.Mock() + mock_worker.execute_unit.return_value = ( + "Transformation-abc:ProtocolUnit-def", + _FailedResultWithDetails(), + ) + + with mock.patch("openfecli.commands.worker._build_worker", return_value=mock_worker): + result = runner.invoke(worker, ["warehouse"]) + + assert result.exit_code == 1 + assert ( + "Failed unit source key: HybridTopologyMultiStateSimulationUnit-deadbeef" + in result.output + ) + assert "Protocol unit exception: RuntimeError: ('simulation blew up',)" in result.output + assert "Protocol unit traceback:" in result.output + + +def test_worker_raises_when_execution_throws(): + runner = CliRunner() + with runner.isolated_filesystem(): + warehouse_path = Path("warehouse") + warehouse_path.mkdir() + (warehouse_path / "tasks.db").touch() + + mock_worker = mock.Mock() + mock_worker.execute_unit.side_effect = RuntimeError("boom") + + with mock.patch("openfecli.commands.worker._build_worker", return_value=mock_worker): + result = runner.invoke(worker, ["warehouse"]) + + assert result.exit_code == 1 + assert "Traceback (most recent call last):" in result.output + assert "RuntimeError: boom" in result.output + assert "Task execution failed: boom" in result.output