lsst-pipe-base 29.2025.4500__py3-none-any.whl → 29.2025.4700__py3-none-any.whl

lsst/pipe/base/quantum_graph/aggregator/_writer.py

@@ -30,7 +30,6 @@ from __future__ import annotations
 __all__ = ("Writer",)
 
 import dataclasses
-import enum
 import itertools
 import logging
 import operator
@@ -62,69 +61,7 @@ from .._provenance import (
     ProvenanceQuantumModel,
 )
 from ._communicators import WriterCommunicator
-from ._structs import ScanResult
-
-
-class _CompressionState(enum.Enum):
-    """Enumeration of the possible states of compression in `_ScanData`."""
-
-    NOT_COMPRESSED = enum.auto()
-    """Nothing is compressed."""
-
-    LOG_AND_METADATA_COMPRESSED = enum.auto()
-    """Only the logs and metadata are compressed."""
-
-    ALL_COMPRESSED = enum.auto()
-    """All `bytes` are compressed."""
-
-
-@dataclasses.dataclass
-class _ScanData:
-    """Information from a quantum scan that has been partially processed for
-    writing.
-    """
-
-    quantum_id: uuid.UUID
-    """Unique ID of the quantum."""
-
-    log_id: uuid.UUID
-    """Unique ID of the log dataset."""
-
-    metadata_id: uuid.UUID
-    """Unique ID of the metadata dataset."""
-
-    quantum: bytes = b""
-    """Possibly-compressed JSON representation of the quantum provenance."""
-
-    datasets: dict[uuid.UUID, bytes] = dataclasses.field(default_factory=dict)
-    """Possibly-compressed JSON representation of output dataset provenance."""
-
-    log: bytes = b""
-    """Possibly-compressed log dataset content."""
-
-    metadata: bytes = b""
-    """Possibly-compressed metadata dataset content."""
-
-    compression: _CompressionState = _CompressionState.NOT_COMPRESSED
-    """Which data is compressed, if any."""
-
-    def compress(self, compressor: Compressor) -> None:
-        """Compress all data in place, if it isn't already.
-
-        Parameters
-        ----------
-        compressor : `Compressor`
-            Object that can compress `bytes`.
-        """
-        if self.compression is _CompressionState.NOT_COMPRESSED:
-            self.metadata = compressor.compress(self.metadata)
-            self.log = compressor.compress(self.log)
-            self.compression = _CompressionState.LOG_AND_METADATA_COMPRESSED
-        if self.compression is _CompressionState.LOG_AND_METADATA_COMPRESSED:
-            self.quantum = compressor.compress(self.quantum)
-            for key in self.datasets.keys():
-                self.datasets[key] = compressor.compress(self.datasets[key])
-        self.compression = _CompressionState.ALL_COMPRESSED
+from ._structs import WriteRequest
 
 
 @dataclasses.dataclass
@@ -267,8 +204,8 @@ class Writer:
     with datasets as well as with quanta.
     """
 
-    pending_compression_training: list[_ScanData] = dataclasses.field(default_factory=list)
-    """Partially processed quantum scans that are being accumulated in order to
+    pending_compression_training: list[WriteRequest] = dataclasses.field(default_factory=list)
+    """Unprocessed quantum scans that are being accumulated in order to
     build a compression dictionary.
     """
 
@@ -299,7 +236,7 @@ class Writer:
         )
 
     def _populate_indices_and_outputs(self) -> None:
