lsst-pipe-base 29.2025.1100__py3-none-any.whl → 29.2025.1300__py3-none-any.whl

lsst/pipe/base/all_dimensions_quantum_graph_builder.py

@@ -213,7 +213,11 @@ class AllDimensionsQuantumGraphBuilder(QuantumGraphBuilder):
                     lines.extend(query.explain_no_results())
                 finally:
                     lines.append("To reproduce this query for debugging purposes, run:")
+                    lines.append("")
                     lines.extend(query_cmd)
+                    lines.append("    print(query.any())")
+                    lines.append("")
+                    lines.append("And then try removing various constraints until query.any() returns True.")
                     # If an exception was raised, write a partial.
                     self.log.error("\n".join(lines))
                 return

lsst/pipe/base/connections.py

@@ -35,6 +35,7 @@ __all__ = [
     "InputQuantizedConnection",
     "OutputQuantizedConnection",
     "PipelineTaskConnections",
+    "QuantaAdjuster",
     "QuantizedConnection",
     "ScalarError",
     "ScalarError",
@@ -45,8 +46,8 @@ import dataclasses
 import itertools
 import string
 import warnings
-from collections import UserDict
-from collections.abc import Collection, Generator, Iterable, Mapping, Sequence, Set
+from collections import UserDict, defaultdict
+from collections.abc import Collection, Generator, Iterable, Iterator, Mapping, Sequence, Set
 from dataclasses import dataclass
 from types import MappingProxyType, SimpleNamespace
 from typing import TYPE_CHECKING, Any
@@ -58,6 +59,8 @@ from .connectionTypes import BaseConnection, BaseInput, Output, PrerequisiteInpu
 
 if TYPE_CHECKING:
     from .config import PipelineTaskConfig
+    from .pipeline_graph import PipelineGraph, TaskNode
+    from .quantum_graph_skeleton import QuantumGraphSkeleton
 
 
 class ScalarError(TypeError):
@@ -999,6 +1002,25 @@ class PipelineTaskConnections(metaclass=PipelineTaskConnectionsMetaclass):
         """
         return ()
 
+    def adjust_all_quanta(self, adjuster: QuantaAdjuster) -> None:
+        """Customize the set of quanta predicted for this task during quantum
+        graph generation.
+
+        Parameters
+        ----------
+        adjuster : `QuantaAdjuster`
+            A helper object that implementations can use to modify the
+            under-construction quantum graph.
+
+        Notes
+        -----
+        This hook is called before `adjustQuantum`, which is where built-in
+        checks for `NoWorkFound` cases and missing prerequisites are handled.
+        This means that the set of preliminary quanta seen by this method could
+        include some that would normally be dropped later.
+        """
+        pass
+
 
 def iterConnections(
     connections: PipelineTaskConnections, connectionType: str | Iterable[str]
@@ -1130,3 +1152,158 @@ class AdjustQuantumHelper:
             self.outputs_adjusted = True
         else:
             self.outputs_adjusted = False
+
+
+class QuantaAdjuster:
+    """A helper class for the `PipelineTaskConnections.adjust_all_quanta` hook.
+
+    Parameters
+    ----------
+    task_label : `str`
+        Label of the task whose quanta are being adjusted.
+    pipeline_graph : `pipeline_graph.PipelineGraph`
+        Pipeline graph the quantum graph is being built from.
+    skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
+        Under-construction quantum graph that will be modified in place.
+    """
+
+    def __init__(self, task_label: str, pipeline_graph: PipelineGraph, skeleton: QuantumGraphSkeleton):
+        self._task_node = pipeline_graph.tasks[task_label]
+        self._pipeline_graph = pipeline_graph
+        self._skeleton = skeleton
+        self._n_removed = 0
+
+    @property
+    def task_label(self) -> str:
+        """The label this task has been configured with."""
+        return self._task_node.label
+
+    @property
+    def task_node(self) -> TaskNode:
+        """The node for this task in the pipeline graph."""
+        return self._task_node
+
+    def iter_data_ids(self) -> Iterator[DataCoordinate]:
+        """Iterate over the data IDs of all quanta for this task."
+
+        Returns
+        -------
+        data_ids : `~collections.abc.Iterator` [ \
+                `~lsst.daf.butler.DataCoordinate` ]
+            Data IDs.  These are minimal data IDs without dimension records or
+            implied values; use `expand_quantum_data_id` to get a full data ID
+            when needed.
+        """
+        for key in self._skeleton.get_quanta(self._task_node.label):
+            yield DataCoordinate.from_required_values(self._task_node.dimensions, key.data_id_values)
+
+    def remove_quantum(self, data_id: DataCoordinate) -> None:
+        """Remove a quantum from the graph.
+
+        Parameters
+        ----------
+        data_id : `~lsst.daf.butler.DataCoordinate`
+            Data ID of the quantum to remove.  All outputs will be removed as
+            well.
+        """
+        from .quantum_graph_skeleton import QuantumKey
+
+        self._skeleton.remove_quantum_node(
+            QuantumKey(self._task_node.label, data_id.required_values), remove_outputs=True
+        )
+        self._n_removed += 1
+
+    def get_inputs(self, quantum_data_id: DataCoordinate) -> dict[str, list[DataCoordinate]]:
+        """Return the data IDs of all regular inputs to a quantum.
+
+        Parameters
+        ----------
+        data_id : `~lsst.daf.butler.DataCoordinate`
+            Data ID of the quantum to get the inputs of.
+
+        Returns
+        -------
+        inputs : `dict` [ `str`, `list` [ `~lsst.daf.butler.DataCoordinate` ] ]
+            Data IDs of inputs, keyed by the connection name (the internal task
+            name, not the dataset type name).  This only contains regular
+            inputs, not init-inputs or prerequisite inputs.
+
+        Notes
+        -----
+        If two connections have the same dataset type, the current
+        implementation assumes the set of datasets is the same for the two
+        connections.  This limitation may be removed in the future.
+        """
+        from .quantum_graph_skeleton import DatasetKey, QuantumKey
+
+        by_dataset_type_name: defaultdict[str, list[DataCoordinate]] = defaultdict(list)
+        quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
+        for dataset_key in self._skeleton.iter_inputs_of(quantum_key):
+            if not isinstance(dataset_key, DatasetKey):
+                continue
+            dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
+            by_dataset_type_name[dataset_key.parent_dataset_type_name].append(
+                DataCoordinate.from_required_values(dataset_type_node.dimensions, dataset_key.data_id_values)
+            )
+        return {
+            edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name]
+            for edge in self._task_node.iter_all_inputs()
+        }
+
+    def add_input(
+        self, quantum_data_id: DataCoordinate, connection_name: str, dataset_data_id: DataCoordinate
+    ) -> None:
+        """Add a new input to a quantum.
+
+        Parameters
+        ----------
+        quantum_data_id : `~lsst.daf.butler.DataCoordinate`
+            Data ID of the quantum to add an input to.
+        connection_name : `str`
+            Name of the connection (the task-internal name, not the butler
+            dataset type name).
+        dataset_data_id : `~lsst.daf.butler.DataCoordinate`
+            Data ID of the input dataset.  Must already exist in the graph
+            as an input to a different quantum of this task, and must be a
+            regular input, not a prerequisite input or init-input.
+
+        Notes
+        -----
+        If two connections have the same dataset type, the current
+        implementation assumes the set of datasets is the same for the two
+        connections.  This limitation may be removed in the future.
+        """
+        from .quantum_graph_skeleton import DatasetKey, QuantumKey
+
+        quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
+        read_edge = self._task_node.inputs[connection_name]
+        dataset_key = DatasetKey(read_edge.parent_dataset_type_name, dataset_data_id.required_values)
+        if dataset_key not in self._skeleton:
+            raise LookupError(
+                f"Dataset {read_edge.parent_dataset_type_name}@{dataset_data_id} is not already in the graph."
+            )
+        self._skeleton.add_input_edge(quantum_key, dataset_key)
+
+    def expand_quantum_data_id(self, data_id: DataCoordinate) -> DataCoordinate:
+        """Expand a quantum data ID to include implied values and dimension
+        records.
+
+        Parameters
+        ----------
+        quantum_data_id : `~lsst.daf.butler.DataCoordinate`
+            A data ID of a quantum already in the graph.
+
+        Returns
+        -------
+        expanded_data_id : `~lsst.daf.butler.DataCoordinate`
+            The same data ID, with implied values included and dimension
+            records attached.
+        """
+        from .quantum_graph_skeleton import QuantumKey
+
+        return self._skeleton.get_data_id(QuantumKey(self._task_node.label, data_id.required_values))
+
+    @property
+    def n_removed(self) -> int:
+        """The number of quanta that have been removed by this helper."""
+        return self._n_removed

lsst/pipe/base/pipeline_graph/visualization/_mermaid.py

@@ -29,7 +29,6 @@ from __future__ import annotations
 __all__ = ("show_mermaid",)
 
 import html
-import importlib.util
 import os
 import sys
 from collections.abc import Mapping
@@ -42,12 +41,14 @@ from ._formatting import NodeKey, format_dimensions, format_task_class
 from ._options import NodeAttributeOptions
 from ._show import parse_display_args
 
-MERMAID_AVAILABLE = importlib.util.find_spec("mermaid") is not None
-
-if MERMAID_AVAILABLE:
+try:
     from mermaid import Mermaid  # type: ignore
     from mermaid.graph import Graph  # type: ignore
 
+    MERMAID_AVAILABLE = True
+except ImportError:
+    MERMAID_AVAILABLE = False
+
 # Configuration constants for label formatting and overflow handling.
 _LABEL_PX_SIZE = 18
 _LABEL_MAX_LINES_SOFT = 10
@@ -237,6 +238,8 @@ def _render_mermaid_image(
 
     Raises
     ------
+    ImportError
+        If `mermaid-py` is not installed.
     ValueError
         If the requested ``output_format`` is not supported.
     RuntimeError
@@ -245,6 +248,9 @@ def _render_mermaid_image(
     if output_format.lower() not in {"svg", "png"}:
         raise ValueError(f"Unsupported format: {output_format}. Use 'svg' or 'png'.")
 
+    if not MERMAID_AVAILABLE:
+        raise ImportError("The `mermaid-py` package is required for rendering images but is not installed.")
+
     # Generate Mermaid graph object.
     graph = Graph(title="Mermaid Diagram", script=mermaid_source)
     diagram = Mermaid(graph, width=width, height=height, scale=scale)

lsst/pipe/base/quantum_graph_builder.py

@@ -62,7 +62,7 @@ from lsst.utils.timer import timeMethod
 from . import automatic_connection_constants as acc
 from ._status import NoWorkFound
 from ._task_metadata import TaskMetadata
-from .connections import AdjustQuantumHelper
+from .connections import AdjustQuantumHelper, QuantaAdjuster
 from .graph import QuantumGraph
 from .pipeline_graph import PipelineGraph, TaskNode
 from .prerequisite_helpers import PrerequisiteInfo, SkyPixBoundsBuilder, TimespanBuilder
@@ -475,7 +475,6 @@ class QuantumGraphBuilder(ABC):
         # Loop over all quanta for this task, remembering the ones we've
         # gotten rid of.
         skipped_quanta = []
-        no_work_quanta = []
         for quantum_key in skeleton.get_quanta(task_node.label):
             if self._skip_quantum_if_metadata_exists(task_node, quantum_key, skeleton):
                 skipped_quanta.append(quantum_key)
@@ -483,17 +482,26 @@ class QuantumGraphBuilder(ABC):
             quantum_data_id = skeleton[quantum_key]["data_id"]
             skypix_bounds_builder = task_prerequisite_info.bounds.make_skypix_bounds_builder(quantum_data_id)
             timespan_builder = task_prerequisite_info.bounds.make_timespan_builder(quantum_data_id)
-            adjusted_outputs = self._gather_quantum_outputs(
-                task_node, quantum_key, skeleton, skypix_bounds_builder, timespan_builder
-            )
-            adjusted_inputs = self._gather_quantum_inputs(
-                task_node,
+            self._update_quantum_for_adjust(
                 quantum_key,
                 skeleton,
                 task_prerequisite_info,
                 skypix_bounds_builder,
                 timespan_builder,
             )
+        for skipped_quantum in skipped_quanta:
+            skeleton.remove_quantum_node(skipped_quantum, remove_outputs=False)
+        # Give the task a chance to adjust all quanta together.  This
+        # operates directly on the skeleton (via a the 'adjuster', which
+        # is just an interface adapter).
+        adjuster = QuantaAdjuster(task_node.label, self._pipeline_graph, skeleton)
+        task_node.get_connections().adjust_all_quanta(adjuster)
+        # Loop over all quanta again, remembering those we get rid of in other
+        # ways.
+        no_work_quanta = []
+        for quantum_key in skeleton.get_quanta(task_node.label):
+            adjusted_outputs = self._adapt_quantum_outputs(task_node, quantum_key, skeleton)
+            adjusted_inputs = self._adapt_quantum_inputs(task_node, quantum_key, skeleton)
             # Give the task's Connections class an opportunity to remove
             # some inputs, or complain if they are unacceptable.  This will
             # raise if one of the check conditions is not met, which is the
@@ -552,8 +560,6 @@ class QuantumGraphBuilder(ABC):
             skeleton[quantum_key]["outputs"] = helper.outputs
         for no_work_quantum in no_work_quanta:
             skeleton.remove_quantum_node(no_work_quantum, remove_outputs=True)
-        for skipped_quantum in skipped_quanta:
-            skeleton.remove_quantum_node(skipped_quantum, remove_outputs=False)
         remaining_quanta = skeleton.get_quanta(task_node.label)
         self._resolve_task_init(task_node, skeleton, bool(skipped_quanta))
         message_terms = []
@@ -561,6 +567,8 @@ class QuantumGraphBuilder(ABC):
             message_terms.append(f"{len(no_work_quanta)} had no work to do")
         if skipped_quanta:
             message_terms.append(f"{len(skipped_quanta)} previously succeeded")
+        if adjuster.n_removed:
+            message_terms.append(f"{adjuster.n_removed} removed by adjust_all_quanta")
         message_parenthetical = f" ({', '.join(message_terms)})" if message_terms else ""
         if remaining_quanta:
             self.log.info(
@@ -634,40 +642,32 @@ class QuantumGraphBuilder(ABC):
         return False
 
     @final
-    def _gather_quantum_outputs(
+    def _update_quantum_for_adjust(
         self,
-        task_node: TaskNode,
         quantum_key: QuantumKey,
         skeleton: QuantumGraphSkeleton,
+        task_prerequisite_info: PrerequisiteInfo,
         skypix_bounds_builder: SkyPixBoundsBuilder,
         timespan_builder: TimespanBuilder,
-    ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
-        """Collect outputs or generate datasets for a preliminary quantum and
-        put them in the form used by `~lsst.daf.butler.Quantum` and
-        `~PipelineTaskConnections.adjustQuantum`.
+    ) -> None:
+        """Update the quantum node in the skeleton by finding remaining
+        prerequisite inputs and dropping regular inputs that we now know will
+        not be produced.
 
         Parameters
         ----------
-        task_node : `pipeline_graph.TaskNode`
-            Node for this task in the pipeline graph.
         quantum_key : `QuantumKey`
             Identifier for this quantum in the graph.
         skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
             Preliminary quantum graph, to be modified in-place.
+        task_prerequisite_info : `~prerequisite_helpers.PrerequisiteInfo`
+            Information about the prerequisite inputs to this task.
         skypix_bounds_builder : `~prerequisite_helpers.SkyPixBoundsBuilder`
             An object that accumulates the appropriate spatial bounds for a
             quantum.
         timespan_builder : `~prerequisite_helpers.TimespanBuilder`
             An object that accumulates the appropriate timespan for a quantum.
 
-        Returns
-        -------
-        outputs : `~lsst.daf.butler.NamedKeyDict` [ \
-                `~lsst.daf.butler.DatasetType`, `list` [ \
-                `~lsst.daf.butler.DatasetRef` ] ]
-            All outputs to the task, using the storage class and components
-            defined by the task's own connections.
-
         Notes
         -----
         This first looks for outputs already present in the `output_run` (i.e.
@@ -680,8 +680,7 @@ class QuantumGraphBuilder(ABC):
         UUID is generated.  In all cases the dataset node in the skeleton is
         associated with a `DatasetRef`.
         """
