From 7292460c82201e73cd9c9fbcb9df8042d0606b11 Mon Sep 17 00:00:00 2001 From: Hsin-Fang Chiang Date: Wed, 6 May 2026 14:50:49 -0700 Subject: [PATCH] Add --retained-dataset-types option to pipetask qgraph and run Expose the retained_dataset_types parameter of QuantumGraphBuilder as a CLI option. Accept a path to a YAML file listing dataset type names or glob-style wildcard patterns for types that are expected to be present in --skip-existing-in when the producing task ran successfully. Has no effect without --skip-existing-in. script/run.py is not modified because commands.run passes all Click options as **kwargs to both script.qgraph (which consumes the new option at graph-build time) and script.run (which absorbs unrecognized kwargs, like the existing dataset_query_constraint). --- doc/changes/DM-54879.feature.md | 1 + .../lsst/ctrl/mpexec/cli/opt/optionGroups.py | 1 + python/lsst/ctrl/mpexec/cli/opt/options.py | 17 ++ python/lsst/ctrl/mpexec/cli/script/qgraph.py | 20 ++ tests/test_run.py | 181 ++++++++++++++++++ 5 files changed, 220 insertions(+) create mode 100644 doc/changes/DM-54879.feature.md diff --git a/doc/changes/DM-54879.feature.md b/doc/changes/DM-54879.feature.md new file mode 100644 index 00000000..ea2a123b --- /dev/null +++ b/doc/changes/DM-54879.feature.md @@ -0,0 +1 @@ +Added `--retained-dataset-types` option to `pipetask qgraph` and `pipetask run`. diff --git a/python/lsst/ctrl/mpexec/cli/opt/optionGroups.py b/python/lsst/ctrl/mpexec/cli/opt/optionGroups.py index d2adef28..1cd88447 100644 --- a/python/lsst/ctrl/mpexec/cli/opt/optionGroups.py +++ b/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(), diff --git a/python/lsst/ctrl/mpexec/cli/opt/options.py b/python/lsst/ctrl/mpexec/cli/opt/options.py index 3b353f20..b8406b72 100644 --- a/python/lsst/ctrl/mpexec/cli/opt/options.py +++ b/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=( diff --git a/python/lsst/ctrl/mpexec/cli/script/qgraph.py b/python/lsst/ctrl/mpexec/cli/script/qgraph.py index 98bd95b0..b5e07aec 100644 --- a/python/lsst/ctrl/mpexec/cli/script/qgraph.py +++ b/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" diff --git a/tests/test_run.py b/tests/test_run.py index 22dc9a25..e093b02e 100644 --- a/tests/test_run.py +++ b/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 . +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."""