DM-54879: Support reprocessing when upstream outputs are selectively retained

doc/changes/DM-54879.feature.md

@@ -0,0 +1 @@
+Added `--retained-dataset-types` option to `pipetask qgraph` and `pipetask run`.

python/lsst/ctrl/mpexec/cli/opt/optionGroups.py

@@ -127,6 +127,7 @@ def __init__(
             ctrlMpExecOpts.qgraph_datastore_records_option(),
             ctrlMpExecOpts.skip_existing_in_option(),
             ctrlMpExecOpts.skip_existing_option(),
+            ctrlMpExecOpts.retained_dataset_types_option(),
             ctrlMpExecOpts.save_qgraph_option(),
             ctrlMpExecOpts.qgraph_dot_option(),
             ctrlMpExecOpts.qgraph_mermaid_option(),

python/lsst/ctrl/mpexec/cli/opt/options.py

@@ -358,6 +358,23 @@
 )
 
 
+retained_dataset_types_option = MWOptionDecorator(
+    "--retained-dataset-types",
+    default=None,
+    metavar="PATH",
+    type=MWPath(file_okay=True, dir_okay=False, readable=True),
+    help=unwrap(
+        """Path to a YAML file listing dataset type names or glob-style wildcard
+        patterns that should exist in --skip-existing-in when the producing
+        task ran successfully.  When a quantum should run, the builder propagates the
+        must-run signal backward through non-retained input datasets, forcing
+        the upstream quanta that need to regenerate those intermediates to also
+        run.  Has no effect without ``skip_existing_in``.
+        """
+    ),
+)
+
+
 clobber_outputs_option = MWOptionDecorator(
     "--clobber-outputs",
     help=(

python/lsst/ctrl/mpexec/cli/script/qgraph.py

@@ -33,6 +33,7 @@
 from collections.abc import Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING
 
+import yaml
 from astropy.table import Table
 
 from lsst.pipe.base import BuildId, QuantumGraph
@@ -68,6 +69,7 @@ def qgraph(
     qgraph_datastore_records: bool,
     skip_existing_in: Iterable[str] | None,
     skip_existing: bool,
+    retained_dataset_types: str | None,
     save_qgraph: ResourcePathExpression | None,
     qgraph_dot: str | None,
     qgraph_mermaid: str | None,
@@ -119,6 +121,13 @@ def qgraph(
         from the QuantumGraph.
     skip_existing : `bool`
         Appends output RUN collection to the ``skip_existing_in`` list.
+    retained_dataset_types : `str` or `None`
+        Path to a YAML file listing dataset type names or glob-style wildcard
+        patterns that should exist in ``skip_existing_in`` when the producing
+        task ran successfully.  When a quantum should run, the builder
+        propagates the must-run signal backward through non-retained input
+        datasets, forcing the upstream quanta that need to regenerate those
+        intermediates to also run.  Has no effect without ``skip_existing_in``.
     save_qgraph : convertible to `lsst.resources.ResourcePath` or `None`
         URI location for saving the quantum graph.
     qgraph_dot : `str` or `None`
@@ -205,6 +214,15 @@ def qgraph(
         skip_existing = True
 
     skip_existing_in = tuple(skip_existing_in) if skip_existing_in is not None else ()
+    retained_dataset_type_patterns: list[str] | None = None
+    if retained_dataset_types is not None:
+        with open(retained_dataset_types) as f:
+            retained_dataset_type_patterns = yaml.safe_load(f)
+        if not isinstance(retained_dataset_type_patterns, list) or not retained_dataset_type_patterns:
+            raise ValueError(
+                f"--retained-dataset-types file {retained_dataset_types!r} must contain "
+                "a non-empty YAML sequence of strings."
+            )
     if data_query is None:
         data_query = ""
     inputs = list(ensure_iterable(input)) if input else []
@@ -303,6 +321,7 @@ def qgraph(
                 butler,
                 where=data_query,
                 skip_existing_in=skip_existing_in,
+                retained_dataset_types=retained_dataset_type_patterns,
                 clobber=clobber_outputs,
                 dataset_query_constraint=DatasetQueryConstraintVariant.fromExpression(
                     dataset_query_constraint
@@ -317,6 +336,7 @@ def qgraph(
                 "extend_run": extend_run,
                 "skip_existing_in": skip_existing_in,
                 "skip_existing": skip_existing,
+                "retained_dataset_types": retained_dataset_types,
                 "data_query": data_query,
             }
             assert run is not None, "Butler output run collection must be defined"

tests/test_run.py

@@ -25,8 +25,10 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+import contextlib
 import logging
 import os
+import tempfile
 import time
 import unittest
 import unittest.mock
@@ -46,6 +48,34 @@
 from lsst.pipe.base.tests.mocks import DirectButlerRepo, DynamicTestPipelineTaskConfig
 
 
+# Copied from test_build.py
+@contextlib.contextmanager
+def make_tmp_file(contents=None, suffix=None):
+    """Context manager for generating temporary file name.
+
+    Temporary file is deleted on exiting context.
+
+    Parameters
+    ----------
+    contents : `bytes` or `None`, optional
+        Data to write into a file.
+    suffix : `str` or `None`, optional
+        Suffix to use for temporary file.
+
+    Yields
+    ------
+    `str`
+        Name of the temporary file.
+    """
+    fd, tmpname = tempfile.mkstemp(suffix=suffix)
+    if contents:
+        os.write(fd, contents)
+    os.close(fd)
+    yield tmpname
+    with contextlib.suppress(OSError):
+        os.remove(tmpname)
+
+
 class RunTestCase(unittest.TestCase):
     """Test pipetask run command-line."""
 
@@ -652,6 +682,157 @@ def test_qg_partial_failure(self):
             with helper.butler.query() as query:
                 self.assertEqual(query.datasets("dataset_auto1", collections=["output"]).count(), 3)
 
+    def test_retained_dataset_types_option(self):
+        """--retained-dataset-types accepts a file path."""
+        with make_tmp_file(b"- '*_metadata'\n- '*_log'\n", suffix=".yaml") as retained_path:
+            kwargs = self._make_run_args(
+                "-b",
+                "fake_repo",
+                "-i",
+                "fake_input",
+                "-o",
+                "fake_output",
+                "--retained-dataset-types",
+                retained_path,
+            )
+            self.assertEqual(kwargs["retained_dataset_types"], retained_path)
+
+    def test_simple_qg_retained_forces_rerun(self):
+        """With --retained-dataset-types listing only metadata types, when
+        task_auto2 has no metadata and must run, task_auto1 is forced to rerun
+        because dataset_auto1 is not retained.
+        """
+        with DirectButlerRepo.make_temporary() as (helper, root):
+            helper.add_task()
+            helper.add_task()
+            helper.insert_datasets("dataset_auto0")
+            kwargs = self._make_run_args(
+                "-b",
+                root,
+                "-i",
+                helper.input_chain,
+                "-o",
+                "output",
+                "--register-dataset-types",
+                pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph),
+            )
+            qg1 = script.qgraph(**kwargs)
+            run1 = qg1.header.output_run
+            kwargs["output_run"] = run1
+            script.run(qg1, **kwargs)
+            # Simulate: task_auto1 ran (metadata kept) but intermediate output
+            # not retained; task_auto2 failed (no metadata, no output).
+            helper.butler.pruneDatasets(
+                helper.butler.query_datasets("dataset_auto1", collections=run1),
+                purge=True,
+                unstore=True,
+                disassociate=True,
+            )
+            helper.butler.pruneDatasets(
+                helper.butler.query_datasets("task_auto2_metadata", collections=run1),
+                purge=True,
+                unstore=True,
+                disassociate=True,
+            )
+            helper.butler.pruneDatasets(
+                helper.butler.query_datasets("dataset_auto2", collections=run1),
+                purge=True,
+                unstore=True,
+                disassociate=True,
+            )
+            time.sleep(1)  # Make sure we don't get the same RUN timestamp.
+            # Only metadata types are retained; dataset_auto1 is not retained.
+            with make_tmp_file(b"- '*_metadata'\n", suffix=".yaml") as retained_path:
+                kwargs = self._make_run_args(
+                    "-b",
+                    root,
+                    "-i",
+                    helper.input_chain,
+                    "-o",
+                    "output",
+                    "--skip-existing-in",
+                    "output",
+                    "--retained-dataset-types",
+                    retained_path,
+                    pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph),
+                )
+                qg2 = script.qgraph(**kwargs)
+            # Both tasks must run: dataset_auto1 is not retained, so
+            # task_auto1 is forced to regenerate it for task_auto2.
+            self.assertEqual(len(qg2.quanta_by_task["task_auto1"]), 1)
+            self.assertEqual(len(qg2.quanta_by_task["task_auto2"]), 1)
+            self.assertEqual(len(qg2), 2)
+
+    def test_simple_qg_retained_both_skipped(self):
+        """When both tasks have metadata, both are skipped regardless of which
+        dataset types are not retained.
+        """
+        with DirectButlerRepo.make_temporary() as (helper, root):
+            helper.add_task()
+            helper.add_task()
+            helper.insert_datasets("dataset_auto0")
+            kwargs = self._make_run_args(
+                "-b",
+                root,
+                "-i",
+                helper.input_chain,
+                "-o",
+                "output",
+                "--register-dataset-types",
+                pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph),
+            )
+            qg1 = script.qgraph(**kwargs)
+            run1 = qg1.header.output_run
+            kwargs["output_run"] = run1
+            script.run(qg1, **kwargs)
+            # Prune only the intermediate; both task metadata are retained.
+            helper.butler.pruneDatasets(
+                helper.butler.query_datasets("dataset_auto1", collections=run1),
+                purge=True,
+                unstore=True,
+                disassociate=True,
+            )
+            time.sleep(1)  # Make sure we don't get the same RUN timestamp.
+            with make_tmp_file(b"- '*_metadata'\n", suffix=".yaml") as retained_path:
+                kwargs = self._make_run_args(
+                    "-b",
+                    root,
+                    "-i",
+                    helper.input_chain,
+                    "-o",
+                    "output",
+                    "--register-dataset-types",
+                    "--skip-existing-in",
+                    "output",
+                    "--retained-dataset-types",
+                    retained_path,
+                    pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph),
+                )
+                qg2 = script.qgraph(**kwargs)
+            # Both tasks have metadata so both are skipped; graph is empty.
+            self.assertIsNone(qg2)
+
+    def test_retained_dataset_types_invalid_yaml_raises(self):
+        """--retained-dataset-types raises ValueError for a non-list YAML
+        or an empty sequence.
+        """
+        with DirectButlerRepo.make_temporary() as (helper, root):
+            helper.add_task()
+            helper.insert_datasets("dataset_auto0")
+            base_kwargs = self._make_run_args(
+                "-b",
+                root,
+                "-i",
+                helper.input_chain,
+                "-o",
+                "output",
+                pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph),
+            )
+            for content in (b"key: value\n", b"[]\n"):
+                with make_tmp_file(content, suffix=".yaml") as retained_path:
+                    with self.assertRaises(ValueError):
+                        script.qgraph(**{**base_kwargs, "retained_dataset_types": retained_path})
+
 
 class CoverageTestCase(unittest.TestCase):
     """Test the coverage context manager."""