-        all_uuids = set(self.predicted.quantum_indices.keys())
+        all_uuids = set(self.predicted.quantum_datasets.keys())
         for quantum in self.comms.periodically_check_for_cancel(
             itertools.chain(
                 self.predicted.init_quanta.root,
@@ -329,13 +266,12 @@ class Writer:
             if not predicted_quantum.task_label:
                 # Skip the 'packages' producer quantum.
                 continue
-            quantum_index = self.indices[predicted_quantum.quantum_id]
             for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
-                self.xgraph.add_edge(self.indices[predicted_input.dataset_id], quantum_index)
+                self.xgraph.add_edge(predicted_input.dataset_id, predicted_quantum.quantum_id)
                 if predicted_input.dataset_id not in self.output_dataset_ids:
                     self.overall_inputs.setdefault(predicted_input.dataset_id, predicted_input)
             for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
-                self.xgraph.add_edge(quantum_index, self.indices[predicted_output.dataset_id])
+                self.xgraph.add_edge(predicted_quantum.quantum_id, predicted_output.dataset_id)
 
     @staticmethod
     def run(predicted_path: str, comms: WriterCommunicator) -> None:
@@ -365,12 +301,11 @@ class Writer:
         self.comms.log.info("Polling for write requests from scanners.")
         for request in self.comms.poll():
             if data_writers is None:
-                self.pending_compression_training.extend(self.make_scan_data(request))
+                self.pending_compression_training.append(request)
                 if len(self.pending_compression_training) >= self.comms.config.zstd_dict_n_inputs:
                     data_writers = self.make_data_writers()
             else:
-                for scan_data in self.make_scan_data(request):
-                    self.write_scan_data(scan_data, data_writers)
+                self.process_request(request, data_writers)
         if data_writers is None:
             data_writers = self.make_data_writers()
         self.write_init_outputs(data_writers)
@@ -398,8 +333,8 @@ class Writer:
         )
         self.comms.check_for_cancel()
         self.comms.log.info("Compressing and writing queued scan requests.")
-        for scan_data in self.pending_compression_training:
-            self.write_scan_data(scan_data, data_writers)
+        for request in self.pending_compression_training:
+            self.process_request(request, data_writers)
         del self.pending_compression_training
         self.comms.check_for_cancel()
         self.write_overall_inputs(data_writers)
@@ -435,11 +370,11 @@ class Writer:
             predicted_quantum.datastore_records.clear()
             training_inputs.append(predicted_quantum.model_dump_json().encode())
         # Add the provenance quanta, metadata, and logs we've accumulated.
-        for scan_data in self.pending_compression_training:
-            assert scan_data.compression is _CompressionState.NOT_COMPRESSED
-            training_inputs.append(scan_data.quantum)
-            training_inputs.append(scan_data.metadata)
-            training_inputs.append(scan_data.log)
+        for write_request in self.pending_compression_training:
+            assert not write_request.is_compressed, "We can't compress without the compression dictionary."
+            training_inputs.append(write_request.quantum)
+            training_inputs.append(write_request.metadata)
+            training_inputs.append(write_request.logs)
         return zstandard.train_dictionary(self.comms.config.zstd_dict_size, training_inputs)
 
     def write_init_outputs(self, data_writers: _DataWriters) -> None:
@@ -458,19 +393,16 @@ class Writer:
                 continue
             existing_outputs = self.existing_init_outputs[predicted_init_quantum.quantum_id]
             for predicted_output in itertools.chain.from_iterable(predicted_init_quantum.outputs.values()):
-                dataset_index = self.indices[predicted_output.dataset_id]
                 provenance_output = ProvenanceDatasetModel.from_predicted(
                     predicted_output,
-                    producer=self.indices[predicted_init_quantum.quantum_id],
-                    consumers=self.xgraph.successors(dataset_index),
+                    producer=predicted_init_quantum.quantum_id,
+                    consumers=self.xgraph.successors(predicted_output.dataset_id),
                 )
-                provenance_output.exists = predicted_output.dataset_id in existing_outputs
+                provenance_output.produced = predicted_output.dataset_id in existing_outputs
                 data_writers.datasets.write_model(
                     provenance_output.dataset_id, provenance_output, data_writers.compressor
                 )
-            init_quanta.root.append(
-                ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum, self.indices)
-            )
+            init_quanta.root.append(ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum))
         data_writers.graph.write_single_model("init_quanta", init_quanta)
 
     def write_overall_inputs(self, data_writers: _DataWriters) -> None:
@@ -484,13 +416,12 @@ class Writer:
         self.comms.log.info("Writing overall inputs.")
         for predicted_input in self.comms.periodically_check_for_cancel(self.overall_inputs.values()):
             if predicted_input.dataset_id not in data_writers.datasets.addresses:
-                dataset_index = self.indices[predicted_input.dataset_id]
                 data_writers.datasets.write_model(
                     predicted_input.dataset_id,
                     ProvenanceDatasetModel.from_predicted(
                         predicted_input,
                         producer=None,
-                        consumers=self.xgraph.successors(dataset_index),
+                        consumers=self.xgraph.successors(predicted_input.dataset_id),
                     ),
                     data_writers.compressor,
                 )
@@ -509,84 +440,61 @@ class Writer:
         data = packages.toBytes("json")
         data_writers.graph.write_single_block("packages", data)
 
-    def make_scan_data(self, request: ScanResult) -> list[_ScanData]:
-        """Process a `ScanResult` into `_ScanData`.
+    def process_request(self, request: WriteRequest, data_writers: _DataWriters) -> None:
+        """Process a `WriteRequest` into `_ScanData`.
 
         Parameters
         ----------
-        request : `ScanResult`
+        request : `WriteRequest`
             Result of a quantum scan.
-
-        Returns
-        -------
-        data : `list` [ `_ScanData` ]
-            A zero- or single-element list of `_ScanData` to write or save for
-            compression-dict training.  A zero-element list is returned if the
-            scan actually represents an init quantum.
+        data_writers : `_DataWriters`
+            Low-level writers struct.
         """
         if (existing_init_outputs := self.existing_init_outputs.get(request.quantum_id)) is not None:
             self.comms.log.debug("Handling init-output scan for %s.", request.quantum_id)
             existing_init_outputs.update(request.existing_outputs)
             self.comms.report_write()
-            return []
+            return
         self.comms.log.debug("Handling quantum scan for %s.", request.quantum_id)
         predicted_quantum = self.predicted.quantum_datasets[request.quantum_id]
-        quantum_index = self.indices[predicted_quantum.quantum_id]
-        (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
-        (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
-        data = _ScanData(
-            request.quantum_id,
-            metadata_id=metadata_output.dataset_id,
-            log_id=log_output.dataset_id,
-            compression=(
-                _CompressionState.LOG_AND_METADATA_COMPRESSED
-                if request.is_compressed
-                else _CompressionState.NOT_COMPRESSED
-            ),
-        )
+        outputs: dict[uuid.UUID, bytes] = {}
         for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
-            dataset_index = self.indices[predicted_output.dataset_id]
             provenance_output = ProvenanceDatasetModel.from_predicted(
                 predicted_output,
-                producer=quantum_index,
-                consumers=self.xgraph.successors(dataset_index),
+                producer=predicted_quantum.quantum_id,
+                consumers=self.xgraph.successors(predicted_output.dataset_id),
             )
-            provenance_output.exists = provenance_output.dataset_id in request.existing_outputs
-            data.datasets[provenance_output.dataset_id] = provenance_output.model_dump_json().encode()
-        provenance_quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum, self.indices)
-        provenance_quantum.status = request.get_run_status()
-        provenance_quantum.caveats = request.caveats
-        provenance_quantum.exception = request.exception
-        provenance_quantum.resource_usage = request.resource_usage
-        data.quantum = provenance_quantum.model_dump_json().encode()
-        data.metadata = request.metadata
-        data.log = request.log
-        return [data]
-
-    def write_scan_data(self, scan_data: _ScanData, data_writers: _DataWriters) -> None:
-        """Write scan data to the provenance graph.
-
-        Parameters
-        ----------
-        scan_data : `_ScanData`
-            Preprocessed information to write.
-        data_writers : `_DataWriters`
-            Low-level writers struct.
-        """
-        self.comms.log.debug("Writing quantum %s.", scan_data.quantum_id)
-        scan_data.compress(data_writers.compressor)
-        data_writers.quanta.write_bytes(scan_data.quantum_id, scan_data.quantum)
-        for dataset_id, dataset_data in scan_data.datasets.items():
+            provenance_output.produced = provenance_output.dataset_id in request.existing_outputs
+            outputs[provenance_output.dataset_id] = data_writers.compressor.compress(
+                provenance_output.model_dump_json().encode()
+            )
+        if not request.quantum:
+            request.quantum = (
+                ProvenanceQuantumModel.from_predicted(predicted_quantum).model_dump_json().encode()
+            )
+            if request.is_compressed:
+                request.quantum = data_writers.compressor.compress(request.quantum)
+        if not request.is_compressed:
+            request.quantum = data_writers.compressor.compress(request.quantum)
+            if request.metadata:
+                request.metadata = data_writers.compressor.compress(request.metadata)
+            if request.logs:
+                request.logs = data_writers.compressor.compress(request.logs)
+        self.comms.log.debug("Writing quantum %s.", request.quantum_id)
+        data_writers.quanta.write_bytes(request.quantum_id, request.quantum)
+        for dataset_id, dataset_data in outputs.items():
             data_writers.datasets.write_bytes(dataset_id, dataset_data)
-        if scan_data.metadata:
-            address = data_writers.metadata.write_bytes(scan_data.quantum_id, scan_data.metadata)
-            data_writers.metadata.addresses[scan_data.metadata_id] = address
-        if scan_data.log:
-            address = data_writers.logs.write_bytes(scan_data.quantum_id, scan_data.log)
-            data_writers.logs.addresses[scan_data.log_id] = address
+        if request.metadata:
+            (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
+            address = data_writers.metadata.write_bytes(request.quantum_id, request.metadata)
+            data_writers.metadata.addresses[metadata_output.dataset_id] = address
+        if request.logs:
+            (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
+            address = data_writers.logs.write_bytes(request.quantum_id, request.logs)
+            data_writers.logs.addresses[log_output.dataset_id] = address
         # We shouldn't need this predicted quantum anymore; delete it in the
         # hopes that'll free up some memory.
-        del self.predicted.quantum_datasets[scan_data.quantum_id]
+        del self.predicted.quantum_datasets[request.quantum_id]
         self.comms.report_write()
 
 

lsst/pipe/base/quantum_graph_builder.py

@@ -1315,7 +1315,6 @@ class QuantumGraphBuilder(ABC):
                     },
                 )
                 components.quantum_datasets[quantum_datasets.quantum_id] = quantum_datasets
-        components.set_quantum_indices()
         components.set_thin_graph()
         components.set_header_counts()
         return components

lsst/pipe/base/quantum_provenance_graph.py

@@ -49,7 +49,7 @@ import threading
 import uuid
 from collections.abc import Callable, Iterator, Mapping, Sequence, Set
 from enum import Enum
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast
+from typing import Any, ClassVar, Literal, TypedDict, cast
 
 import astropy.table
 import networkx
@@ -72,7 +72,7 @@ from lsst.daf.butler import (
 from lsst.resources import ResourcePathExpression
 from lsst.utils.logging import PeriodicLogger, getLogger
 
-from ._status import QuantumSuccessCaveats
+from ._status import ExceptionInfo, QuantumSuccessCaveats
 from .automatic_connection_constants import (
     LOG_OUTPUT_CONNECTION_NAME,
     LOG_OUTPUT_TEMPLATE,
@@ -82,9 +82,6 @@ from .automatic_connection_constants import (
 )
 from .graph import QuantumGraph, QuantumNode
 
-if TYPE_CHECKING:
-    from ._task_metadata import TaskMetadata
-
 _LOG = getLogger(__name__)
 
 
@@ -188,45 +185,6 @@ class QuantumRunStatus(Enum):
     SUCCESSFUL = 1
 
 
-class ExceptionInfo(pydantic.BaseModel):
-    """Information about an exception that was raised."""
-
-    type_name: str
-    """Fully-qualified Python type name for the exception raised."""
-
-    message: str
-    """String message included in the exception."""
-
-    metadata: dict[str, float | int | str | bool | None]
-    """Additional metadata included in the exception."""
-
-    @classmethod
-    def _from_metadata(cls, md: TaskMetadata) -> ExceptionInfo:
-        """Construct from task metadata.
-
-        Parameters
-        ----------
-        md : `TaskMetadata`
-            Metadata about the error, as written by
-            `AnnotatedPartialOutputsError`.
-
-        Returns
-        -------
-        info : `ExceptionInfo`
-            Information about the exception.
-        """
-        result = cls(type_name=md["type"], message=md["message"], metadata={})
-        if "metadata" in md:
-            raw_err_metadata = md["metadata"].to_dict()
-            for k, v in raw_err_metadata.items():
-                # Guard against error metadata we couldn't serialize later
-                # via Pydantic; don't want one weird value bringing down our
-                # ability to report on an entire run.
-                if isinstance(v, float | int | str | bool):
-                    result.metadata[k] = v
-        return result
-
-
 class QuantumRun(pydantic.BaseModel):
     """Information about a quantum in a given run collection."""
 

lsst/pipe/base/single_quantum_executor.py

@@ -44,12 +44,19 @@ from lsst.daf.butler import (
     NamedKeyDict,
     Quantum,
 )
+from lsst.utils.introspection import get_full_type_name
 from lsst.utils.timer import logInfo
 
 from ._quantumContext import ExecutionResources, QuantumContext
-from ._status import AnnotatedPartialOutputsError, InvalidQuantumError, NoWorkFound, QuantumSuccessCaveats
+from ._status import (
+    AnnotatedPartialOutputsError,
+    ExceptionInfo,
+    InvalidQuantumError,
+    NoWorkFound,
+    QuantumSuccessCaveats,
+)
 from .connections import AdjustQuantumHelper
-from .log_capture import LogCapture
+from .log_capture import LogCapture, _ExecutionLogRecordsExtra
 from .pipeline_graph import TaskNode
 from .pipelineTask import PipelineTask
 from .quantum_graph_executor import QuantumExecutor
@@ -147,6 +154,7 @@ class SingleQuantumExecutor(QuantumExecutor):
             self._skip_existing = self._butler.run in self._butler.collections.query(
                 skip_existing_in, flatten_chains=True
             )
+        self._previous_process_quanta: list[uuid.UUID] = []
 
     def execute(
         self, task_node: TaskNode, /, quantum: Quantum, quantum_id: uuid.UUID | None = None
@@ -196,7 +204,7 @@ class SingleQuantumExecutor(QuantumExecutor):
             # or raises an exception do not try to store logs, as they may be
             # already in butler.
             captureLog.store = False
-            if self._check_existing_outputs(quantum, task_node, limited_butler):
+            if self._check_existing_outputs(quantum, task_node, limited_butler, captureLog.extra):
                 _LOG.info(
                     "Skipping already-successful quantum for label=%s dataId=%s.",
                     task_node.label,
@@ -205,6 +213,9 @@ class SingleQuantumExecutor(QuantumExecutor):
                 return quantum
             captureLog.store = True
 
+            captureLog.extra.previous_process_quanta.extend(self._previous_process_quanta)
+            if quantum_id is not None:
+                self._previous_process_quanta.append(quantum_id)
             try:
                 quantum = self._updated_quantum_inputs(quantum, task_node, limited_butler)
             except NoWorkFound as exc:
@@ -261,6 +272,11 @@ class SingleQuantumExecutor(QuantumExecutor):
                     e.__class__.__name__,
                     str(e),
                 )
+                captureLog.extra.exception = ExceptionInfo(
+                    type_name=get_full_type_name(e),
+                    message=str(e),
+                    metadata={},
+                )
                 raise
             else:
                 quantumMetadata["butler_metrics"] = butler_metrics.model_dump()
@@ -268,11 +284,13 @@ class SingleQuantumExecutor(QuantumExecutor):
                 # Stringify the UUID for easier compatibility with
                 # PropertyList.
                 quantumMetadata["outputs"] = [str(output) for output in outputsPut]
-            logInfo(None, "end", metadata=quantumMetadata)  # type: ignore[arg-type]
-            fullMetadata = task.getFullMetadata()
-            fullMetadata["quantum"] = quantumMetadata
-            if self._job_metadata is not None:
-                fullMetadata["job"] = self._job_metadata
+            finally:
+                logInfo(None, "end", metadata=quantumMetadata)  # type: ignore[arg-type]
+                fullMetadata = task.getFullMetadata()
+                fullMetadata["quantum"] = quantumMetadata
+                if self._job_metadata is not None:
+                    fullMetadata["job"] = self._job_metadata
+                captureLog.extra.metadata = fullMetadata
             self._write_metadata(quantum, fullMetadata, task_node, limited_butler)
             stopTime = time.time()
             _LOG.info(
@@ -284,7 +302,12 @@ class SingleQuantumExecutor(QuantumExecutor):
         return quantum
 
     def _check_existing_outputs(
-        self, quantum: Quantum, task_node: TaskNode, /, limited_butler: LimitedButler
+        self,
+        quantum: Quantum,
+        task_node: TaskNode,
+        /,
+        limited_butler: LimitedButler,
+        log_extra: _ExecutionLogRecordsExtra,
     ) -> bool:
         """Decide whether this quantum needs to be executed.
 
