lsst-pipe-base 29.2025.3000__py3-none-any.whl → 29.2025.3200__py3-none-any.whl

lsst/pipe/base/quantum_graph_executor.py

@@ -0,0 +1,125 @@
+# 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__ = ["QuantumExecutor", "QuantumGraphExecutor"]
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from .quantum_reports import QuantumReport, Report
+
+if TYPE_CHECKING:
+    import uuid
+
+    from lsst.daf.butler import Quantum
+
+    from .graph import QuantumGraph
+    from .pipeline_graph import TaskNode
+
+
+class QuantumExecutor(ABC):
+    """Class which abstracts execution of a single Quantum.
+
+    In general implementation should not depend on execution model and
+    execution should always happen in-process. Main reason for existence
+    of this class is to provide do-nothing implementation that can be used
+    in the unit tests.
+    """
+
+    @abstractmethod
+    def execute(
+        self, task_node: TaskNode, /, quantum: Quantum, quantum_id: uuid.UUID | None = None
+    ) -> tuple[Quantum, QuantumReport | None]:
+        """Execute single quantum.
+
+        Parameters
+        ----------
+        task_node : `~.pipeline_graph.TaskNode`
+            Task definition structure.
+        quantum : `~lsst.daf.butler.Quantum`
+            Quantum for this execution.
+        quantum_id : `uuid.UUID` or `None`, optional
+            The ID of the quantum to be executed.
+
+        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.
+
+        Notes
+        -----
+        Any exception raised by the task or code that wraps task execution is
+        propagated to the caller of this method.
+        """
+        raise NotImplementedError()
+
+
+class QuantumGraphExecutor(ABC):
+    """Class which abstracts QuantumGraph execution.
+
+    Any specific execution model is implemented in sub-class by overriding
+    the `execute` method.
+    """
+
+    @abstractmethod
+    def execute(self, graph: QuantumGraph) -> None:
+        """Execute whole graph.
+
+        Implementation of this method depends on particular execution model
+        and it has to be provided by a subclass. Execution model determines
+        what happens here; it can be either actual running of the task or,
+        for example, generation of the scripts for delayed batch execution.
+
+        Parameters
+        ----------
+        graph : `.QuantumGraph`
+            Execution graph.
+        """
+        raise NotImplementedError()
+
+    def getReport(self) -> Report | None:
+        """Return execution report from last call to `execute`.
+
+        Returns
+        -------
+        report : `~.quantum_reports.Report`, optional
+            Structure describing the status of the execution of a quantum
+            graph. `None` is returned if implementation does not support
+            this feature.
+
+        Raises
+        ------
+        RuntimeError
+            Raised if this method is called before `execute`.
+        """
+        return None

lsst/pipe/base/quantum_graph_skeleton.py

@@ -40,12 +40,21 @@ __all__ = (
 )
 
 import dataclasses
+from collections import defaultdict
 from collections.abc import Iterable, Iterator, MutableMapping, Set
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias
 
 import networkx
 
-from lsst.daf.butler import DataCoordinate, DataIdValue, DatasetRef
+from lsst.daf.butler import (
+    Butler,
+    DataCoordinate,
+    DataIdValue,
+    DatasetRef,
+    DimensionDataAttacher,
+    DimensionGroup,
+    DimensionRecordSet,
+)
 from lsst.utils.logging import getLogger
 
 if TYPE_CHECKING:
@@ -170,6 +179,7 @@ class QuantumGraphSkeleton:
         self._tasks: dict[str, tuple[TaskInitKey, set[QuantumKey]]] = {}
         self._xgraph: networkx.DiGraph = networkx.DiGraph()
         self._global_init_outputs: set[DatasetKey] = set()
+        self._dimension_data: dict[str, DimensionRecordSet] = {}
         for task_label in task_labels:
             task_init_key = TaskInitKey(task_label)
             self._tasks[task_label] = (task_init_key, set())
@@ -310,6 +320,10 @@ class QuantumGraphSkeleton:
         for task_label, (_, quanta) in other._tasks.items():
             self._tasks[task_label][1].update(quanta)
         self._xgraph.update(other._xgraph)
