lsst-pipe-base 29.2025.3900__py3-none-any.whl → 29.2025.4100__py3-none-any.whl

lsst/pipe/base/_task_metadata.py

@@ -686,6 +686,21 @@ class TaskMetadata(BaseModel):
             """See `pydantic.BaseModel.model_json_schema`."""
             return super().model_json_schema(*args, **kwargs)
 
+        @classmethod
+        def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate`."""
+            return super().model_validate(*args, **kwargs)
+
+        @classmethod
+        def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_json`."""
+            return super().model_validate_json(*args, **kwargs)
+
+        @classmethod
+        def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_strings`."""
+            return super().model_validate_strings(*args, **kwargs)
+
 
 # Needed because a TaskMetadata can contain a TaskMetadata.
 TaskMetadata.model_rebuild()

lsst/pipe/base/dot_tools.py

@@ -33,147 +33,27 @@ from __future__ import annotations
 
 __all__ = ["graph2dot", "pipeline2dot"]
 
-# -------------------------------
-#  Imports of standard modules --
-# -------------------------------
-import html
-import io
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
-# -----------------------------
-#  Imports for other modules --
-# -----------------------------
 from .pipeline import Pipeline
 
 if TYPE_CHECKING:
-    from lsst.daf.butler import DatasetRef
-    from lsst.pipe.base import QuantumGraph, QuantumNode, TaskDef
+    from .graph import QuantumGraph
+    from .pipeline import TaskDef
+    from .quantum_graph import PredictedQuantumGraph
 
-# ----------------------------------
-#  Local non-exported definitions --
-# ----------------------------------
 
