Merge pull request #367 from lsst/tickets/DM-55320

DM-55320: Add a feature to prune unanchored quanta in qgraph building

Author: hsinfang

doc/changes/DM-55320.feature.md

@@ -0,0 +1 @@
+Add --prune-unanchored-quanta option to pipetask qgraph

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(),

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=(

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

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."""