+        for record_set in other._dimension_data.values():
+            self._dimension_data.setdefault(
+                record_set.element.name, DimensionRecordSet(record_set.element)
+            ).update(record_set)
 
     def add_quantum_node(self, task_label: str, data_id: DataCoordinate, **attrs: Any) -> QuantumKey:
         """Add a new node representing a quantum.
@@ -710,3 +724,48 @@ class QuantumGraphSkeleton:
             Raised if this node does not have an expanded data ID.
         """
         return self._xgraph.nodes[key]["data_id"]
+
+    def attach_dimension_records(
+        self, butler: Butler, dimensions: DimensionGroup, dimension_records: Iterable[DimensionRecordSet]
+    ) -> None:
+        """Attach dimension records to the data IDs in the skeleton.
+
+        Parameters
+        ----------
+        butler : `lsst.daf.butler.Butler`
+            Butler to use to query for missing dimension records.
+        dimensions : `lsst.daf.butler.DimensionGroup`
+            Superset of all of the dimensions of all data IDs.
+        dimension_records : `~collections.abc.Iterable` [ \
+                `lsst.daf.butler.DimensionRecordSet` ]
+            Iterable of sets of dimension records to attach.
+        """
+        for record_set in dimension_records:
+            self._dimension_data.setdefault(
+                record_set.element.name, DimensionRecordSet(record_set.element)
+            ).update(record_set)
+        # Group all nodes by data ID (and dimensions of data ID).
+        data_ids_to_expand: defaultdict[DimensionGroup, defaultdict[DataCoordinate, list[Key]]] = defaultdict(
+            lambda: defaultdict(list)
+        )
+        data_id: DataCoordinate | None
+        for node_key in self:
+            if data_id := self[node_key].get("data_id"):
+                data_ids_to_expand[data_id.dimensions][data_id].append(node_key)
+        attacher = DimensionDataAttacher(records=self._dimension_data.values(), dimensions=dimensions)
+        for dimensions, data_ids in data_ids_to_expand.items():
+            with butler.query() as query:
+                # Butler query will be used as-needed to get dimension records
+                # (from prerequisites) we didn't fetch in advance.  These are
+                # cached in the attacher so we don't look them up multiple
+                # times.
+                expanded_data_ids = attacher.attach(dimensions, data_ids.keys(), query=query)
+            for expanded_data_id, node_keys in zip(expanded_data_ids, data_ids.values()):
+                for node_key in node_keys:
+                    self.set_data_id(node_key, expanded_data_id)
+        # Hold on to any records that we had to query for.
+        self._dimension_data = attacher.records
+
+    def get_dimension_data(self) -> list[DimensionRecordSet]:
+        """Return the dimension records attached to data IDs."""
+        return list(self._dimension_data.values())

lsst/pipe/base/quantum_reports.py