-# Attributes applied to directed graph objects.
-_NODELABELPOINTSIZE = "18"
-_ATTRIBS = dict(
-    defaultGraph=dict(splines="ortho", nodesep="0.5", ranksep="0.75", pad="0.5"),
-    defaultNode=dict(shape="box", fontname="Monospace", fontsize="14", margin="0.2,0.1", penwidth="3"),
-    defaultEdge=dict(color="black", arrowsize="1.5", penwidth="1.5"),
-    task=dict(style="filled", color="black", fillcolor="#B1F2EF"),
-    quantum=dict(style="filled", color="black", fillcolor="#B1F2EF"),
-    dsType=dict(style="rounded,filled,bold", color="#00BABC", fillcolor="#F5F5F5"),
-    dataset=dict(style="rounded,filled,bold", color="#00BABC", fillcolor="#F5F5F5"),
-)
-
-
-def _renderDefault(type: str, attribs: dict[str, str], file: io.TextIOBase) -> None:
-    """Set default attributes for a given type."""
-    default_attribs = ", ".join([f'{key}="{val}"' for key, val in attribs.items()])
-    print(f"{type} [{default_attribs}];", file=file)
-
-
-def _renderNode(file: io.TextIOBase, nodeName: str, style: str, labels: list[str]) -> None:
-    """Render GV node"""
-    label = r"</TD></TR><TR><TD>".join(labels)
-    attrib_dict = dict(_ATTRIBS[style], label=label)
-    pre = '<<TABLE BORDER="0" CELLPADDING="5"><TR><TD>'
-    post = "</TD></TR></TABLE>>"
-    attrib = ", ".join(
-        [
-            f'{key}="{val}"' if key != "label" else f"{key}={pre}{val}{post}"
-            for key, val in attrib_dict.items()
-        ]
-    )
-    print(f'"{nodeName}" [{attrib}];', file=file)
-
-
-def _renderTaskNode(nodeName: str, taskDef: TaskDef, file: io.TextIOBase, idx: Any = None) -> None:
-    """Render GV node for a task"""
-    labels = [
-        f'<B><FONT POINT-SIZE="{_NODELABELPOINTSIZE}">' + html.escape(taskDef.label) + "</FONT></B>",
-        html.escape(taskDef.taskName),
-    ]
-    if idx is not None:
-        labels.append(f"<I>index:</I>&nbsp;{idx}")
-    if taskDef.connections:
-        # don't print collection of str directly to avoid visually noisy quotes
-        dimensions_str = ", ".join(sorted(taskDef.connections.dimensions))
-        labels.append(f"<I>dimensions:</I>&nbsp;{html.escape(dimensions_str)}")
-    _renderNode(file, nodeName, "task", labels)
-
-
-def _renderQuantumNode(
-    nodeName: str, taskDef: TaskDef, quantumNode: QuantumNode, file: io.TextIOBase
-) -> None:
-    """Render GV node for a quantum"""
-    labels = [f"{quantumNode.nodeId}", html.escape(taskDef.label)]
-    dataId = quantumNode.quantum.dataId
-    assert dataId is not None, "Quantum DataId cannot be None"
-    labels.extend(f"{key} = {dataId[key]}" for key in sorted(dataId.required.keys()))
-    _renderNode(file, nodeName, "quantum", labels)
-
-
-def _renderDSTypeNode(name: str, dimensions: list[str], file: io.TextIOBase) -> None:
-    """Render GV node for a dataset type"""
-    labels = [f'<B><FONT POINT-SIZE="{_NODELABELPOINTSIZE}">' + html.escape(name) + "</FONT></B>"]
-    if dimensions:
-        labels.append("<I>dimensions:</I>&nbsp;" + html.escape(", ".join(sorted(dimensions))))
-    _renderNode(file, name, "dsType", labels)
-
-
-def _renderDSNode(nodeName: str, dsRef: DatasetRef, file: io.TextIOBase) -> None:
-    """Render GV node for a dataset"""
-    labels = [html.escape(dsRef.datasetType.name), f"run: {dsRef.run!r}"]
-    labels.extend(f"{key} = {dsRef.dataId[key]}" for key in sorted(dsRef.dataId.required.keys()))
-    _renderNode(file, nodeName, "dataset", labels)
-
-
-def _renderEdge(fromName: str, toName: str, file: io.TextIOBase, **kwargs: Any) -> None:
-    """Render GV edge"""
-    if kwargs:
-        attrib = ", ".join([f'{key}="{val}"' for key, val in kwargs.items()])
-        print(f'"{fromName}" -> "{toName}" [{attrib}];', file=file)
-    else:
-        print(f'"{fromName}" -> "{toName}";', file=file)
-
-
-def _datasetRefId(dsRef: DatasetRef) -> str:
-    """Make an identifying string for given ref"""
-    dsId = [dsRef.datasetType.name]
-    dsId.extend(f"{key} = {dsRef.dataId[key]}" for key in sorted(dsRef.dataId.required.keys()))
-    return ":".join(dsId)
-
-
-def _makeDSNode(dsRef: DatasetRef, allDatasetRefs: dict[str, str], file: io.TextIOBase) -> str:
-    """Make new node for dataset if  it does not exist.
-
-    Returns node name.
-    """
-    dsRefId = _datasetRefId(dsRef)
-    nodeName = allDatasetRefs.get(dsRefId)
-    if nodeName is None:
-        idx = len(allDatasetRefs)
-        nodeName = f"dsref_{idx}"
-        allDatasetRefs[dsRefId] = nodeName
-        _renderDSNode(nodeName, dsRef, file)
-    return nodeName
-
-
-# ------------------------
-#  Exported definitions --
-# ------------------------
-
-
-def graph2dot(qgraph: QuantumGraph, file: Any) -> None:
+def graph2dot(qgraph: QuantumGraph | PredictedQuantumGraph, file: Any) -> None:
     """Convert QuantumGraph into GraphViz digraph.
 
     This method is mostly for documentation/presentation purposes.
 
     Parameters
     ----------
-    qgraph : `lsst.pipe.base.QuantumGraph`
-        QuantumGraph instance.
+    qgraph : `lsst.pipe.base.QuantumGraph` or \
+            `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`
+        Quantum graph object.
     file : `str` or file object
         File where GraphViz graph (DOT language) is written, can be a file name
         or file object.
@@ -185,38 +65,20 @@ def graph2dot(qgraph: QuantumGraph, file: Any) -> None:
     ImportError
         Raised if the task class cannot be imported.
     """
+    from .quantum_graph import PredictedQuantumGraph, visualization
+
+    if not isinstance(qgraph, PredictedQuantumGraph):
+        qgraph = PredictedQuantumGraph.from_old_quantum_graph(qgraph)
+
     # open a file if needed
     close = False
     if not hasattr(file, "write"):
         file = open(file, "w")
         close = True
 
