Coverage for tests/test_execution_storage_class_conversion.py: 24%
111 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-14 01:20 -0800
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/>.
28from __future__ import annotations
30import os
31import shutil
32import tempfile
33import unittest
35import lsst.daf.butler
36import lsst.utils.tests
37from lsst.pipe.base import PipelineGraph
38from lsst.pipe.base.pipeline_graph import IncompatibleDatasetTypeError
39from lsst.pipe.base.simple_pipeline_executor import SimplePipelineExecutor
40from lsst.pipe.base.tests.mocks import (
41 DynamicConnectionConfig,
42 DynamicTestPipelineTask,
43 DynamicTestPipelineTaskConfig,
44 MockStorageClass,
45 get_mock_name,
46)
48TESTDIR = os.path.abspath(os.path.dirname(__file__))
51class TestExecutionStorageClassConversion(lsst.utils.tests.TestCase):
52 """Test storage class conversions during execution.
54 Task connection declarations should always define which storage class they
55 see, while data repository registrations should always define what is
56 stored.
58 This test uses mock storage classes for intermediate and output datasets,
59 which let us load the dataset to see what storage class the task saw when
60 it was running. These storage class names need to be wrapped in
61 get_mock_name calls to get what the butler actually sees. Overall input
62 datasets are not declared with mock datasets, so we can `put` them directly
63 in test code.
64 """
66 def setUp(self):
67 self.path = tempfile.mkdtemp()
68 # standalone parameter forces the returned config to also include
69 # the information from the search paths.
70 config = lsst.daf.butler.Butler.makeRepo(
71 self.path, standalone=True, searchPaths=[os.path.join(TESTDIR, "config")]
72 )
73 self.butler = SimplePipelineExecutor.prep_butler(config, [], "fake")
74 self.butler.registry.registerDatasetType(
75 lsst.daf.butler.DatasetType(
76 "input",
77 dimensions=self.butler.dimensions.empty,
78 storageClass="StructuredDataDict",
79 )
80 )
81 self.butler.put({"zero": 0}, "input")
82 MockStorageClass.get_or_register_mock("StructuredDataDict")
83 MockStorageClass.get_or_register_mock("TaskMetadataLike")
85 def tearDown(self):
86 shutil.rmtree(self.path, ignore_errors=True)
88 def _make_config(
89 self,
90 input_storage_class="StructuredDataDict",
91 output_storage_class="StructuredDataDict",
92 input_name="input",
93 output_name="output",
94 ):
95 """Create configuration for a test task with a single input and single
96 output of the given storage classes and dataset type names.
97 """
98 config = DynamicTestPipelineTaskConfig()
99 config.inputs["i"] = DynamicConnectionConfig(
100 dataset_type_name=input_name,
101 storage_class=input_storage_class,
102 # Since the overall input is special, we only use a mock storage
103 # class for it when there's a storage class conversion.
104 mock_storage_class=(input_name != "input" or (input_storage_class != "StructuredDataDict")),
105 )
106 config.outputs["o"] = DynamicConnectionConfig(
107 dataset_type_name=output_name, storage_class=output_storage_class
108 )
109 return config
111 def _make_executor(
112 self,
113 a_i_storage_class="StructuredDataDict",
114 a_o_storage_class="StructuredDataDict",
115 b_i_storage_class="StructuredDataDict",
116 b_o_storage_class="StructuredDataDict",
117 ):
118 """Configure a SimplePipelineExecutor with tasks with the given
119 storage classes as inputs and outputs.
121 This sets up a simple pipeline with two tasks ('a' and 'b') where the
122 second task's only input is the first task's only output.
123 """
124 config_a = self._make_config(a_i_storage_class, a_o_storage_class, output_name="intermediate")
125 config_b = self._make_config(b_i_storage_class, b_o_storage_class, input_name="intermediate")
126 pipeline_graph = PipelineGraph()
127 pipeline_graph.add_task("a", DynamicTestPipelineTask, config_a)
128 pipeline_graph.add_task("b", DynamicTestPipelineTask, config_b)
129 executor = SimplePipelineExecutor.from_pipeline_graph(pipeline_graph, butler=self.butler)
130 return executor
132 def _assert_datasets(
133 self,
134 a_i_storage_class="StructuredDataDict",
135 a_o_storage_class="StructuredDataDict",
136 b_i_storage_class="StructuredDataDict",
137 b_o_storage_class="StructuredDataDict",
138 stored_intermediate_storage_class="StructuredDataDict",
139 stored_output_storage_class="StructuredDataDict",
140 butler: lsst.daf.butler.Butler | None = None,
141 ):
142 """Check that a butler repository's contents are consistent with
143 running a pipeline created by _make_executor.
144 """
145 if butler is None:
146 butler = self.butler
147 # Read input and output datasets from butler, inspect their storage
148 # classes directly.
149 stored_intermediate = butler.get("intermediate")
150 stored_output = butler.get("output")
151 self.assertEqual(
152 butler.get_dataset_type("intermediate").storageClass_name,
153 get_mock_name(stored_intermediate_storage_class),
154 )
155 self.assertEqual(stored_output.storage_class, get_mock_name(stored_output_storage_class))
156 self.assertEqual(
157 butler.get_dataset_type("output").storageClass_name,
158 get_mock_name(stored_output_storage_class),
159 )
160 # Since we didn't tell the butler to convert storage classes on read,
161 # they'll remember their last conversion (on write).
162 if a_o_storage_class != stored_intermediate_storage_class:
163 self.assertEqual(
164 stored_intermediate.converted_from.storage_class,
165 get_mock_name(a_o_storage_class),
166 )
167 else:
168 self.assertIsNone(stored_intermediate.converted_from)
169 if b_o_storage_class != stored_output_storage_class:
170 self.assertEqual(
171 stored_output.converted_from.storage_class,
172 get_mock_name(b_o_storage_class),
173 )
174 else:
175 self.assertIsNone(stored_output.converted_from)
176 # Extract the inputs as seen by the tasks from those stored outputs.
177 quantum_a = stored_intermediate.quantum
178 quantum_b = stored_output.quantum
179 b_input = quantum_b.inputs["i"][0]
180 a_input = quantum_a.inputs["i"][0]
181 if a_i_storage_class == "StructuredDataDict":
182 self.assertIsNone(a_input.converted_from, None)
183 else:
184 self.assertEqual(a_input.original_type, "dict")
185 self.assertEqual(b_input.storage_class, get_mock_name(b_i_storage_class))
187 def test_no_conversions(self):
188 """Test execution with no storage class conversions as a baseline."""
189 executor = self._make_executor()
190 quanta = executor.run(register_dataset_types=True, save_versions=False)
191 self.assertEqual(len(quanta), 2)
192 self._assert_datasets()
194 def test_intermediate_registration_differs(self):
195 """Test execution where an intermediate is registered to be different
196 from both the producing and consuming task.
197 """
198 self.butler.registry.registerDatasetType(
199 lsst.daf.butler.DatasetType(
200 "intermediate",
201 dimensions=self.butler.dimensions.empty,
202 storageClass=get_mock_name("TaskMetadataLike"),
203 )
204 )
205 executor = self._make_executor()
206 quanta = executor.run(register_dataset_types=True, save_versions=False)
207 self.assertEqual(len(quanta), 2)
208 self._assert_datasets(stored_intermediate_storage_class="TaskMetadataLike")
210 def test_intermediate_producer_differs(self):
211 """Test execution where an intermediate is registered to be consistent
212 with the consumer but different from its producer.
213 """
214 self.butler.registry.registerDatasetType(
215 lsst.daf.butler.DatasetType(
216 "intermediate",
217 dimensions=self.butler.dimensions.empty,
218 storageClass=get_mock_name("TaskMetadataLike"),
219 )
220 )
221 executor = self._make_executor(b_i_storage_class="TaskMetadataLike")
222 quanta = executor.run(register_dataset_types=True, save_versions=False)
223 self.assertEqual(len(quanta), 2)
224 self._assert_datasets(
225 stored_intermediate_storage_class="TaskMetadataLike", b_i_storage_class="TaskMetadataLike"
226 )
228 def test_intermediate_consumer_differs(self):
229 """Test execution where an intermediate is registered to be consistent
230 with its producer but different from its consumer.
231 """
232 executor = self._make_executor(a_o_storage_class="TaskMetadataLike")
233 quanta = executor.run(register_dataset_types=True, save_versions=False)
234 self.assertEqual(len(quanta), 2)
235 self._assert_datasets(
236 stored_intermediate_storage_class="TaskMetadataLike", a_o_storage_class="TaskMetadataLike"
237 )
239 def test_output_differs(self):
240 """Test execution where an overall output is registered to be
241 different from the producing task.
242 """
243 self.butler.registry.registerDatasetType(
244 lsst.daf.butler.DatasetType(
245 "output",
246 dimensions=self.butler.dimensions.empty,
247 storageClass=get_mock_name("TaskMetadataLike"),
248 )
249 )
250 executor = self._make_executor()
251 quanta = executor.run(register_dataset_types=True, save_versions=False)
252 self.assertEqual(len(quanta), 2)
253 self._assert_datasets(stored_output_storage_class="TaskMetadataLike")
255 def test_input_differs(self):
256 """Test execution where an overall input's storage class is different
257 from the consuming task.
258 """
259 executor = self._make_executor(a_i_storage_class="TaskMetadataLike")
260 quanta = executor.run(register_dataset_types=True, save_versions=False)
261 self.assertEqual(len(quanta), 2)
262 self._assert_datasets(a_i_storage_class="TaskMetadataLike")
264 def test_input_differs_use_local_butler(self):
265 """Test execution where an overall input's storage class is different
266 from the consuming task, and we use a local butler.
267 """
268 executor = self._make_executor(a_i_storage_class="TaskMetadataLike")
269 with tempfile.TemporaryDirectory() as tempdir:
270 root = os.path.join(tempdir, "butler_root")
271 executor.use_local_butler(root)
272 quanta = executor.run(register_dataset_types=True, save_versions=False)
273 self.assertEqual(len(quanta), 2)
274 self._assert_datasets(a_i_storage_class="TaskMetadataLike", butler=executor.butler)
276 def test_incompatible(self):
277 """Test that we cannot make a QG if the registry and pipeline have
278 incompatible storage classes for a dataset type.
279 """
280 # Incompatible output dataset type.
281 self.butler.registry.registerDatasetType(
282 lsst.daf.butler.DatasetType(
283 "output",
284 dimensions=self.butler.dimensions.empty,
285 storageClass="StructuredDataList",
286 )
287 )
288 with self.assertRaisesRegex(
289 IncompatibleDatasetTypeError, "Incompatible definition.*StructuredDataDict.*StructuredDataList.*"
290 ):
291 self._make_executor()
293 def test_registry_changed(self):
294 """Run pipeline, but change registry dataset types between making the
295 QG and executing it.
297 This only fails with full-butler execution; we don't have a way to
298 prevent it with QBB.
299 """
300 executor = self._make_executor()
301 self.butler.registry.registerDatasetType(
302 lsst.daf.butler.DatasetType(
303 "output",
304 dimensions=self.butler.dimensions.empty,
305 storageClass="TaskMetadataLike", # even compatible is not okay
306 )
307 )
308 with self.assertRaisesRegex(
309 lsst.daf.butler.registry.ConflictingDefinitionError,
310 ".*_mock_StructuredDataDict.*is inconsistent with.*TaskMetadataLike.*",
311 ):
312 executor.run(register_dataset_types=True, save_versions=False)
315class MemoryTester(lsst.utils.tests.MemoryTestCase):
316 """Generic tests for file leaks."""
319def setup_module(module):
320 """Set up the module for pytest.
322 Parameters
323 ----------
324 module : `~types.ModuleType`
325 Module to set up.
326 """
327 lsst.utils.tests.init()
330if __name__ == "__main__":
331 lsst.utils.tests.init()
332 unittest.main()