From a2b1723bd2ad5d130c988f794f7a6f90ebf6e53b Mon Sep 17 00:00:00 2001 From: Hsin-Fang Chiang Date: Thu, 25 Jun 2026 13:23:57 -0700 Subject: [PATCH] Add prune_unanchored_quanta to QuantumGraphBuilder This absorbs the functionality of ap_pipe's prune_orphan_preloads.py post-processing script into qgraph generation. --- doc/changes/DM-55320.feature.md | 1 + .../lsst/pipe/base/quantum_graph_builder.py | 35 ++++++++ .../lsst/pipe/base/quantum_graph_skeleton.py | 55 +++++++++++++ tests/test_graphBuilder.py | 46 +++++++++++ tests/test_graphSkeleton.py | 79 +++++++++++++++++++ 5 files changed, 216 insertions(+) create mode 100644 doc/changes/DM-55320.feature.md create mode 100644 tests/test_graphSkeleton.py diff --git a/doc/changes/DM-55320.feature.md b/doc/changes/DM-55320.feature.md new file mode 100644 index 000000000..65605d54a --- /dev/null +++ b/doc/changes/DM-55320.feature.md @@ -0,0 +1 @@ +Add prune_unanchored_quanta parameter to QuantumGraphBuilder diff --git a/python/lsst/pipe/base/quantum_graph_builder.py b/python/lsst/pipe/base/quantum_graph_builder.py index 5e7c4c8dc..88b792f25 100644 --- a/python/lsst/pipe/base/quantum_graph_builder.py +++ b/python/lsst/pipe/base/quantum_graph_builder.py @@ -138,6 +138,12 @@ class QuantumGraphBuilder(ABC): the upstream quanta that need to regenerate those intermediates to also run. Has no effect without ``skip_existing_in``. ``["*"]`` means retaining all datasets, equivalent to not providing this option. + prune_unanchored_quanta : `tuple` [ `str`, `str` ], optional + A ``(source_label, anchor_label)`` pair of task labels triggering + unanchored quanta pruning after the skeleton is assembled. A + ``source_label`` quantum is removed along with its entire downstream + chain if no ``anchor_label`` quantum is reachable from it along + directed graph edges. clobber : `bool`, optional Whether to raise if predicted outputs already exist in ``output_run`` (not including those quanta that would be skipped because they've @@ -182,6 +188,7 @@ def __init__( output_run: str | None = None, skip_existing_in: Sequence[str] = (), retained_dataset_types: Sequence[str] | None = None, + prune_unanchored_quanta: tuple[str, str] | None = None, clobber: bool = False, ): self.log = getLogger(__name__) @@ -204,6 +211,7 @@ def __init__( ) if self._retained_dataset_type_patterns is not None and not skip_existing_in: raise ValueError("retained_dataset_types has no effect without skip_existing_in.") + self._prune_unanchored_quanta = prune_unanchored_quanta self.empty_data_id = DataCoordinate.make_empty(butler.dimensions) self.clobber = clobber # See whether the output run already exists. @@ -249,6 +257,21 @@ def __init__( task_node.label: PrerequisiteInfo(task_node, self._pipeline_graph) for task_node in pipeline_graph.tasks.values() } + if self._prune_unanchored_quanta is not None: + source_label, anchor_label = self._prune_unanchored_quanta + if source_label not in self._pipeline_graph.tasks: + self.log.warning( + "prune_unanchored_quanta source label %s is not present in the pipeline; " + "pruning will have no effect.", + source_label, + ) + elif anchor_label not in self._pipeline_graph.tasks: + self.log.warning( + "prune_unanchored_quanta anchor label %s is not present in the pipeline; " + "all %s quanta will be treated as unanchored and removed.", + anchor_label, + source_label, + ) log: LsstLogAdapter """Logger to use for all quantum-graph generation messages. @@ -470,6 +493,18 @@ def _build_skeleton(self, attach_datastore_records: bool = True) -> QuantumGraph # with the quanta because no quantum knows if its the only # consumer). full_skeleton.remove_orphan_datasets() + if self._prune_unanchored_quanta is not None: + source_label, anchor_label = self._prune_unanchored_quanta + removed = full_skeleton.remove_unanchored_quanta(source_label, anchor_label) + if removed: + for task_label, n_removed in removed.items(): + self.log.info( + "Pruned %d unanchored or downstream %s quanta (anchor: %s).", + n_removed, + task_label, + anchor_label, + ) + full_skeleton.remove_orphan_datasets() if attach_datastore_records: self._attach_datastore_records(full_skeleton) return full_skeleton diff --git a/python/lsst/pipe/base/quantum_graph_skeleton.py b/python/lsst/pipe/base/quantum_graph_skeleton.py index 386b3f7c3..4cf5e0384 100644 --- a/python/lsst/pipe/base/quantum_graph_skeleton.py +++ b/python/lsst/pipe/base/quantum_graph_skeleton.py @@ -562,6 +562,61 @@ def remove_orphan_datasets(self) -> None: if not orphan.is_task and orphan not in self._global_init_outputs: self._xgraph.remove_node(orphan) + def remove_unanchored_quanta(self, source_label: str, anchor_label: str) -> dict[str, int]: + """Remove unanchored source quanta and their entire downstream chain. + + A source quantum is considered unanchored if no quantum with + ``anchor_label`` is reachable along directed edges from it. + Unanchored source quanta and every descendant reachable from them are + removed unconditionally, regardless of what else may reach it. + + Parameters + ---------- + source_label : `str` + Task label of the source task whose unanchored quanta to remove. + anchor_label : `str` + Task label that must appear downstream of a source quantum for that + quantum to be considered anchored. + + Returns + ------- + removed : `dict` [`str`, `int`] + Mapping of task label to number of quanta removed for that task. + Empty if nothing was removed. + """ + if not self.has_task(source_label): + return {} + source_quanta = set(self.get_quanta(source_label)) + anchor_quanta = set(self.get_quanta(anchor_label)) if self.has_task(anchor_label) else set() + reachable = set() + for quantum in anchor_quanta: + reachable.update(networkx.ancestors(self._xgraph, quantum)) + + unanchored = source_quanta - reachable + if not unanchored: + return {} + + # to_remove collects unanchored source quanta and all their + # descendants. It has both QuantumKey and DatasetKey nodes. + to_remove = set(unanchored) + for quantum in unanchored: + to_remove.update(networkx.descendants(self._xgraph, quantum)) + removed: dict[str, int] = {} + for node in to_remove: + if isinstance(node, QuantumKey): + removed[node.task_label] = removed.get(node.task_label, 0) + 1 + _, quanta = self._tasks[node.task_label] + quanta.remove(node) + self._xgraph.remove_nodes_from(to_remove) + # For any task with no quanta remaining, remove its TaskInitKey and + # any init-output dataset nodes attached to it, then drop the task. + for label in removed: + task_init_key, remaining = self._tasks[label] + if not remaining: + self._xgraph.remove_nodes_from(list(self._xgraph.successors(task_init_key))) + self.remove_task(label) + return removed + def extract_overall_inputs(self) -> dict[DatasetKey | PrerequisiteDatasetKey, DatasetRef]: """Find overall input datasets. diff --git a/tests/test_graphBuilder.py b/tests/test_graphBuilder.py index 6bc764af4..69196295c 100644 --- a/tests/test_graphBuilder.py +++ b/tests/test_graphBuilder.py @@ -465,6 +465,52 @@ def test_full_chain_unskipped_when_none_retained(self): self.assertEqual(len(qgraph), 3) +class PruneUnanchoredQuantaTestCase(unittest.TestCase): + """Tests for the prune_unanchored_quanta behavior of QuantumGraphBuilder. + + Graph: auto0 -> source -> auto1 -> anchor -> auto2 + + All tasks are dimensionless so each is one quantum. + """ + + def setUp(self): + self.helper = InMemoryRepo() + self.enterContext(self.helper) + self.helper.add_task("source") + self.helper.add_task("anchor") + self.helper.make_quantum_graph_builder(output_run="output_run") + + def _build(self, **kwargs): + return AllDimensionsQuantumGraphBuilder( + self.helper.pipeline_graph, + self.helper.butler, + input_collections=[self.helper.input_chain], + output_run="output_run", + **kwargs, + ).build(attach_datastore_records=False) + + def test_no_effect_without_parameter(self): + """Without prune_unanchored_quanta, all quanta are kept.""" + qg = self._build() + self.assertEqual(len(qg), 2) + + def test_no_pruning_when_anchor_reachable(self): + """Anchor reachable from source quantum: nothing is pruned.""" + qg = self._build(prune_unanchored_quanta=("source", "anchor")) + self.assertEqual(len(qg), 2) + + def test_all_pruned_when_anchor_label_absent(self): + """Anchor is absent: all source quanta and task removed.""" + qg = self._build(prune_unanchored_quanta=("source", "no_such_task")) + self.assertEqual(len(qg), 0) + self.assertNotIn("source", {td.label for td in qg.iterTaskGraph()}) + + def test_noop_when_source_label_absent(self): + """source_label not in pipeline: nothing happens.""" + qg = self._build(prune_unanchored_quanta=("no_such_task", "anchor")) + self.assertEqual(len(qg), 2) + + if __name__ == "__main__": lsst.utils.tests.init() unittest.main() diff --git a/tests/test_graphSkeleton.py b/tests/test_graphSkeleton.py new file mode 100644 index 000000000..5c49fa9bb --- /dev/null +++ b/tests/test_graphSkeleton.py @@ -0,0 +1,79 @@ +# This file is part of pipe_base. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This software is dual licensed under the GNU General Public License and also +# under a 3-clause BSD license. Recipients may choose which of these licenses +# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, +# respectively. If you choose the GPL option then the following text applies +# (but note that there is still no warranty even if you opt for BSD instead): +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Unit tests for QuantumGraphSkeleton.""" + +import unittest + +from lsst.daf.butler import DataCoordinate, DimensionUniverse +from lsst.pipe.base.quantum_graph_skeleton import DatasetKey, QuantumGraphSkeleton + + +class RemoveUnanchoredQuantaTestCase(unittest.TestCase): + """Tests for ``QuantumGraphSkeleton.remove_unanchored_quanta``. + + Graph: + source1 -> d1 -> anchor1 (band 1: anchored) + source2 -> d2 -> anchor2 (band 2: anchored) + source3 -> d3 (band 3: unanchored) + source4 -> d4 (band 4: unanchored) + """ + + def setUp(self): + universe = DimensionUniverse() + self.skeleton = QuantumGraphSkeleton(["source", "anchor"]) + for i in range(1, 5): + data_id = DataCoordinate.standardize({"band": i}, universe=universe) + source_quantum = self.skeleton.add_quantum_node("source", data_id) + dataset = self.skeleton.add_dataset_node(f"d{i}", data_id) + self.skeleton.add_output_edge(source_quantum, dataset) + if i <= 2: + anchor_quantum = self.skeleton.add_quantum_node("anchor", data_id) + self.skeleton.add_input_edge(anchor_quantum, dataset) + + def test_remove_unanchored(self): + """Unanchored source quanta and their descendants are removed.""" + removed = self.skeleton.remove_unanchored_quanta("source", "anchor") + self.assertEqual(removed, {"source": 2}) + self.assertEqual(len(self.skeleton.get_quanta("source")), 2) + self.assertEqual(len(self.skeleton.get_quanta("anchor")), 2) + self.assertNotIn(DatasetKey("d3", (3,)), self.skeleton) + self.assertNotIn(DatasetKey("d4", (4,)), self.skeleton) + + # Second call on an already-pruned skeleton is a no-op. + removed = self.skeleton.remove_unanchored_quanta("source", "anchor") + self.assertEqual(removed, {}) + + def test_task_dropped_when_all_unanchored(self): + """Task is dropped when all its quanta are removed.""" + removed = self.skeleton.remove_unanchored_quanta("source", "nonexistent") + self.assertEqual(removed["source"], 4) + self.assertFalse(self.skeleton.has_task("source")) + + +if __name__ == "__main__": + unittest.main()