-    print("digraph QuantumGraph {", file=file)
-    _renderDefault("graph", _ATTRIBS["defaultGraph"], file)
-    _renderDefault("node", _ATTRIBS["defaultNode"], file)
-    _renderDefault("edge", _ATTRIBS["defaultEdge"], file)
-
-    allDatasetRefs: dict[str, str] = {}
-    for taskId, taskDef in enumerate(qgraph.taskGraph):
-        quanta = qgraph.getNodesForTask(taskDef)
-        for qId, quantumNode in enumerate(quanta):
-            # node for a task
-            taskNodeName = f"task_{taskId}_{qId}"
-            _renderQuantumNode(taskNodeName, taskDef, quantumNode, file)
-
-            # quantum inputs
-            for dsRefs in quantumNode.quantum.inputs.values():
-                for dsRef in dsRefs:
-                    nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
-                    _renderEdge(nodeName, taskNodeName, file)
-
-            # quantum outputs
-            for dsRefs in quantumNode.quantum.outputs.values():
-                for dsRef in dsRefs:
-                    nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
-                    _renderEdge(taskNodeName, nodeName, file)
+    v = visualization.QuantumGraphDotVisualizer()
+    v.write_bipartite(qgraph, file)
 
-    print("}", file=file)
     if close:
         file.close()
 

lsst/pipe/base/exec_fixup_data_id.py

@@ -27,16 +27,17 @@
 
 __all__ = ["ExecutionGraphFixup"]
 
-import contextlib
 import itertools
+import uuid
 from collections import defaultdict
-from collections.abc import Sequence
-from typing import Any
+from collections.abc import Mapping, Sequence
 
 import networkx as nx
 
+from lsst.daf.butler import DataCoordinate, DataIdValue
+
 from .execution_graph_fixup import ExecutionGraphFixup
-from .graph import QuantumGraph, QuantumNode
+from .graph import QuantumGraph
 
 
 class ExecFixupDataId(ExecutionGraphFixup):
@@ -88,44 +89,16 @@ class ExecFixupDataId(ExecutionGraphFixup):
         else:
             self.dimensions = tuple(self.dimensions)
 
-    def _key(self, qnode: QuantumNode) -> tuple[Any, ...]:
-        """Produce comparison key for quantum data.
-
-        Parameters
-        ----------
-        qnode : `QuantumNode`
-            An individual node in a `~lsst.pipe.base.QuantumGraph`
-
-        Returns
-        -------
-        key : `tuple`
-        """
-        dataId = qnode.quantum.dataId
-        assert dataId is not None, "Quantum DataId cannot be None"
-        key = tuple(dataId[dim] for dim in self.dimensions)
-        return key
-
     def fixupQuanta(self, graph: QuantumGraph) -> QuantumGraph:
-        taskDef = graph.findTaskDefByLabel(self.taskLabel)
-        if taskDef is None:
-            raise ValueError(f"Cannot find task with label {self.taskLabel}")
-        quanta = list(graph.getNodesForTask(taskDef))
-        keyQuanta = defaultdict(list)
-        for q in quanta:
-            key = self._key(q)
-            keyQuanta[key].append(q)
-        keys = sorted(keyQuanta.keys(), reverse=self.reverse)
-        networkGraph = graph.graph
-
-        for prev_key, key in itertools.pairwise(keys):
-            for prev_node in keyQuanta[prev_key]:
-                for node in keyQuanta[key]:
-                    # remove any existing edges between the two nodes, but
-                    # don't fail if there are not any. Both directions need
-                    # tried because in a directed graph, order maters
-                    for edge in ((node, prev_node), (prev_node, node)):
-                        with contextlib.suppress(nx.NetworkXException):
-                            networkGraph.remove_edge(*edge)
-
-                    networkGraph.add_edge(prev_node, node)
-        return graph
+        raise NotImplementedError()
+
+    def fixup_graph(
+        self, xgraph: nx.DiGraph, quanta_by_task: Mapping[str, Mapping[DataCoordinate, uuid.UUID]]
+    ) -> None:
+        quanta_by_sort_key: defaultdict[tuple[DataIdValue, ...], list[uuid.UUID]] = defaultdict(list)
+        for data_id, quantum_id in quanta_by_task[self.taskLabel].items():
+            key = tuple(data_id[dim] for dim in self.dimensions)
+            quanta_by_sort_key[key].append(quantum_id)
+        sorted_keys = sorted(quanta_by_sort_key.keys(), reverse=self.reverse)
+        for prev_key, key in itertools.pairwise(sorted_keys):
+            xgraph.add_edges_from(itertools.product(quanta_by_sort_key[prev_key], quanta_by_sort_key[key]))

lsst/pipe/base/execution_graph_fixup.py

@@ -27,7 +27,13 @@
 
 __all__ = ["ExecutionGraphFixup"]
 
