Merge pull request #556 from lsst/tickets/DM-54162

DM-54162: write in-process-gathered provenance when exceptions are raised and handled

Author: TallJimbo

doc/changes/DM-54162.feature.md

@@ -0,0 +1,3 @@
+Write provenance gathered from in-process execution (i.e. `SeparablePipelineExecutor` and `MPGraphExecutor`) when a task raises an exception and the executor successfully catches it and continues with other tasks.
+
+This also includes task construction in the quantum executor's error-handling logic, so exceptions raised at that point do not break provenance writing or require any other special-casing at higher levels.

python/lsst/pipe/base/mp_graph_executor.py

@@ -531,6 +531,7 @@ def execute(
             new_graph = graph
         xgraph = self._make_xgraph(new_graph, old_graph)
         self._report = Report(qgraphSummary=new_graph._make_summary())
+        err: MPGraphExecutorError | None = None
         with ExitStack() as exit_stack:
             provenance_writer: ProvenanceQuantumGraphWriter | None = None
             if provenance_graph_file is not None:
@@ -549,13 +550,21 @@ def execute(
                     self._execute_quanta_mp(xgraph, self._report)
                 else:
                     self._execute_quanta_in_process(xgraph, self._report, provenance_writer)
+            except MPGraphExecutorError as exc:
+                self._report.set_exception(exc)
+                err = exc
+                # Defer re-raising this exception only to let provenance writes
+                # finish as the ExitStack closes. The original traceback for
+                # this exception isn't useful anyway.
             except Exception as exc:
                 self._report.set_exception(exc)
                 raise
             if provenance_writer is not None:
                 provenance_writer.write_overall_inputs()
                 provenance_writer.write_packages()
                 provenance_writer.write_init_outputs(assume_existence=True)
+        if err is not None:
+            raise err
 
     def _make_xgraph(
         self, new_graph: PredictedQuantumGraph, old_graph: QuantumGraph | None
@@ -724,6 +733,8 @@ def tiebreaker_sort_key(quantum_id: uuid.UUID) -> tuple:
                             taskLabel=downstream_node_state["task_label"],
                         )
                         report.quantaReports.append(failed_quantum_report)
+                        if provenance_writer is not None:
+                            provenance_writer.write_blocked_quantum_provenance(downstream_quantum_id)
                         _LOG.error(
                             "Upstream job failed for task %s (%s@%s), skipping this quantum.",
                             downstream_quantum_id,

python/lsst/pipe/base/quantum_graph/_provenance.py

@@ -2040,6 +2040,17 @@ def write_quantum_provenance(
         scan_data = provenance_models.to_scan_data(predicted_quantum, compressor=self.compressor)
         self.write_scan_data(scan_data)
 
+    def write_blocked_quantum_provenance(self, quantum_id: uuid.UUID) -> None:
+        """Gather and write provenance for a quantum that was blocked by an
+        upstream failure.
+
+        Parameters
+        ----------
+        quantum_id : `uuid.UUID`
+            Unique ID for the quantum.
+        """
+        self.write_scan_data(ProvenanceQuantumScanData.make_blocked(quantum_id))
+
     def write_scan_data(self, scan_data: ProvenanceQuantumScanData) -> None:
         """Write the output of a quantum provenance scan to disk.
 
@@ -2436,6 +2447,27 @@ class ProvenanceQuantumScanData:
     compressed.
     """
 
+    @classmethod
+    def make_blocked(cls, quantum_id: uuid.UUID) -> ProvenanceQuantumScanData:
+        """Construct provenance information for a quantum blocked by an
+        upstream failure.
+
+        Parameters
+        ----------
+        quantum_id : `uuid.UUID`
+            Unique ID of the quantum
+
+        Returns
+        -------
+        scan_data : `ProvenanceQuantumScanData`
+            Struct with ready-to-write provenance data.
+        """
+        return ProvenanceQuantumScanData(
+            quantum_id,
+            status=ProvenanceQuantumScanStatus.BLOCKED,
+            is_compressed=True,  # nothing to compress
+        )
+
     def compress(self, compressor: Compressor) -> None:
         """Compress the data in this struct if it has not been compressed
         already.

python/lsst/pipe/base/separable_pipeline_executor.py

@@ -46,7 +46,7 @@
 from ._quantumContext import ExecutionResources
 from .all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
 from .graph import QuantumGraph
-from .mp_graph_executor import MPGraphExecutor
+from .mp_graph_executor import MPGraphExecutor, MPGraphExecutorError
 from .pipeline import Pipeline
 from .quantum_graph import PredictedQuantumGraph
 from .quantum_graph_builder import QuantumGraphBuilder
@@ -390,11 +390,12 @@ def run_pipeline(
         provenance_dataset_ref : `lsst.daf.butler.DatasetRef`, optional
             Dataset that should be used to save provenance.  Provenance is only
             supported when running in a single process (at least for the
-            default quantum executor), and should not be used with
-            ``skip_existing_in=[output_run]`` when retrying a previous
-            execution attempt. The caller is responsible for registering the
-            dataset type and for ensuring that the dimensions of this dataset
-            do not lead to uniqueness conflicts.
+            default quantum executor), and should not be enabled in contexts
+            where a quantum might be executed more than once (i.e. retried)
+            within the same `~lsst.daf.butler.CollectionType.RUN` collection.
+            The caller is responsible for registering the dataset type and for
+            ensuring that the dimensions of this dataset do not lead to
+            uniqueness conflicts.
         """
         if not graph_executor:
             quantum_executor = SingleQuantumExecutor(
@@ -417,7 +418,16 @@ def run_pipeline(
 
         if provenance_dataset_ref is not None:
             with TemporaryForIngest(self._butler, provenance_dataset_ref) as temporary:
-                graph_executor.execute(graph, provenance_graph_file=temporary.ospath)
-                temporary.ingest()
+                try:
+                    graph_executor.execute(graph, provenance_graph_file=temporary.ospath)
+                    temporary.ingest()
+                except MPGraphExecutorError:
+                    # If the graph executor itself raised, it will have
+                    # finished the provenance rewrite.  In other cases the
+                    # temporary file might be incomplete or corrupted and we
+                    # can't roll the dice on ingesting it.
+                    temporary.ingest()
+                    raise
+
         else:
             graph_executor.execute(graph)

python/lsst/pipe/base/single_quantum_executor.py

@@ -293,10 +293,10 @@ def _execute_with_limited_butler(
                 task_node.label,
                 quantum.dataId,
             )
-            task = self._task_factory.makeTask(task_node, limited_butler, init_input_refs)
-            logInfo(None, "start", metadata=quantumMetadata)  # type: ignore[arg-type]
-            outputs_put: list[uuid.UUID] = []
             try:
+                task = self._task_factory.makeTask(task_node, limited_butler, init_input_refs)
+                logInfo(None, "start", metadata=quantumMetadata)  # type: ignore[arg-type]
+                outputs_put: list[uuid.UUID] = []
                 with limited_butler.record_metrics() as butler_metrics:
                     caveats = self._run_quantum(
                         task, quantum, task_node, limited_butler, quantum_id=quantum_id, ids_put=outputs_put

tests/test_separable_pipeline_executor.py

@@ -50,6 +50,8 @@
     PROVENANCE_DATASET_TYPE_NAME,
     PROVENANCE_STORAGE_CLASS,
 )
+from lsst.pipe.base.mp_graph_executor import MPGraphExecutorError
+from lsst.pipe.base.quantum_graph import ProvenanceQuantumGraph
 from lsst.pipe.base.quantum_graph_builder import OutputExistsError
 from lsst.pipe.base.separable_pipeline_executor import SeparablePipelineExecutor
 from lsst.pipe.base.tests.mocks import (
@@ -1153,60 +1155,106 @@ class SeparablePipelineExecutorMockTests(lsst.utils.tests.TestCase):
     the lsst.pipe.base.tests.mocks system to define complex pipelines.
     """
 
+    def setUp(self):
+        # 'base.yaml' adds an instrument, 'Cam1', with four detectors and
+        # two physical filters.
+        self.helper, _ = self.enterContext(DirectButlerRepo.make_temporary("base.yaml"))
+
+    def run_base_test(
+        self, b_config: DynamicTestPipelineTaskConfig, expected_error: type[Exception] | None
+    ) -> ProvenanceQuantumGraph:
+        """Build and run a quantum graph with three tasks and four data IDs,
+        with customization of the middle task.
+        """
+        self.helper.add_task("a", dimensions=["detector"])
+        self.helper.add_task("b", dimensions=["detector"], config=b_config)
+        self.helper.add_task("c", dimensions=["detector"])
+        qg = self.helper.make_quantum_graph()
+        self.helper.butler.collections.register(qg.header.output_run)
+        qg.init_output_run(self.helper.butler, existing=False)
+        executor = SeparablePipelineExecutor(
+            self.helper.butler.clone(collections=qg.header.inputs, run=qg.header.output_run)
+        )
+        provenance_type = lsst.daf.butler.DatasetType(
+            PROVENANCE_DATASET_TYPE_NAME,
+            self.helper.butler.dimensions.empty,
+            PROVENANCE_STORAGE_CLASS,
+        )
+        self.helper.butler.registry.registerDatasetType(provenance_type)
+        provenance_ref = lsst.daf.butler.DatasetRef(
+            provenance_type,
+            lsst.daf.butler.DataCoordinate.make_empty(self.helper.butler.dimensions),
+            run=qg.header.output_run,
+        )
+        if expected_error is None:
+            executor.run_pipeline(qg, provenance_dataset_ref=provenance_ref)
+        else:
+            with self.assertRaises(expected_error):
+                executor.run_pipeline(qg, provenance_dataset_ref=provenance_ref)
+        provenance_graph = self.helper.butler.get(provenance_ref)
+        self.assertEqual(len(provenance_graph.quanta_by_task), 3)
+        self.assertEqual(len(provenance_graph.quanta_by_task["a"]), 4)
+        self.assertEqual(len(provenance_graph.quanta_by_task["b"]), 4)
+        self.assertEqual(len(provenance_graph.quanta_by_task["c"]), 4)
+        return provenance_graph
+
     def test_no_work_chain_provenance(self):
         """Test provenance recording when a NoWorkFound error chains to
         downstream tasks during execution.
         """
+        b_config = DynamicTestPipelineTaskConfig()
+        b_config.fail_exception = "lsst.pipe.base.NoWorkFound"
+        b_config.fail_condition = "detector=2"
+        provenance_graph = self.run_base_test(b_config, expected_error=None)
+        xgraph = provenance_graph.quantum_only_xgraph
+        for quantum_id in provenance_graph.quanta_by_task["a"].values():
+            self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.SUCCESSFUL)
+            self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
+        for data_id, quantum_id in provenance_graph.quanta_by_task["b"].items():
+            self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.SUCCESSFUL)
+            if data_id["detector"] == 2:
+                self.assertTrue(xgraph.nodes[quantum_id]["caveats"] & QuantumSuccessCaveats.NO_WORK)
+            else:
+                self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
+        for data_id, quantum_id in provenance_graph.quanta_by_task["c"].items():
+            self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.SUCCESSFUL)
+            if data_id["detector"] == 2:
+                self.assertTrue(
+                    xgraph.nodes[quantum_id]["caveats"] & QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED
+                )
+            else:
+                self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
+
+    def test_failure_block_provenance(self):
+        """Test provenance recording when an exception blocks one branch of a
+        QG but not another.
+        """
         # 'base.yaml' adds an instrument, 'Cam1', with four detectors and
         # two physical filters.
-        with DirectButlerRepo.make_temporary("base.yaml") as (helper, _):
-            helper.add_task("a", dimensions=["detector"])
-            b_config = DynamicTestPipelineTaskConfig()
-            b_config.fail_exception = "lsst.pipe.base.NoWorkFound"
-            b_config.fail_condition = "detector=2"
-            helper.add_task("b", dimensions=["detector"], config=b_config)
-            helper.add_task("c", dimensions=["detector"])
-            qg = helper.make_quantum_graph()
-            helper.butler.collections.register(qg.header.output_run)
-            qg.init_output_run(helper.butler, existing=False)
-            provenance_type = lsst.daf.butler.DatasetType(
-                PROVENANCE_DATASET_TYPE_NAME,
-                helper.butler.dimensions.empty,
-                PROVENANCE_STORAGE_CLASS,
-            )
-            helper.butler.registry.registerDatasetType(provenance_type)
-            provenance_ref = lsst.daf.butler.DatasetRef(
-                provenance_type,
-                lsst.daf.butler.DataCoordinate.make_empty(helper.butler.dimensions),
-                run=qg.header.output_run,
-            )
-            executor = SeparablePipelineExecutor(
-                helper.butler.clone(collections=qg.header.inputs, run=qg.header.output_run)
-            )
-            executor.run_pipeline(qg, provenance_dataset_ref=provenance_ref)
-            provenance_graph = helper.butler.get(provenance_ref)
-            self.assertEqual(len(provenance_graph.quanta_by_task), 3)
-            self.assertEqual(len(provenance_graph.quanta_by_task["a"]), 4)
-            self.assertEqual(len(provenance_graph.quanta_by_task["b"]), 4)
-            self.assertEqual(len(provenance_graph.quanta_by_task["c"]), 4)
-            xgraph = provenance_graph.quantum_only_xgraph
-            for quantum_id in provenance_graph.quanta_by_task["a"].values():
+        b_config = DynamicTestPipelineTaskConfig()
+        b_config.fail_exception = "builtins.RuntimeError"
+        b_config.fail_condition = "detector=2"
+        provenance_graph = self.run_base_test(b_config, MPGraphExecutorError)
+        self.assertEqual(len(provenance_graph.quanta_by_task), 3)
+        self.assertEqual(len(provenance_graph.quanta_by_task["a"]), 4)
+        self.assertEqual(len(provenance_graph.quanta_by_task["b"]), 4)
+        self.assertEqual(len(provenance_graph.quanta_by_task["c"]), 4)
+        xgraph = provenance_graph.quantum_only_xgraph
+        for quantum_id in provenance_graph.quanta_by_task["a"].values():
+            self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.SUCCESSFUL)
+            self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
+        for data_id, quantum_id in provenance_graph.quanta_by_task["b"].items():
+            if data_id["detector"] == 2:
+                self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.FAILED)
+            else:
                 self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.SUCCESSFUL)
                 self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
-            for data_id, quantum_id in provenance_graph.quanta_by_task["b"].items():
+        for data_id, quantum_id in provenance_graph.quanta_by_task["c"].items():
+            if data_id["detector"] == 2:
+                self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.BLOCKED)
+            else:
                 self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.SUCCESSFUL)
-                if data_id["detector"] == 2:
-                    self.assertTrue(xgraph.nodes[quantum_id]["caveats"] & QuantumSuccessCaveats.NO_WORK)
-                else:
-                    self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
-            for data_id, quantum_id in provenance_graph.quanta_by_task["c"].items():
-                self.assertEqual(xgraph.nodes[quantum_id]["status"], QuantumAttemptStatus.SUCCESSFUL)
-                if data_id["detector"] == 2:
-                    self.assertTrue(
-                        xgraph.nodes[quantum_id]["caveats"] & QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED
-                    )
-                else:
-                    self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
+                self.assertEqual(xgraph.nodes[quantum_id]["caveats"], QuantumSuccessCaveats.NO_CAVEATS)
 
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):