lsst-pipe-base 30.0.0rc2__py3-none-any.whl → 30.0.1rc1__py3-none-any.whl

lsst/pipe/base/quantum_graph_executor.py

@@ -27,23 +27,113 @@
 
 from __future__ import annotations
 
-__all__ = ["QuantumExecutor", "QuantumGraphExecutor"]
+__all__ = ["QuantumExecutionResult", "QuantumExecutor", "QuantumGraphExecutor"]
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Self
+
+from lsst.daf.butler import Quantum
 
 from .quantum_reports import QuantumReport, Report
 
 if TYPE_CHECKING:
     import uuid
 
-    from lsst.daf.butler import Quantum
+    from lsst.daf.butler.logging import ButlerLogRecords
 
+    from ._task_metadata import TaskMetadata
     from .graph import QuantumGraph
     from .pipeline_graph import TaskNode
     from .quantum_graph import PredictedQuantumGraph
 
 
+class QuantumExecutionResult(tuple[Quantum, QuantumReport | None]):
+    """A result struct that captures information about a single quantum's
+    execution.
+
+    Parameters
+    ----------
+    quantum : `lsst.daf.butler.Quantum`
+        Quantum that was executed.
+    report : `.quantum_reports.QuantumReport`
+        Report with basic information about the execution.
+    task_metadata : `TaskMetadata`, optional
+        Metadata saved by the task and executor during execution.
+    skipped_existing : `bool`, optional
+        If `True`, this quantum was not executed because it appeared to have
+        already been executed successfully.
+    adjusted_no_work : `bool`, optional
+        If `True`, this quantum was not executed because the
+        `PipelineTaskConnections.adjustQuanta` hook raised `NoWorkFound`.
+
+    Notes
+    -----
+    For backwards compatibility, this class is a two-element tuple that allows
+    the ``quantum`` and ``report`` attributes to be unpacked.  Additional
+    regular attributes may be added by executors (but the tuple must remain
+    only two elements to enable the current unpacking interface).
+    """
+
+    def __new__(
+        cls,
+        quantum: Quantum,
+        report: QuantumReport | None,
+        *,
+        task_metadata: TaskMetadata | None = None,
+        skipped_existing: bool | None = None,
+        adjusted_no_work: bool | None = None,
+    ) -> Self:
+        return super().__new__(cls, (quantum, report))
+
+    # We need to define both __init__ and __new__ because tuple inheritance
+    # requires __new__ and numpydoc requires __init__.
+
+    def __init__(
+        self,
+        quantum: Quantum,
+        report: QuantumReport | None,
+        *,
+        task_metadata: TaskMetadata | None = None,
+        skipped_existing: bool | None = None,
+        adjusted_no_work: bool | None = None,
+    ):
+        self._task_metadata = task_metadata
+        self._skipped_existing = skipped_existing
+        self._adjusted_no_work = adjusted_no_work
+
+    @property
+    def quantum(self) -> Quantum:
+        """The quantum actually executed."""
+        return self[0]
+
+    @property
+    def report(self) -> QuantumReport | None:
+        """Structure describing the status of the execution of a quantum.
+
+        This is `None` if the implementation does not support this feature.
+        """
+        return self[1]
+
+    @property
+    def task_metadata(self) -> TaskMetadata | None:
+        """Metadata saved by the task and executor during execution."""
+        return self._task_metadata
+
+    @property
+    def skipped_existing(self) -> bool | None:
+        """If `True`, this quantum was not executed because it appeared to have
+        already been executed successfully.
+        """
+        return self._skipped_existing
+
+    @property
+    def adjusted_no_work(self) -> bool | None:
+        """If `True`, this quantum was not executed because the
+        `PipelineTaskConnections.adjustQuanta` hook raised `NoWorkFound`.
+        """
+        return self._adjusted_no_work
+
+
 class QuantumExecutor(ABC):
     """Class which abstracts execution of a single Quantum.
 
@@ -55,8 +145,14 @@ class QuantumExecutor(ABC):
 
     @abstractmethod
     def execute(
-        self, task_node: TaskNode, /, quantum: Quantum, quantum_id: uuid.UUID | None = None
-    ) -> tuple[Quantum, QuantumReport | None]:
+        self,
+        task_node: TaskNode,
+        /,
+        quantum: Quantum,
+        quantum_id: uuid.UUID | None = None,
+        *,
+        log_records: ButlerLogRecords | None = None,
+    ) -> QuantumExecutionResult:
         """Execute single quantum.
 
         Parameters
