From f278d4fb8a56af128219983c787a4038bd0f3d28 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:58:29 -0700 Subject: [PATCH 1/3] epic: add warehouse (#1864) * refactor: remove result_server references in favor of gufe ExternalStorage (#1632) * feat: remove result_server references in favor of gufe external_storage Signed-off-by: Ethan Holz * fix: remove extra external_store * test: change from result_server to external_storage * test: remove external_store reference * test: change extension to item per class changes * refactor: remove result server import * chore: remove metadata store from resultclient Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * refactor: remove result server * refactor: remove metadata store --------- Signed-off-by: Ethan Holz Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * refactor: rename to WarehouseBaseClass (#1660) * refactor(warehouse): rename to warehouse * refactor: rename _ResultContainerTest to _ContainerTest * refactor: rename _Container instances to _DataContainer * feat/refactor: add setup store (#1671) * refactor: remove extra implementations We may end up adding these back later but for the time being, we are going to remove them to keep life easy. * refactor: remove load_* and store_* functions Signed-off-by: Ethan Holz * refactor: keys are being stored really strangely * refactor(warehouse): huge refactor to simplify working with Warehouse * refactor(warehouse): remove dead code * fix: delete failed due to incorrect object * test: add a test for the FileSystemWarehouse * refactor: remove dead code * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * docs: update warehouse docstrings * feat(warehouse): leverage keyed_chain for object dedup * Update openfe/storage/warehouse.py * import Literal * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * docs(warehouse): change from return, to raises * refactor(warehouse): move _key_exists to exists * chore(warehouse): remove extra todo * chore(warehouse): remove _load_stream * fix(warehouse): deduplicate objects on the filesystem * chore(warehouse): add typehint to exists --------- Signed-off-by: Ethan Holz Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> Co-authored-by: Alyssa Travitz * Temporarily build pooch from main w/ hotfix (#1806) * build with pooch@main to see if hotfix works * add link * Add support for result tokenizables to warehouse (#1763) * feat(warehouse): add result tokenizable store * docs(warehouse): add docs on the result store * chore: add property for result store * fix: use the correct function signature for handling setup --------- Co-authored-by: Alyssa Travitz * test: cleanup warehouse test to be more modular (#1809) * fix dict type * add news item * run mypy CI on PRs into all branches (not just main) (#2112) * ignore a couple type check violations for now --------- Signed-off-by: Ethan Holz Co-authored-by: Ethan Holz --- news/warehouse.rst | 23 ++ src/openfe/storage/metadatastore.py | 100 ------ src/openfe/storage/resultclient.py | 287 ---------------- src/openfe/storage/resultserver.py | 60 ---- src/openfe/storage/warehouse.py | 321 ++++++++++++++++++ .../tests/storage/test_metadatastore.py | 147 -------- src/openfe/tests/storage/test_resultclient.py | 304 ----------------- src/openfe/tests/storage/test_resultserver.py | 117 ------- src/openfe/tests/storage/test_warehouse.py | 177 ++++++++++ 9 files changed, 521 insertions(+), 1015 deletions(-) create mode 100644 news/warehouse.rst delete mode 100644 src/openfe/storage/metadatastore.py delete mode 100644 src/openfe/storage/resultclient.py delete mode 100644 src/openfe/storage/resultserver.py create mode 100644 src/openfe/storage/warehouse.py delete mode 100644 src/openfe/tests/storage/test_metadatastore.py delete mode 100644 src/openfe/tests/storage/test_resultclient.py delete mode 100644 src/openfe/tests/storage/test_resultserver.py create mode 100644 src/openfe/tests/storage/test_warehouse.py diff --git a/news/warehouse.rst b/news/warehouse.rst new file mode 100644 index 000000000..8e7d5dc20 --- /dev/null +++ b/news/warehouse.rst @@ -0,0 +1,23 @@ +**Added:** + +* Added ``openfe.storage.warehouse``, which includes ``WarehouseBaseClass`` and the example implementation ``FileSystemWarehouse``. + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* Removed the unused methods ``metadatastore``, ``resultclient``, and ``resultserver`` from ``openfe.storage``. + +**Fixed:** + +* + +**Security:** + +* 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..eb91efc0e --- /dev/null +++ b/src/openfe/storage/warehouse.py @@ -0,0 +1,321 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/gufe +import json +import re +from typing import Literal, TypedDict + +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 : ExternalStorage + Storage location for setup-related objects and configurations. + result : ExternalStorage + Storage location for result-related object. + + Notes + ----- + Additional stores for results and tasks may be added in future versions. + """ + + setup: ExternalStorage + result: 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): + self.stores = stores + + 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 + Thrown if the object you are trying to delete, can't delete from the store + """ + store: ExternalStorage = self.stores[store_name] + store.delete(location) + + 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: + """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: + """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 exists(self, key: GufeKey) -> bool: + """Check if an object with the given key exists in any store. + + 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"], obj: GufeTokenizable): + """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" + ) + 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) + + @property + def setup_store(self): + """Get the setup store + + Returns + ------- + ExternalStorage + The setup storage location + """ + return self.stores["setup"] + + @property + def result_store(self): + """Get the result store. + + Returns + ------- + ExternalStorage + The result storage location + """ + return self.stores["result"] + + +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) 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..572769d70 --- /dev/null +++ b/src/openfe/tests/storage/test_warehouse.py @@ -0,0 +1,177 @@ +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 _test_store_load_same_process( + obj, store_func_name, load_func_name, store_name: Literal["setup", "result"] + ): + 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 setup_store._data == {} + assert result_store._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"], + ): + 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 setup_store._data == {} + assert result_store._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 + + @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", + ) + + @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", + ) From 1f6ae4d496f387ff6583df43e37986804c138857 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Thu, 30 Jul 2026 17:13:11 -0600 Subject: [PATCH 2/3] feat: filesystem exorcist worker (#1833) * refactor: remove result_server references in favor of gufe ExternalStorage (#1632) * feat: remove result_server references in favor of gufe external_storage Signed-off-by: Ethan Holz * fix: remove extra external_store * test: change from result_server to external_storage * test: remove external_store reference * test: change extension to item per class changes * refactor: remove result server import * chore: remove metadata store from resultclient Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * refactor: remove result server * refactor: remove metadata store --------- Signed-off-by: Ethan Holz Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * refactor: rename to WarehouseBaseClass (#1660) * refactor(warehouse): rename to warehouse * refactor: rename _ResultContainerTest to _ContainerTest * refactor: rename _Container instances to _DataContainer * feat/refactor: add setup store (#1671) * refactor: remove extra implementations We may end up adding these back later but for the time being, we are going to remove them to keep life easy. * refactor: remove load_* and store_* functions Signed-off-by: Ethan Holz * refactor: keys are being stored really strangely * refactor(warehouse): huge refactor to simplify working with Warehouse * refactor(warehouse): remove dead code * fix: delete failed due to incorrect object * test: add a test for the FileSystemWarehouse * refactor: remove dead code * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * docs: update warehouse docstrings * feat(warehouse): leverage keyed_chain for object dedup * Update openfe/storage/warehouse.py * import Literal * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * Update openfe/storage/warehouse.py Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> * docs(warehouse): change from return, to raises * refactor(warehouse): move _key_exists to exists * chore(warehouse): remove extra todo * chore(warehouse): remove _load_stream * fix(warehouse): deduplicate objects on the filesystem * chore(warehouse): add typehint to exists --------- Signed-off-by: Ethan Holz Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> Co-authored-by: Alyssa Travitz * Temporarily build pooch from main w/ hotfix (#1806) * build with pooch@main to see if hotfix works * add link * Add support for result tokenizables to warehouse (#1763) * feat(warehouse): add result tokenizable store * docs(warehouse): add docs on the result store * chore: add property for result store * fix: use the correct function signature for handling setup --------- Co-authored-by: Alyssa Travitz * test: cleanup warehouse test to be more modular (#1809) * feat: add warehouse primitives for handling protocol units * feat: inital worker for exorcist * test: add tests for warehouse * fix: can now return protocol unit * refactor: make things more consistent * test: initial test setup for orchestration subpackage * test: initial exorcist utility testing * refactor: provide a root path to the exorcist DB * test: inital worker testing * feat: add shared_store * feat: add better handling for CLI application Signed-off-by: Ethan Holz * test: add new worker tests * feat: add exorcist worker to CLI * test: add for worker CLI command * docs: add numpy docstrings * feat: add support for planning an RBFE to a Warehouse * fix: correct edge direction for task graph * refactor: remove extra debugging from warehouse * refactor: cleanup handling of tasks for worker * refactor: fix issues in the CLI for running the worker * Update src/openfe/orchestration/__init__.py * fix dict type * add news item * ignore a couple type check violations for now * fix type checking * add exorcist to docs build --------- Signed-off-by: Ethan Holz Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> Co-authored-by: Alyssa Travitz --- docs/environment.yaml | 2 + environment.yml | 1 + src/openfe/orchestration/__init__.py | 244 ++++++++++++++++ src/openfe/orchestration/exorcist_utils.py | 95 ++++++ src/openfe/storage/warehouse.py | 45 ++- src/openfe/tests/orchestration/__init__.py | 2 + src/openfe/tests/orchestration/conftest.py | 118 ++++++++ .../orchestration/test_exorcist_utils.py | 226 +++++++++++++++ src/openfe/tests/orchestration/test_worker.py | 274 ++++++++++++++++++ src/openfe/tests/storage/test_warehouse.py | 96 +++++- src/openfecli/commands/plan_rbfe_network.py | 10 +- src/openfecli/commands/worker.py | 133 +++++++++ src/openfecli/parameters/__init__.py | 1 + src/openfecli/parameters/warehouse.py | 4 + .../plan_alchemical_networks_utils.py | 50 ++-- src/openfecli/tests/commands/test_worker.py | 143 +++++++++ 16 files changed, 1407 insertions(+), 37 deletions(-) create mode 100644 src/openfe/orchestration/exorcist_utils.py create mode 100644 src/openfe/tests/orchestration/__init__.py create mode 100644 src/openfe/tests/orchestration/conftest.py create mode 100644 src/openfe/tests/orchestration/test_exorcist_utils.py create mode 100644 src/openfe/tests/orchestration/test_worker.py create mode 100644 src/openfecli/commands/worker.py create mode 100644 src/openfecli/parameters/warehouse.py create mode 100644 src/openfecli/tests/commands/test_worker.py 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/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/src/openfe/orchestration/__init__.py b/src/openfe/orchestration/__init__.py index e69de29bb..a290e945c 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 "Transformation-:ProtocolUnit-" + taskid = db.check_out_task() + if taskid is None: + return None + + _, protocol_unit_key = taskid.split(":", maxsplit=1) + 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..414f5da8c --- /dev/null +++ b/src/openfe/orchestration/exorcist_utils.py @@ -0,0 +1,95 @@ +"""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 +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 = f"{str(transformation.key)}:{str(unit.key)}" + global_dag.add_node( + node_id, + ) + warehouse.store_task(unit) + for dependent_unit, dependency_unit in dag.graph.edges: + upstream_id = f"{str(transformation.key)}:{str(dependency_unit.key)}" + downstream_id = f"{str(transformation.key)}:{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 + + +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 ``Path("tasks.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("tasks.db") + + global_dag = 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/warehouse.py b/src/openfe/storage/warehouse.py index eb91efc0e..445a77144 100644 --- a/src/openfe/storage/warehouse.py +++ b/src/openfe/storage/warehouse.py @@ -4,6 +4,8 @@ import re from typing import 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, @@ -34,6 +36,8 @@ class WarehouseStores(TypedDict): setup: ExternalStorage result: ExternalStorage + shared: ExternalStorage + tasks: ExternalStorage class WarehouseBaseClass: @@ -82,6 +86,15 @@ def delete(self, store_name: Literal["setup", "result"], location: str): 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. @@ -133,7 +146,7 @@ def load_result_tokenizable(self, obj: GufeKey) -> GufeTokenizable: return self._load_gufe_tokenizable(gufe_key=obj) def exists(self, key: GufeKey) -> bool: - """Check if an object with the given key exists in any store. + """Check if an object with the given key exists in any store that holds tokenizables. Parameters ---------- @@ -173,7 +186,12 @@ def _get_store_for_key(self, key: GufeKey) -> ExternalStorage: return self.stores[name] # type: ignore raise ValueError(f"GufeKey {key} is not stored") - def _store_gufe_tokenizable(self, store_name: Literal["setup", "result"], obj: GufeTokenizable): + def _store_gufe_tokenizable( + self, + store_name: Literal["setup", "result", "tasks"], + obj: GufeTokenizable, + name: str | None = None, + ): """Store a GufeTokenizable object with deduplication. Parameters @@ -199,7 +217,10 @@ def _store_gufe_tokenizable(self, store_name: Literal["setup", "result"], obj: G data = json.dumps(keyed_dict, cls=JSON_HANDLER.encoder, sort_keys=True).encode( "utf-8" ) - target.store_bytes(gufe_key, data) + 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. @@ -295,6 +316,17 @@ def result_store(self): """ 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. @@ -315,7 +347,12 @@ class FileSystemWarehouse(WarehouseBaseClass): """ def __init__(self, root_dir: str = "warehouse"): + self.root_dir = root_dir setup_store = FileStorage(f"{root_dir}/setup") result_store = FileStorage(f"{root_dir}/result") - stores = WarehouseStores(setup=setup_store, result=result_store) + shared_store = FileStorage(f"{root_dir}/shared") + tasks_store = FileStorage(f"{root_dir}/tasks") + stores = WarehouseStores( + setup=setup_store, result=result_store, shared=shared_store, tasks=tasks_store + ) super().__init__(stores) 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..22413c771 --- /dev/null +++ b/src/openfe/tests/orchestration/test_exorcist_utils.py @@ -0,0 +1,226 @@ +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 _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)) + + transformation_keys = {str(transformation.key) for transformation in network.edges} + expected_protocol_unit_keys = sorted(str(unit.key) for unit in warehouse.stored_tasks) + observed_protocol_unit_keys = [] + + for node in graph.nodes: + transformation_key, protocol_unit_key = node.split(":", maxsplit=1) + assert transformation_key in transformation_keys + 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: + _, upstream_key = upstream_id.split(":", maxsplit=1) + _, downstream_key = downstream_id.split(":", maxsplit=1) + 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.split(":", maxsplit=1) + 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("tasks.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..98edcc8e2 --- /dev/null +++ b/src/openfe/tests/orchestration/test_worker.py @@ -0,0 +1,274 @@ +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 = f"{absolute_transformation.key}:{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.split(":", maxsplit=1)[1] 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.endswith(f":{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.split(":", maxsplit=1)[1] 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.endswith(f":{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) + + transformation_key = "Transformation-toy" + first_taskid = f"{transformation_key}:{first_unit.key}" + second_taskid = f"{transformation_key}:{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 = f"Transformation-toy:{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 = f"Transformation-toy:{first_unit.key}" + second_taskid = f"Transformation-toy:{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_warehouse.py b/src/openfe/tests/storage/test_warehouse.py index 572769d70..d113cf1d7 100644 --- a/src/openfe/tests/storage/test_warehouse.py +++ b/src/openfe/tests/storage/test_warehouse.py @@ -19,18 +19,35 @@ 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(), + ) + + @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"] + obj, + store_func_name, + load_func_name, + store_name: Literal["setup", "result", "tasks"], ): - setup_store = MemoryStorage() - result_store = MemoryStorage() - stores = WarehouseStores(setup=setup_store, result=result_store) + stores = TestWarehouseBaseClass._build_stores() client = WarehouseBaseClass(stores) store_func = getattr(client, store_func_name) load_func = getattr(client, load_func_name) - assert setup_store._data == {} - assert result_store._data == {} + 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 != {} @@ -43,16 +60,16 @@ def _test_store_load_different_process( obj: GufeTokenizable, store_func_name, load_func_name, - store_name: Literal["setup", "result"], + store_name: Literal["setup", "result", "tasks"], ): - setup_store = MemoryStorage() - result_store = MemoryStorage() - stores = WarehouseStores(setup=setup_store, result=result_store) + stores = TestWarehouseBaseClass._build_stores() client = WarehouseBaseClass(stores) store_func = getattr(client, store_func_name) load_func = getattr(client, load_func_name) - assert setup_store._data == {} - assert result_store._data == {} + 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 != {} @@ -65,6 +82,45 @@ def _test_store_load_different_process( 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) + 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) + + 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) + + 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"], @@ -164,6 +220,22 @@ def test_store_load_transformation_same_process(self, request, fixture): "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"], 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 From cc8edd500077932e1685c01edfdd4644dd7576f5 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz <31974495+atravitz@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:28:12 -0700 Subject: [PATCH 3/3] updates to filesystem warehouse worker-based execution (#2153) * only use protocolunit as key * add specific protocol dag store * add get_protocol_dags * fix protocol dag store deduplication order * type checking and require a name * default task db to use warehouse name * update docstrings * update type hint for python 3.12 --- src/openfe/orchestration/__init__.py | 4 +- src/openfe/orchestration/exorcist_utils.py | 21 +++-- src/openfe/storage/warehouse.py | 91 ++++++++++++++++--- .../orchestration/test_exorcist_utils.py | 21 +++-- src/openfe/tests/orchestration/test_worker.py | 21 ++--- src/openfe/tests/storage/test_warehouse.py | 11 ++- 6 files changed, 123 insertions(+), 46 deletions(-) diff --git a/src/openfe/orchestration/__init__.py b/src/openfe/orchestration/__init__.py index a290e945c..14999e2f2 100644 --- a/src/openfe/orchestration/__init__.py +++ b/src/openfe/orchestration/__init__.py @@ -160,12 +160,12 @@ def _checkout_task(self) -> tuple[TaskStatusDB, str, ProtocolUnit] | None: """ db: TaskStatusDB = TaskStatusDB.from_filename(self.task_db_path) - # The format for the taskid is "Transformation-:ProtocolUnit-" + # The format for the taskid is "ProtocolUnit-" taskid = db.check_out_task() if taskid is None: return None - _, protocol_unit_key = taskid.split(":", maxsplit=1) + protocol_unit_key = taskid unit = self.warehouse.load_task(GufeKey(protocol_unit_key)) return db, taskid, unit diff --git a/src/openfe/orchestration/exorcist_utils.py b/src/openfe/orchestration/exorcist_utils.py index 414f5da8c..459877b9c 100644 --- a/src/openfe/orchestration/exorcist_utils.py +++ b/src/openfe/orchestration/exorcist_utils.py @@ -8,6 +8,7 @@ import exorcist import networkx as nx +import pandas as pd from gufe import AlchemicalNetwork from openfe.storage.warehouse import WarehouseBaseClass @@ -43,14 +44,15 @@ def alchemical_network_to_task_graph( for transformation in alchemical_network.edges: dag = transformation.create() for unit in dag.protocol_units: - node_id = f"{str(transformation.key)}:{str(unit.key)}" - global_dag.add_node( - node_id, - ) + 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 = f"{str(transformation.key)}:{str(dependency_unit.key)}" - downstream_id = f"{str(transformation.key)}:{str(dependent_unit.key)}" + 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): @@ -59,6 +61,7 @@ def alchemical_network_to_task_graph( 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, @@ -75,7 +78,7 @@ def build_task_db_from_alchemical_network( 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 ``Path("tasks.db")`` in the current working directory. + 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``. @@ -87,9 +90,9 @@ def build_task_db_from_alchemical_network( edges derived from ``alchemical_network``. """ if db_path is None: - db_path = Path("tasks.db") + db_path = Path(f"{warehouse.name}.db") - global_dag = alchemical_network_to_task_graph(alchemical_network, warehouse) + 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/warehouse.py b/src/openfe/storage/warehouse.py index 445a77144..6f8acb537 100644 --- a/src/openfe/storage/warehouse.py +++ b/src/openfe/storage/warehouse.py @@ -1,8 +1,9 @@ # 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 Literal, TypedDict +from typing import Generator, Literal, TypedDict from gufe.protocols.protocoldag import ProtocolDAG from gufe.protocols.protocolunit import ProtocolUnit @@ -28,6 +29,12 @@ class WarehouseStores(TypedDict): Storage location for setup-related objects and configurations. result : 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 ----- @@ -38,6 +45,7 @@ class WarehouseStores(TypedDict): result: ExternalStorage shared: ExternalStorage tasks: ExternalStorage + protocol_dags: ExternalStorage class WarehouseBaseClass: @@ -58,8 +66,11 @@ class WarehouseBaseClass: The storage locations managed by this warehouse instance. """ - def __init__(self, stores: WarehouseStores): + 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 @@ -106,6 +117,7 @@ def store_setup_tokenizable(self, obj: GufeTokenizable): 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 @@ -131,6 +143,7 @@ def store_result_tokenizable(self, obj: GufeTokenizable): 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 @@ -145,6 +158,38 @@ def load_result_tokenizable(self, obj: GufeKey) -> GufeTokenizable: """ 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. @@ -188,7 +233,7 @@ def _get_store_for_key(self, key: GufeKey) -> ExternalStorage: def _store_gufe_tokenizable( self, - store_name: Literal["setup", "result", "tasks"], + store_name: Literal["setup", "result", "tasks", "protocol_dags"], obj: GufeTokenizable, name: str | None = None, ): @@ -294,6 +339,23 @@ def recursive_build_object_cache(key: GufeKey) -> GufeTokenizable: 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): """Get the setup store @@ -346,13 +408,20 @@ class FileSystemWarehouse(WarehouseBaseClass): for results and other data types. """ - def __init__(self, root_dir: str = "warehouse"): - self.root_dir = root_dir - setup_store = FileStorage(f"{root_dir}/setup") - result_store = FileStorage(f"{root_dir}/result") - shared_store = FileStorage(f"{root_dir}/shared") - tasks_store = FileStorage(f"{root_dir}/tasks") + 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 + setup=setup_store, + result=result_store, + shared=shared_store, + tasks=tasks_store, + protocol_dags=protocol_dag_store, ) - super().__init__(stores) + super().__init__(stores, name) diff --git a/src/openfe/tests/orchestration/test_exorcist_utils.py b/src/openfe/tests/orchestration/test_exorcist_utils.py index 22413c771..aea4df87f 100644 --- a/src/openfe/tests/orchestration/test_exorcist_utils.py +++ b/src/openfe/tests/orchestration/test_exorcist_utils.py @@ -22,6 +22,13 @@ def __init__(self): 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 = [] @@ -35,7 +42,6 @@ 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] @@ -52,13 +58,11 @@ def test_alchemical_network_to_task_graph_uses_canonical_task_ids(request, fixtu graph = alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) - transformation_keys = {str(transformation.key) for transformation in network.edges} expected_protocol_unit_keys = sorted(str(unit.key) for unit in warehouse.stored_tasks) observed_protocol_unit_keys = [] for node in graph.nodes: - transformation_key, protocol_unit_key = node.split(":", maxsplit=1) - assert transformation_key in transformation_keys + protocol_unit_key = node observed_protocol_unit_keys.append(protocol_unit_key) assert sorted(observed_protocol_unit_keys) == expected_protocol_unit_keys @@ -86,8 +90,9 @@ def test_alchemical_network_to_task_graph_edge_direction_matches_dependencies(re units_by_key = {str(unit.key): unit for unit in warehouse.stored_tasks} for upstream_id, downstream_id in graph.edges: - _, upstream_key = upstream_id.split(":", maxsplit=1) - _, downstream_key = downstream_id.split(":", maxsplit=1) + # 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 @@ -148,7 +153,7 @@ def test_build_task_db_checkout_order_is_dependency_safe(tmp_path, request, fixt break checkout_order.append(taskid) - _, protocol_unit_key = taskid.split(":", maxsplit=1) + 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) @@ -190,7 +195,7 @@ def test_build_task_db_default_path(request, fixture): 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("tasks.db")) + 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 diff --git a/src/openfe/tests/orchestration/test_worker.py b/src/openfe/tests/orchestration/test_worker.py index 98edcc8e2..221e8de54 100644 --- a/src/openfe/tests/orchestration/test_worker.py +++ b/src/openfe/tests/orchestration/test_worker.py @@ -71,7 +71,7 @@ def worker_with_executable_task_db(tmp_path, absolute_transformation): unit = _get_dependency_free_unit(absolute_transformation) warehouse.store_task(unit) - taskid = f"{absolute_transformation.key}:{unit.key}" + taskid = unit.key task_graph = nx.DiGraph() task_graph.add_node(taskid) @@ -94,24 +94,24 @@ def test_get_task_uses_default_db_path_without_patching( worker = Worker(warehouse=warehouse) taskid, loaded = worker._get_task() - expected_keys = {task_row.taskid.split(":", maxsplit=1)[1] for task_row in db.get_all_tasks()} + 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.endswith(f":{loaded.key}") + 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.split(":", maxsplit=1)[1] for task_id in task_ids} + 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.endswith(f":{loaded.key}") + assert taskid == loaded.key def test_execute_unit_stores_real_result(worker_with_executable_task_db, tmp_path): @@ -185,9 +185,8 @@ def test_execute_unit_resolves_dependency_results(tmp_path): warehouse.store_task(first_unit) warehouse.store_task(second_unit) - transformation_key = "Transformation-toy" - first_taskid = f"{transformation_key}:{first_unit.key}" - second_taskid = f"{transformation_key}:{second_unit.key}" + first_taskid = first_unit.key + second_taskid = second_unit.key task_graph = nx.DiGraph() task_graph.add_edge(first_taskid, second_taskid) @@ -220,7 +219,7 @@ def test_execute_unit_marks_missing_dependency_as_failed(tmp_path): dependent_unit = _ToyProtocolUnit(name="dependent", upstream=missing_upstream, increment=2) warehouse.store_task(dependent_unit) - taskid = f"Transformation-toy:{dependent_unit.key}" + taskid = dependent_unit.key task_graph = nx.DiGraph() task_graph.add_node(taskid) @@ -247,8 +246,8 @@ def test_execute_unit_uses_isolated_shared_workspace_per_task(tmp_path): warehouse.store_task(first_unit) warehouse.store_task(second_unit) - first_taskid = f"Transformation-toy:{first_unit.key}" - second_taskid = f"Transformation-toy:{second_unit.key}" + first_taskid = first_unit.key + second_taskid = second_unit.key task_graph = nx.DiGraph() task_graph.add_node(first_taskid) diff --git a/src/openfe/tests/storage/test_warehouse.py b/src/openfe/tests/storage/test_warehouse.py index d113cf1d7..d55709b5b 100644 --- a/src/openfe/tests/storage/test_warehouse.py +++ b/src/openfe/tests/storage/test_warehouse.py @@ -26,6 +26,7 @@ def _build_stores() -> WarehouseStores: result=MemoryStorage(), shared=MemoryStorage(), tasks=MemoryStorage(), + protocol_dags=MemoryStorage(), ) @staticmethod @@ -41,7 +42,7 @@ def _test_store_load_same_process( store_name: Literal["setup", "result", "tasks"], ): stores = TestWarehouseBaseClass._build_stores() - client = WarehouseBaseClass(stores) + client = WarehouseBaseClass(stores, "test_warehouse") store_func = getattr(client, store_func_name) load_func = getattr(client, load_func_name) assert stores["setup"]._data == {} @@ -63,7 +64,7 @@ def _test_store_load_different_process( store_name: Literal["setup", "result", "tasks"], ): stores = TestWarehouseBaseClass._build_stores() - client = WarehouseBaseClass(stores) + client = WarehouseBaseClass(stores, "test_warehouse") store_func = getattr(client, store_func_name) load_func = getattr(client, load_func_name) assert stores["setup"]._data == {} @@ -93,7 +94,7 @@ def test_store_load_task_different_process(self, absolute_transformation): 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) + client = WarehouseBaseClass(stores, name="test_warehouse") client.store_task(unit) assert stores["tasks"]._data != {} @@ -104,7 +105,7 @@ def test_store_task_writes_to_tasks_store(self, absolute_transformation): def test_exists_finds_task_key(self, absolute_transformation): unit = self._get_protocol_unit(absolute_transformation) stores = self._build_stores() - client = WarehouseBaseClass(stores) + client = WarehouseBaseClass(stores, "test_warehouse") client.store_task(unit) @@ -113,7 +114,7 @@ def test_exists_finds_task_key(self, absolute_transformation): def test_load_task_returns_object(self, absolute_transformation): unit = self._get_protocol_unit(absolute_transformation) stores = self._build_stores() - client = WarehouseBaseClass(stores) + client = WarehouseBaseClass(stores, name="test_warehouse") client.store_task(unit) loaded = client.load_task(unit.key)