diff --git a/doc/changes/DM-55320.feature.md b/doc/changes/DM-55320.feature.md new file mode 100644 index 00000000..b33c6204 --- /dev/null +++ b/doc/changes/DM-55320.feature.md @@ -0,0 +1 @@ +Add --prune-unanchored-quanta option to pipetask qgraph diff --git a/python/lsst/ctrl/mpexec/cli/opt/optionGroups.py b/python/lsst/ctrl/mpexec/cli/opt/optionGroups.py index 1cd88447..708c5e52 100644 --- a/python/lsst/ctrl/mpexec/cli/opt/optionGroups.py +++ b/python/lsst/ctrl/mpexec/cli/opt/optionGroups.py @@ -128,6 +128,7 @@ def __init__( ctrlMpExecOpts.skip_existing_in_option(), ctrlMpExecOpts.skip_existing_option(), ctrlMpExecOpts.retained_dataset_types_option(), + ctrlMpExecOpts.prune_unanchored_quanta_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 b8406b72..4f5d155a 100644 --- a/python/lsst/ctrl/mpexec/cli/opt/options.py +++ b/python/lsst/ctrl/mpexec/cli/opt/options.py @@ -375,6 +375,56 @@ ) +def parse_prune_unanchored_quanta( + ctx: click.Context, param: click.Option, value: str | None +) -> tuple[str, str] | None: + """Parse the --prune-unanchored-quanta option value into a tuple. + + Parameters + ---------- + ctx : `click.Context` + Context provided by Click. + param : `click.Option` + Click option. + value : `str` or `None` + Value from option, expected to be ``SOURCE:ANCHOR`` or `None` if the + option was not provided. + + Returns + ------- + result : `tuple` [`str`, `str`] or `None` + A ``(source_label, anchor_label)`` tuple, or `None` if ``value`` is + `None`. + + Raises + ------ + click.UsageError + Raised if ``value`` is not `None` and does not match the + ``SOURCE:ANCHOR`` format. + """ + if value is None: + return None + parts = value.split(":", 1) + if len(parts) != 2 or not parts[0] or not parts[1]: + raise click.UsageError( + f"Invalid value for --prune-unanchored-quanta: {value!r}; expected SOURCE:ANCHOR." + ) + return (parts[0], parts[1]) + + +prune_unanchored_quanta_option = MWOptionDecorator( + "--prune-unanchored-quanta", + callback=parse_prune_unanchored_quanta, + default=None, + metavar="SOURCE:ANCHOR", + help=unwrap( + """Remove source quanta that have no reachable anchor quantum downstream, + along with their entire downstream chain. Specify as a colon-separated + pair of task labels SOURCE:ANCHOR.""" + ), +) + + 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 b5e07aec..ce40cb9d 100644 --- a/python/lsst/ctrl/mpexec/cli/script/qgraph.py +++ b/python/lsst/ctrl/mpexec/cli/script/qgraph.py @@ -70,6 +70,7 @@ def qgraph( skip_existing_in: Iterable[str] | None, skip_existing: bool, retained_dataset_types: str | None, + prune_unanchored_quanta: tuple[str, str] | None = None, save_qgraph: ResourcePathExpression | None, qgraph_dot: str | None, qgraph_mermaid: str | None, @@ -128,6 +129,10 @@ def qgraph( 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``. + prune_unanchored_quanta : `tuple` [ `str`, `str` ] or `None`, optional + A ``(source_label, anchor_label)`` pair of task labels. If not `None`, + source quanta with no reachable anchor quantum downstream are removed, + along with their entire downstream chain. save_qgraph : convertible to `lsst.resources.ResourcePath` or `None` URI location for saving the quantum graph. qgraph_dot : `str` or `None` @@ -322,6 +327,7 @@ def qgraph( where=data_query, skip_existing_in=skip_existing_in, retained_dataset_types=retained_dataset_type_patterns, + prune_unanchored_quanta=prune_unanchored_quanta, clobber=clobber_outputs, dataset_query_constraint=DatasetQueryConstraintVariant.fromExpression( dataset_query_constraint diff --git a/tests/test_run.py b/tests/test_run.py index e093b02e..0727567a 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -833,6 +833,53 @@ def test_retained_dataset_types_invalid_yaml_raises(self): with self.assertRaises(ValueError): script.qgraph(**{**base_kwargs, "retained_dataset_types": retained_path}) + def test_simple_qg_prune_unanchored_anchor_absent(self): + """With --prune-unanchored-quanta SOURCE:ANCHOR where ANCHOR does not + exist in the pipeline, all source quanta are pruned and the graph is + empty. + """ + with DirectButlerRepo.make_temporary() as (helper, root): + helper.add_task("source") + helper.add_task("anchor") + helper.insert_datasets("dataset_auto0") + kwargs = self._make_run_args( + "-b", + root, + "-i", + helper.input_chain, + "-o", + "output", + "--register-dataset-types", + "--prune-unanchored-quanta", + "source:no_such_task", + pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), + ) + qg = script.qgraph(**kwargs) + self.assertIsNone(qg) + + def test_simple_qg_prune_unanchored_anchor_reachable(self): + """With --prune-unanchored-quanta SOURCE:ANCHOR where every source + quantum has an anchor quantum downstream, nothing is pruned. + """ + with DirectButlerRepo.make_temporary() as (helper, root): + helper.add_task("source") + helper.add_task("anchor") + helper.insert_datasets("dataset_auto0") + kwargs = self._make_run_args( + "-b", + root, + "-i", + helper.input_chain, + "-o", + "output", + "--register-dataset-types", + "--prune-unanchored-quanta", + "source:anchor", + pipeline_graph_factory=PipelineGraphFactory(pipeline_graph=helper.pipeline_graph), + ) + qg = script.qgraph(**kwargs) + self.assertEqual(len(qg), 2) + class CoverageTestCase(unittest.TestCase): """Test the coverage context manager."""