lsst-pipe-base 30.0.1rc1__py3-none-any.whl → 30.2025.5100__py3-none-any.whl

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

@@ -30,14 +30,130 @@ from __future__ import annotations
 __all__ = ("Writer",)
 
 import dataclasses
+import itertools
+import logging
+import operator
+import uuid
+from typing import TypeVar
 
+import networkx
 import zstandard
 
-from ...log_on_close import LogOnClose
+from lsst.utils.packages import Packages
+
+from ... import automatic_connection_constants as acc
 from ...pipeline_graph import TaskImportMode
-from .._predicted import PredictedQuantumGraphComponents, PredictedQuantumGraphReader
-from .._provenance import ProvenanceQuantumGraphWriter, ProvenanceQuantumScanData
+from .._common import BaseQuantumGraphWriter
+from .._multiblock import Compressor, MultiblockWriter
+from .._predicted import PredictedDatasetModel, PredictedQuantumGraphComponents, PredictedQuantumGraphReader
+from .._provenance import (
+    DATASET_ADDRESS_INDEX,
+    DATASET_MB_NAME,
+    LOG_ADDRESS_INDEX,
+    LOG_MB_NAME,
+    METADATA_ADDRESS_INDEX,
+    METADATA_MB_NAME,
+    QUANTUM_ADDRESS_INDEX,
+    QUANTUM_MB_NAME,
+    ProvenanceDatasetModel,
+    ProvenanceInitQuantaModel,
+    ProvenanceInitQuantumModel,
+    ProvenanceQuantumModel,
+)
 from ._communicators import WriterCommunicator
+from ._structs import WriteRequest
+
+
+@dataclasses.dataclass
+class _DataWriters:
+    """A struct of low-level writer objects for the main components of a
+    provenance quantum graph.
+
+    Parameters
+    ----------
+    comms : `WriterCommunicator`
+        Communicator helper object for the writer.
+    predicted : `.PredictedQuantumGraphComponents`
+        Components of the predicted graph.
+    indices : `dict` [ `uuid.UUID`, `int` ]
+        Mapping from UUID to internal integer ID, including both quanta and
+        datasets.
+    compressor : `Compressor`
+        Object that can compress `bytes`.
+    cdict_data : `bytes` or `None`, optional
+        Bytes representation of the compression dictionary used by the
+        compressor.
+    """
+
+    def __init__(
+        self,
+        comms: WriterCommunicator,
+        predicted: PredictedQuantumGraphComponents,
+        indices: dict[uuid.UUID, int],
+        compressor: Compressor,
+        cdict_data: bytes | None = None,
+    ) -> None:
+        assert comms.config.output_path is not None
+        header = predicted.header.model_copy()
+        header.graph_type = "provenance"
+        self.graph = comms.enter(
+            BaseQuantumGraphWriter.open(
+                comms.config.output_path,
+                header,
+                predicted.pipeline_graph,
+                indices,
+                address_filename="nodes",
+                compressor=compressor,
+                cdict_data=cdict_data,
+            ),
+            on_close="Finishing writing provenance quantum graph.",
+            is_progress_log=True,
+        )
+        self.graph.address_writer.addresses = [{}, {}, {}, {}]
+        self.logs = comms.enter(
+            MultiblockWriter.open_in_zip(self.graph.zf, LOG_MB_NAME, header.int_size, use_tempfile=True),
+            on_close="Copying logs into zip archive.",
+            is_progress_log=True,
+        )
+        self.graph.address_writer.addresses[LOG_ADDRESS_INDEX] = self.logs.addresses
+        self.metadata = comms.enter(
+            MultiblockWriter.open_in_zip(self.graph.zf, METADATA_MB_NAME, header.int_size, use_tempfile=True),
+            on_close="Copying metadata into zip archive.",
+            is_progress_log=True,
+        )
+        self.graph.address_writer.addresses[METADATA_ADDRESS_INDEX] = self.metadata.addresses
+        self.datasets = comms.enter(
+            MultiblockWriter.open_in_zip(self.graph.zf, DATASET_MB_NAME, header.int_size, use_tempfile=True),
+            on_close="Copying dataset provenance into zip archive.",
+            is_progress_log=True,
+        )
+        self.graph.address_writer.addresses[DATASET_ADDRESS_INDEX] = self.datasets.addresses
+        self.quanta = comms.enter(
+            MultiblockWriter.open_in_zip(self.graph.zf, QUANTUM_MB_NAME, header.int_size, use_tempfile=True),
+            on_close="Copying quantum provenance into zip archive.",
+            is_progress_log=True,
+        )
+        self.graph.address_writer.addresses[QUANTUM_ADDRESS_INDEX] = self.quanta.addresses
+
+    graph: BaseQuantumGraphWriter
+    """The parent graph writer."""
+
+    datasets: MultiblockWriter
+    """A writer for dataset provenance."""
+
+    quanta: MultiblockWriter
+    """A writer for quantum provenance."""
+
+    metadata: MultiblockWriter
+    """A writer for metadata content."""
+
+    logs: MultiblockWriter
+    """A writer for log content."""
+
+    @property
+    def compressor(self) -> Compressor:
+        """Object that should be used to compress all JSON blocks."""
+        return self.graph.compressor
 
 
 @dataclasses.dataclass
