add results handling for new worker-based execution

src/openfe/orchestration/__init__.py

@@ -20,6 +20,7 @@
 from .exorcist_utils import (
     alchemical_network_to_task_graph,
     build_task_db_from_alchemical_network,
+    get_task_df,
 )
 
 
@@ -158,7 +159,6 @@ def _checkout_task(self) -> tuple[TaskStatusDB, str, ProtocolUnit] | None:
             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 "ProtocolUnit-<HASH>"
         taskid = db.check_out_task()

src/openfe/orchestration/exorcist_utils.py

@@ -4,6 +4,7 @@
 structures and can initialize an Exorcist task database from that graph.
 """
 
+import sys
 from pathlib import Path
 
 import exorcist
@@ -15,17 +16,22 @@
 
 
 def alchemical_network_to_task_graph(
-    alchemical_network: AlchemicalNetwork, warehouse: WarehouseBaseClass
+    alchemical_network: AlchemicalNetwork,
+    warehouse: WarehouseBaseClass,
 ) -> nx.DiGraph:
-    """Build a global task DAG from an alchemical network.
+    """Build a global task DAG from `alchemical_network` and store its relevant data
+    in `warehouse` the following warehouse stores:
+        - 'setup': The AlchemicalNetwork, deduplicated on disk
+        - 'tasks': The ProtocolUnits to be executed as tasks
+        - 'protocol_dags': The ProtocolDAGs that the ProtocolUnits belong to.
+                           Used to gather results after execution.
 
     Parameters
     ----------
     alchemical_network : AlchemicalNetwork
-        Network containing transformations to execute.
+        Network containing alchemical transformations to be executed.
     warehouse : WarehouseBaseClass
-        Warehouse used to persist protocol units as tasks while the graph is
-        constructed.
+        Warehouse used to store data used by the execution and simulation engines.
 
     Returns
     -------
@@ -37,11 +43,18 @@ def alchemical_network_to_task_graph(
     Raises
     ------
     ValueError
-        Raised if the assembled task graph is not acyclic.
+        If the assembled task graph is not acyclic.
+        If the input `alchemical_network` is not a valid openfe.AlchemicalNetwork
     """
 
+    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_dag = nx.DiGraph()
-    for transformation in alchemical_network.edges:
+    for transformation in alchemical_network.edges:  # TODO: skip edges that already have units?
         dag = transformation.create()
         for unit in dag.protocol_units:
             node_id = str(unit.key)
@@ -91,8 +104,47 @@ 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"{db_path} already exists.")  # TODO: add more user flexibility here
+        sys.exit()
 
     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
+
+
+def get_task_df(task_db: exorcist.TaskStatusDB) -> pd.DataFrame:
+    """Create a pandas Dataframe from task_db.
+
+    Parameters
+    ----------
+    task_db : exorcist.TaskStatusDB
+        A task database.
+
+    Returns
+    -------
+    pd.DataFrame
+        A dataframe of the tasks and their statuses
+    """
+    status_name_encoding = {e.value: e.name for e in exorcist.TaskStatus}
+    task_table = pd.read_sql_table("tasks", task_db.engine)
+    task_table.replace({"status": status_name_encoding}, inplace=True)
+    return task_table
+
+
+def get_dependency_df(task_db: exorcist.TaskStatusDB) -> pd.DataFrame:
+    """Create a pandas Dataframe from task_db.
+
+    Parameters
+    ----------
+    task_db : exorcist.TaskStatusDB
+        A task database.
+
+    Returns
+    -------
+    pd.DataFrame
+        A dataframe of the tasks and their dependencies.
+
+    """
+    return pd.read_sql_table("dependencies", task_db.engine)

src/openfe/storage/warehouse.py

@@ -1,12 +1,15 @@
 # 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
-from typing import Generator, Literal, TypedDict
+from typing import Generator, Iterable, Literal, TypedDict
 