@@ -0,0 +1,334 @@
+# 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__ = ["ExceptionInfo", "ExecutionStatus", "QuantumReport", "Report"]
+
+import enum
+import sys
+from typing import Any
+
+import pydantic
+
+from lsst.daf.butler import DataCoordinate, DataId, DataIdValue
+from lsst.utils.introspection import get_full_type_name
+
+from .graph import QgraphSummary
+
+
+def _serializeDataId(dataId: DataId) -> dict[str, DataIdValue]:
+    if isinstance(dataId, DataCoordinate):
+        return dict(dataId.required)
+    else:
+        return dataId  # type: ignore
+
+
+class ExecutionStatus(enum.Enum):
+    """Possible values for job execution status.
+
+    Status `FAILURE` is set if one or more tasks failed. Status `TIMEOUT` is
+    set if there are no failures but one or more tasks timed out. Timeouts can
+    only be detected in multi-process mode, child task is killed on timeout
+    and usually should have non-zero exit code.
+    """
+
+    SUCCESS = "success"
+    FAILURE = "failure"
+    TIMEOUT = "timeout"
+    SKIPPED = "skipped"
+
+
+class ExceptionInfo(pydantic.BaseModel):
+    """Information about exception."""
+
+    className: str
+    """Name of the exception class if exception was raised."""
+
+    message: str
+    """Exception message for in-process quantum execution, None if
+    quantum was executed in sub-process.
+    """
+
+    @classmethod
+    def from_exception(cls, exception: Exception) -> ExceptionInfo:
+        """Construct instance from an exception.
+
+        Parameters
+        ----------
+        exception : `Exception`
+            Exception to wrap.
+
+        Returns
+        -------
+        info : `ExceptionInfo`
+            Information about the exception.
+        """
+        return cls(className=get_full_type_name(exception), message=str(exception))
+
+    # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
+    # when we inherit those docstrings in our public classes.
+    if "sphinx" in sys.modules:
+
+        def copy(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.copy`."""
+            return super().copy(*args, **kwargs)
+
+        def model_dump(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_dump`."""
+            return super().model_dump(*args, **kwargs)
+
+        def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_dump_json`."""
+            return super().model_dump(*args, **kwargs)
+
+        def model_copy(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_copy`."""
+            return super().model_copy(*args, **kwargs)
+
+        @classmethod
+        def model_construct(cls, *args: Any, **kwargs: Any) -> Any:  # type: ignore[misc, override]
+            """See `pydantic.BaseModel.model_construct`."""
+            return super().model_construct(*args, **kwargs)
+
+        @classmethod
+        def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_json_schema`."""
+            return super().model_json_schema(*args, **kwargs)
+
+
+class QuantumReport(pydantic.BaseModel):
+    """Task execution report for a single Quantum.
+
+    Parameters
+    ----------
+    dataId : `~lsst.daf.butler.DataId`
+        Quantum data ID.
+    taskLabel : `str`
+        Label for task executing this Quantum.
+    status : `ExecutionStatus`
+        Status of this quantum execution.
+    exitCode : `int` or `None`, optional
+        Exit code for sub-process executing this Quantum. `None` for
+        in-process execution. Negative if process was killed by a signal.
+    exceptionInfo : `ExceptionInfo` or `None`, optional
+        Exception information if an exception was raised.
+    """
+
+    status: ExecutionStatus = ExecutionStatus.SUCCESS
+    """Execution status, one of the values in `ExecutionStatus` enum."""
+
+    dataId: dict[str, DataIdValue]
+    """Quantum DataId."""
+
+    taskLabel: str | None
+    """Label for a task executing this Quantum."""
+
+    exitCode: int | None = None
+    """Exit code for a sub-process executing Quantum, None for in-process
+    Quantum execution. Negative if process was killed by a signal.
+    """
+
+    exceptionInfo: ExceptionInfo | None = None
+    """Exception information if exception was raised."""
+
+    def __init__(
+        self,
+        dataId: DataId,
+        taskLabel: str,
+        status: ExecutionStatus = ExecutionStatus.SUCCESS,
+        exitCode: int | None = None,
+        exceptionInfo: ExceptionInfo | None = None,
+    ):
+        super().__init__(
+            status=status,
+            dataId=_serializeDataId(dataId),
+            taskLabel=taskLabel,
+            exitCode=exitCode,
+            exceptionInfo=exceptionInfo,
+        )
+
+    @classmethod
+    def from_exception(
+        cls,
+        exception: Exception,
+        dataId: DataId,
+        taskLabel: str,
+        *,
+        exitCode: int | None = None,
+    ) -> QuantumReport:
+        """Construct report instance from an exception and other pieces of
+        data.
+
+        Parameters
+        ----------
+        exception : `Exception`
+            Exception caught from processing quantum.
+        dataId : `~lsst.daf.butler.DataId`
+            Data ID of quantum.
+        taskLabel : `str`
+            Label of task.
+        exitCode : `int`, optional
+            Exit code for the process, used when it is known that the process
+            will exit with that exit code.
+        """
+        return cls(
+            status=ExecutionStatus.FAILURE,
+            dataId=dataId,
+            taskLabel=taskLabel,
+            exitCode=exitCode,
+            exceptionInfo=ExceptionInfo.from_exception(exception),
+        )
+
+    @classmethod
+    def from_exit_code(
+        cls,
+        exitCode: int,
+        dataId: DataId,
+        taskLabel: str,
+    ) -> QuantumReport:
+        """Construct report instance from an exit code and other pieces of
+        data.
+
+        Parameters
+        ----------
+        exitCode : `int`
+            The exit code of the subprocess.
+        dataId : `~lsst.daf.butler.DataId`
+            The quantum Data ID.
+        taskLabel : `str`
+            The task label.
+        """
+        return cls(
+            status=ExecutionStatus.SUCCESS if exitCode == 0 else ExecutionStatus.FAILURE,
+            dataId=dataId,
+            taskLabel=taskLabel,
+            exitCode=exitCode,
+        )
+
+    # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
+    # when we inherit those docstrings in our public classes.
+    if "sphinx" in sys.modules:
+
+        def copy(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.copy`."""
+            return super().copy(*args, **kwargs)
+
+        def model_dump(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_dump`."""
+            return super().model_dump(*args, **kwargs)
+
+        def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_dump_json`."""
+            return super().model_dump(*args, **kwargs)
+
+        def model_copy(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_copy`."""
+            return super().model_copy(*args, **kwargs)
+
+        @classmethod
+        def model_construct(cls, *args: Any, **kwargs: Any) -> Any:  # type: ignore[misc, override]
+            """See `pydantic.BaseModel.model_construct`."""
+            return super().model_construct(*args, **kwargs)
+
+        @classmethod
+        def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_json_schema`."""
+            return super().model_json_schema(*args, **kwargs)
+
+
+class Report(pydantic.BaseModel):
+    """Execution report for the whole job with one or few quanta."""
+
+    qgraphSummary: QgraphSummary
+    """Summary report about QuantumGraph."""
+
+    status: ExecutionStatus = ExecutionStatus.SUCCESS
+    """Job status."""
+
+    cmdLine: list[str] | None = None
+    """Command line for the whole job."""
+
+    exitCode: int | None = None
+    """Job exit code, this obviously cannot be set in pipetask."""
+
+    exceptionInfo: ExceptionInfo | None = None
+    """Exception information if exception was raised."""
+
+    quantaReports: list[QuantumReport] = []
+    """List of per-quantum reports, ordering is not specified. Some or all
+    quanta may not produce a report.
+    """
+
+    # Always want to validate the default value for cmdLine so
+    # use a model_validator.
+    @pydantic.model_validator(mode="before")
+    @classmethod
+    def _set_cmdLine(cls, data: Any) -> Any:
+        if data.get("cmdLine") is None:
+            data["cmdLine"] = sys.argv
+        return data
+
+    def set_exception(self, exception: Exception) -> None:
+        """Update exception information from an exception object.
+
+        Parameters
+        ----------
+        exception : `Exception`
+            Exception to use to extract information from.
+        """
+        self.exceptionInfo = ExceptionInfo.from_exception(exception)
+
+    # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
+    # when we inherit those docstrings in our public classes.
+    if "sphinx" in sys.modules:
+
+        def copy(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.copy`."""
+            return super().copy(*args, **kwargs)
+
+        def model_dump(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_dump`."""
+            return super().model_dump(*args, **kwargs)
+
+        def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_dump_json`."""
+            return super().model_dump(*args, **kwargs)
+
+        def model_copy(self, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_copy`."""
+            return super().model_copy(*args, **kwargs)
+
+        @classmethod
+        def model_construct(cls, *args: Any, **kwargs: Any) -> Any:  # type: ignore[misc, override]
+            """See `pydantic.BaseModel.model_construct`."""
+            return super().model_construct(*args, **kwargs)
+
+        @classmethod
+        def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
+            """See `pydantic.BaseModel.model_json_schema`."""
+            return super().model_json_schema(*args, **kwargs)

lsst/pipe/base/script/transfer_from_graph.py

@@ -167,10 +167,13 @@ def _update_chain(butler: Butler, output_chain: str, output_run: str, inputs: li
     if created_now:
         _LOG.verbose("Registered chain collection: %s", output_chain)
         if inputs:
+            # First must flatten any input chains
+            flattened = butler.collections.query(inputs, flatten_chains=True)
+
             # Add input collections to chain collection just made.  Using
             # extend instead of prepend in case of race condition where another
             # execution adds a run before this adds the inputs to the chain.
-            butler.collections.extend_chain(output_chain, inputs)
+            butler.collections.extend_chain(output_chain, flattened)
     _LOG.verbose(
         "Prepending output chain collection (%s) with output RUN collection (%s)", output_chain, output_run
     )