@@ -67,15 +163,18 @@ class QuantumExecutor(ABC):
             Quantum for this execution.
         quantum_id : `uuid.UUID` or `None`, optional
             The ID of the quantum to be executed.
+        log_records : `lsst.daf.butler.ButlerLogRecords`, optional
+            Container that should be used to store logs in memory before
+            writing them to the butler.  This disables streaming log (since
+            we'd have to store them in memory anyway), but it permits the
+            caller to prepend logs to be stored in the butler and allows task
+            logs to be inspected by the caller after execution is complete.
 
         Returns
         -------
-        quantum : `~lsst.daf.butler.Quantum`
-            The quantum actually executed.
-        report : `~.quantum_reports.QuantumReport`
-            Structure describing the status of the execution of a quantum.
-            `None` is returned if implementation does not support this
-            feature.
+        result : `QuantumExecutionResult`
+            Result struct.  May also be unpacked as a 2-tuple (see type
+            documentation).
 
         Notes
         -----
@@ -93,7 +192,9 @@ class QuantumGraphExecutor(ABC):
     """
 
     @abstractmethod
-    def execute(self, graph: QuantumGraph | PredictedQuantumGraph) -> None:
+    def execute(
+        self, graph: QuantumGraph | PredictedQuantumGraph, *, provenance_graph_file: str | None = None
+    ) -> None:
         """Execute whole graph.
 
         Implementation of this method depends on particular execution model
@@ -103,8 +204,10 @@ class QuantumGraphExecutor(ABC):
 
         Parameters
         ----------
-        graph : `.QuantumGraph`
+        graph : `.QuantumGraph` or `.quantum_graph.PredictedQuantumGraph`
             Execution graph.
+        provenance_graph_file : `str`, optional
+            A filename to write provenance to.
         """
         raise NotImplementedError()
 

lsst/pipe/base/quantum_graph_skeleton.py

@@ -42,7 +42,7 @@ __all__ = (
 import dataclasses
 from collections import defaultdict
 from collections.abc import Iterable, Iterator, MutableMapping, Set
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 import networkx
 
@@ -145,7 +145,7 @@ class PrerequisiteDatasetKey:
     is_prerequisite: ClassVar[Literal[True]] = True
 
 
-Key: TypeAlias = QuantumKey | TaskInitKey | DatasetKey | PrerequisiteDatasetKey
+type Key = QuantumKey | TaskInitKey | DatasetKey | PrerequisiteDatasetKey
 
 
 class QuantumGraphSkeleton:
@@ -383,12 +383,6 @@ class QuantumGraphSkeleton:
             The dataset ref of the prerequisite.
         **attrs : `~typing.Any`
             Additional attributes for the node.
-
-        Notes
-        -----
-        This automatically sets the 'existing_input' ref attribute (see
-        `set_existing_input_ref`), since prerequisites are always overall
-        inputs.
         """
         key = PrerequisiteDatasetKey(ref.datasetType.name, ref.id.bytes)
         self._xgraph.add_node(key, data_id=ref.dataId, ref=ref, **attrs)
@@ -577,12 +571,13 @@ class QuantumGraphSkeleton:
     def set_dataset_ref(
         self, ref: DatasetRef, key: DatasetKey | PrerequisiteDatasetKey | None = None
     ) -> None:
-        """Associate a dataset node with a `DatasetRef` instance.
+        """Associate a dataset node with a `~lsst.daf.butler.DatasetRef`
+        instance.
 
         Parameters
         ----------
-        ref : `DatasetRef`
-            `DatasetRef` to associate with the node.
+        ref : `~lsst.daf.butler.DatasetRef`
+            `~lsst.daf.butler.DatasetRef` to associate with the node.
         key : `DatasetKey` or `PrerequisiteDatasetKey`, optional
             Identifier for the graph node.  If not provided, a `DatasetKey`
             is constructed from the dataset type name and data ID of ``ref``.
@@ -592,32 +587,33 @@ class QuantumGraphSkeleton:
         self._xgraph.nodes[key]["ref"] = ref
 
     def set_output_for_skip(self, ref: DatasetRef) -> None:
-        """Associate a dataset node with a `DatasetRef` that represents an
-        existing output in a collection where such outputs can cause a quantum
-        to be skipped.
+        """Associate a dataset node with a `~lsst.daf.butler.DatasetRef` that
+        represents an existing output in a collection where such outputs can
+        cause a quantum to be skipped.
 
         Parameters
         ----------
-        ref : `DatasetRef`
-            `DatasetRef` to associate with the node.
+        ref : `~lsst.daf.butler.DatasetRef`
+            `~lsst.daf.butler.DatasetRef` to associate with the node.
         """
         key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
         self._xgraph.nodes[key]["output_for_skip"] = ref
 
     def set_output_in_the_way(self, ref: DatasetRef) -> None:
-        """Associate a dataset node with a `DatasetRef` that represents an
-        existing output in the output RUN collectoin.
+        """Associate a dataset node with a `~lsst.daf.butler.DatasetRef` that
+        represents an existing output in the output RUN collection.
 
         Parameters
         ----------
-        ref : `DatasetRef`
-            `DatasetRef` to associate with the node.
+        ref : `~lsst.daf.butler.DatasetRef`
+            `~lsst.daf.butler.DatasetRef` to associate with the node.
         """
         key = DatasetKey(ref.datasetType.name, ref.dataId.required_values)
         self._xgraph.nodes[key]["output_in_the_way"] = ref
 
     def get_dataset_ref(self, key: DatasetKey | PrerequisiteDatasetKey) -> DatasetRef | None:
-        """Return the `DatasetRef` associated with the given node.
+        """Return the `~lsst.daf.butler.DatasetRef` associated with the given
+        node.
 
         This does not return "output for skip" and "output in the way"
         datasets.
@@ -629,14 +625,14 @@ class QuantumGraphSkeleton:
 
         Returns
         -------
-        ref : `DatasetRef` or `None`
+        ref : `~lsst.daf.butler.DatasetRef` or `None`
             Dataset reference associated with the node.
         """
         return self._xgraph.nodes[key].get("ref")
 
     def get_output_for_skip(self, key: DatasetKey) -> DatasetRef | None:
-        """Return the `DatasetRef` associated with the given node in a
-        collection where it could lead to a quantum being skipped.
+        """Return the `~lsst.daf.butler.DatasetRef` associated with the given
+        node in a collection where it could lead to a quantum being skipped.
 
         Parameters
         ----------
@@ -645,14 +641,14 @@ class QuantumGraphSkeleton:
 
         Returns
         -------
-        ref : `DatasetRef` or `None`
+        ref : `~lsst.daf.butler.DatasetRef` or `None`
             Dataset reference associated with the node.
         """
         return self._xgraph.nodes[key].get("output_for_skip")
 
     def get_output_in_the_way(self, key: DatasetKey) -> DatasetRef | None:
-        """Return the `DatasetRef` associated with the given node in the
-        output RUN collection.
+        """Return the `~lsst.daf.butler.DatasetRef` associated with the given
+        node in the output RUN collection.
 
         Parameters
         ----------
@@ -661,16 +657,16 @@ class QuantumGraphSkeleton:
 
         Returns
         -------
-        ref : `DatasetRef` or `None`
+        ref : `~lsst.daf.butler.DatasetRef` or `None`
             Dataset reference associated with the node.
         """
         return self._xgraph.nodes[key].get("output_in_the_way")
 
     def discard_output_in_the_way(self, key: DatasetKey) -> None:
-        """Drop any `DatasetRef` associated with this node in the output RUN
-        collection.
+        """Drop any `~lsst.daf.butler.DatasetRef` associated with this node in
+        the output RUN collection.
 
-        Does nothing if there is no such `DatasetRef`.
+        Does nothing if there is no such `~lsst.daf.butler.DatasetRef`.
 
         Parameters
         ----------
@@ -682,8 +678,8 @@ class QuantumGraphSkeleton:
     def set_data_id(self, key: Key, data_id: DataCoordinate) -> None:
         """Set the data ID associated with a node.
 
-        This updates the data ID in any `DatasetRef` objects associated with
-        the node via `set_ref`, `set_output_for_skip`, or
+        This updates the data ID in any `~lsst.daf.butler.DatasetRef` objects
+        associated with the node via `set_ref`, `set_output_for_skip`, or
         `set_output_in_the_way` as well, assuming it is an expanded version
         of the original data ID.
 
@@ -691,7 +687,7 @@ class QuantumGraphSkeleton:
         ----------
         key : `Key`
             Identifier for the graph node.
-        data_id : `DataCoordinate`
+        data_id : `~lsst.daf.butler.DataCoordinate`
             Data ID for the node.
         """
         state: MutableMapping[str, Any] = self._xgraph.nodes[key]
@@ -716,7 +712,7 @@ class QuantumGraphSkeleton:
 
         Returns
         -------
-        data_id : `DataCoordinate`
+        data_id : `~lsst.daf.butler.DataCoordinate`
             Expanded data ID for the node, if one is available.
 
         Raises

lsst/pipe/base/quantum_provenance_graph.py

@@ -79,6 +79,7 @@ from .automatic_connection_constants import (
     METADATA_OUTPUT_CONNECTION_NAME,
     METADATA_OUTPUT_STORAGE_CLASS,
     METADATA_OUTPUT_TEMPLATE,
+    PROVENANCE_DATASET_TYPE_NAME,
 )
 from .graph import QuantumGraph, QuantumNode
 
@@ -586,8 +587,8 @@ class TaskSummary(pydantic.BaseModel):
 
         Unpack the `QuantumInfo` object, sorting quanta of each status into
         the correct place in the `TaskSummary`. If looking for error messages
-        in the `Butler` logs is desired, take special care to catch issues
-        with missing logs.
+        in the `lsst.daf.butler.Butler` logs is desired, take special care to
+        catch issues with missing logs.
 
         Parameters
         ----------
@@ -866,7 +867,7 @@ class DatasetTypeSummary(pydantic.BaseModel):
 class Summary(pydantic.BaseModel):
     """A summary of the contents of the QuantumProvenanceGraph, including
     all information on the quanta for each task and the datasets of each
-    `DatasetType`.
+    `~lsst.daf.butler.DatasetType`.
     """
 
     tasks: dict[str, TaskSummary] = pydantic.Field(default_factory=dict)
@@ -885,7 +886,7 @@ class Summary(pydantic.BaseModel):
 
         Parameters
         ----------
-        summaries : `Sequence[Summary]`
+        summaries : `~collections.abc.Sequence` [`Summary`]
             Sequence of all `Summary` objects to aggregate.
         """
         result = cls()
@@ -1245,8 +1246,8 @@ class QuantumProvenanceGraph:
         Returns
         -------
         dataset_info : `DatasetInfo`
-            The `TypedDict` with information about the `DatasetType`-dataID
-            pair across all runs.
+            The `TypedDict` with information about the
+            `~lsst.daf.butler.DatasetType`-dataID pair across all runs.
         """
         return self._xgraph.nodes[key]
 
@@ -1262,6 +1263,7 @@ class QuantumProvenanceGraph:
         do_store_logs : `bool`
             Store the logs in the summary dictionary.
         n_cores : `int`, optional
+            Number of cores to use.
 
         Returns
         -------
@@ -1513,8 +1515,22 @@ class QuantumProvenanceGraph:
                     len(self._datasets.keys()),
                 )
         if use_qbb:
-            _LOG.verbose("Using quantum-backed butler for metadata loads.")
-            self._butler_wrappers[output_run] = _ThreadLocalButlerWrapper.wrap_qbb(butler, qgraph)
+            provenance_graph_ref: DatasetRef | None = None
+            try:
+                provenance_graph_ref = butler.find_dataset(
+                    PROVENANCE_DATASET_TYPE_NAME, collections=output_run
+                )
+            except MissingDatasetTypeError:
+                pass
+            if provenance_graph_ref is not None:
+                _LOG.warning(
+                    "Cannot use QBB for metadata/log reads after provenance has been ingested; "
+                    "falling back to full butler."
+                )
+                self._butler_wrappers[output_run] = _ThreadLocalButlerWrapper.wrap_full(butler)
+            else:
+                _LOG.verbose("Using quantum-backed butler for metadata loads.")
+                self._butler_wrappers[output_run] = _ThreadLocalButlerWrapper.wrap_qbb(butler, qgraph)
         else:
             _LOG.verbose("Using full butler for metadata loads.")
             self._butler_wrappers[output_run] = _ThreadLocalButlerWrapper.wrap_full(butler)
@@ -1777,9 +1793,10 @@ class QuantumProvenanceGraph:
             successes. If "exhaustive", all metadata files will be read.  If
             "lazy", only metadata files where at least one predicted output is
             missing will be read.
-        butler : `lsst.daf.butler.Butler`
-            The Butler used for this report. This should match the Butler
-            used for the run associated with the executed quantum graph.
+        executor : `concurrent.futures.Executor`
+            The futures executor to use.
+        futures : `list` [ `concurrent.futures.Future` ]
+            Current list of futures. Will be modified.
         """
         if read_caveats == "lazy" and all(
             self.get_dataset_info(dataset_key)["runs"][output_run].produced
@@ -1994,7 +2011,7 @@ class _ThreadLocalButlerWrapper:
         full_butler : `~lsst.daf.butler.Butler`
             Full butler to draw datastore and dimension configuration from.
         qg : `QuantumGraph`
-            Quantum graph,
+            Quantum graph.
 
         Returns
         -------

lsst/pipe/base/separable_pipeline_executor.py

@@ -40,7 +40,8 @@ from collections.abc import Iterable
 from typing import Any
 
 import lsst.resources
-from lsst.daf.butler import Butler
+from lsst.daf.butler import Butler, DatasetRef
+from lsst.daf.butler._rubin.temporary_for_ingest import TemporaryForIngest
 
 from ._quantumContext import ExecutionResources
 from .all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
@@ -78,7 +79,7 @@ class SeparablePipelineExecutor:
     clobber_output : `bool`, optional
         If set, the pipeline execution overwrites existing output files.
         Otherwise, any conflict between existing and new outputs is an error.
-    skip_existing_in : iterable [`str`], optional
+    skip_existing_in : `~collections.abc.Iterable` [`str`], optional
         If not empty, the pipeline execution searches the listed collections
         for existing outputs, and skips any quanta that have run to completion
         (or have no work to do). Otherwise, all tasks are attempted (subject to
@@ -362,6 +363,8 @@ class SeparablePipelineExecutor:
         fail_fast: bool = False,
         graph_executor: QuantumGraphExecutor | None = None,
         num_proc: int = 1,
+        *,
+        provenance_dataset_ref: DatasetRef | None = None,
     ) -> None:
         """Run a pipeline in the form of a prepared quantum graph.
 
@@ -384,6 +387,14 @@ class SeparablePipelineExecutor:
             The number of processes that can be used to run the pipeline. The
             default value ensures that no subprocess is created. Only used with
             the default graph executor.
+        provenance_dataset_ref : `lsst.daf.butler.DatasetRef`, optional
+            Dataset that should be used to save provenance.  Provenance is only
+            supported when running in a single process (at least for the
+            default quantum executor), and should not be used with
+            ``skip_existing_in=[output_run]`` when retrying a previous
+            execution attempt. The caller is responsible for registering the
+            dataset type and for ensuring that the dimensions of this dataset
+            do not lead to uniqueness conflicts.
         """
         if not graph_executor:
             quantum_executor = SingleQuantumExecutor(
@@ -404,4 +415,9 @@ class SeparablePipelineExecutor:
             # forked processes.
             self._butler.registry.resetConnectionPool()
 
-        graph_executor.execute(graph)
+        if provenance_dataset_ref is not None:
+            with TemporaryForIngest(self._butler, provenance_dataset_ref) as temporary:
+                graph_executor.execute(graph, provenance_graph_file=temporary.ospath)
+                temporary.ingest()
+        else:
+            graph_executor.execute(graph)