docs: add numpy docstrings

Author: ethanholz

src/openfe/orchestration/__init__.py

@@ -1,3 +1,5 @@
+"""Task orchestration utilities backed by Exorcist and a warehouse."""
+
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -21,10 +23,29 @@
 
 @dataclass
 class Worker:
+    """Execute protocol units from an Exorcist task database.
+
+    Parameters
+    ----------
+    warehouse : FileSystemWarehouse
+        Warehouse used to load queued tasks and store execution results.
+    task_db_path : pathlib.Path, default=Path("./warehouse/tasks.db")
+        Path to the Exorcist SQLite task database.
+    """
+
     warehouse: FileSystemWarehouse
     task_db_path: Path = Path("./warehouse/tasks.db")
 
     def _checkout_task(self) -> tuple[str, ProtocolUnit] | None:
+        """Check out one available task and load its protocol unit.
+
+        Returns
+        -------
+        tuple[str, ProtocolUnit] or None
+            The checked-out task ID and corresponding protocol unit, or
+            ``None`` if no task is currently available.
+        """
+
         db: TaskStatusDB = TaskStatusDB.from_filename(self.task_db_path)
         # The format for the taskid is "Transformation-<HASH>:ProtocolUnit-<HASH>"
         taskid = db.check_out_task()
@@ -36,13 +57,46 @@ def _checkout_task(self) -> tuple[str, ProtocolUnit] | None:
         return taskid, unit
 
     def _get_task(self) -> ProtocolUnit:
+        """Return the next available protocol unit.
+
+        Returns
+        -------
+        ProtocolUnit
+            A protocol unit loaded from the warehouse.
+
+        Raises
+        ------
+        RuntimeError
+            Raised when no task is available in the task database.
+        """
+
         task = self._checkout_task()
         if task is None:
             raise RuntimeError("No AVAILABLE tasks found in the task database.")
         _, unit = task
         return unit
 
     def execute_unit(self, scratch: Path) -> tuple[str, ProtocolUnitResult] | None:
+        """Execute one checked-out protocol unit and persist its result.
+
+        Parameters
+        ----------
+        scratch : pathlib.Path
+            Scratch directory passed to the protocol execution context.
+
+        Returns
+        -------
+        tuple[str, ProtocolUnitResult] or None
+            The task ID and execution result for the processed task, or
+            ``None`` if no task is currently available.
+
+        Raises
+        ------
+        Exception
+            Re-raises any exception thrown during protocol unit execution after
+            marking the task as failed.
+        """
+
         # 1. Get task/unit
         task = self._checkout_task()
         if task is None:

src/openfe/orchestration/exorcist_utils.py

@@ -1,4 +1,8 @@
-"""Utilities for building Exorcist task graphs and task databases."""
+"""Utilities for building Exorcist task graphs and task databases.
+
+This module translates an :class:`gufe.AlchemicalNetwork` into Exorcist task
+structures and can initialize an Exorcist task database from that graph.
+"""
 
 from pathlib import Path
 
@@ -12,7 +16,28 @@
 def alchemical_network_to_task_graph(
     alchemical_network: AlchemicalNetwork, warehouse: WarehouseBaseClass
 ) -> nx.DiGraph:
-    """Build a global task DAG from an AlchemicalNetwork."""
+    """Build a global task DAG from an alchemical network.
+
+    Parameters
+    ----------
+    alchemical_network : AlchemicalNetwork
+        Network containing transformations to execute.
+    warehouse : WarehouseBaseClass
+        Warehouse used to persist protocol units as tasks while the graph is
+        constructed.
+
+    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.
+
+    Raises
+    ------
+    ValueError
+        Raised if the assembled task graph is not acyclic.
+    """
 
     global_dag = nx.DiGraph()
     for transformation in alchemical_network.edges:
@@ -40,7 +65,27 @@ def build_task_db_from_alchemical_network(
     db_path: Path | None = None,
     max_tries: int = 1,
 ) -> exorcist.TaskStatusDB:
-    """Create an Exorcist TaskStatusDB from an AlchemicalNetwork."""
+    """Create and populate a task database from an alchemical network.
+
+    Parameters
+    ----------
+    alchemical_network : AlchemicalNetwork
+        Network containing transformations to convert into task records.
+    warehouse : WarehouseBaseClass
+        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.
+    max_tries : int, default=1
+        Maximum number of retries for each task before Exorcist marks it as
+        ``TOO_MANY_RETRIES``.
+
+    Returns
+    -------
+    exorcist.TaskStatusDB
+        Initialized task database populated with graph nodes and dependency
+        edges derived from ``alchemical_network``.
+    """
     if db_path is None:
         db_path = Path("tasks.db")