lsst-pipe-base 29.2025.2900__py3-none-any.whl → 29.2025.3100__py3-none-any.whl

lsst/pipe/base/_datasetQueryConstraints.py

@@ -26,7 +26,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 """Symbols defined in this package should be imported from
-`all_dimensions_quantum_graph_builder` instead; it only appears in the docs
+`.all_dimensions_quantum_graph_builder` instead; it only appears in the docs
 due to limitations in Sphinx.
 """
 

lsst/pipe/base/all_dimensions_quantum_graph_builder.py

@@ -65,13 +65,14 @@ if TYPE_CHECKING:
 
 @final
 class AllDimensionsQuantumGraphBuilder(QuantumGraphBuilder):
-    """An implementation of `QuantumGraphBuilder` that uses a single large
-    query for data IDs covering all dimensions in the pipeline.
+    """An implementation of `.quantum_graph_builder.QuantumGraphBuilder` that
+    uses a single large query for data IDs covering all dimensions in the
+    pipeline.
 
     Parameters
     ----------
     pipeline_graph : `.pipeline_graph.PipelineGraph`
-        Pipeline to build a `QuantumGraph` from, as a graph.  Will be resolved
+        Pipeline to build a `.QuantumGraph` from, as a graph.  Will be resolved
         in-place with the given butler (any existing resolution is ignored).
     butler : `lsst.daf.butler.Butler`
         Client for the data repository.  Should be read-only.
@@ -92,7 +93,8 @@ class AllDimensionsQuantumGraphBuilder(QuantumGraphBuilder):
         are constrained by the ``where`` argument or pipeline data ID will be
         filled in automatically.
     **kwargs
-        Additional keyword arguments forwarded to `QuantumGraphBuilder`.
+        Additional keyword arguments forwarded to
+        `.quantum_graph_builder.QuantumGraphBuilder`.
 
     Notes
     -----

lsst/pipe/base/connectionTypes.py

@@ -26,7 +26,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 """Module defining connection types to be used within a
-`PipelineTaskConnections` class.
+`.PipelineTaskConnections` class.
 """
 
 __all__ = ["BaseConnection", "InitInput", "InitOutput", "Input", "Output", "PrerequisiteInput"]
@@ -53,7 +53,7 @@ class BaseConnection:
         Indicates if this connection should expect to contain multiple objects
         of the given dataset type.  Tasks with more than one connection with
         ``multiple=True`` with the same dimensions may want to implement
-        `PipelineTaskConnections.adjustQuantum` to ensure those datasets are
+        `.PipelineTaskConnections.adjustQuantum` to ensure those datasets are
         consistent (i.e. zip-iterable) in `PipelineTask.runQuantum()` and
         notify the execution system as early as possible of outputs that will
         not be produced because the corresponding input is missing.
@@ -121,7 +121,7 @@ class DimensionedConnection(BaseConnection):
         Indicates if this connection should expect to contain multiple objects
         of the given dataset type.  Tasks with more than one connection with
         ``multiple=True`` with the same dimensions may want to implement
-        `PipelineTaskConnections.adjustQuantum` to ensure those datasets are
+        `.PipelineTaskConnections.adjustQuantum` to ensure those datasets are
         consistent (i.e. zip-iterable) in `PipelineTask.runQuantum` and notify
         the execution system as early as possible of outputs that will not be
         produced because the corresponding input is missing.
@@ -161,7 +161,7 @@ class BaseInput(DimensionedConnection):
         Indicates if this connection should expect to contain multiple objects
         of the given dataset type.  Tasks with more than one connection with
         ``multiple=True`` with the same dimensions may want to implement
-        `PipelineTaskConnections.adjustQuantum` to ensure those datasets are
+        `.PipelineTaskConnections.adjustQuantum` to ensure those datasets are
         consistent (i.e. zip-iterable) in `PipelineTask.runQuantum` and notify
         the execution system as early as possible of outputs that will not be
         produced because the corresponding input is missing.
@@ -175,14 +175,14 @@ class BaseInput(DimensionedConnection):
     minimum : `bool`
         Minimum number of datasets required for this connection, per quantum.
         This is checked in the base implementation of
-        `PipelineTaskConnections.adjustQuantum`, which raises `NoWorkFound` if
+        `.PipelineTaskConnections.adjustQuantum`, which raises `NoWorkFound` if
         the minimum is not met for `Input` connections (causing the quantum to
         be pruned, skipped, or never created, depending on the context), and
         `FileNotFoundError` for `PrerequisiteInput` connections (causing
         QuantumGraph generation to fail).  `PipelineTask` implementations may
