Coverage for tests/test_graphBuilder.py: 100%
207 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-25 08:09 +0000
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <https://www.gnu.org/licenses/>.
28"""Tests of things related to the GraphBuilder class."""
30import io
31import logging
32import unittest
34import lsst.utils.tests
35from lsst.daf.butler import Butler, DataCoordinate, DatasetType
36from lsst.daf.butler.registry import UserExpressionError
37from lsst.pipe.base import PipelineGraph, QuantumGraph
38from lsst.pipe.base.all_dimensions_quantum_graph_builder import (
39 AllDimensionsQuantumGraphBuilder,
40 DatasetQueryConstraintVariant,
41)
42from lsst.pipe.base.tests import simpleQGraph
43from lsst.pipe.base.tests.mocks import (
44 DynamicConnectionConfig,
45 DynamicTestPipelineTask,
46 DynamicTestPipelineTaskConfig,
47 InMemoryRepo,
48 MockDataset,
49 MockStorageClass,
50)
51from lsst.utils.tests import temporaryDirectory
53_LOG = logging.getLogger(__name__)
56class GraphBuilderTestCase(unittest.TestCase):
57 """Test graph building."""
59 def _assertGraph(self, graph: QuantumGraph) -> None:
60 """Check basic structure of the graph."""
61 for taskDef in graph.iterTaskGraph():
62 refs = graph.initOutputRefs(taskDef)
63 # task has one initOutput, second ref is for config dataset
64 self.assertEqual(len(refs), 2)
66 self.assertEqual(len(list(graph.inputQuanta)), 1)
68 # This includes only "packages" dataset for now.
69 refs = graph.globalInitOutputRefs()
70 self.assertEqual(len(refs), 1)
72 def testDefault(self):
73 """Simple test to verify makeSimpleQGraph can be used to make a Quantum
74 Graph.
75 """
76 with temporaryDirectory() as root:
77 # makeSimpleQGraph calls GraphBuilder.
78 butler, qgraph = simpleQGraph.makeSimpleQGraph(root=root)
79 self.enterContext(butler)
80 # by default makeSimpleQGraph makes a graph with 5 nodes
81 self.assertEqual(len(qgraph), 5)
82 self._assertGraph(qgraph)
83 constraint = DatasetQueryConstraintVariant.OFF
84 _, qgraph2 = simpleQGraph.makeSimpleQGraph(
85 butler=butler, datasetQueryConstraint=constraint, callPopulateButler=False
86 )
87 # When all outputs are random resolved refs, direct comparison
88 # of graphs does not work because IDs are different. Can only
89 # verify the number of quanta in the graph without doing something
90 # terribly complicated.
91 self.assertEqual(len(qgraph2), 5)
92 constraint = DatasetQueryConstraintVariant.fromExpression("add_dataset0")
93 _, qgraph3 = simpleQGraph.makeSimpleQGraph(
94 butler=butler, datasetQueryConstraint=constraint, callPopulateButler=False
95 )
96 self.assertEqual(len(qgraph3), 5)
98 def test_empty_qg(self):
99 """Test that making an empty QG doesn't raise exceptions."""
100 config = DynamicTestPipelineTaskConfig()
101 config.inputs["i"] = DynamicConnectionConfig(
102 dataset_type_name="input",
103 storage_class="StructuredDataDict",
104 dimensions=["detector"],
105 )
106 config.init_inputs["ii"] = DynamicConnectionConfig(
107 dataset_type_name="init_input",
108 storage_class="StructuredDataDict",
109 )
110 config.outputs["o"] = DynamicConnectionConfig(
111 dataset_type_name="output",
112 storage_class="StructuredDataDict",
113 dimensions=["detector"],
114 )
115 config.init_outputs["io"] = DynamicConnectionConfig(
116 dataset_type_name="init_output",
117 storage_class="StructuredDataDict",
118 )
119 pipeline_graph = PipelineGraph()
120 pipeline_graph.add_task("a", DynamicTestPipelineTask, config)
121 with temporaryDirectory() as repo_path:
122 Butler.makeRepo(repo_path)
123 butler = Butler.from_config(repo_path, writeable=True, run="test_empty_qg")
124 self.enterContext(butler)
125 MockStorageClass.get_or_register_mock("StructuredDataDict")
126 butler.registry.registerDatasetType(
127 DatasetType(
128 "input",
129 dimensions=butler.dimensions.conform(["detector"]),
130 storageClass="_mock_StructuredDataDict",
131 )
132 )
133 init_input_dataset_type = DatasetType(
134 "init_input",
135 dimensions=butler.dimensions.empty,
136 storageClass="_mock_StructuredDataDict",
137 )
138 butler.registry.registerDatasetType(init_input_dataset_type)
139 # Init-input initially exists, but input does not (hence empty QG).
140 butler.put(
141 MockDataset(
142 dataset_id=None,
143 dataset_type=init_input_dataset_type.to_simple(),
144 data_id={},
145 run=butler.run,
146 ),
147 "init_input",
148 )
149 # Attempt to make QG; should just be empty, with no exceptions.
150 self.assertFalse(AllDimensionsQuantumGraphBuilder(pipeline_graph, butler).build())
151 # Initialize the output run, try again, with same expected result.
152 pipeline_graph.register_dataset_types(butler)
153 pipeline_graph.init_output_run(butler)
154 self.assertFalse(AllDimensionsQuantumGraphBuilder(pipeline_graph, butler).build())
156 # Inconsistent governor dimensions are no longer an error, so this test
157 # fails with the new query system. We should probably check instead that
158 # logging includes an explanation for the empty QG, but it might not
159 # because explain_no_results isn't good enough yet.
160 @unittest.expectedFailure
161 def testAddInstrumentMismatch(self):
162 """Verify that a RuntimeError is raised if the instrument in the user
163 query does not match the instrument in the pipeline.
164 """
165 with temporaryDirectory() as root:
166 pipeline = simpleQGraph.makeSimplePipeline(
167 nQuanta=5, instrument="lsst.pipe.base.tests.simpleQGraph.SimpleInstrument"
168 )
169 with self.assertRaises(UserExpressionError):
170 butler, _ = simpleQGraph.makeSimpleQGraph(
171 root=root, pipeline=pipeline, userQuery="instrument = 'foo'"
172 )
173 self.enterContext(butler)
175 def testUserQueryBind(self):
176 """Verify that bind values work for user query."""
177 pipeline = simpleQGraph.makeSimplePipeline(
178 nQuanta=5, instrument="lsst.pipe.base.tests.simpleQGraph.SimpleInstrument"
179 )
180 instr = simpleQGraph.SimpleInstrument.getName()
181 # With a literal in the user query
182 with temporaryDirectory() as root:
183 butler, _ = simpleQGraph.makeSimpleQGraph(
184 root=root, pipeline=pipeline, userQuery=f"instrument = '{instr}'"
185 )
186 self.enterContext(butler)
187 # With a bind for the user query
188 with temporaryDirectory() as root:
189 butler, _ = simpleQGraph.makeSimpleQGraph(
190 root=root, pipeline=pipeline, userQuery="instrument = :instr", bind={"instr": instr}
191 )
192 self.enterContext(butler)
194 def test_datastore_records(self):
195 """Test for generating datastore records."""
196 with temporaryDirectory() as root:
197 # need FileDatastore for this tests
198 butler, qgraph1 = simpleQGraph.makeSimpleQGraph(
199 root=root, inMemory=False, makeDatastoreRecords=True
200 )
201 self.enterContext(butler)
203 # save and reload
204 buffer = io.BytesIO()
205 qgraph1.save(buffer)
206 buffer.seek(0)
207 qgraph2 = QuantumGraph.load(buffer, universe=butler.dimensions)
208 del buffer
210 for qgraph in (qgraph1, qgraph2):
211 self.assertEqual(len(qgraph), 5)
212 for i, qnode in enumerate(qgraph):
213 quantum = qnode.quantum
214 self.assertIsNotNone(quantum.datastore_records)
215 # only the first quantum has a pre-existing input
216 if i == 0:
217 datastore_name = "FileDatastore@<butlerRoot>"
218 self.assertEqual(set(quantum.datastore_records.keys()), {datastore_name})
219 records_data = quantum.datastore_records[datastore_name]
220 records = dict(records_data.records)
221 self.assertEqual(len(records), 1)
222 _, records = records.popitem()
223 records = records["file_datastore_records"]
224 self.assertEqual(
225 [record.path for record in records],
226 ["test/add_dataset0/add_dataset0_INSTR_det0_test.pickle"],
227 )
228 else:
229 self.assertEqual(quantum.datastore_records, {})
232class SkipExistingInTestCase(unittest.TestCase):
233 """Tests for the skip_existing_in behavior of QuantumGraphBuilder."""
235 def setUp(self):
236 self.helper = InMemoryRepo()
237 self.enterContext(self.helper)
238 self.helper.add_task()
239 self.helper.make_quantum_graph_builder(output_run="new_run")
240 self.helper.butler.collections.register("prior_run")
241 self._task_node = self.helper.pipeline_graph.tasks["task_auto1"]
242 self._empty_data_id = DataCoordinate.make_empty(self.helper.butler.dimensions)
244 def _insert(self, *names, run="prior_run"):
245 """Register datasets with empty data IDs into a run collection."""
246 for name in names:
247 dt = self.helper.pipeline_graph.dataset_types[name].dataset_type
248 self.helper.butler.registry.insertDatasets(dt, [self._empty_data_id], run=run)
250 def _build(self, *, output_run="new_run", **kwargs):
251 return AllDimensionsQuantumGraphBuilder(
252 self.helper.pipeline_graph,
253 self.helper.butler,
254 input_collections=[self.helper.input_chain],
255 output_run=output_run,
256 **kwargs,
257 ).build(attach_datastore_records=False)
259 def test_not_skipped_without_skip_existing_in(self):
260 """Without skip_existing_in, a quantum is never skipped even if
261 metadata exists in an input collection.
262 """
263 self._insert(self._task_node.metadata_output.parent_dataset_type_name)
264 qgraph = self._build()
265 self.assertEqual(len(qgraph), 1)
267 def test_skipped_when_metadata_exists(self):
268 """With skip_existing_in, a quantum is skipped when its metadata
269 dataset is present in the specified collections.
270 """
271 self._insert(self._task_node.metadata_output.parent_dataset_type_name)
272 # Init-outputs required, otherwise InitInputMissingError.
273 for edge in self._task_node.init.iter_all_outputs():
274 self._insert(edge.parent_dataset_type_name)
275 qgraph = self._build(skip_existing_in=["prior_run"])
276 self.assertEqual(len(qgraph), 0)
278 def test_not_skipped_when_metadata_absent(self):
279 """With skip_existing_in, a quantum is not skipped when its metadata
280 dataset is absent from the specified collections.
281 """
282 qgraph = self._build(skip_existing_in=["prior_run"])
283 self.assertEqual(len(qgraph), 1)
286class RetainedDatasetTypesTestCase(unittest.TestCase):
287 """Tests for QuantumGraphBuilder.retained_dataset_types.
289 dataset_auto0 -> task_auto1 -> dataset_auto1 -> task_auto2
290 """
292 def setUp(self):
293 self.helper = InMemoryRepo()
294 self.enterContext(self.helper)
295 self.helper.add_task()
296 self.helper.add_task()
297 self.helper.make_quantum_graph_builder(output_run="new_run")
298 self.helper.butler.collections.register("prior_run")
299 self._task1 = self.helper.pipeline_graph.tasks["task_auto1"]
300 self._task2 = self.helper.pipeline_graph.tasks["task_auto2"]
301 self._empty_data_id = DataCoordinate.make_empty(self.helper.butler.dimensions)
303 def _insert(self, *names, run="prior_run"):
304 """Register datasets with empty data IDs into a run collection."""
305 for name in names:
306 dt = self.helper.pipeline_graph.dataset_types[name].dataset_type
307 self.helper.butler.registry.insertDatasets(dt, [self._empty_data_id], run=run)
309 def _build(self, *, output_run="new_run", **kwargs):
310 return AllDimensionsQuantumGraphBuilder(
311 self.helper.pipeline_graph,
312 self.helper.butler,
313 input_collections=[self.helper.input_chain],
314 output_run=output_run,
315 **kwargs,
316 ).build(attach_datastore_records=False)
318 def test_raises_without_skip_existing_in(self):
319 """retained_dataset_types invalid without skip_existing_in."""
320 with self.assertRaises(ValueError):
321 self._build(retained_dataset_types=["dataset_auto1"])
323 def test_ancestor_unskipped_when_output_not_retained(self):
324 """task1 ran (metadata present) but did not retain its output;
325 task2 must run. Because dataset_auto1 is not retained, task1
326 is unskipped to regenerate it.
327 """
328 # task1 succeeded previously, but dataset_auto1 not retained.
329 self._insert(self._task1.metadata_output.parent_dataset_type_name)
330 qgraph = self._build(
331 skip_existing_in=["prior_run"],
332 retained_dataset_types=["*_metadata"],
333 )
334 # Both tasks run: task1 regenerate dataset_auto1 for task2.
335 self.assertEqual(len(qgraph), 2)
337 def test_ancestor_not_unskipped_when_output_retained(self):
338 """When the intermediate output is declared retained and is present in
339 skip_existing_in, unskipping stops there and task1 remains skipped.
340 """
341 # task1 metadata and its output dataset_auto1 both present in
342 # prior_run.
343 self._insert(self._task1.metadata_output.parent_dataset_type_name)
344 self._insert("dataset_auto1")
345 for edge in self._task1.init.iter_all_outputs():
346 self._insert(edge.parent_dataset_type_name)
347 qgraph = self._build(
348 skip_existing_in=["prior_run"],
349 retained_dataset_types=["dataset_auto1", "*_metadata"],
350 )
351 # Only task2 runs.
352 self.assertEqual(len(qgraph), 1)
354 def test_both_skipped_when_both_have_metadata(self):
355 """When both tasks have metadata, both remain skipped regardless of
356 which outputs are not retained.
357 """
358 self._insert(self._task1.metadata_output.parent_dataset_type_name)
359 self._insert(self._task2.metadata_output.parent_dataset_type_name)
360 for edge in self._task1.init.iter_all_outputs():
361 self._insert(edge.parent_dataset_type_name)
362 for edge in self._task2.init.iter_all_outputs():
363 self._insert(edge.parent_dataset_type_name)
364 qgraph = self._build(
365 skip_existing_in=["prior_run"],
366 retained_dataset_types=["*_metadata"],
367 )
368 self.assertEqual(len(qgraph), 0)
370 def test_unrecognised_pattern_warns(self):
371 """Literal names and wildcard patterns that match nothing in the
372 pipeline emit a WARNING log message.
373 """
374 with self.assertLogs("lsst.pipe.base.quantum_graph_builder", level="WARNING") as cm:
375 self._build(
376 skip_existing_in=["prior_run"],
377 retained_dataset_types=["no_such_dataset_type", "no_such_*"],
378 )
379 self.assertTrue(any("no_such_dataset_type" in msg for msg in cm.output))
380 self.assertTrue(any("no_such_*" in msg for msg in cm.output))
382 def test_no_unskipping_when_all_retained(self):
383 """'*' matches all dataset types; no ancestor unskipping occurs,
384 equivalent to not providing retained_dataset_types.
385 """
386 # task1 ran; metadata and dataset_auto1 present.
387 self._insert(self._task1.metadata_output.parent_dataset_type_name)
388 self._insert("dataset_auto1")
389 for edge in self._task1.init.iter_all_outputs():
390 self._insert(edge.parent_dataset_type_name)
391 qgraph = self._build(
392 skip_existing_in=["prior_run"],
393 retained_dataset_types=["*"],
394 )
395 # All types retained -> no unskipping -> task1 stays skipped,
396 # only task2 runs.
397 self.assertEqual(len(qgraph), 1)
400class RetainedDatasetTypesThreeTaskTestCase(unittest.TestCase):
401 """Tests for retained_dataset_types with a 3-task chain.
403 Pipeline: dataset_auto0 -> task_auto1 -> dataset_auto1
404 -> task_auto2 -> dataset_auto2
405 -> task_auto3 -> dataset_auto3
406 """
408 def setUp(self):
409 self.helper = InMemoryRepo()
410 self.enterContext(self.helper)
411 self.helper.add_task()
412 self.helper.add_task()
413 self.helper.add_task()
414 self.helper.make_quantum_graph_builder(output_run="new_run")
415 self.helper.butler.collections.register("prior_run")
416 self._task1 = self.helper.pipeline_graph.tasks["task_auto1"]
417 self._task2 = self.helper.pipeline_graph.tasks["task_auto2"]
418 self._task3 = self.helper.pipeline_graph.tasks["task_auto3"]
419 self._empty_data_id = DataCoordinate.make_empty(self.helper.butler.dimensions)
421 def _insert(self, *names, run="prior_run"):
422 """Register datasets with empty data IDs into a run collection."""
423 for name in names:
424 dt = self.helper.pipeline_graph.dataset_types[name].dataset_type
425 self.helper.butler.registry.insertDatasets(dt, [self._empty_data_id], run=run)
427 def _build(self, *, output_run="new_run", **kwargs):
428 return AllDimensionsQuantumGraphBuilder(
429 self.helper.pipeline_graph,
430 self.helper.butler,
431 input_collections=[self.helper.input_chain],
432 output_run=output_run,
433 **kwargs,
434 ).build(attach_datastore_records=False)
436 def test_unskipping_stops_at_retained_intermediate(self):
437 """task2's output is retained and present in skip_existing_in.
438 Only task3 runs, task1 and task2 remain skipped.
439 """
440 self._insert(self._task1.metadata_output.parent_dataset_type_name)
441 self._insert(self._task2.metadata_output.parent_dataset_type_name)
442 self._insert("dataset_auto2")
443 for edge in self._task1.init.iter_all_outputs():
444 self._insert(edge.parent_dataset_type_name)
445 for edge in self._task2.init.iter_all_outputs():
446 self._insert(edge.parent_dataset_type_name)
447 qgraph = self._build(
448 skip_existing_in=["prior_run"],
449 retained_dataset_types=["dataset_auto2", "*_metadata"],
450 )
451 self.assertEqual(len(qgraph), 1)
453 def test_full_chain_unskipped_when_none_retained(self):
454 """task3 needs to run. Unskipping walks back through
455 dataset_auto2 (not retained) to unskip task2, then through
456 dataset_auto1 (not retained) to unskip task1. All three tasks
457 run to regenerate the non-retained datasets.
458 """
459 self._insert(self._task1.metadata_output.parent_dataset_type_name)
460 self._insert(self._task2.metadata_output.parent_dataset_type_name)
461 qgraph = self._build(
462 skip_existing_in=["prior_run"],
463 retained_dataset_types=["*_metadata"],
464 )
465 self.assertEqual(len(qgraph), 3)
468if __name__ == "__main__":
469 lsst.utils.tests.init()
470 unittest.main()