-from gufe.protocols.protocoldag import ProtocolDAG
-from gufe.protocols.protocolunit import ProtocolUnit
+from gufe.protocols import ProtocolResult
+from gufe.protocols.protocoldag import ProtocolDAG, ProtocolDAGResult
+from gufe.protocols.protocolunit import ProtocolUnit, ProtocolUnitResult
 from gufe.storage.externalresource import ExternalStorage, FileStorage
 from gufe.tokenization import (
     JSON_HANDLER,
@@ -356,6 +359,108 @@ def get_protocol_dags(self) -> Generator[ProtocolDAG, None, None]:
                 dag = self.load_protocol_dag(item)
                 yield dag
 
+    def get_unit_results(self) -> Generator[ProtocolUnitResult]:
+        """Yield all ProtocolUnitResult(s) stored in the Warehouse's 'result' store.
+
+        Yields
+        ------
+        Generator[ProtocolUnitResult]
+            The ProtocolUnitResults found in this Warehouse's 'result' store
+
+        Raises
+        ------
+        RuntimeError
+            If any object in the result store is not a ProtocolUnitResult
+        """
+        for i in self.stores["result"]:
+            obj = self.load_result_tokenizable(i)
+            if isinstance(obj, ProtocolUnitResult):
+                yield obj
+            else:
+                raise RuntimeError(
+                    f"gufe tokenizable {obj} found in result store, but is not a ProtocolUnitResult."
+                )
+
+    def gather_all_results(self) -> list[tuple[ProtocolResult, ProtocolDAGResult]]:
+        """From this warehouse, gather all ProtocolDAGResults corresponding to the recorded
+        ProtocolDAGs, and return all (ProtocolResult, ProtocolDAGResult) pairs.
+
+        Note: this requires the Warehouse to explicitly have stored the ProtocolDAGs and
+        their ProtocolUnits when constructing the task graph.
+
+        Returns
+        -------
+        list[tuple[ProtocolResult, ProtocolDAGResult]]
+            ProtocolResults and their corresponding ProtocolDAGResults
+        """
+
+        # construct a map of all the ProtocolDAGs and their corresponding ProtocolUnitResults
+        dags_to_purs = self._construct_dags_to_purs(
+            self.get_protocol_dags(), self.get_unit_results()
+        )
+        # load all dags
+        dags_with_results = [
+            self.load_protocol_dag(d) for d in dags_to_purs if dags_to_purs[d] != []
+        ]
+
+        result_edges: list[tuple[ProtocolResult, ProtocolDAGResult]] = []
+        for dag in dags_with_results:
+            prot_dag_result: ProtocolDAGResult = self._construct_protocol_dag_result(
+                protocol_dag=dag, dags_to_purs=dags_to_purs
+            )
+            prot_result: ProtocolResult = self.gather_result(
+                protocol_dag=dag, dags_to_purs=dags_to_purs
+            )
+            result_edges.append((prot_result, prot_dag_result))
+
+        return result_edges
+
+    @staticmethod
+    def _construct_dags_to_purs(dags: Iterable[ProtocolDAG], purs: Iterable[ProtocolUnitResult]):
+        """Creating a mapping of protocolDAGs to their corresponding ProtocolUnitResults"""
+        pur_pu_keys = {str(pur.source_key): pur for pur in purs}
+        dag_map = {}
+        for dag in dags:
+            dag_purs = []
+            for unit in dag.protocol_units:
+                if unit.key in pur_pu_keys:
+                    dag_purs.append(pur_pu_keys[unit.key])
+            dag_map[str(dag.key)] = dag_purs
+        return dag_map
+
+    @staticmethod
+    def _construct_protocol_dag_result(
+        protocol_dag: ProtocolDAG,
+        dags_to_purs: dict[str, list[ProtocolUnitResult]],
+    ) -> ProtocolDAGResult:
+        """Create a ProtocolDAGResult from the ProtocolDAG and its corresponding ProtocolUnitResults
+
+        Parameters
+        ----------
+        protocol_dag : ProtocolDAG
+            The ProtocolDAG to construct a ProtocolDAGResult for.
+        dags_to_purs : dict[str, list[ProtocolUnitResult]]
+            Mapping of all ProtocolDAG keys and their ProtocolUnitResults
+
+        Returns
+        -------
+        ProtocolDAGResult
+        """
+        purs = dags_to_purs[str(protocol_dag.key)]
+        dag_result = ProtocolDAGResult(
+            protocol_units=protocol_dag.protocol_units,
+            protocol_unit_results=purs,
+            transformation_key=protocol_dag.transformation_key,
+            extends_key=protocol_dag.extends_key,
+        )
+        return dag_result
+
+    def gather_result(self, protocol_dag, dags_to_purs) -> ProtocolResult:
+        protocol_dag_result = self._construct_protocol_dag_result(protocol_dag, dags_to_purs)
+        transformation = self.load_setup_tokenizable(protocol_dag.transformation_key)
+        result = transformation.gather([protocol_dag_result])
+        return result
+
     @property
     def setup_store(self):
         """Get the setup store
@@ -398,8 +503,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 +513,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 +531,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/data/warehouse/mc1_campaign/protocol_dags/ProtocolDAG-32c19e7ec6864e13658a8d3ca48e582d

@@ -0,0 +1 @@
+{":version:": 1, "__module__": "gufe.protocols.protocoldag", "__qualname__": "ProtocolDAG", "extends_key": null, "name": null, "protocol_units": [{":gufe-key:": "HybridTopologySetupUnit-d568ebe569b445c7875d9b6d5ce04cd0"}, {":gufe-key:": "HybridTopologyMultiStateSimulationUnit-2821f959834844a79ff406338c9e4101"}, {":gufe-key:": "HybridTopologyMultiStateAnalysisUnit-a7252f67bcd541378e9846b01ad21a5e"}], "transformation_key": "Transformation-a8ccb7f841aafdc860143dad3f2bd82a"}

src/openfe/tests/data/warehouse/mc1_campaign/protocol_dags/ProtocolDAG-616ec7a89d894b5f430007cd3a3b5d49

@@ -0,0 +1 @@
+{":version:": 1, "__module__": "gufe.protocols.protocoldag", "__qualname__": "ProtocolDAG", "extends_key": null, "name": null, "protocol_units": [{":gufe-key:": "HybridTopologySetupUnit-78321eda905c4c2c9f75eb69c36496da"}, {":gufe-key:": "HybridTopologyMultiStateSimulationUnit-f94c4264ca8f483095da93f626892c1a"}, {":gufe-key:": "HybridTopologyMultiStateAnalysisUnit-7e9620845eff4575a757b1993e88d517"}], "transformation_key": "Transformation-d51c5a0397bfe45865a78edd1935e3b0"}

src/openfe/tests/data/warehouse/mc1_campaign/protocol_dags/ProtocolDAG-72133e3079ecad750b71ca7994237228

@@ -0,0 +1 @@
+{":version:": 1, "__module__": "gufe.protocols.protocoldag", "__qualname__": "ProtocolDAG", "extends_key": null, "name": null, "protocol_units": [{":gufe-key:": "HybridTopologySetupUnit-0238f65b55044b1ea751447537a28f2b"}, {":gufe-key:": "HybridTopologyMultiStateSimulationUnit-cf34ef156b844592b4e860d9f47197a8"}, {":gufe-key:": "HybridTopologyMultiStateAnalysisUnit-72c71598dc75432a90aed831a4e4cb24"}], "transformation_key": "Transformation-622f1d19bea91cd47a7e478e42376305"}

src/openfe/tests/data/warehouse/mc1_campaign/protocol_dags/ProtocolDAG-b51ae0bc4944b776b4fe63f17158bac1

@@ -0,0 +1 @@
+{":version:": 1, "__module__": "gufe.protocols.protocoldag", "__qualname__": "ProtocolDAG", "extends_key": null, "name": null, "protocol_units": [{":gufe-key:": "HybridTopologySetupUnit-17dc05e0d79747e79d3e93e55383c9bd"}, {":gufe-key:": "HybridTopologyMultiStateSimulationUnit-15972a17ad31428280da2cc53c8a6cdc"}, {":gufe-key:": "HybridTopologyMultiStateAnalysisUnit-e44e96bb39c947ed9d68f5c1ada6defc"}], "transformation_key": "Transformation-ae340a619b2794a9a3da4e72ef977885"}

src/openfe/tests/data/warehouse/mc1_campaign/result/ProtocolUnitFailure-64e7d042322d4495a096d379660dbfeb

@@ -0,0 +1 @@
+{":version:": 1, "__module__": "gufe.protocols.protocolunit", "__qualname__": "ProtocolUnitFailure", "_key": "ProtocolUnitFailure-64e7d042322d4495a096d379660dbfeb", "end_time": {":is_custom:": true, "__class__": "datetime", "__module__": "datetime", "isotime": "2026-08-14T11:25:53.636128"}, "exception": ["SimulationNaNError", ["Propagating replica 0 at state 0 resulted in a NaN!\nThe state of the system and integrator before the error were saved in mc1_campaign/shared/task_workdirs/HybridTopologyMultiStateSimulationUnit-cf34ef156b844592b4e860d9f47197a8/nan-error-logs"]], "inputs": {"generation": 0, "protocol": {":gufe-key:": "RelativeHybridTopologyProtocol-59d604ce336f12d44f6e4647a6f7e616"}, "repeat_id": 284817259896887492684538463771894716234, "setup_results": {":gufe-key:": "ProtocolUnitResult-4388526f592f4586a4f0cb62c6a71ef2"}}, "name": "HybridTopology Simulation: ligand_2 to ligand_3 repeat 0 generation 0", "outputs": {}, "source_key": "HybridTopologyMultiStateSimulationUnit-cf34ef156b844592b4e860d9f47197a8", "start_time": {":is_custom:": true, "__class__": "datetime", "__module__": "datetime", "isotime": "2026-08-14T11:25:34.771776"}, "stderr": {}, "stdout": {}, "traceback": "Traceback (most recent call last):\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/openmmtools/multistate/multistatesampler.py\", line 1323, in _propagate_replica\n    mcmc_move.apply(thermodynamic_state, sampler_state, context_cache=self.sampler_context_cache)\n    ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/openmmtools/mcmc.py\", line 1151, in apply\n    super().apply(thermodynamic_state, sampler_state,\n    ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n                                            context_cache=context_cache)\n                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/openmmtools/mcmc.py\", line 755, in apply\n    raise IntegratorMoveError(err_msg, self, context)\nopenmmtools.mcmc.IntegratorMoveError: Potential energy is NaN after 20 attempts of integration with move LangevinDynamicsMove\n\nDuring handling of the above exception, another exception occurred:\n\nTraceback (most recent call last):\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/gufe/protocols/protocolunit.py\", line 367, in execute\n    outputs = self._execute(context, **inputs)\n  File \"/Users/atravitz/software/openfe/src/openfe/protocols/openmm_rfe/hybridtop_units.py\", line 1444, in _execute\n    outputs = self.run(\n        system=system,\n    ...<3 lines>...\n        shared_basepath=ctx.shared,\n    )\n  File \"/Users/atravitz/software/openfe/src/openfe/protocols/openmm_rfe/hybridtop_units.py\", line 1380, in run\n    self._run_simulation(\n    ~~~~~~~~~~~~~~~~~~~~^\n        sampler=sampler,\n        ^^^^^^^^^^^^^^^^\n    ...<4 lines>...\n        dry=dry,\n        ^^^^^^^^\n    )\n    ^\n  File \"/Users/atravitz/software/openfe/src/openfe/protocols/openmm_rfe/hybridtop_units.py\", line 1248, in _run_simulation\n    sampler.equilibrate(int(equil_steps / mc_steps))\n    ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/openmmtools/multistate/multistatesampler.py\", line 693, in equilibrate\n    self._propagate_replicas()\n    ~~~~~~~~~~~~~~~~~~~~~~~~^^\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/openmmtools/utils/utils.py\", line 95, in _wrapper\n    return func(*args, **kwargs)\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/openmmtools/multistate/multistatesampler.py\", line 1296, in _propagate_replicas\n    propagated_states, replica_ids = mpiplus.distribute(self._propagate_replica, range(self.n_replicas),\n                                     ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n                                                    send_results_to=0)\n                                                    ^^^^^^^^^^^^^^^^^^\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/mpiplus/mpiplus.py\", line 523, in distribute\n    all_results = [task(job_args, *other_args, **kwargs) for job_args in distributed_args]\n                   ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File \"/Users/atravitz/micromamba/envs/openfe-exorcist/lib/python3.14/site-packages/openmmtools/multistate/multistatesampler.py\", line 1334, in _propagate_replica\n    raise SimulationNaNError(message)\nopenmmtools.multistate.utils.SimulationNaNError: Propagating replica 0 at state 0 resulted in a NaN!\nThe state of the system and integrator before the error were saved in mc1_campaign/shared/task_workdirs/HybridTopologyMultiStateSimulationUnit-cf34ef156b844592b4e860d9f47197a8/nan-error-logs\n"}