From b960536264300a7282619060a487023431a99ee8 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 17 Dec 2025 15:24:50 -0700 Subject: [PATCH 1/5] feat(warehouse): add result tokenizable store --- openfe/storage/warehouse.py | 48 +++++++++++++++++++++++++++---------- 1 file changed, 36 insertions(+), 12 deletions(-) diff --git a/openfe/storage/warehouse.py b/openfe/storage/warehouse.py index 02a81d5c4..b80ab6121 100644 --- a/openfe/storage/warehouse.py +++ b/openfe/storage/warehouse.py @@ -32,7 +32,7 @@ class WarehouseStores(TypedDict): """ setup: ExternalStorage - # We will add a result and task store here in the future. + result: ExternalStorage class WarehouseBaseClass: @@ -63,7 +63,7 @@ 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"], location: str): + def delete(self, store_name: Literal["setup", "result"], location: str): """Delete an object from a specific store. Parameters @@ -106,6 +106,31 @@ def load_setup_tokenizable(self, obj: GufeKey) -> GufeTokenizable: """ return self._load_gufe_tokenizable(gufe_key=obj) + def store_result_tokenizable(self, obj: GufeKey) -> 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. @@ -144,15 +169,15 @@ def _get_store_for_key(self, key: GufeKey) -> ExternalStorage: return self.stores[name] raise ValueError(f"GufeKey {key} is not stored") - def _store_gufe_tokenizable(self, store_name: Literal["setup"], obj: GufeTokenizable): + 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. + Parameters + ---------- + store_name : Literal["setup"] + Name of the store to store the object in. + obj : GufeTokenizable + The object to store. Notes ----- @@ -276,7 +301,6 @@ class FileSystemWarehouse(WarehouseBaseClass): def __init__(self, root_dir: str = "warehouse"): setup_store = FileStorage(f"{root_dir}/setup") - # When we add a result store it will look like this - # result_store = FileStorage(f"{root_dir}/results") - stores = WarehouseStores(setup=setup_store) + result_store = FileStorage(f"{root_dir}/result") + stores = WarehouseStores(setup=setup_store, result=result_store) super().__init__(stores) From dbff24f8b2f6a68faf5441a8c1b9ef0a8c5d3538 Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Wed, 17 Dec 2025 15:26:52 -0700 Subject: [PATCH 2/5] docs(warehouse): add docs on the result store --- openfe/storage/warehouse.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/openfe/storage/warehouse.py b/openfe/storage/warehouse.py index b80ab6121..45c311da9 100644 --- a/openfe/storage/warehouse.py +++ b/openfe/storage/warehouse.py @@ -25,6 +25,8 @@ class WarehouseStores(TypedDict): ---------- setup : ExternalStorage Storage location for setup-related objects and configurations. + result : ExternalStorage + Storage location for result-related object. Notes ----- From 10dae747914fb1a3ee7c42ef82c3973fef3c7baa Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Fri, 19 Dec 2025 14:58:38 -0700 Subject: [PATCH 3/5] test: cleanup warehouse test to be more modular --- openfe/tests/storage/test_warehouse.py | 101 ++++++++++++++----------- 1 file changed, 56 insertions(+), 45 deletions(-) diff --git a/openfe/tests/storage/test_warehouse.py b/openfe/tests/storage/test_warehouse.py index 850ae172a..572769d70 100644 --- a/openfe/tests/storage/test_warehouse.py +++ b/openfe/tests/storage/test_warehouse.py @@ -1,6 +1,7 @@ import os import tempfile from pathlib import Path +from typing import Literal from unittest import mock import pytest @@ -19,28 +20,42 @@ 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() - stores = WarehouseStores(setup=store) + 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 store._data == {} + assert setup_store._data == {} + assert result_store._data == {} store_func(obj) - assert store._data != {} - reloaded = load_func(obj.key) + 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 = MemoryStorage() - stores = WarehouseStores(setup=store) + 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 store._data == {} + assert setup_store._data == {} + assert result_store._data == {} store_func(obj) - assert store._data != {} + 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 @@ -54,60 +69,56 @@ def _test_store_load_different_process(obj: GufeTokenizable, store_func_name, lo "fixture", ["absolute_transformation", "complex_equilibrium"], ) - def test_store_load_transformation_same_process(self, request, fixture): + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_store_load_transformation_same_process(self, request, fixture, store): transformation = request.getfixturevalue(fixture) - self._test_store_load_same_process( - transformation, - "store_setup_tokenizable", - "load_setup_tokenizable", - ) + 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"], ) - def test_store_load_transformation_different_process(self, request, fixture): + @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_setup_tokenizable", - "load_setup_tokenizable", + transformation, store_func_name, load_func_name, store ) # @pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) - def test_store_load_network_same_process(self, request, fixture): + @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) - self._test_store_load_same_process( - network, "store_setup_tokenizable", "load_setup_tokenizable" - ) + 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"]) - def test_store_load_network_different_process(self, request, fixture): + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_store_load_network_different_process(self, request, fixture, store): network = request.getfixturevalue(fixture) - self._test_store_load_different_process( - network, "store_setup_tokenizable", "load_setup_tokenizable" - ) + 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"]) - def test_delete(self, request, fixture): - store = MemoryStorage() - stores = WarehouseStores(setup=store) - client = WarehouseBaseClass(stores) - + @pytest.mark.parametrize("store", ["setup", "result"]) + def test_delete(self, request, fixture, store): network = request.getfixturevalue(fixture) - assert store._data == {} - client.store_setup_tokenizable(network) - assert store._data != {} - key = network.key - loaded = client.load_setup_tokenizable(key) - assert loaded is network - assert client.setup_store.exists(key) - client.delete("setup", key) - assert not client.exists(key) + 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: From 47b49abb83ee0c7c837bfd229471d97bd684fb9b Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Fri, 19 Dec 2025 15:05:30 -0700 Subject: [PATCH 4/5] chore: add property for result store --- openfe/storage/warehouse.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/openfe/storage/warehouse.py b/openfe/storage/warehouse.py index 45c311da9..f0fb04b65 100644 --- a/openfe/storage/warehouse.py +++ b/openfe/storage/warehouse.py @@ -273,15 +273,26 @@ def recursive_build_object_cache(key: GufeKey) -> GufeTokenizable: @property def setup_store(self): - """Get the setup store. + """Get the setup store Returns ------- ExternalStorage - The setup storage location. + 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. From 09f546c385c03eaf2c45f29e3b573c24b5e319bd Mon Sep 17 00:00:00 2001 From: Ethan Holz Date: Mon, 29 Dec 2025 10:58:53 -0700 Subject: [PATCH 5/5] fix: use the correct function signature for handling setup --- openfe/storage/warehouse.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openfe/storage/warehouse.py b/openfe/storage/warehouse.py index f0fb04b65..c10baa12c 100644 --- a/openfe/storage/warehouse.py +++ b/openfe/storage/warehouse.py @@ -108,7 +108,7 @@ def load_setup_tokenizable(self, obj: GufeKey) -> GufeTokenizable: """ return self._load_gufe_tokenizable(gufe_key=obj) - def store_result_tokenizable(self, obj: GufeKey) -> GufeTokenizable: + def store_result_tokenizable(self, obj: GufeTokenizable): """Store a GufeTokenizable object from the result store. Parameters