-from abc import ABC, abstractmethod
+import uuid
+from abc import ABC
+from collections.abc import Mapping
+
+import networkx
+
+from lsst.daf.butler import DataCoordinate
 
 from .graph import QuantumGraph
 
@@ -35,27 +41,26 @@ from .graph import QuantumGraph
 class ExecutionGraphFixup(ABC):
     """Interface for classes which update quantum graphs before execution.
 
-    Primary goal of this class is to modify quanta dependencies which may not
-    be possible to reflect in a quantum graph using standard tools. One known
-    use case for that is to guarantee particular execution order of visits in
-    CI jobs for cases when outcome depends on the processing order of visits
-    (e.g. AP association pipeline).
-
-    Instances of this class receive pre-ordered sequence of quanta
-    (`.QuantumGraph` instances) and they are allowed to modify quanta data in
-    place, for example update ``dependencies`` field to add additional
-    dependencies. Returned list of quanta will be re-ordered once again by the
-    graph executor to reflect new dependencies.
+    Notes
+    -----
+    The primary goal of this class is to modify quanta dependencies which may
+    not be possible to reflect in a quantum graph using standard tools. One
+    known use case for that is to guarantee particular execution order of
+    visits in CI jobs for cases when outcome depends on the processing order of
+    visits (e.g. AP association pipeline).
+
+    Instances of this class receive a preliminary graph and are allowed to
+    add edges, as long as those edges do not result in a cycle.  Edges and
+    nodes may not be removed.
+
+    New subclasses should implement only `fixup_graph`, which will always be
+    called first.  `fixupQuanta` is only called if `fixup_graph` raises
+    `NotImplementedError`.
     """
 
-    @abstractmethod
     def fixupQuanta(self, graph: QuantumGraph) -> QuantumGraph:
         """Update quanta in a graph.
 
-        Potentially anything in the graph could be changed if it does not
-        break executor assumptions. If modifications result in a dependency
-        cycle the executor will raise an exception.
-
         Parameters
         ----------
         graph : `.QuantumGraph`
@@ -65,5 +70,31 @@ class ExecutionGraphFixup(ABC):
         -------
         graph : `.QuantumGraph`
             Modified graph.
+
+        Notes
+        -----
+        This hook is provided for backwards compatibility only.
+        """
+        raise NotImplementedError()
+
+    def fixup_graph(
+        self, xgraph: networkx.DiGraph, quanta_by_task: Mapping[str, Mapping[DataCoordinate, uuid.UUID]]
+    ) -> None:
+        """Update a networkx graph of quanta in place by adding edges to
+        further constrain the ordering.
+
+        Parameters
+        ----------
+        xgraph : `networkx.DiGraph`
+            A directed acyclic graph of quanta to modify in place.  Node keys
+            are quantum UUIDs, and attributes include ``task_label`` (`str`)
+            and ``data_id`` (a full `lsst.daf.butler.DataCoordinate`, without
+            dimension records attached).  Edges may be added, but not removed.
+            Nodes may not be modified.
+        quanta_by_task : `~collections.abc.Mapping` [ `str`,\
+                `~collections.abc.Mapping` [ `lsst.daf.butler.DataCoordinate`,\
+                `uuid.UUID` ] ]
+            All quanta in the graph, grouped first by task label and then by
+            data ID.
         """
-        raise NotImplementedError
+        raise NotImplementedError()

lsst/pipe/base/graph/_versionDeserializers.py

@@ -38,7 +38,7 @@ from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, ClassVar, cast
+from typing import TYPE_CHECKING, ClassVar
 
 import networkx as nx
 
@@ -50,6 +50,7 @@ from lsst.daf.butler import (
     Quantum,
     SerializedDimensionRecord,
 )
+from lsst.daf.butler._rubin import generate_uuidv7
 from lsst.utils import doImportType
 
 from ..config import PipelineTaskConfig
@@ -242,7 +243,7 @@ class DeserializerV1(DeserializerBase):
 
             # reconstruct node
             qNode = pickle.loads(dump)
-            object.__setattr__(qNode, "nodeId", uuid.uuid4())
+            object.__setattr__(qNode, "nodeId", generate_uuidv7())
 
             # read the saved node, name. If it has been loaded, attach it, if
             # not read in the taskDef first, and then load it
@@ -376,7 +377,7 @@ class DeserializerV2(DeserializerBase):
 
             # reconstruct node
             qNode = pickle.loads(dump)