@@ -302,6 +325,8 @@ class SingleQuantumExecutor(QuantumExecutor):
             Task definition structure.
         limited_butler : `~lsst.daf.butler.LimitedButler`
             Butler to use for querying and clobbering.
+        log_extra : `.log_capture.TaskLogRecordsExtra`
+            Extra information to attach to log records.
 
         Returns
         -------
@@ -337,6 +362,15 @@ class SingleQuantumExecutor(QuantumExecutor):
             "Looking for existing outputs in the way for label=%s dataId=%s.", task_node.label, quantum.dataId
         )
         ref_dict = limited_butler.stored_many(chain.from_iterable(quantum.outputs.values()))
+        if task_node.log_output is not None:
+            (log_ref,) = quantum.outputs[task_node.log_output.dataset_type_name]
+            if ref_dict[log_ref]:
+                _LOG.debug(
+                    "Attaching logs from previous attempt on label=%s dataId=%s.",
+                    task_node.label,
+                    quantum.dataId,
+                )
+                log_extra.attach_previous_attempt(limited_butler.get(log_ref))
         existingRefs = [ref for ref, exists in ref_dict.items() if exists]
         missingRefs = [ref for ref, exists in ref_dict.items() if not exists]
         if existingRefs:

lsst/pipe/base/tests/mocks/_data_id_match.py

