Enable provenance recording.

Author: TallJimbo

python/activator/exception.py

@@ -23,7 +23,7 @@
 __all__ = ["NonRetriableError", "RetriableError", "GracefulShutdownInterrupt", "TimeoutInterrupt",
            "InvalidVisitError", "IgnorableVisit",
            "InvalidPipelineError", "NoGoodPipelinesError",
-           "PipelinePreExecutionError", "PipelineExecutionError",
+           "PipelinePreExecutionError", "PipelineExecutionError", "ProvenanceDimensionsError",
            ]
 
 
@@ -144,3 +144,9 @@ class PipelineExecutionError(RuntimeError):
 
     Usually chained to an internal exception.
     """
+
+
+class ProvenanceDimensionsError(RuntimeError):
+    """Exception raised if a viable data ID for provenance could not be
+    defined.
+    """

python/activator/middleware_interface.py

@@ -62,7 +62,8 @@
 from shared.visit import FannedOutVisit
 from .caching import DatasetCache
 from .exception import GracefulShutdownInterrupt, TimeoutInterrupt, NonRetriableError, RetriableError, \
-    InvalidPipelineError, NoGoodPipelinesError, PipelinePreExecutionError, PipelineExecutionError
+    InvalidPipelineError, NoGoodPipelinesError, PipelinePreExecutionError, PipelineExecutionError, \
+    ProvenanceDimensionsError
 from .timer import enforce_schema, time_this_to_bundle
 
 _log = logging.getLogger("lsst." + __name__)
@@ -363,6 +364,7 @@ def __init__(self, read_butler: Butler, butler_writer: ButlerWriter, image_bucke
         self._define_dimensions()
         self._init_ingester()
         self._init_visit_definer()
+        self._init_provenance_dataset_type()
 
         # How much to pad the spatial region we will copy over.
         self.padding = padding*lsst.geom.arcseconds
@@ -439,6 +441,18 @@ def _init_governor_datasets(self, timestamp, skymap):
         self.skymap = self.read_central_butler.get("skyMap", skymap=self.skymap_name,
                                                    collections=self._collection_skymap)
 
+    def _init_provenance_dataset_type(self):
+        """Register the dataset types used to store provenance information.
+
+        ``self._init_local_butler`` must have already been run.
+        """
+        self._provenance_dataset_type = DatasetType(
+            "prompt_provenance",
+            self.butler.dimensions.conform(["group", "detector"]),
+            "ProvenanceQuantumGraph",
+        )
+        self.butler.registry.registerDatasetType(self._provenance_dataset_type)
+
     def _define_dimensions(self):
         """Define any dimensions that must be computed from this object's visit.
 
@@ -1292,7 +1306,7 @@ def _get_graph_executor(self, butler, factory):
         )
         graph_executor = MPGraphExecutor(
             # TODO: re-enable parallel execution once we can log as desired with CliLog or a successor
-            # (see issues linked from DM-42063)
+            # (see issues linked from DM-42063) AND once provenance is supported with multiprocessing.
             num_proc=1,  # Avoid spawning processes, because they bypass our logger
             timeout=2_592_000.0,  # In practice, timeout is never helpful; set to 30 days.
             quantum_executor=quantum_executor,
@@ -1370,6 +1384,11 @@ def _try_pipelines(self, pipelines, in_collections, data_ids, *, label):
                 # Diagnostic logs are the responsibility of GraphBuilder.
                 _log.error(f"Empty quantum graph for {pipeline_file}; see previous logs for details.")
                 continue
+            try:
+                provenance_ref = self._make_provenance_ref(data_ids, output_run)
+            except ProvenanceDimensionsError:
+                _log.exception(f"Failed to determine data ID for provenance for {pipeline_file}.")
+                continue
             # Past this point, partial execution creates datasets.
             # Don't retry -- either fail (raise) or break.
 
@@ -1384,7 +1403,8 @@ def _try_pipelines(self, pipelines, in_collections, data_ids, *, label):
                         _log, msg=f"executor.run_pipeline ({label})", level=logging.DEBUG):
                     executor.run_pipeline(
                         qgraph,
-                        graph_executor=self._get_graph_executor(exec_butler, factory)
+                        graph_executor=self._get_graph_executor(exec_butler, factory),
+                        provenance_dataset_ref=provenance_ref,
                     )
                     _log.info(f"{label.capitalize()} pipeline successfully run.")
                     return output_run
@@ -1397,6 +1417,33 @@ def _try_pipelines(self, pipelines, in_collections, data_ids, *, label):
         else:
             raise NoGoodPipelinesError(f"No {label} pipeline graph could be built.")
 
+    def _make_provenance_ref(self, where, output_run):
+        """Make the provenance DatasetRef for a quantum graph.
+
+        Parameters
+        ----------
+        where : `str`
+            Butler query expression that can be related to a single
+            ``{group, detector}`` data ID.
+        output_run : `str`
+            Output RUN collection.
+
+        Returns
+        -------
+        ref : `lsst.daf.butler.DatasetRef`
+            A reference to a to-be-written provenance dataset in ``output_run``.
+        """
+        query_results = self.butler.query_data_ids(
+            self._provenance_dataset_type.dimensions, where=where, explain=False
+        )
+        try:
+            (data_id,) = query_results
+        except ValueError:
+            raise ProvenanceDimensionsError(
+                f"Expected exactly one data ID for {self._provenance_dataset_type}; got {query_results}."
+            ) from None
+        return DatasetRef(self._provenance_dataset_type, data_id, run=output_run)
+
     def _run_preprocessing(self) -> None:
         """Preprocess a visit ahead of incoming image(s).
 

tests/test_middleware_interface.py

@@ -41,7 +41,9 @@
 import lsst.afw.image
 import lsst.afw.table
 from lsst.dax.apdb import ApdbSql
-from lsst.daf.butler import Butler, CollectionType, DataCoordinate, DimensionUniverse, EmptyQueryResultError
+from lsst.daf.butler import (
+    Butler, CollectionType, DataCoordinate, DatasetType, DimensionUniverse, EmptyQueryResultError
+)
 import lsst.daf.butler.tests as butler_tests
 from lsst.obs.base.formatters.fitsExposure import FitsImageFormatter
 from lsst.obs.base.ingest import RawFileDatasetInfo, RawFileData
@@ -717,6 +719,13 @@ def _check_run_pipeline_fallback(self, callable, pipe_files, graphs, final_label
             The description of the pipeline that should be run, given
             ``pipe_files`` and ``graphs``.
         """
+        test_provenance_dataset_type = DatasetType(
+            "test_provenance",
+            # Mocked QGs do not have realistic dimensions, and provenance
+            # dataset types need to have the same dimensions.
+            self.interface.butler.dimensions.conform(["detector"]),
+            "ProvenanceQuantumGraph"
+        )
         with (
             unittest.mock.patch(
                 "activator.middleware_interface.MiddlewareInterface.get_pre_pipeline_files",
@@ -736,6 +745,11 @@ def _check_run_pipeline_fallback(self, callable, pipe_files, graphs, final_label
             unittest.mock.patch(
                 "activator.middleware_interface.SeparablePipelineExecutor.run_pipeline"
             ) as mock_run,
+            unittest.mock.patch.object(
+                self.interface,
+                "_provenance_dataset_type",
+                test_provenance_dataset_type
+            ),
             self.assertLogs(self.logger_name, level="INFO") as logs,
         ):
             callable()