-            object.__setattr__(qNode, "nodeId", uuid.uuid4())
+            object.__setattr__(qNode, "nodeId", generate_uuidv7())
 
             # read the saved node, name. If it has been loaded, attach it, if
             # not read in the taskDef first, and then load it
@@ -599,11 +600,11 @@ class DeserializerV3(DeserializerBase):
                 # initInputRefs and initOutputRefs are optional
                 if (refs := taskDefDump.get("initInputRefs")) is not None:
                     initInputRefs[recreatedTaskDef.label] = [
-                        cast(DatasetRef, DatasetRef.from_json(ref, universe=universe)) for ref in refs
+                        DatasetRef.from_json(ref, universe=universe) for ref in refs
                     ]
                 if (refs := taskDefDump.get("initOutputRefs")) is not None:
                     initOutputRefs[recreatedTaskDef.label] = [
-                        cast(DatasetRef, DatasetRef.from_json(ref, universe=universe)) for ref in refs
+                        DatasetRef.from_json(ref, universe=universe) for ref in refs
                     ]
 
                 # rebuild the mappings that associate dataset type names with

lsst/pipe/base/graph/graph.py

@@ -59,6 +59,7 @@ from lsst.daf.butler import (
     Quantum,
     QuantumBackedButler,
 )
+from lsst.daf.butler._rubin import generate_uuidv7
 from lsst.daf.butler.datastore.record_data import DatastoreRecordData
 from lsst.daf.butler.persistence_context import PersistenceContextVars
 from lsst.daf.butler.registry import ConflictingDefinitionError
@@ -246,7 +247,7 @@ class QuantumGraph:
                             "associated value in the mapping"
                         )
                 else:
-                    nodeId = uuid.uuid4()
+                    nodeId = generate_uuidv7()
 
                 inits = quantum.initInputs.values()
                 inputs = quantum.inputs.values()
@@ -806,11 +807,18 @@ class QuantumGraph:
         uri : convertible to `~lsst.resources.ResourcePath`
             URI to where the graph should be saved.
         """
-        buffer = self._buildSaveObject()
         path = ResourcePath(uri)
-        if path.getExtension() not in (".qgraph"):
-            raise TypeError(f"Can currently only save a graph in qgraph format not {uri}")
-        path.write(buffer)  # type: ignore  # Ignore because bytearray is safe to use in place of bytes
+        match path.getExtension():
+            case ".qgraph":
+                buffer = self._buildSaveObject()
+                path.write(buffer)  # type: ignore  # Ignore because bytearray is safe to use in place of bytes
+            case ".qg":
+                from ..quantum_graph import PredictedQuantumGraphComponents
+
+                pqg = PredictedQuantumGraphComponents.from_old_quantum_graph(self)
+                pqg.write(path)
+            case ext:
+                raise TypeError(f"Can currently only save a graph in .qgraph or .qg format, not {ext!r}.")
 
     @property
     def metadata(self) -> MappingProxyType[str, Any]:
@@ -962,11 +970,23 @@ class QuantumGraph:
             the graph.
         """
         uri = ResourcePath(uri)
-        if uri.getExtension() in {".qgraph"}:
-            with LoadHelper(uri, minimumVersion, fullRead=(nodes is None)) as loader:
-                qgraph = loader.load(universe, nodes, graphID)
-        else:
-            raise ValueError(f"Only know how to handle files saved as `.qgraph`, not {uri}")
+        match uri.getExtension():
+            case ".qgraph":
+                with LoadHelper(uri, minimumVersion, fullRead=(nodes is None)) as loader:
+                    qgraph = loader.load(universe, nodes, graphID)
+            case ".qg":
+                from ..quantum_graph import PredictedQuantumGraphReader
+
+                with PredictedQuantumGraphReader.open(uri, page_size=100000) as qgr:
+                    quantum_ids = (
+                        [uuid.UUID(q) if not isinstance(q, uuid.UUID) else q for q in nodes]
+                        if nodes is not None
+                        else None
+                    )
+                    qgr.read_execution_quanta(quantum_ids)
+                    qgraph = qgr.finish().to_old_quantum_graph()
+            case _:
+                raise ValueError(f"Only know how to handle files saved as `.qgraph`, not {uri}")
         if not isinstance(qgraph, QuantumGraph):
             raise TypeError(f"QuantumGraph file {uri} contains unexpected object type: {type(qgraph)}")
         return qgraph

lsst/pipe/base/graph/graphSummary.py