-        provide custom `~PipelineTaskConnections.adjustQuantum` implementations
-        for more fine-grained or configuration-driven constraints, as long as
-        they are compatible with this minium.
+        provide custom `~.PipelineTaskConnections.adjustQuantum`
+        implementations for more fine-grained or configuration-driven
+        constraints, as long as they are compatible with this minium.
 
     Raises
     ------
@@ -216,7 +216,7 @@ class Input(BaseInput):
         Indicates if this connection should expect to contain multiple objects
         of the given dataset type.  Tasks with more than one connection with
         ``multiple=True`` with the same dimensions may want to implement
-        `PipelineTaskConnections.adjustQuantum` to ensure those datasets are
+        `.PipelineTaskConnections.adjustQuantum` to ensure those datasets are
         consistent (i.e. zip-iterable) in `PipelineTask.runQuantum` and notify
         the execution system as early as possible of outputs that will not be
         produced because the corresponding input is missing.
@@ -230,14 +230,14 @@ class Input(BaseInput):
     minimum : `bool`
         Minimum number of datasets required for this connection, per quantum.
         This is checked in the base implementation of
-        `PipelineTaskConnections.adjustQuantum`, which raises `NoWorkFound` if
+        `.PipelineTaskConnections.adjustQuantum`, which raises `NoWorkFound` if
         the minimum is not met for `Input` connections (causing the quantum to
         be pruned, skipped, or never created, depending on the context), and
         `FileNotFoundError` for `PrerequisiteInput` connections (causing
         QuantumGraph generation to fail).  `PipelineTask` implementations may
-        provide custom `~PipelineTaskConnections.adjustQuantum` implementations
-        for more fine-grained or configuration-driven constraints, as long as
-        they are compatible with this minium.
+        provide custom `~.PipelineTaskConnections.adjustQuantum`
+        implementations for more fine-grained or configuration-driven
+        constraints, as long as they are compatible with this minium.
     deferGraphConstraint : `bool`, optional
         If `True`, do not include this dataset type's existence in the initial
         query that starts the QuantumGraph generation process.  This can be
@@ -286,7 +286,7 @@ class PrerequisiteInput(BaseInput):
         Indicates if this connection should expect to contain multiple objects
         of the given dataset type.  Tasks with more than one connection with
         ``multiple=True`` with the same dimensions may want to implement
-        `PipelineTaskConnections.adjustQuantum` to ensure those datasets are
+        `.PipelineTaskConnections.adjustQuantum` to ensure those datasets are
         consistent (i.e. zip-iterable) in `PipelineTask.runQuantum` and notify
         the execution system as early as possible of outputs that will not be
         produced because the corresponding input is missing.
@@ -296,12 +296,12 @@ class PrerequisiteInput(BaseInput):
     minimum : `bool`
         Minimum number of datasets required for this connection, per quantum.
         This is checked in the base implementation of
-        `PipelineTaskConnections.adjustQuantum`, which raises
+        `.PipelineTaskConnections.adjustQuantum`, which raises
         `FileNotFoundError` (causing QuantumGraph generation to fail).
-        `PipelineTask` implementations may
-        provide custom `~PipelineTaskConnections.adjustQuantum` implementations
-        for more fine-grained or configuration-driven constraints, as long as
-        they are compatible with this minium.
+        `PipelineTask` implementations may provide custom
+        `~.PipelineTaskConnections.adjustQuantum` implementations for more
+        fine-grained or configuration-driven constraints, as long as they are
+        compatible with this minium.
     lookupFunction : `typing.Callable`, optional
         An optional callable function that will look up PrerequisiteInputs
         using the DatasetType, registry, quantum dataId, and input collections

lsst/pipe/base/connections.py