@@ -37,7 +37,7 @@ from uuid import UUID
 import astropy.time
 
 from lsst.daf.butler import DataId
-from lsst.daf.butler.registry.queries.expressions.parser import Node, TreeVisitor, parse_expression
+from lsst.daf.butler.queries.expressions.parser import Node, TreeVisitor, parse_expression
 
 
 class _DataIdMatchTreeVisitor(TreeVisitor):

lsst/pipe/base/tests/mocks/_pipeline_task.py

@@ -95,7 +95,7 @@ class ForcedFailure:
 
     memory_required: Quantity | None = None
     """If not `None`, this failure simulates an out-of-memory failure by
-    raising only if this value exceeds `ExecutionResources.max_mem`.f
+    raising only if this value exceeds `ExecutionResources.max_mem`.
     """
 
     def set_config(self, config: MockPipelineTaskConfig) -> None:

lsst/pipe/base/version.py

@@ -1,2 +1,2 @@
 __all__ = ["__version__"]
-__version__ = "29.2025.4500"
+__version__ = "29.2025.4700"

{lsst_pipe_base-29.2025.4500.dist-info → lsst_pipe_base-29.2025.4700.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lsst-pipe-base
-Version: 29.2025.4500
+Version: 29.2025.4700
 Summary: Pipeline infrastructure for the Rubin Science Pipelines.
 Author-email: Rubin Observatory Data Management <dm-admin@lists.lsst.org>
 License-Expression: BSD-3-Clause OR GPL-3.0-or-later