Coverage for tests / test_quantumBackedButler.py: 9%
290 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-24 01:10 -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 registry = _RegistryFactory(registryConfig).create_from_config(butlerRoot=self.root)
71 registry.close()
72 butler = Butler.from_config(self.config, writeable=True, run="RUN", metrics=self.metrics)
73 self.enterContext(butler)
74 assert isinstance(butler, DirectButler)
75 self.butler = butler
76 self.butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml")
78 # make all dataset types
79 graph = self.universe.conform(("instrument", "detector"))
80 storageClass = StorageClass("StructuredDataDict")
81 self.datasetTypeInit = DatasetType("test_ds_init", graph, storageClass)
82 self.datasetTypeInput = DatasetType("test_ds_input", graph, storageClass)
83 self.datasetTypeOutput = DatasetType("test_ds_output", graph, storageClass)
84 self.datasetTypeOutput2 = DatasetType("test_ds_output2", graph, storageClass)
85 self.datasetTypeExtra = DatasetType("test_ds_extra", graph, storageClass)
87 self.dataset_types: dict[str, DatasetType] = {}
88 dataset_types = (
89 self.datasetTypeInit,
90 self.datasetTypeInput,
91 self.datasetTypeOutput,
92 self.datasetTypeOutput2,
93 self.datasetTypeExtra,
94 )
95 for dataset_type in dataset_types:
96 self.butler.registry.registerDatasetType(dataset_type)
97 self.dataset_types[dataset_type.name] = dataset_type
99 dataIds = [
100 self.butler.registry.expandDataId(dict(instrument="Cam1", detector=detector_id))
101 for detector_id in (1, 2, 3, 4)
102 ]
104 # make actual input datasets
105 self.input_refs = [
106 self.butler.put({"data": dataId["detector"]}, self.datasetTypeInput, dataId) for dataId in dataIds
107 ]
108 self.init_inputs_refs = [self.butler.put({"data": -1}, self.datasetTypeInit, dataIds[0])]
109 self.all_input_refs = self.input_refs + self.init_inputs_refs
111 # generate dataset refs for outputs
112 self.output_refs = [DatasetRef(self.datasetTypeOutput, dataId, run="RUN") for dataId in dataIds]
113 self.output_refs2 = [DatasetRef(self.datasetTypeOutput2, dataId, run="RUN") for dataId in dataIds]
115 self.missing_refs = [DatasetRef(self.datasetTypeExtra, dataId, run="RUN") for dataId in dataIds]
117 def tearDown(self) -> None:
118 removeTestTempDir(self.root)
120 def make_quantum(self, step: int = 1) -> Quantum:
121 """Make a Quantum which includes datastore records."""
122 if step == 1:
123 datastore_records = self.butler._datastore.export_records(self.all_input_refs)
124 predictedInputs = {self.datasetTypeInput: self.input_refs}
125 outputs = {self.datasetTypeOutput: self.output_refs}
126 initInputs = {self.datasetTypeInit: self.init_inputs_refs[0]}
127 elif step == 2:
128 # The result should be empty, this is just to test that it works.
129 datastore_records = self.butler._datastore.export_records(self.output_refs)
130 predictedInputs = {self.datasetTypeInput: self.output_refs}
131 outputs = {self.datasetTypeOutput2: self.output_refs2}
132 initInputs = {}
133 else:
134 raise ValueError(f"unexpected {step} value")
136 return Quantum(
137 taskName="some.task.name",
138 inputs=predictedInputs,
139 outputs=outputs,
140 initInputs=initInputs,
141 datastore_records=datastore_records,
142 )
144 def test_initialize(self) -> None:
145 """Test for initialize factory method"""
146 quantum = self.make_quantum()
147 qbb = QuantumBackedButler.initialize(
148 config=self.config,
149 quantum=quantum,
150 dimensions=self.universe,
151 dataset_types=self.dataset_types,
152 metrics=self.metrics,
153 )
154 self.addCleanup(qbb.close)
155 self._test_factory(qbb)
157 def test_initialize_repo_index(self) -> None:
158 """Test for initialize using config file and repo index."""
159 # Store config to a file.
160 self.config.dumpToUri(self.root)
162 butler_index = Config()
163 butler_index["label"] = self.root
164 with ResourcePath.temporary_uri(suffix=".yaml") as index_path:
165 butler_index.dumpToUri(index_path)
167 with unittest.mock.patch.dict(os.environ, {"DAF_BUTLER_REPOSITORY_INDEX": str(index_path)}):
168 quantum = self.make_quantum()
169 qbb = QuantumBackedButler.initialize(
170 config="label",
171 quantum=quantum,
172 dimensions=self.universe,
173 dataset_types=self.dataset_types,
174 metrics=self.metrics,
175 )
176 self.addCleanup(qbb.close)
177 self._test_factory(qbb)
179 def test_from_predicted(self) -> None:
180 """Test for from_predicted factory method"""
181 datastore_records = self.butler._datastore.export_records(self.all_input_refs)
182 qbb = QuantumBackedButler.from_predicted(
183 config=self.config,
184 predicted_inputs=[ref.id for ref in self.all_input_refs],
185 predicted_outputs=[ref.id for ref in self.output_refs],
186 dimensions=self.universe,
187 datastore_records=datastore_records,
188 dataset_types=self.dataset_types,
189 )
190 self.addCleanup(qbb.close)
191 self._test_factory(qbb)
193 def _test_factory(self, qbb: QuantumBackedButler) -> None:
194 """Test state immediately after construction."""
195 self.assertTrue(qbb.isWriteable())
196 self.assertEqual(qbb._predicted_inputs, {ref.id for ref in self.all_input_refs})
197 self.assertEqual(qbb._predicted_outputs, {ref.id for ref in self.output_refs})
198 self.assertEqual(qbb._available_inputs, set())
199 self.assertEqual(qbb._unavailable_inputs, set())
200 self.assertEqual(qbb._actual_inputs, set())
201 self.assertEqual(qbb._actual_output_refs, set())
203 def test_getput(self) -> None:
204 """Test for get/put methods"""
205 quantum = self.make_quantum()
206 qbb = QuantumBackedButler.initialize(
207 config=self.config,
208 quantum=quantum,
209 dimensions=self.universe,
210 dataset_types=self.dataset_types,
211 metrics=self.metrics,
212 )
213 self.addCleanup(qbb.close)
215 # Verify all input data are readable.
216 for ref in self.input_refs:
217 data = qbb.get(ref)
218 self.assertEqual(data, {"data": ref.dataId["detector"]})
219 for ref in self.init_inputs_refs:
220 data = qbb.get(ref)
221 self.assertEqual(data, {"data": -1})
222 for ref in self.missing_refs:
223 with self.assertRaises(FileNotFoundError):
224 data = qbb.get(ref)
226 self.assertEqual(qbb._available_inputs, qbb._predicted_inputs)
227 self.assertEqual(qbb._actual_inputs, qbb._predicted_inputs)
228 self.assertEqual(qbb._unavailable_inputs, {ref.id for ref in self.missing_refs})
230 self.metrics.reset()
232 # Write all expected outputs.
233 for ref in self.output_refs:
234 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
235 self.assertEqual(self.metrics.n_put, len(self.output_refs))
237 # Must be able to read them back
238 for ref in self.output_refs:
239 data = qbb.get(ref)
240 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 2})
242 self.assertEqual(qbb._actual_output_refs, set(self.output_refs))
243 self.assertEqual(self.metrics.n_get, len(self.output_refs))
245 # Retrieve them as a Zip artifact.
246 with tempfile.TemporaryDirectory() as tmpdir:
247 zip = qbb.retrieve_artifacts_zip(self.output_refs, destination=tmpdir)
249 index = ZipIndex.from_zip_file(zip)
250 zip_refs = index.refs.to_refs(universe=qbb.dimensions)
251 self.assertEqual(len(zip_refs), 4)
252 self.assertEqual(set(zip_refs), set(self.output_refs))
253 self.assertEqual(len(index.artifact_map), 4) # Count number of artifacts in Zip.
255 # Retrieve them to a directory.
256 with tempfile.TemporaryDirectory() as tmpdir:
257 retrieved = qbb.retrieve_artifacts(self.output_refs, destination=tmpdir, preserve_path=False)
258 self.assertEqual(len(retrieved), 4)
259 self.assertTrue(retrieved[0].exists())
260 with open(os.path.join(tmpdir, ZipIndex.index_name)) as zf:
261 index = ZipIndex.model_validate_json(zf.read())
262 self.assertEqual(len(index.artifact_map), len(retrieved))
264 def test_getDeferred(self) -> None:
265 """Test for getDeferred method"""
266 quantum = self.make_quantum()
267 qbb = QuantumBackedButler.initialize(
268 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
269 )
270 self.addCleanup(qbb.close)
272 # get some input data
273 input_refs = self.input_refs[:2]
274 for ref in input_refs:
275 data = qbb.getDeferred(ref)
276 self.assertEqual(data.get(), {"data": ref.dataId["detector"]})
277 for ref in self.init_inputs_refs:
278 data = qbb.getDeferred(ref)
279 self.assertEqual(data.get(), {"data": -1})
280 for ref in self.missing_refs:
281 data = qbb.getDeferred(ref)
282 with self.assertRaises(FileNotFoundError):
283 data.get()
285 # _avalable_inputs is not
286 self.assertEqual(qbb._available_inputs, {ref.id for ref in input_refs + self.init_inputs_refs})
287 self.assertEqual(qbb._actual_inputs, {ref.id for ref in input_refs + self.init_inputs_refs})
288 self.assertEqual(qbb._unavailable_inputs, {ref.id for ref in self.missing_refs})
290 def test_stored(self) -> None:
291 """Test for dataset existence method"""
292 quantum = self.make_quantum()
293 qbb = QuantumBackedButler.initialize(
294 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
295 )
296 self.addCleanup(qbb.close)
298 # get some input data
299 input_refs = self.input_refs[:2]
300 for ref in input_refs:
301 exists = qbb.stored(ref)
302 self.assertTrue(exists)
303 for ref in self.init_inputs_refs:
304 exists = qbb.stored(ref)
305 self.assertTrue(exists)
306 for ref in self.missing_refs:
307 exists = qbb.stored(ref)
308 self.assertFalse(exists)
310 # Now do the same checks in bulk.
311 missing_set = set(self.missing_refs)
312 refs = input_refs + self.init_inputs_refs + self.missing_refs
313 stored_many = qbb.stored_many(refs)
314 for ref, stored in stored_many.items():
315 if ref in missing_set:
316 self.assertFalse(stored)
317 else:
318 self.assertTrue(stored)
320 # _available_inputs is not
321 self.assertEqual(qbb._available_inputs, {ref.id for ref in input_refs + self.init_inputs_refs})
322 self.assertEqual(qbb._actual_inputs, set())
323 self.assertEqual(qbb._unavailable_inputs, set()) # this is not consistent with get?
325 def test_markInputUnused(self) -> None:
326 """Test for markInputUnused method"""
327 quantum = self.make_quantum()
328 qbb = QuantumBackedButler.initialize(
329 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
330 )
331 self.addCleanup(qbb.close)
333 # get some input data
334 for ref in self.input_refs:
335 data = qbb.get(ref)
336 self.assertEqual(data, {"data": ref.dataId["detector"]})
337 for ref in self.init_inputs_refs:
338 data = qbb.get(ref)
339 self.assertEqual(data, {"data": -1})
341 self.assertEqual(qbb._available_inputs, qbb._predicted_inputs)
342 self.assertEqual(qbb._actual_inputs, qbb._predicted_inputs)
344 qbb.markInputUnused(self.input_refs[0])
345 self.assertEqual(qbb._actual_inputs, {ref.id for ref in self.input_refs[1:] + self.init_inputs_refs})
347 def test_pruneDatasets(self) -> None:
348 """Test for pruneDatasets methods"""
349 quantum = self.make_quantum()
350 qbb = QuantumBackedButler.initialize(
351 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
352 )
353 self.addCleanup(qbb.close)
355 # Write all expected outputs.
356 for ref in self.output_refs:
357 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
359 # Must be able to read them back
360 for ref in self.output_refs:
361 data = qbb.get(ref)
362 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 2})
364 # Check for invalid arguments.
365 with self.assertRaisesRegex(TypeError, "Cannot pass purge=True without disassociate=True"):
366 qbb.pruneDatasets(self.output_refs, disassociate=False, unstore=True, purge=True)
367 with self.assertRaisesRegex(TypeError, "Cannot pass purge=True without unstore=True"):
368 qbb.pruneDatasets(self.output_refs, disassociate=True, unstore=False, purge=True)
369 with self.assertRaisesRegex(TypeError, "Cannot pass disassociate=True without purge=True"):
370 qbb.pruneDatasets(self.output_refs, disassociate=True, unstore=True, purge=False)
372 # Disassociate only.
373 ref = self.output_refs[0]
374 qbb.pruneDatasets([ref], disassociate=False, unstore=True, purge=False)
375 self.assertFalse(qbb.stored(ref))
376 with self.assertRaises(FileNotFoundError):
377 data = qbb.get(ref)
379 # can store it again
380 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
381 self.assertTrue(qbb.stored(ref))
383 # Purge completely.
384 ref = self.output_refs[1]
385 qbb.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
386 self.assertFalse(qbb.stored(ref))
387 with self.assertRaises(FileNotFoundError):
388 data = qbb.get(ref)
389 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
390 self.assertTrue(qbb.stored(ref))
392 def test_extract_provenance_data(self) -> None:
393 """Test for extract_provenance_data method"""
394 quantum = self.make_quantum()
395 qbb = QuantumBackedButler.initialize(
396 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
397 )
398 self.addCleanup(qbb.close)
400 # read/store everything
401 for ref in self.input_refs:
402 qbb.get(ref)
403 for ref in self.init_inputs_refs:
404 qbb.get(ref)
405 for ref in self.output_refs:
406 qbb.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
408 provenance1 = qbb.extract_provenance_data()
409 prov_json = provenance1.model_dump_json()
410 provenance2 = QuantumProvenanceData.direct(**json.loads(prov_json))
411 for provenance in (provenance1, provenance2):
412 input_ids = {ref.id for ref in self.input_refs + self.init_inputs_refs}
413 self.assertEqual(provenance.predicted_inputs, input_ids)
414 self.assertEqual(provenance.available_inputs, input_ids)
415 self.assertEqual(provenance.actual_inputs, input_ids)
416 output_ids = {ref.id for ref in self.output_refs}
417 self.assertEqual(provenance.predicted_outputs, output_ids)
418 self.assertEqual(provenance.actual_outputs, output_ids)
419 datastore_name = "FileDatastore@<butlerRoot>/datastore"
420 self.assertEqual(set(provenance.datastore_records.keys()), {datastore_name})
421 datastore_records = provenance.datastore_records[datastore_name]
422 self.assertEqual(set(datastore_records.dataset_ids), output_ids)
423 class_name = "lsst.daf.butler.datastore.stored_file_info.StoredFileInfo"
424 self.assertEqual(set(datastore_records.records.keys()), {class_name})
425 self.assertEqual(set(datastore_records.records[class_name].keys()), {id.hex for id in output_ids})
426 table_name = "file_datastore_records"
427 for dataset_data in datastore_records.records[class_name].values():
428 self.assertEqual(set(dataset_data), {table_name})
430 def test_export_predicted_datastore_records(self) -> None:
431 """Test for export_predicted_datastore_records method"""
432 quantum = self.make_quantum()
433 qbb = QuantumBackedButler.initialize(
434 config=self.config, quantum=quantum, dimensions=self.universe, dataset_types=self.dataset_types
435 )
436 self.addCleanup(qbb.close)
438 records = qbb.export_predicted_datastore_records(self.output_refs)
439 self.assertEqual(len(records["FileDatastore@<butlerRoot>/datastore"].records), len(self.output_refs))
441 with self.assertRaises(ValueError):
442 qbb.export_predicted_datastore_records(self.missing_refs)
444 def test_collect_and_transfer(self) -> None:
445 """Test for collect_and_transfer method"""
446 quantum1 = self.make_quantum(1)
447 qbb1 = QuantumBackedButler.initialize(
448 config=self.config, quantum=quantum1, dimensions=self.universe, dataset_types=self.dataset_types
449 )
450 self.addCleanup(qbb1.close)
452 quantum2 = self.make_quantum(2)
453 qbb2 = QuantumBackedButler.initialize(
454 config=self.config, quantum=quantum2, dimensions=self.universe, dataset_types=self.dataset_types
455 )
456 self.addCleanup(qbb2.close)
458 # read/store everything
459 for ref in self.input_refs:
460 qbb1.get(ref)
461 for ref in self.init_inputs_refs:
462 qbb1.get(ref)
463 for ref in self.output_refs:
464 qbb1.put({"data": cast(int, ref.dataId["detector"]) ** 2}, ref)
466 for ref in self.output_refs:
467 qbb2.get(ref)
468 for ref in self.output_refs2:
469 qbb2.put({"data": cast(int, ref.dataId["detector"]) ** 3}, ref)
471 QuantumProvenanceData.collect_and_transfer(
472 self.butler,
473 [quantum1, quantum2],
474 [qbb1.extract_provenance_data(), qbb2.extract_provenance_data()],
475 )
477 for ref in self.output_refs:
478 data = self.butler.get(ref)
479 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 2})
481 for ref in self.output_refs2:
482 data = self.butler.get(ref)
483 self.assertEqual(data, {"data": cast(int, ref.dataId["detector"]) ** 3})
485 def test_record_data_merge_mappings(self):
486 """Test DatastoreRecordData.merge_mappings."""
487 r1 = self.butler._datastore.export_records(self.input_refs)
488 r2 = self.butler._datastore.export_records(self.init_inputs_refs)
489 r3 = self.butler._datastore.export_records(self.input_refs) # intentional duplicate
490 merged = DatastoreRecordData.merge_mappings(r1, r2, r3)
491 self.assertEqual(merged.keys(), r1.keys() | r2.keys())
492 for datastore_name in merged.keys():
493 if datastore_name in r1:
494 self.assertGreaterEqual(
495 merged[datastore_name].records.keys(), r1[datastore_name].records.keys()
496 )
497 if datastore_name in r2:
498 self.assertGreaterEqual(
499 merged[datastore_name].records.keys(), r2[datastore_name].records.keys()
500 )
503if __name__ == "__main__":
504 unittest.main()