@@ -55,13 +171,46 @@ class Writer:
     predicted: PredictedQuantumGraphComponents = dataclasses.field(init=False)
     """Components of the predicted quantum graph."""
 
-    pending_compression_training: list[ProvenanceQuantumScanData] = dataclasses.field(default_factory=list)
+    existing_init_outputs: dict[uuid.UUID, set[uuid.UUID]] = dataclasses.field(default_factory=dict)
+    """Mapping that tracks which init-outputs exist.
+
+    This mapping is updated as scanners inform the writer about init-output
+    existence, since we want to write that provenance information out only at
+    the end.
+    """
+
+    indices: dict[uuid.UUID, int] = dataclasses.field(default_factory=dict)
+    """Mapping from UUID to internal integer ID, including both quanta and
+    datasets.
+
+    This is fully initialized at construction.
+    """
+
+    output_dataset_ids: set[uuid.UUID] = dataclasses.field(default_factory=set)
+    """The IDs of all datasets that are produced by this graph.
+
+    This is fully initialized at construction.
+    """
+
+    overall_inputs: dict[uuid.UUID, PredictedDatasetModel] = dataclasses.field(default_factory=dict)
+    """All datasets that are not produced by any quantum in this graph."""
+
+    xgraph: networkx.DiGraph = dataclasses.field(default_factory=networkx.DiGraph)
+    """A bipartite NetworkX graph linking datasets to quanta and quanta to
+    datasets.
+
+    This is fully initialized at construction.  There are no node or edge
+    attributes in this graph; we only need it to store adjacency information
+    with datasets as well as with quanta.
+    """
+
+    pending_compression_training: list[WriteRequest] = dataclasses.field(default_factory=list)
     """Unprocessed quantum scans that are being accumulated in order to
     build a compression dictionary.
     """
 
     def __post_init__(self) -> None:
-        assert self.comms.config.is_writing_provenance, "Writer should not be used if writing is disabled."
+        assert self.comms.config.output_path is not None, "Writer should not be used if writing is disabled."
         self.comms.log.info("Reading predicted quantum graph.")
         with PredictedQuantumGraphReader.open(
             self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT
@@ -71,6 +220,58 @@ class Writer:
             self.comms.check_for_cancel()
             reader.read_quantum_datasets()
             self.predicted = reader.components
+        for predicted_init_quantum in self.predicted.init_quanta.root:
+            self.existing_init_outputs[predicted_init_quantum.quantum_id] = set()
+        self.comms.check_for_cancel()
+        self.comms.log.info("Generating integer indexes and identifying outputs.")
+        self._populate_indices_and_outputs()
+        self.comms.check_for_cancel()
+        self._populate_xgraph_and_inputs()
+        self.comms.check_for_cancel()
+        self.comms.log_progress(
+            # We add one here for 'packages', which we do ingest but don't
+            # record provenance for.
+            logging.INFO,
+            f"Graph has {len(self.output_dataset_ids) + 1} predicted output dataset(s).",
+        )
+
+    def _populate_indices_and_outputs(self) -> None:
+        all_uuids = set(self.predicted.quantum_datasets.keys())
+        for quantum in self.comms.periodically_check_for_cancel(
+            itertools.chain(
+                self.predicted.init_quanta.root,
+                self.predicted.quantum_datasets.values(),
+            )
+        ):
+            if not quantum.task_label:
+                # Skip the 'packages' producer quantum.
+                continue
+            all_uuids.update(quantum.iter_input_dataset_ids())
+            self.output_dataset_ids.update(quantum.iter_output_dataset_ids())
+        all_uuids.update(self.output_dataset_ids)
+        self.indices = {
+            node_id: node_index
+            for node_index, node_id in self.comms.periodically_check_for_cancel(
+                enumerate(sorted(all_uuids, key=operator.attrgetter("int")))
+            )
+        }
+
+    def _populate_xgraph_and_inputs(self) -> None:
+        for predicted_quantum in self.comms.periodically_check_for_cancel(
+            itertools.chain(
+                self.predicted.init_quanta.root,
+                self.predicted.quantum_datasets.values(),
+            )
+        ):
+            if not predicted_quantum.task_label:
+                # Skip the 'packages' producer quantum.
+                continue
+            for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
+                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(predicted_quantum.quantum_id, predicted_output.dataset_id)
 
     @staticmethod
     def run(predicted_path: str, comms: WriterCommunicator) -> None:
@@ -86,7 +287,7 @@ class Writer:
         Notes
         -----
         This method is designed to run as the ``target`` in
-        `WorkerFactory.make_worker`.
+        `WorkerContext.make_worker`.
         """
         with comms:
             writer = Writer(predicted_path, comms)
@@ -94,59 +295,52 @@ class Writer:
 
     def loop(self) -> None:
         """Run the main loop for the writer."""
-        qg_writer: ProvenanceQuantumGraphWriter | None = None
+        data_writers: _DataWriters | None = None
         if not self.comms.config.zstd_dict_size:
-            qg_writer = self.make_qg_writer()
+            data_writers = self.make_data_writers()
         self.comms.log.info("Polling for write requests from scanners.")
         for request in self.comms.poll():
-            if qg_writer is None:
+            if data_writers is None:
                 self.pending_compression_training.append(request)
                 if len(self.pending_compression_training) >= self.comms.config.zstd_dict_n_inputs:
-                    qg_writer = self.make_qg_writer()
+                    data_writers = self.make_data_writers()
             else:
-                qg_writer.write_scan_data(request)
-                self.comms.report_write()
-        if qg_writer is None:
-            qg_writer = self.make_qg_writer()
-        self.comms.log.info("Writing init outputs.")
-        qg_writer.write_init_outputs(assume_existence=False)
+                self.process_request(request, data_writers)
+        if data_writers is None:
+            data_writers = self.make_data_writers()
+        self.write_init_outputs(data_writers)
 
-    def make_qg_writer(self) -> ProvenanceQuantumGraphWriter:
+    def make_data_writers(self) -> _DataWriters:
         """Make a compression dictionary, open the low-level writers, and
         write any accumulated scans that were needed to make the compression
         dictionary.
 
         Returns
         -------
-        qg_writer : `ProvenanceQuantumGraphWriter`
+        data_writers : `_DataWriters`
             Low-level writers struct.
         """
         cdict = self.make_compression_dictionary()
         self.comms.send_compression_dict(cdict.as_bytes())
-        assert self.comms.config.is_writing_provenance and self.comms.config.output_path is not None
-        self.comms.log.info("Opening output files and processing predicted graph.")
-        qg_writer = ProvenanceQuantumGraphWriter(
-            self.comms.config.output_path,
-            exit_stack=self.comms.exit_stack,
-            log_on_close=LogOnClose(self.comms.log_progress),
-            predicted=self.predicted,
-            zstd_level=self.comms.config.zstd_level,
+        assert self.comms.config.output_path is not None
+        self.comms.log.info("Opening output files.")
+        data_writers = _DataWriters(
+            self.comms,
+            self.predicted,
+            self.indices,
+            compressor=zstandard.ZstdCompressor(self.comms.config.zstd_level, cdict),
             cdict_data=cdict.as_bytes(),
-            loop_wrapper=self.comms.periodically_check_for_cancel,
-            log=self.comms.log,
         )
         self.comms.check_for_cancel()
         self.comms.log.info("Compressing and writing queued scan requests.")
         for request in self.pending_compression_training:
-            qg_writer.write_scan_data(request)
-            self.comms.report_write()
+            self.process_request(request, data_writers)
         del self.pending_compression_training
         self.comms.check_for_cancel()
-        self.comms.log.info("Writing overall inputs.")
-        qg_writer.write_overall_inputs(self.comms.periodically_check_for_cancel)
-        qg_writer.write_packages()
+        self.write_overall_inputs(data_writers)
+        self.write_packages(data_writers)
         self.comms.log.info("Returning to write request loop.")
-        return qg_writer
+        return data_writers
 
     def make_compression_dictionary(self) -> zstandard.ZstdCompressionDict:
         """Make the compression dictionary.
@@ -182,3 +376,126 @@ class Writer:
             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:
+        """Write provenance for init-output datasets and init-quanta.
+
+        Parameters
+        ----------
+        data_writers : `_DataWriters`
+            Low-level writers struct.
+        """
+        self.comms.log.info("Writing init outputs.")
+        init_quanta = ProvenanceInitQuantaModel()
+        for predicted_init_quantum in self.predicted.init_quanta.root:
+            if not predicted_init_quantum.task_label:
+                # Skip the 'packages' producer quantum.
+                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()):
+                provenance_output = ProvenanceDatasetModel.from_predicted(
+                    predicted_output,
+                    producer=predicted_init_quantum.quantum_id,
+                    consumers=self.xgraph.successors(predicted_output.dataset_id),
+                )
+                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))
+        data_writers.graph.write_single_model("init_quanta", init_quanta)
+
+    def write_overall_inputs(self, data_writers: _DataWriters) -> None:
+        """Write provenance for overall-input datasets.
+
+        Parameters
+        ----------
+        data_writers : `_DataWriters`
+            Low-level writers struct.
+        """
+        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:
+                data_writers.datasets.write_model(
+                    predicted_input.dataset_id,
+                    ProvenanceDatasetModel.from_predicted(
+                        predicted_input,
+                        producer=None,
+                        consumers=self.xgraph.successors(predicted_input.dataset_id),
+                    ),
+                    data_writers.compressor,
+                )
+        del self.overall_inputs
+
+    @staticmethod
+    def write_packages(data_writers: _DataWriters) -> None:
+        """Write package version information to the provenance graph.
+
+        Parameters
+        ----------
+        data_writers : `_DataWriters`
+            Low-level writers struct.
+        """
+        packages = Packages.fromSystem(include_all=True)
+        data = packages.toBytes("json")
+        data_writers.graph.write_single_block("packages", data)
+
+    def process_request(self, request: WriteRequest, data_writers: _DataWriters) -> None:
+        """Process a `WriteRequest` into `_ScanData`.
+
+        Parameters
+        ----------
+        request : `WriteRequest`
+            Result of a quantum scan.
+        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
+        self.comms.log.debug("Handling quantum scan for %s.", request.quantum_id)
+        predicted_quantum = self.predicted.quantum_datasets[request.quantum_id]
+        outputs: dict[uuid.UUID, bytes] = {}
+        for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
+            provenance_output = ProvenanceDatasetModel.from_predicted(
+                predicted_output,
+                producer=predicted_quantum.quantum_id,
+                consumers=self.xgraph.successors(predicted_output.dataset_id),
+            )
+            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 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[request.quantum_id]
+        self.comms.report_write()
+
+
+_T = TypeVar("_T")

lsst/pipe/base/quantum_graph/visualization.py

@@ -37,16 +37,12 @@ from typing import IO, ClassVar, Generic, TypeVar
 from ..pipeline_graph import NodeType
 from ._common import BaseQuantumGraph, BipartiteEdgeInfo, DatasetInfo, QuantumInfo
 
-# We use the old generic syntax in this module because for some reason the new
-# one confused Sphinx (or one of its plugins), even though it seems fine with
-# it in other places.  We can try again when we're ready to remove types from
-# the docstrings of annotated functions, in case that matters.
 _G = TypeVar("_G", bound=BaseQuantumGraph, contravariant=True)
 _Q = TypeVar("_Q", bound=QuantumInfo, contravariant=True)
 _D = TypeVar("_D", bound=DatasetInfo, contravariant=True)
 
 
-class QuantumGraphVisualizer(Generic[_G, _Q, _D]):  # noqa: UP046
+class QuantumGraphVisualizer(Generic[_G, _Q, _D]):
     """A base class for exporting quantum graphs to graph-visualization
     languages.
 