@@ -1063,8 +1063,8 @@ def iterConnections(
 class AdjustQuantumHelper:
     """Helper class for calling `PipelineTaskConnections.adjustQuantum`.
 
-    This class holds `input` and `output` mappings in the form used by
-    `Quantum` and execution harness code, i.e. with
+    This class holds `inputs` and `outputs` mappings in the form used by
+    `lsst.daf.butler.Quantum` and execution harness code, i.e. with
     `~lsst.daf.butler.DatasetType` keys, translating them to and from the
     connection-oriented mappings used inside `PipelineTaskConnections`.
     """

lsst/pipe/base/exec_fixup_data_id.py

@@ -0,0 +1,131 @@
+# 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/>.
+
+__all__ = ["ExecutionGraphFixup"]
+
+import contextlib
+import itertools
+from collections import defaultdict
+from collections.abc import Sequence
+from typing import Any
+
+import networkx as nx
+
+from .execution_graph_fixup import ExecutionGraphFixup
+from .graph import QuantumGraph, QuantumNode
+
+
+class ExecFixupDataId(ExecutionGraphFixup):
+    """Implementation of ExecutionGraphFixup for ordering of tasks based
+    on DataId values.
+
+    This class is a trivial implementation mostly useful as an example,
+    though it can be used to make actual fixup instances by defining
+    a method that instantiates it, e.g.::
+
+        # lsst/ap/verify/ci_fixup.py
+
+        from lsst.pipe.base.exec_fixup_data_id import ExecFixupDataId
+
+
+        def assoc_fixup():
+            return ExecFixupDataId(
+                taskLabel="ap_assoc", dimensions=("visit", "detector")
+            )
+
+    and then executing pipetask::
+
+        pipetask run --graph-fixup=lsst.ap.verify.ci_fixup.assoc_fixup ...
+
+    This will add new dependencies between quanta executed by the task with
+    label "ap_assoc". Quanta with higher visit number will depend on quanta
+    with lower visit number and their execution will wait until lower visit
+    number finishes.
+
+    Parameters
+    ----------
+    taskLabel : `str`
+        The label of the task for which to add dependencies.
+    dimensions : `str` or sequence [`str`]
+        One or more dimension names, quanta execution will be ordered
+        according to values of these dimensions.
+    reverse : `bool`, optional
+        If `False` (default) then quanta with higher values of dimensions
+        will be executed after quanta with lower values, otherwise the order
+        is reversed.
+    """
+
+    def __init__(self, taskLabel: str, dimensions: str | Sequence[str], reverse: bool = False):
+        self.taskLabel = taskLabel
+        self.dimensions = dimensions
+        self.reverse = reverse
+        if isinstance(self.dimensions, str):
+            self.dimensions = (self.dimensions,)
+        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

lsst/pipe/base/execution_graph_fixup.py

@@ -0,0 +1,69 @@
+# 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/>.
+
+__all__ = ["ExecutionGraphFixup"]
+
+from abc import ABC, abstractmethod
+
+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.
+    """
+
+    @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`
+            Quantum Graph that will be executed by the executor.
+
+        Returns
+        -------
+        graph : `.QuantumGraph`
+            Modified graph.
+        """
+        raise NotImplementedError

lsst/pipe/base/log_capture.py

@@ -0,0 +1,227 @@
+# 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__ = ["LogCapture"]
+
+import logging
+import os
+import shutil
+import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager, suppress
+from logging import FileHandler
+
+from lsst.daf.butler import Butler, FileDataset, LimitedButler, Quantum
+from lsst.daf.butler.logging import ButlerLogRecordHandler, ButlerLogRecords, ButlerMDC, JsonLogFormatter
+
+from ._status import InvalidQuantumError
+from .pipeline_graph import TaskNode
+
+_LOG = logging.getLogger(__name__)
+
+
+class _LogCaptureFlag:
+    """Simple flag to enable/disable log-to-butler saving."""
+
+    store: bool = True
+
+
+class LogCapture:
+    """Class handling capture of logging messages and their export to butler.
+
+    Parameters
+    ----------
+    butler : `~lsst.daf.butler.LimitedButler`
+        Data butler with limited API.
+    full_butler : `~lsst.daf.butler.Butler` or `None`
+        Data butler with full API, or `None` if full Butler is not available.
+        If not none, then this must be the same instance as ``butler``.
+    """
+
+    stream_json_logs = True
+    """If True each log record is written to a temporary file and ingested
+    when quantum completes. If False the records are accumulated in memory
+    and stored in butler on quantum completion. If full butler is not available
+    then temporary file is not used."""
+
+    def __init__(
+        self,
+        butler: LimitedButler,
+        full_butler: Butler | None,
+    ):
+        self.butler = butler
+        self.full_butler = full_butler
+
+    @classmethod
+    def from_limited(cls, butler: LimitedButler) -> LogCapture:
+        return cls(butler, None)
+
+    @classmethod
+    def from_full(cls, butler: Butler) -> LogCapture:
+        return cls(butler, butler)
+
+    @contextmanager
+    def capture_logging(self, task_node: TaskNode, /, quantum: Quantum) -> Iterator[_LogCaptureFlag]:
+        """Configure logging system to capture logs for execution of this task.
+
+        Parameters
+        ----------
+        task_node : `~lsst.pipe.base.pipeline_graph.TaskNode`
+            The task definition.
+        quantum : `~lsst.daf.butler.Quantum`
+            Single Quantum instance.
+
+        Notes
+        -----
+        Expected to be used as a context manager to ensure that logging
+        records are inserted into the butler once the quantum has been
+        executed:
+
+        .. code-block:: py
+
+           with self.capture_logging(task_node, quantum):
+               # Run quantum and capture logs.
+
+        Ths method can also setup logging to attach task- or
+        quantum-specific information to log messages. Potentially this can
+        take into account some info from task configuration as well.
+        """
+        # include quantum dataId and task label into MDC
+        mdc = {"LABEL": task_node.label, "RUN": ""}
+        if quantum.dataId:
+            mdc["LABEL"] += f":{quantum.dataId}"
+        if self.full_butler is not None:
+            mdc["RUN"] = self.full_butler.run or ""
+        ctx = _LogCaptureFlag()
+        log_dataset_name = (
+            task_node.log_output.dataset_type_name if task_node.log_output is not None else None
+        )
+
+        # Add a handler to the root logger to capture execution log output.
+        if log_dataset_name is not None:
+            # Either accumulate into ButlerLogRecords or stream JSON records to
+            # file and ingest that (ingest is possible only with full butler).
+            if self.stream_json_logs and self.full_butler is not None:
+                # Create the log file in a temporary directory rather than
+                # creating a temporary file. This is necessary because
+                # temporary files are created with restrictive permissions
+                # and during file ingest these permissions persist in the
+                # datastore. Using a temp directory allows us to create
+                # a file with umask default permissions.
+                tmpdir = tempfile.mkdtemp(prefix="butler-temp-logs-")
+
+                # Construct a file to receive the log records and "touch" it.
+                log_file = os.path.join(tmpdir, f"butler-log-{task_node.label}.json")
+                with open(log_file, "w"):
+                    pass
+                log_handler_file = FileHandler(log_file)
+                log_handler_file.setFormatter(JsonLogFormatter())
+                logging.getLogger().addHandler(log_handler_file)
+
+                try:
+                    with ButlerMDC.set_mdc(mdc):
+                        yield ctx
+                finally:
+                    # Ensure that the logs are stored in butler.
+                    logging.getLogger().removeHandler(log_handler_file)
+                    log_handler_file.close()
+                    if ctx.store:
+                        self._ingest_log_records(quantum, log_dataset_name, log_file)
+                    shutil.rmtree(tmpdir, ignore_errors=True)
+
+            else:
+                log_handler_memory = ButlerLogRecordHandler()
+                logging.getLogger().addHandler(log_handler_memory)
+
+                try:
+                    with ButlerMDC.set_mdc(mdc):
+                        yield ctx
+                finally:
+                    # Ensure that the logs are stored in butler.
+                    logging.getLogger().removeHandler(log_handler_memory)
+                    if ctx.store:
+                        self._store_log_records(quantum, log_dataset_name, log_handler_memory)
+                    log_handler_memory.records.clear()
+
+        else:
+            with ButlerMDC.set_mdc(mdc):
+                yield ctx
+
+    def _store_log_records(
+        self, quantum: Quantum, dataset_type: str, log_handler: ButlerLogRecordHandler
+    ) -> None:
+        # DatasetRef has to be in the Quantum outputs, can lookup by name.
+        try:
+            [ref] = quantum.outputs[dataset_type]
+        except LookupError as exc:
+            raise InvalidQuantumError(
+                f"Quantum outputs is missing log output dataset type {dataset_type};"
+                " this could happen due to inconsistent options between QuantumGraph generation"
+                " and execution"
+            ) from exc
+
+        self.butler.put(log_handler.records, ref)
+
+    def _ingest_log_records(self, quantum: Quantum, dataset_type: str, filename: str) -> None:
+        # If we are logging to an external file we must always try to
+        # close it.
+        assert self.full_butler is not None, "Expected to have full butler for ingest"
+        ingested = False
+        try:
+            # DatasetRef has to be in the Quantum outputs, can lookup by name.
+            try:
+                [ref] = quantum.outputs[dataset_type]
+            except LookupError as exc:
+                raise InvalidQuantumError(
+                    f"Quantum outputs is missing log output dataset type {dataset_type};"
+                    " this could happen due to inconsistent options between QuantumGraph generation"
+                    " and execution"
+                ) from exc
+
+            # Need to ingest this file directly into butler.
+            dataset = FileDataset(path=filename, refs=ref)
+            try:
+                self.full_butler.ingest(dataset, transfer="move")
+                ingested = True
+            except NotImplementedError:
+                # Some datastores can't receive files (e.g. in-memory datastore
+                # when testing), we store empty list for those just to have a
+                # dataset. Alternative is to read the file as a
+                # ButlerLogRecords object and put it.
+                _LOG.info(
+                    "Log records could not be stored in this butler because the"
+                    " datastore can not ingest files, empty record list is stored instead."
+                )
+                records = ButlerLogRecords.from_records([])
+                self.full_butler.put(records, ref)
+        finally:
+            # remove file if it is not ingested
+            if not ingested:
+                with suppress(OSError):
+                    os.remove(filename)