From 158fe31697c520c9e51e65299efe3104a3bb0495 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 13 Aug 2026 14:21:50 -0700 Subject: [PATCH 1/4] warehouse loading --- src/openfe/storage/warehouse.py | 28 ++++++++++++++++------ src/openfe/tests/storage/test_warehouse.py | 19 ++++++++------- 2 files changed, 32 insertions(+), 15 deletions(-) diff --git a/src/openfe/storage/warehouse.py b/src/openfe/storage/warehouse.py index 6f8acb537..87ca01620 100644 --- a/src/openfe/storage/warehouse.py +++ b/src/openfe/storage/warehouse.py @@ -1,5 +1,7 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/gufe +from __future__ import annotations + import json import pathlib import re @@ -398,8 +400,8 @@ class FileSystemWarehouse(WarehouseBaseClass): Parameters ---------- - root_dir : str, optional - Root directory for the warehouse storage, by default "warehouse". + root_dir : pathlib.Path, optional + Root directory in which to create the warehouse storage. Notes ----- @@ -408,14 +410,16 @@ class FileSystemWarehouse(WarehouseBaseClass): for results and other data types. """ - def __init__(self, name): - # TODO: should name and location be different? - self.root_dir = pathlib.Path(f"{name}") + def __init__(self, root_dir: pathlib.Path): + self.root_dir = pathlib.Path(root_dir) + if self.root_dir.is_dir(): + raise ValueError( + "`root_dir` already exists. To load an existing Warehouse, use FileSystemWarehouse.load(`root_dir`)" + ) 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, @@ -424,4 +428,14 @@ def __init__(self, name): tasks=tasks_store, protocol_dags=protocol_dag_store, ) - super().__init__(stores, name) + name = self.root_dir.resolve().name + super().__init__(stores=stores, name=name) + + @classmethod + def load(cls, root_dir: pathlib.Path) -> FileSystemWarehouse: + root_dir = pathlib.Path(root_dir) + if not root_dir.is_dir(): + raise ValueError( + "`root_dir` must be an existing filepath. To create a new Warehouse, use FileSystemWarehouse(`root_dir`)" + ) + return cls(root_dir=root_dir) diff --git a/src/openfe/tests/storage/test_warehouse.py b/src/openfe/tests/storage/test_warehouse.py index d55709b5b..acb671507 100644 --- a/src/openfe/tests/storage/test_warehouse.py +++ b/src/openfe/tests/storage/test_warehouse.py @@ -1,4 +1,3 @@ -import os import tempfile from pathlib import Path from typing import Literal @@ -42,7 +41,7 @@ def _test_store_load_same_process( store_name: Literal["setup", "result", "tasks"], ): stores = TestWarehouseBaseClass._build_stores() - client = WarehouseBaseClass(stores, "test_warehouse") + client = WarehouseBaseClass(stores=stores, name="test_warehouse") store_func = getattr(client, store_func_name) load_func = getattr(client, load_func_name) assert stores["setup"]._data == {} @@ -64,7 +63,7 @@ def _test_store_load_different_process( store_name: Literal["setup", "result", "tasks"], ): stores = TestWarehouseBaseClass._build_stores() - client = WarehouseBaseClass(stores, "test_warehouse") + client = WarehouseBaseClass(stores=stores, name="test_warehouse") store_func = getattr(client, store_func_name) load_func = getattr(client, load_func_name) assert stores["setup"]._data == {} @@ -182,24 +181,26 @@ class TestFileSystemWarehouse: @staticmethod def _test_store_load_same_process(obj, store_func_name, load_func_name): with tempfile.TemporaryDirectory() as tmpdir: - client = FileSystemWarehouse(tmpdir) + wh_dir = Path(tmpdir) / "warehouse_name" + client = FileSystemWarehouse(wh_dir) 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()) + assert any(Path(f"{wh_dir}").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) + wh_dir = Path(tmpdir) / "warehouse_name" + client = FileSystemWarehouse(root_dir=wh_dir) 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()) + assert any(Path(f"{wh_dir}").iterdir()) # make it look like we have an empty cache, as if this was a # different process key = obj.key @@ -225,7 +226,9 @@ def test_filesystemwarehouse_has_shared_and_tasks_stores(self, absolute_transfor unit = TestWarehouseBaseClass._get_protocol_unit(absolute_transformation) with tempfile.TemporaryDirectory() as tmpdir: - client = FileSystemWarehouse(tmpdir) + wh_dir = Path(tmpdir) / "warehouse_name" + client = FileSystemWarehouse(wh_dir) + assert client.name == "warehouse_name" assert "shared" in client.stores assert "tasks" in client.stores From 9a4b3e05748ef48179a7160045c6d5f4e4e17dbd Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Thu, 13 Aug 2026 14:32:36 -0700 Subject: [PATCH 2/4] separate creating new filesystem warehouse from loading an existing one --- src/openfe/storage/warehouse.py | 6 +++--- src/openfe/tests/storage/test_warehouse.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/openfe/storage/warehouse.py b/src/openfe/storage/warehouse.py index 87ca01620..8d59b701a 100644 --- a/src/openfe/storage/warehouse.py +++ b/src/openfe/storage/warehouse.py @@ -410,9 +410,9 @@ class FileSystemWarehouse(WarehouseBaseClass): for results and other data types. """ - def __init__(self, root_dir: pathlib.Path): + def __init__(self, root_dir: pathlib.Path, exist_okay=False): self.root_dir = pathlib.Path(root_dir) - if self.root_dir.is_dir(): + if self.root_dir.is_dir() and not exist_okay: raise ValueError( "`root_dir` already exists. To load an existing Warehouse, use FileSystemWarehouse.load(`root_dir`)" ) @@ -438,4 +438,4 @@ def load(cls, root_dir: pathlib.Path) -> FileSystemWarehouse: raise ValueError( "`root_dir` must be an existing filepath. To create a new Warehouse, use FileSystemWarehouse(`root_dir`)" ) - return cls(root_dir=root_dir) + return cls(root_dir=root_dir, exist_okay=True) diff --git a/src/openfe/tests/storage/test_warehouse.py b/src/openfe/tests/storage/test_warehouse.py index acb671507..352941509 100644 --- a/src/openfe/tests/storage/test_warehouse.py +++ b/src/openfe/tests/storage/test_warehouse.py @@ -240,6 +240,18 @@ def test_filesystemwarehouse_has_shared_and_tasks_stores(self, absolute_transfor client.store_task(unit) assert client.exists(unit.key) + def test_filesystem_warehouse_exists_error(self): + with tempfile.TemporaryDirectory() as tmpdir: + wh_dir = Path(tmpdir) / "warehouse_name" + client = FileSystemWarehouse(wh_dir) + # store some data so files are created + client.stores["shared"].store_bytes("sentinel", b"shared-data") + + with pytest.raises(ValueError, match="already exists"): + _ = FileSystemWarehouse(wh_dir) + + reloaded_client = FileSystemWarehouse.load(root_dir=wh_dir) + @pytest.mark.parametrize( "fixture", ["absolute_transformation", "complex_equilibrium"], From a670034e9a577fbd54ebe1da555f98393dabf309 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Fri, 14 Aug 2026 13:39:02 -0700 Subject: [PATCH 3/4] update tests --- src/openfe/tests/orchestration/test_worker.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/openfe/tests/orchestration/test_worker.py b/src/openfe/tests/orchestration/test_worker.py index 221e8de54..c6cd6a1f0 100644 --- a/src/openfe/tests/orchestration/test_worker.py +++ b/src/openfe/tests/orchestration/test_worker.py @@ -156,7 +156,7 @@ 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)) + warehouse = FileSystemWarehouse.load(str(warehouse_root)) exorcist.TaskStatusDB.from_filename(db_path) worker = Worker(warehouse=warehouse, task_db_path=db_path) @@ -167,7 +167,7 @@ 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)) + warehouse = FileSystemWarehouse.load(str(warehouse_root)) exorcist.TaskStatusDB.from_filename(db_path) worker = Worker(warehouse=warehouse, task_db_path=db_path) From 86ebbf7e16e17b231e4c6ceba3b6c0de3e58a390 Mon Sep 17 00:00:00 2001 From: Alyssa Travitz Date: Fri, 14 Aug 2026 17:33:43 -0700 Subject: [PATCH 4/4] make alchemical_network_to_task_graph private --- src/openfe/orchestration/__init__.py | 2 +- src/openfe/orchestration/exorcist_utils.py | 54 +++++++++++-------- .../orchestration/test_exorcist_utils.py | 21 ++++---- 3 files changed, 43 insertions(+), 34 deletions(-) diff --git a/src/openfe/orchestration/__init__.py b/src/openfe/orchestration/__init__.py index 14999e2f2..77105b24f 100644 --- a/src/openfe/orchestration/__init__.py +++ b/src/openfe/orchestration/__init__.py @@ -18,7 +18,7 @@ from openfe.storage.warehouse import FileSystemWarehouse from .exorcist_utils import ( - alchemical_network_to_task_graph, + _alchemical_network_to_task_graph, build_task_db_from_alchemical_network, ) diff --git a/src/openfe/orchestration/exorcist_utils.py b/src/openfe/orchestration/exorcist_utils.py index 459877b9c..0d82783c4 100644 --- a/src/openfe/orchestration/exorcist_utils.py +++ b/src/openfe/orchestration/exorcist_utils.py @@ -4,18 +4,19 @@ structures and can initialize an Exorcist task database from that graph. """ +import sys from pathlib import Path import exorcist import networkx as nx -import pandas as pd -from gufe import AlchemicalNetwork +from gufe import AlchemicalNetwork, ProtocolDAG from openfe.storage.warehouse import WarehouseBaseClass -def alchemical_network_to_task_graph( - alchemical_network: AlchemicalNetwork, warehouse: WarehouseBaseClass +def _alchemical_network_to_task_graph( + alchemical_network: AlchemicalNetwork, + warehouse: WarehouseBaseClass, ) -> nx.DiGraph: """Build a global task DAG from an alchemical network. @@ -30,9 +31,8 @@ def alchemical_network_to_task_graph( Returns ------- nx.DiGraph - A directed acyclic graph where each node is a task ID in the form - ``":"`` and edges encode - protocol-unit dependencies. + A directed acyclic graph where each node is a task ID with the ProtocolUnit key as a name + and edges encode protocol-unit dependencies. Raises ------ @@ -40,35 +40,40 @@ def alchemical_network_to_task_graph( Raised if the assembled task graph is not acyclic. """ - global_dag = nx.DiGraph() + if not isinstance(alchemical_network, AlchemicalNetwork): + raise ValueError( + f"alchemical_network must be an AlchemicalNetwork, not {type(alchemical_network)}." + ) + + warehouse.store_setup_tokenizable(alchemical_network) + + global_task_dag = nx.DiGraph() for transformation in alchemical_network.edges: - dag = transformation.create() + dag: ProtocolDAG = transformation.create() for unit in dag.protocol_units: - node_id = str(unit.key) - global_dag.add_node(node_id) + global_task_dag.add_node(str(unit.key)) warehouse.store_task(unit) - # store the protocol_dag as a shallow dict, since all its units are - # already written to disk - warehouse.store_protocol_dag(dag) for dependent_unit, dependency_unit in dag.graph.edges: upstream_id = str(dependency_unit.key) downstream_id = str(dependent_unit.key) - global_dag.add_edge(upstream_id, downstream_id) + global_task_dag.add_edge(upstream_id, downstream_id) + + # at this point, stored as a shallow dict since all its units are already stored + warehouse.store_protocol_dag(dag) - if not nx.is_directed_acyclic_graph(global_dag): + if not nx.is_directed_acyclic_graph(global_task_dag): raise ValueError("AlchemicalNetwork produced a task graph that is not a DAG.") - return global_dag + return global_task_dag -# TODO: do we test adding a multiple alchemical networks to the same task graph? def build_task_db_from_alchemical_network( alchemical_network: AlchemicalNetwork, warehouse: WarehouseBaseClass, db_path: Path | None = None, max_tries: int = 1, ) -> exorcist.TaskStatusDB: - """Create and populate a task database from an alchemical network. + """Create and populate a task database and warehouse from an alchemical network. Parameters ---------- @@ -80,8 +85,8 @@ def build_task_db_from_alchemical_network( Location of the SQLite-backed Exorcist database. If ``None``, defaults to {warehouse.name}.db in the current working directory. max_tries : int, default=1 - Maximum number of retries for each task before Exorcist marks it as - ``TOO_MANY_RETRIES``. + Maximum number of times a task will attempt to be submitted before it + is labelled ``TOO_MANY_RETRIES``. Returns ------- @@ -91,8 +96,11 @@ def build_task_db_from_alchemical_network( """ if db_path is None: db_path = Path(f"{warehouse.name}.db") + if db_path.exists(): + print(f"Error: {db_path} already exists.") # TODO: add more user flexibility here + sys.exit() - global_dag: nx.DiGraph = alchemical_network_to_task_graph(alchemical_network, warehouse) + global_task_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) + db.add_task_network(global_task_dag, max_tries) return db diff --git a/src/openfe/tests/orchestration/test_exorcist_utils.py b/src/openfe/tests/orchestration/test_exorcist_utils.py index aea4df87f..6f72132bf 100644 --- a/src/openfe/tests/orchestration/test_exorcist_utils.py +++ b/src/openfe/tests/orchestration/test_exorcist_utils.py @@ -6,10 +6,11 @@ import networkx as nx import pytest import sqlalchemy as sqla +from gufe import AlchemicalNetwork from gufe.tokenization import GufeKey -from openfe.orchestration.exorcist_utils import ( - alchemical_network_to_task_graph, +from openfe.orchestration import ( + _alchemical_network_to_task_graph, build_task_db_from_alchemical_network, ) from openfe.storage.warehouse import FileSystemWarehouse, WarehouseBaseClass @@ -42,7 +43,7 @@ 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)) + _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] @@ -56,7 +57,7 @@ def test_alchemical_network_to_task_graph_uses_canonical_task_ids(request, fixtu warehouse = _RecordingWarehouse() network = request.getfixturevalue(fixture) - graph = alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) + graph = _alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) expected_protocol_unit_keys = sorted(str(unit.key) for unit in warehouse.stored_tasks) observed_protocol_unit_keys = [] @@ -73,7 +74,7 @@ def test_alchemical_network_to_task_graph_edges_reference_existing_nodes(request warehouse = _RecordingWarehouse() network = request.getfixturevalue(fixture) - graph = alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) + graph = _alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) assert len(graph.edges) > 0 for u, v in graph.edges: @@ -86,7 +87,7 @@ def test_alchemical_network_to_task_graph_edge_direction_matches_dependencies(re warehouse = _RecordingWarehouse() network = request.getfixturevalue(fixture) - graph = alchemical_network_to_task_graph(network, cast(WarehouseBaseClass, warehouse)) + 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: @@ -118,12 +119,12 @@ def create(self): dag.graph.add_edges_from([(unit_a, unit_b), (unit_b, unit_a)]) return dag - network = mock.Mock() + network = mock.Mock(AlchemicalNetwork) network.edges = [_Transformation()] warehouse = mock.Mock() with pytest.raises(ValueError, match="not a DAG"): - alchemical_network_to_task_graph(network, warehouse) + _alchemical_network_to_task_graph(network, warehouse) @pytest.mark.parametrize("fixture", ["benzene_variants_star_map"]) @@ -184,7 +185,7 @@ def test_build_task_db_default_path(request, fixture): with ( mock.patch( - "openfe.orchestration.exorcist_utils.alchemical_network_to_task_graph", + "openfe.orchestration.exorcist_utils._alchemical_network_to_task_graph", return_value=fake_graph, ) as task_graph_mock, mock.patch( @@ -210,7 +211,7 @@ def test_build_task_db_forwards_graph_and_max_tries(request, tmp_path, fixture): with ( mock.patch( - "openfe.orchestration.exorcist_utils.alchemical_network_to_task_graph", + "openfe.orchestration.exorcist_utils._alchemical_network_to_task_graph", return_value=fake_graph, ) as task_graph_mock, mock.patch(