lsst/pipe/base/quantum_graph_builder.py

@@ -380,6 +380,8 @@ class QuantumGraphBuilder(ABC):
 
         Parameters
         ----------
+        metadata : `~collections.abc.Mapping`, optional
+            Flexible metadata to add to the quantum graph.
         attach_datastore_records : `bool`, optional
             Whether to include datastore records in the graph.  Required for
             `lsst.daf.butler.QuantumBackedButler` execution.
@@ -678,26 +680,6 @@ class QuantumGraphBuilder(ABC):
                 "Dropping task %s because no quanta remain%s.", task_node.label, message_parenthetical
             )
             skeleton.remove_task(task_node.label)
-        if len(no_work_quanta) > len(remaining_quanta):
-            only_overall_inputs = self._get_task_inputs_if_overall_only(task_node)
-            self.log.warning(
-                "More than half of %s quanta had no work to do given available inputs.\n"
-                "A query constraint on one of %s may yield a much faster build.",
-                task_node.label,
-                only_overall_inputs,
-            )
-
-    def _get_task_inputs_if_overall_only(self, task_node: TaskNode) -> list[str] | None:
-        """If the given task consumes only overall-inputs, return their names.
-        Otherwise return `None`.
-        """
-        result: list[str] = []
-        for read_edge in task_node.inputs.values():
-            if self._pipeline_graph.producer_of(read_edge.parent_dataset_type_name) is None:
-                result.append(read_edge.parent_dataset_type_name)
-            else:
-                return None
-        return result
 
     def _skip_quantum_if_metadata_exists(
         self, task_node: TaskNode, quantum_key: QuantumKey, skeleton: QuantumGraphSkeleton
@@ -905,6 +887,11 @@ class QuantumGraphBuilder(ABC):
             Identifier for this quantum in the graph.
         skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
             Preliminary quantum graph, to be modified in-place.
+        skypix_bounds_builder : `~prerequisite_helpers.SkyPixBoundsBuilder`
+            An object that accumulates the appropriate spatial bounds for a
+            quantum.
+        timespan_builder : `~prerequisite_helpers.TimespanBuilder`
+            An object that accumulates the appropriate timespan for a quantum.
 
         Returns
         -------
@@ -1157,7 +1144,7 @@ class QuantumGraphBuilder(ABC):
             "outputs" attributes on all quantum nodes, as added by
             `_resolve_task_quanta`, as well as a "datastore_records" attribute
             as added by `_attach_datastore_records`.
-        metadata : `~collections.abc.Mapping`
+        metadata : `Mapping`
             Flexible metadata to add to the graph.
 
         Returns