filesystem warehouse usability improvements

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,
 )
 

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,45 +31,49 @@ def alchemical_network_to_task_graph(
     Returns
     -------
     nx.DiGraph
-        A directed acyclic graph where each node is a task ID in the form
-        ``"<transformation_key>:<protocol_unit_key>"`` 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
     ------
     ValueError
         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

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, exist_okay=False):
+        self.root_dir = pathlib.Path(root_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`)"
+            )
         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, exist_okay=True)

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(

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)
 

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
@@ -237,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"],