-        outputs_by_type: dict[str, list[DatasetRef]] = {}
-        dataset_key: DatasetKey
+        dataset_key: DatasetKey | PrerequisiteDatasetKey
         for dataset_key in skeleton.iter_outputs_of(quantum_key):
             dataset_data_id = skeleton[dataset_key]["data_id"]
             dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
@@ -699,6 +698,66 @@ class QuantumGraphBuilder(ABC):
             skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
             timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
             skeleton.set_dataset_ref(ref, dataset_key)
+        quantum_data_id = skeleton[quantum_key]["data_id"]
+        # Process inputs already present in the skeleton - this should include
+        # all regular inputs (including intermediates) and may include some
+        # prerequisites.
+        for dataset_key in list(skeleton.iter_inputs_of(quantum_key)):
+            if (ref := skeleton.get_dataset_ref(dataset_key)) is None:
+                # If the dataset ref hasn't been set either as an existing
+                # input or as an output of an already-processed upstream
+                # quantum, it's not going to be produced; remove it.
+                skeleton.remove_dataset_nodes([dataset_key])
+                continue
+            skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
+            timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
+        # Query for any prerequisites not handled by process_subgraph.  Note
+        # that these were not already in the skeleton graph, so we add them
+        # now.
+        skypix_bounds = skypix_bounds_builder.finish()
+        timespan = timespan_builder.finish()
+        for finder in task_prerequisite_info.finders.values():
+            dataset_keys = []
+            for ref in finder.find(
+                self.butler, self.input_collections, quantum_data_id, skypix_bounds, timespan
+            ):
+                dataset_key = skeleton.add_prerequisite_node(ref)
+                dataset_keys.append(dataset_key)
+            skeleton.add_input_edges(quantum_key, dataset_keys)
+
+    @final
+    def _adapt_quantum_outputs(
+        self,
+        task_node: TaskNode,
+        quantum_key: QuantumKey,
+        skeleton: QuantumGraphSkeleton,
+    ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
+        """Adapt outputs for a preliminary quantum and put them into the form
+        used by `~lsst.daf.butler.Quantum` and
+        `~PipelineTaskConnections.adjustQuantum`.
+
+        Parameters
+        ----------
+        task_node : `pipeline_graph.TaskNode`
+            Node for this task in the pipeline graph.
+        quantum_key : `QuantumKey`
+            Identifier for this quantum in the graph.
+        skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
+            Preliminary quantum graph, to be modified in-place.
+
+        Returns
+        -------
+        outputs : `~lsst.daf.butler.NamedKeyDict` [ \
+                `~lsst.daf.butler.DatasetType`, `list` [ \
+                `~lsst.daf.butler.DatasetRef` ] ]
+            All outputs to the task, using the storage class and components
+            defined by the task's own connections.
+        """
+        outputs_by_type: dict[str, list[DatasetRef]] = {}
+        dataset_key: DatasetKey
+        for dataset_key in skeleton.iter_outputs_of(quantum_key):
+            ref = skeleton.get_dataset_ref(dataset_key)
+            assert ref is not None, "Should have been added (or the node removed) in a previous pass."
             outputs_by_type.setdefault(dataset_key.parent_dataset_type_name, []).append(ref)
         adapted_outputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
         for write_edge in task_node.iter_all_outputs():
@@ -711,17 +770,14 @@ class QuantumGraphBuilder(ABC):
         return adapted_outputs
 
     @final
-    def _gather_quantum_inputs(
+    def _adapt_quantum_inputs(
         self,
         task_node: TaskNode,
         quantum_key: QuantumKey,
         skeleton: QuantumGraphSkeleton,
-        task_prerequisite_info: PrerequisiteInfo,
-        skypix_bounds_builder: SkyPixBoundsBuilder,
-        timespan_builder: TimespanBuilder,
     ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
-        """Collect input datasets for a preliminary quantum and put them in the
-        form used by `~lsst.daf.butler.Quantum` and
+        """Adapt input datasets for a preliminary quantum into the form used by
+        `~lsst.daf.butler.Quantum` and
         `~PipelineTaskConnections.adjustQuantum`.
 
         Parameters
@@ -752,37 +808,12 @@ class QuantumGraphBuilder(ABC):
         with a `DatasetRef`, and queries for prerequisite input nodes that do
         not exist.
         """
-        quantum_data_id = skeleton[quantum_key]["data_id"]
         inputs_by_type: dict[str, set[DatasetRef]] = {}
         dataset_key: DatasetKey | PrerequisiteDatasetKey
-        # Process inputs already present in the skeleton - this should include
-        # all regular inputs (including intermediates) and may include some
-        # prerequisites.
         for dataset_key in list(skeleton.iter_inputs_of(quantum_key)):
-            if (ref := skeleton.get_dataset_ref(dataset_key)) is None:
-                # If the dataset ref hasn't been set either as an existing
-                # input or as an output of an already-processed upstream
-                # quantum, it's not going to be produced; remove it.
-                skeleton.remove_dataset_nodes([dataset_key])
-                continue
+            ref = skeleton.get_dataset_ref(dataset_key)
+            assert ref is not None, "Should have been added (or the node removed) in a previous pass."
             inputs_by_type.setdefault(dataset_key.parent_dataset_type_name, set()).add(ref)
-            skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
-            timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
-        # Query for any prerequisites not handled by process_subgraph.  Note
-        # that these were not already in the skeleton graph, so we add them
-        # now.
-        skypix_bounds = skypix_bounds_builder.finish()
-        timespan = timespan_builder.finish()
-        for finder in task_prerequisite_info.finders.values():
-            inputs_for_type = inputs_by_type.setdefault(finder.dataset_type_node.name, set())
-            dataset_keys = []
-            for ref in finder.find(
-                self.butler, self.input_collections, quantum_data_id, skypix_bounds, timespan
-            ):
-                dataset_key = skeleton.add_prerequisite_node(ref)
-                dataset_keys.append(dataset_key)
-                inputs_for_type.add(ref)
-            skeleton.add_input_edges(quantum_key, dataset_keys)
         adapted_inputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
         for read_edge in task_node.iter_all_inputs():
             dataset_type_node = self._pipeline_graph.dataset_types[read_edge.parent_dataset_type_name]

lsst/pipe/base/quantum_graph_skeleton.py

@@ -669,3 +669,23 @@ class QuantumGraphSkeleton:
         output_in_the_way: DatasetRef | None
         if (output_in_the_way := state.get("output_in_the_way")) is not None:
             state["output_in_the_way"] = output_in_the_way.expanded(data_id)
+
+    def get_data_id(self, key: Key) -> DataCoordinate:
+        """Return the full data ID for a quantum or dataset, if available.
+
+        Parameters
+        ----------
+        key : `Key`
+            Identifier for the graph node.
+
+        Returns
+        -------
+        data_id : `DataCoordinate`
+            Expanded data ID for the node, if one is available.
+
+        Raises
+        ------
+        KeyError
+            Raised if this node does not have an expanded data ID.
+        """
+        return self._xgraph.nodes[key]["data_id"]