Coverage for tests/test_quantumBackedButler.py: 9%
276 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-04 01:36 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://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 <http://www.gnu.org/licenses/>.
28import json
29import os
30import tempfile
31import unittest
32import unittest.mock
33from typing import cast
35from lsst.daf.butler import (
36 Butler,
37 ButlerMetrics,
38 Config,
39 DatasetRef,
40 DatasetType,
41 DimensionUniverse,
42 Quantum,
43 QuantumBackedButler,
44 QuantumProvenanceData,
45 RegistryConfig,
46 StorageClass,
47)
48from lsst.daf.butler.datastore.record_data import DatastoreRecordData
49from lsst.daf.butler.datastores.file_datastore.retrieve_artifacts import ZipIndex
50from lsst.daf.butler.direct_butler import DirectButler
51from lsst.daf.butler.registry import _RegistryFactory
52from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir
53from lsst.resources import ResourcePath
55TESTDIR = os.path.abspath(os.path.dirname(__file__))
58class QuantumBackedButlerTestCase(unittest.TestCase):
59 """Test case for QuantumBackedButler."""
61 def setUp(self) -> None:
62 self.root = makeTestTempDir(TESTDIR)
63 self.config = Config()
64 self.config["root"] = self.root
65 self.universe = DimensionUniverse()
66 self.metrics = ButlerMetrics()
68 # Make a butler and import dimension definitions.
69 registryConfig = RegistryConfig(self.config.get("registry"))
70 _RegistryFactory(registryConfig).create_from_config(butlerRoot=self.root)
71 butler = Butler.from_config(self.config, writeable=True, run="RUN", metrics=self.metrics)
72 assert isinstance(butler, DirectButler)
73 self.butler = butler
74 self.butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml")
76 # make all dataset types
77 graph = self.universe.conform(("instrument", "detector"))
78 storageClass = StorageClass("StructuredDataDict")
79 self.datasetTypeInit = DatasetType("test_ds_init", graph, storageClass)
80 self.datasetTypeInput = DatasetType("test_ds_input", graph, storageClass)
81 self.datasetTypeOutput = DatasetType("test_ds_output", graph, storageClass)
82 self.datasetTypeOutput2 = DatasetType("test_ds_output2", graph, storageClass)
83 self.datasetTypeExtra = DatasetType("test_ds_extra", graph, storageClass)
85 self.dataset_types: dict[str, DatasetType] = {}
86 dataset_types = (
87 self.datasetTypeInit,
88 self.datasetTypeInput,
89 self.datasetTypeOutput,
90 self.datasetTypeOutput2,
91 self.datasetTypeExtra,
92 )
93 for dataset_type in dataset_types:
94 self.butler.registry.registerDatasetType(dataset_type)
95 self.dataset_types[dataset_type.name] = dataset_type
97 dataIds = [
98 self.butler.registry.expandDataId(dict(instrument="Cam1", detector=detector_id))
99 for detector_id in (1, 2, 3, 4)
100 ]
102 # make actual input datasets
103 self.input_refs = [
104 self.butler.put({"data": dataId["detector"]}, self.datasetTypeInput, dataId) for dataId in dataIds
105 ]
106 self.init_inputs_refs = [self.butler.put({"data": -1}, self.datasetTypeInit, dataIds[0])]
107 self.all_input_refs = self.input_refs + self.init_inputs_refs
109 # generate dataset refs for outputs
110 self.output_refs = [DatasetRef(self.datasetTypeOutput, dataId, run="RUN") for dataId in dataIds]
111 self.output_refs2 = [DatasetRef(self.datasetTypeOutput2, dataId, run="RUN") for dataId in dataIds]
113 self.missing_refs = [DatasetRef(self.datasetTypeExtra, dataId, run="RUN") for dataId in dataIds]
115 def tearDown(self) -> None:
116 removeTestTempDir(self.root)
118 def make_quantum(self, step: int = 1) -> Quantum:
119 """Make a Quantum which includes datastore records."""
120 if step == 1:
121 datastore_records = self.butler._datastore.export_records(self.all_input_refs)
122 predictedInputs = {self.datasetTypeInput: self.input_refs}
123 outputs = {self.datasetTypeOutput: self.output_refs}
124 initInputs = {self.datasetTypeInit: self.init_inputs_refs[0]}
125 elif step == 2:
126 # The result should be empty, this is just to test that it works.
127 datastore_records = self.butler._datastore.export_records(self.output_refs)
128 predictedInputs = {self.datasetTypeInput: self.output_refs}
129 outputs = {self.datasetTypeOutput2: self.output_refs2}
130 initInputs = {}
131 else:
132 raise ValueError(f"unexpected {step} value")
134 return Quantum(
135 taskName="some.task.name",
136 inputs=predictedInputs,
137 outputs=outputs,
138 initInputs=initInputs,
139 datastore_records=datastore_records,
140 )
142 def test_initialize(self) -> None:
143 """Test for initialize factory method"""
144 quantum = self.make_quantum()
145 qbb = QuantumBackedButler.initialize(
146 config=self.config,
147 quantum=quantum,
148 dimensions=self.universe,
149 dataset_types=self.dataset_types,
150 metrics=self.metrics,
151 )
152 self._test_factory(qbb)
154 def test_initialize_repo_index(self) -> None:
155 """Test for initialize using config file and repo index."""
156 # Store config to a file.
157 self.config.dumpToUri(self.root)
159 butler_index = Config()
160 butler_index["label"] = self.root
161 with ResourcePath.temporary_uri(suffix=".yaml") as index_path:
162 butler_index.dumpToUri(index_path)
164 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(index_path)}):
165 quantum = self.make_quantum()
166 qbb = QuantumBackedButler.initialize(
167 config="label",
168 quantum=quantum,
169 dimensions=self.universe,
170 dataset_types=self.dataset_types,
171 metrics=self.metrics,
172 )
173 self._test_factory(qbb)
175 def test_from_predicted(self) -> None:
176 """Test for from_predicted factory method"""
177 datastore_records = self.butler._datastore.export_records(self.all_input_refs)
178 qbb = QuantumBackedButler.from_predicted(
179 config=self.config,
180 predicted_inputs=[ref.id for ref in self.all_input_refs],
181 predicted_outputs=[ref.id for ref in self.output_refs],
182 dimensions=self.universe,
183 datastore_records=datastore_records,
184 dataset_types=self.dataset_types,
185 )
186 self._test_factory(qbb)
188 def _test_factory(self, qbb: QuantumBackedButler) -> None:
189 """Test state immediately after construction."""
190 self.assertTrue(qbb.isWriteable())
191 self.assertEqual(qbb._predicted_inputs, {ref.id for ref in self.all_input_refs})
192 self.assertEqual(qbb._predicted_outputs, {ref.id for ref in self.output_refs})
193 self.assertEqual(qbb._available_inputs, set())
194 self.assertEqual(qbb._unavailable_inputs, set())
195 self.assertEqual(qbb._actual_inputs, set())
196 self.assertEqual(qbb._actual_output_refs, set())
198 def test_getput(self) -> None:
199 """Test for get/put methods"""
200 quantum = self.make_quantum()
201 qbb = QuantumBackedButler.initialize(
202 config=self.config,
203 quantum=quantum,
204 dimensions=self.universe,
205 dataset_types=self.dataset_types,
206 metrics=self.metrics,
207 )
209 # Verify all input data are readable.
210 for ref in self.input_refs:
211 data = qbb.get(ref)
212 self.assertEqual(data, {"data": ref.dataId["detector"]})
213 for ref in self.init_inputs_refs:
214 data = qbb.get(ref)
215 self.assertEqual(data, {"data": -1})
216 for ref in self.missing_refs:
217 with self.assertRaises(FileNotFoundError):
218 data = qbb.get(ref)
220 self.assertEqual(qbb._available_inputs, qbb._predicted_inputs)
221 self.assertEqual(qbb._actual_inputs, qbb._predicted_inputs)
222 self.assertEqual(qbb._unavailable_inputs, {ref.id for ref in self.missing_refs})
224 self.metrics.reset()
226 # Write all expected outputs.
227 for ref in self.output_refs:
228 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
229 self.assertEqual(self.metrics.n_put, len(self.output_refs))
231 # Must be able to read them back
232 for ref in self.output_refs:
233 data = qbb.get(ref)
234 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 2})
236 self.assertEqual(qbb._actual_output_refs, set(self.output_refs))
237 self.assertEqual(self.metrics.n_get, len(self.output_refs))
239 # Retrieve them as a Zip artifact.
240 with tempfile.TemporaryDirectory() as tmpdir:
241 zip = qbb.retrieve_artifacts_zip(self.output_refs, destination=tmpdir)
243 index = ZipIndex.from_zip_file(zip)
244 zip_refs = index.refs.to_refs(universe=qbb.dimensions)
245 self.assertEqual(len(zip_refs), 4)
246 self.assertEqual(set(zip_refs), set(self.output_refs))
247 self.assertEqual(len(index.artifact_map), 4) # Count number of artifacts in Zip.
249 # Retrieve them to a directory.
250 with tempfile.TemporaryDirectory() as tmpdir:
251 retrieved = qbb.retrieve_artifacts(self.output_refs, destination=tmpdir, preserve_path=False)
252 self.assertEqual(len(retrieved), 4)
253 self.assertTrue(retrieved[0].exists())
254 with open(os.path.join(tmpdir, ZipIndex.index_name)) as zf:
255 index = ZipIndex.model_validate_json(zf.read())
256 self.assertEqual(len(index.artifact_map), len(retrieved))
258 def test_getDeferred(self) -> None:
259 """Test for getDeferred method"""
260 quantum = self.make_quantum()
261 qbb = QuantumBackedButler.initialize(
262 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
263 )
265 # get some input data
266 input_refs = self.input_refs[:2]
267 for ref in input_refs:
268 data = qbb.getDeferred(ref)
269 self.assertEqual(data.get(), {"data": ref.dataId["detector"]})
270 for ref in self.init_inputs_refs:
271 data = qbb.getDeferred(ref)
272 self.assertEqual(data.get(), {"data": -1})
273 for ref in self.missing_refs:
274 data = qbb.getDeferred(ref)
275 with self.assertRaises(FileNotFoundError):
276 data.get()
278 # _avalable_inputs is not
279 self.assertEqual(qbb._available_inputs, {ref.id for ref in input_refs + self.init_inputs_refs})
280 self.assertEqual(qbb._actual_inputs, {ref.id for ref in input_refs + self.init_inputs_refs})
281 self.assertEqual(qbb._unavailable_inputs, {ref.id for ref in self.missing_refs})
283 def test_stored(self) -> None:
284 """Test for dataset existence method"""
285 quantum = self.make_quantum()
286 qbb = QuantumBackedButler.initialize(
287 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
288 )
290 # get some input data
291 input_refs = self.input_refs[:2]
292 for ref in input_refs:
293 exists = qbb.stored(ref)
294 self.assertTrue(exists)
295 for ref in self.init_inputs_refs:
296 exists = qbb.stored(ref)
297 self.assertTrue(exists)
298 for ref in self.missing_refs:
299 exists = qbb.stored(ref)
300 self.assertFalse(exists)
302 # Now do the same checks in bulk.
303 missing_set = set(self.missing_refs)
304 refs = input_refs + self.init_inputs_refs + self.missing_refs
305 stored_many = qbb.stored_many(refs)
306 for ref, stored in stored_many.items():
307 if ref in missing_set:
308 self.assertFalse(stored)
309 else:
310 self.assertTrue(stored)
312 # _available_inputs is not
313 self.assertEqual(qbb._available_inputs, {ref.id for ref in input_refs + self.init_inputs_refs})
314 self.assertEqual(qbb._actual_inputs, set())
315 self.assertEqual(qbb._unavailable_inputs, set()) # this is not consistent with get?
317 def test_markInputUnused(self) -> None:
318 """Test for markInputUnused method"""
319 quantum = self.make_quantum()
320 qbb = QuantumBackedButler.initialize(
321 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
322 )
324 # get some input data
325 for ref in self.input_refs:
326 data = qbb.get(ref)
327 self.assertEqual(data, {"data": ref.dataId["detector"]})
328 for ref in self.init_inputs_refs:
329 data = qbb.get(ref)
330 self.assertEqual(data, {"data": -1})
332 self.assertEqual(qbb._available_inputs, qbb._predicted_inputs)
333 self.assertEqual(qbb._actual_inputs, qbb._predicted_inputs)
335 qbb.markInputUnused(self.input_refs[0])
336 self.assertEqual(qbb._actual_inputs, {ref.id for ref in self.input_refs[1:] + self.init_inputs_refs})
338 def test_pruneDatasets(self) -> None:
339 """Test for pruneDatasets methods"""
340 quantum = self.make_quantum()
341 qbb = QuantumBackedButler.initialize(
342 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
343 )
345 # Write all expected outputs.
346 for ref in self.output_refs:
347 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
349 # Must be able to read them back
350 for ref in self.output_refs:
351 data = qbb.get(ref)
352 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 2})
354 # Check for invalid arguments.
355 with self.assertRaisesRegex(TypeError, "Cannot pass purge=True without disassociate=True"):
356 qbb.pruneDatasets(self.output_refs, disassociate=False, unstore=True, purge=True)
357 with self.assertRaisesRegex(TypeError, "Cannot pass purge=True without unstore=True"):
358 qbb.pruneDatasets(self.output_refs, disassociate=True, unstore=False, purge=True)
359 with self.assertRaisesRegex(TypeError, "Cannot pass disassociate=True without purge=True"):
360 qbb.pruneDatasets(self.output_refs, disassociate=True, unstore=True, purge=False)
362 # Disassociate only.
363 ref = self.output_refs[0]
364 qbb.pruneDatasets([ref], disassociate=False, unstore=True, purge=False)
365 self.assertFalse(qbb.stored(ref))
366 with self.assertRaises(FileNotFoundError):
367 data = qbb.get(ref)
369 # can store it again
370 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
371 self.assertTrue(qbb.stored(ref))
373 # Purge completely.
374 ref = self.output_refs[1]
375 qbb.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
376 self.assertFalse(qbb.stored(ref))
377 with self.assertRaises(FileNotFoundError):
378 data = qbb.get(ref)
379 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
380 self.assertTrue(qbb.stored(ref))
382 def test_extract_provenance_data(self) -> None:
383 """Test for extract_provenance_data method"""
384 quantum = self.make_quantum()
385 qbb = QuantumBackedButler.initialize(
386 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
387 )
389 # read/store everything
390 for ref in self.input_refs:
391 qbb.get(ref)
392 for ref in self.init_inputs_refs:
393 qbb.get(ref)
394 for ref in self.output_refs:
395 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
397 provenance1 = qbb.extract_provenance_data()
398 prov_json = provenance1.model_dump_json()
399 provenance2 = QuantumProvenanceData.direct(**json.loads(prov_json))
400 for provenance in (provenance1, provenance2):
401 input_ids = {ref.id for ref in self.input_refs + self.init_inputs_refs}
402 self.assertEqual(provenance.predicted_inputs, input_ids)
403 self.assertEqual(provenance.available_inputs, input_ids)
404 self.assertEqual(provenance.actual_inputs, input_ids)
405 output_ids = {ref.id for ref in self.output_refs}
406 self.assertEqual(provenance.predicted_outputs, output_ids)
407 self.assertEqual(provenance.actual_outputs, output_ids)
408 datastore_name = "FileDatastore@<butlerRoot>/datastore"
409 self.assertEqual(set(provenance.datastore_records.keys()), {datastore_name})
410 datastore_records = provenance.datastore_records[datastore_name]
411 self.assertEqual(set(datastore_records.dataset_ids), output_ids)
412 class_name = "lsst.daf.butler.datastore.stored_file_info.StoredFileInfo"
413 self.assertEqual(set(datastore_records.records.keys()), {class_name})
414 self.assertEqual(set(datastore_records.records[class_name].keys()), {id.hex for id in output_ids})
415 table_name = "file_datastore_records"
416 for dataset_data in datastore_records.records[class_name].values():
417 self.assertEqual(set(dataset_data), {table_name})
419 def test_export_predicted_datastore_records(self) -> None:
420 """Test for export_predicted_datastore_records method"""
421 quantum = self.make_quantum()
422 qbb = QuantumBackedButler.initialize(
423 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
424 )
426 records = qbb.export_predicted_datastore_records(self.output_refs)
427 self.assertEqual(len(records["FileDatastore@<butlerRoot>/datastore"].records), len(self.output_refs))
429 with self.assertRaises(ValueError):
430 qbb.export_predicted_datastore_records(self.missing_refs)
432 def test_collect_and_transfer(self) -> None:
433 """Test for collect_and_transfer method"""
434 quantum1 = self.make_quantum(1)
435 qbb1 = QuantumBackedButler.initialize(
436 config=self.config, quantum=quantum1, dimensions=self.universe, dataset_types=self.dataset_types
437 )
439 quantum2 = self.make_quantum(2)
440 qbb2 = QuantumBackedButler.initialize(
441 config=self.config, quantum=quantum2, dimensions=self.universe, dataset_types=self.dataset_types
442 )
444 # read/store everything
445 for ref in self.input_refs:
446 qbb1.get(ref)
447 for ref in self.init_inputs_refs:
448 qbb1.get(ref)
449 for ref in self.output_refs:
450 qbb1.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
452 for ref in self.output_refs:
453 qbb2.get(ref)
454 for ref in self.output_refs2:
455 qbb2.put({"data": cast(int, ref.dataId["detector"]) ** 3}, ref)
457 QuantumProvenanceData.collect_and_transfer(
458 self.butler,
459 [quantum1, quantum2],
460 [qbb1.extract_provenance_data(), qbb2.extract_provenance_data()],
461 )
463 for ref in self.output_refs:
464 data = self.butler.get(ref)
465 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 2})
467 for ref in self.output_refs2:
468 data = self.butler.get(ref)
469 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 3})
471 def test_record_data_merge_mappings(self):
472 """Test DatastoreRecordData.merge_mappings."""
473 r1 = self.butler._datastore.export_records(self.input_refs)
474 r2 = self.butler._datastore.export_records(self.init_inputs_refs)
475 r3 = self.butler._datastore.export_records(self.input_refs) # intentional duplicate
476 merged = DatastoreRecordData.merge_mappings(r1, r2, r3)
477 self.assertEqual(merged.keys(), r1.keys() | r2.keys())
478 for datastore_name in merged.keys():
479 if datastore_name in r1:
480 self.assertGreaterEqual(
481 merged[datastore_name].records.keys(), r1[datastore_name].records.keys()
482 )
483 if datastore_name in r2:
484 self.assertGreaterEqual(
485 merged[datastore_name].records.keys(), r2[datastore_name].records.keys()
486 )
489if __name__ == "__main__":
490 unittest.main()