@@ -75,6 +75,21 @@ class QgraphTaskSummary(pydantic.BaseModel):
             """See `pydantic.BaseModel.model_json_schema`."""
             return super().model_json_schema(*args, **kwargs)
 
+        @classmethod
+        def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate`."""
+            return super().model_validate(*args, **kwargs)
+
+        @classmethod
+        def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_json`."""
+            return super().model_validate_json(*args, **kwargs)
+
+        @classmethod
+        def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_strings`."""
+            return super().model_validate_strings(*args, **kwargs)
+
 
 class QgraphSummary(pydantic.BaseModel):
     """Report for the QuantumGraph creation or reading."""
@@ -129,3 +144,18 @@ class QgraphSummary(pydantic.BaseModel):
         def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
             """See `pydantic.BaseModel.model_json_schema`."""
             return super().model_json_schema(*args, **kwargs)
+
+        @classmethod
+        def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate`."""
+            return super().model_validate(*args, **kwargs)
+
+        @classmethod
+        def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_json`."""
+            return super().model_validate_json(*args, **kwargs)
+
+        @classmethod
+        def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_validate_strings`."""
+            return super().model_validate_strings(*args, **kwargs)

lsst/pipe/base/graph_walker.py

@@ -0,0 +1,119 @@
+# This file is part of pipe_base.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This software is dual licensed under the GNU General Public License and also
+# under a 3-clause BSD license. Recipients may choose which of these licenses
+# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
+# respectively.  If you choose the GPL option then the following text applies
+# (but note that there is still no warranty even if you opt for BSD instead):
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from __future__ import annotations
+
+__all__ = ("GraphWalker",)
+
+from typing import Generic, Self, TypeVar
+
+import networkx
+
+_T = TypeVar("_T")
+
+
+class GraphWalker(Generic[_T]):
+    """A helper for traversing directed acyclic graphs.
+
+    Parameters
+    ----------
+    xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
+        Networkx graph to process.  Will be consumed during iteration, so this
+        should often be a copy.
+
+    Notes
+    -----
+    Each iteration yields a `frozenset` of nodes, which may be empty if there
+    are no nodes ready for processing.  A node is only considered ready if all
+    of its predecessor nodes have been marked as complete with `finish`.
+    Iteration only completes when all nodes have been finished or failed.
+
+    `GraphWalker` is not thread-safe; calling one `GraphWalker` method while
+    another is in progress is undefined behavior.  It is designed to be used
+    in the management thread or process in a parallel traversal.
+    """
+
+    def __init__(self, xgraph: networkx.DiGraph | networkx.MultiDiGraph):
+        self._xgraph = xgraph
+        self._ready: set[_T] = set(next(iter(networkx.dag.topological_generations(self._xgraph)), []))
+        self._active: set[_T] = set()
+        self._incomplete: set[_T] = set(self._xgraph)
+
+    def __iter__(self) -> Self:
+        return self
+
+    def __next__(self) -> frozenset[_T]:
+        if not self._incomplete:
+            raise StopIteration()
+        new_active = frozenset(self._ready)
+        self._active.update(new_active)
+        self._ready.clear()
+        return new_active
+
+    def finish(self, key: _T) -> None:
+        """Mark a node as successfully processed, unblocking (at least in part)
+        iteration over successor nodes.
+
+        Parameters
+        ----------
+        key : unspecified
+            NetworkX key of the node to mark finished.
+        """
+        self._active.remove(key)
+        self._incomplete.remove(key)
+        successors = list(self._xgraph.successors(key))
+        for successor in successors:
+            assert successor not in self._active, (
+                "A node downstream of an active one should not have been yielded yet."
+            )
+            if all(
+                predecessor not in self._incomplete for predecessor in self._xgraph.predecessors(successor)
+            ):
+                self._ready.add(successor)
+
+    def fail(self, key: _T) -> list[_T]:
+        """Mark a node as unsuccessfully processed, permanently blocking all
+        recursive descendants.
+
+        Parameters
+        ----------
+        key : unspecified
+            NetworkX key of the node to mark as a failure.
+
+        Returns
+        -------
+        blocked : `list`
+            NetworkX keys of nodes that were recursive descendants of the
+            failed node, and will hence never be yielded by the iterator.
+        """
+        self._active.remove(key)
+        self._incomplete.remove(key)
+        descendants = list(networkx.dag.descendants(self._xgraph, key))
+        self._xgraph.remove_node(key)
+        self._xgraph.remove_nodes_from(descendants)
+        self._incomplete.difference_update(descendants)
+        return descendants