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

Author: atravitz

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-<HASH>:ProtocolUnit-<HASH>"
+        # The format for the taskid is "ProtocolUnit-<HASH>"
         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
 

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

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)

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
 

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)