Coverage for python/lsst/pipe/base/tests/mocks/_repo.py: 20%
145 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-21 01:24 -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# (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/>.
28from __future__ import annotations
30__all__ = ("DirectButlerRepo", "InMemoryRepo", "MockRepo")
32import tempfile
33from abc import ABC, abstractmethod
34from collections.abc import Iterable, Iterator, Mapping
35from contextlib import contextmanager
36from typing import Any
38from lsst.daf.butler import (
39 Butler,
40 CollectionType,
41 DataCoordinate,
42 DatasetRef,
43 DatasetType,
44 DimensionConfig,
45 LimitedButler,
46 RegistryConfig,
47)
48from lsst.daf.butler.tests.utils import create_populated_sqlite_registry
49from lsst.resources import ResourcePath, ResourcePathExpression
50from lsst.sphgeom import RangeSet
52from ...all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
53from ...pipeline_graph import PipelineGraph
54from ...quantum_graph import PredictedQuantumGraph
55from ...single_quantum_executor import SingleQuantumExecutor
56from ..in_memory_limited_butler import InMemoryLimitedButler
57from ._pipeline_task import (
58 DynamicConnectionConfig,
59 DynamicTestPipelineTask,
60 DynamicTestPipelineTaskConfig,
61)
62from ._storage_class import MockDataset, is_mock_name
65class MockRepo(ABC):
66 """A test helper that populates a butler repository for task execution.
68 Parameters
69 ----------
70 butler : `lsst.daf.butler.Butler`
71 Butler to use for at least quantum graph building. Must be writeable.
72 input_run : `str`, optional
73 Name of a `~lsst.daf.butler.CollectionType.RUN` collection that will be
74 used as an input to quantum graph generation. Input datasets created
75 by the helper are added to this collection.
76 input_chain : `str`, optional
77 Name of a `~lsst.daf.butler.CollectionType.CHAINED` collection that
78 will be the direct input to quantum graph generation. This always
79 includes ``input_run``.
80 input_children : `str` or `~collections.abc.Iterable` [ `str`], optional
81 Additional collections to include in ``input_chain``.
82 """
84 def __init__(
85 self,
86 butler: Butler,
87 input_run: str = "input_run",
88 input_chain: str = "input_chain",
89 input_children: Iterable[str] = (),
90 ):
91 self.butler = butler
92 input_chain_definition = [input_run]
93 input_chain_definition.extend(input_children)
94 self.input_run = input_run
95 self.input_chain = input_chain
96 self.butler.collections.register(self.input_run)
97 self.butler.collections.register(self.input_chain, CollectionType.CHAINED)
98 self.butler.collections.redefine_chain(self.input_chain, input_chain_definition)
99 self.pipeline_graph = PipelineGraph()
100 self.last_auto_dataset_type_index = 0
101 self.last_auto_task_index = 0
103 def add_task(
104 self,
105 label: str | None = None,
106 *,
107 task_class: type[DynamicTestPipelineTask] = DynamicTestPipelineTask,
108 config: DynamicTestPipelineTaskConfig | None = None,
109 dimensions: Iterable[str] | None = None,
110 inputs: Mapping[str, DynamicConnectionConfig] | None = None,
111 outputs: Mapping[str, DynamicConnectionConfig] | None = None,
112 prerequisite_inputs: Mapping[str, DynamicConnectionConfig] | None = None,
113 init_inputs: Mapping[str, DynamicConnectionConfig] | None = None,
114 init_outputs: Mapping[str, DynamicConnectionConfig] | None = None,
115 ) -> None:
116 """Add a task to the helper's pipeline graph.
118 Parameters
119 ----------
120 label : `str`, optional
121 Label for the task. If not provided, the task name will be
122 ``task_auto{self.last_auto_task_index}``, with that variable
123 incremented.
124 task_class : `type`, optional
125 Subclass of `DynamicTestPipelineTask` to use.
126 config : `DynamicTestPipelineTaskConfig`, optional
127 Task configuration to use. Note that the dimensions are always
128 overridden by the ``dimensions`` argument and ``inputs`` and
129 ``outputs`` are updated by those arguments unless they are
130 explicitly set to empty dictionaries.
131 dimensions : `~collections.abc.Iterable` [ `str` ], optional
132 Dimensions of the task and any automatically-added input or output
133 connection.
134 inputs : `~collections.abc.Mapping` [ `str`, \
135 `DynamicConnectionConfig` ], optional
136 Input connections to add. If not provided, a single connection is
137 added with the same dimensions as the task and dataset type name
138 ``dataset_auto{self.last_auto_dataset_type_index}``.
139 outputs : `~collections.abc.Mapping` [ `str`, \
140 `DynamicConnectionConfig` ], optional
141 Output connections to add. If not provided, a single connection is
142 added with the same dimensions as the task and dataset type name
143 ``dataset_auto{self.last_auto_dataset_type_index}``, with that
144 variable incremented first.
145 prerequisite_inputs : `~collections.abc.Mapping` [ `str`, \
146 `DynamicConnectionConfig` ], optional
147 Prerequisite input connections to add. Defaults to an empty
148 mapping.
149 init_inputs : `~collections.abc.Mapping` [ `str`, \
150 `DynamicConnectionConfig` ], optional
151 Init input connections to add. Defaults to an empty mapping.
152 init_outputs : `~collections.abc.Mapping` [ `str`, \
153 `DynamicConnectionConfig` ], optional
154 Init output connections to add. Defaults to an empty mapping.
156 Notes
157 -----
158 The defaults for this method's arguments are designed to allow it to be
159 called in succession to create a sequence of "one-to-one" tasks in
160 which each consumes the output of the last.
161 """
162 if config is None:
163 config = DynamicTestPipelineTaskConfig()
164 if dimensions is not None:
165 config.dimensions = list(dimensions)
166 if inputs is not None:
167 config.inputs.update(inputs)
168 else:
169 config.inputs["input_connection"] = DynamicConnectionConfig(
170 dataset_type_name=f"dataset_auto{self.last_auto_dataset_type_index}",
171 dimensions=list(config.dimensions),
172 )
173 if outputs is not None:
174 config.outputs.update(outputs)
175 else:
176 self.last_auto_dataset_type_index += 1
177 config.outputs["output_connection"] = DynamicConnectionConfig(
178 dataset_type_name=f"dataset_auto{self.last_auto_dataset_type_index}",
179 dimensions=list(config.dimensions),
180 )
181 if prerequisite_inputs is not None:
182 config.prerequisite_inputs.update(prerequisite_inputs)
183 if init_inputs is not None:
184 config.init_inputs.update(init_inputs)
185 if init_outputs is not None:
186 config.init_outputs.update(init_outputs)
187 if label is None:
188 self.last_auto_task_index += 1
189 label = f"task_auto{self.last_auto_task_index}"
190 self.pipeline_graph.add_task(label, task_class=task_class, config=config)
192 def make_quantum_graph(
193 self,
194 *,
195 output: str | None = None,
196 output_run: str = "output_run",
197 insert_mocked_inputs: bool = True,
198 register_output_dataset_types: bool = True,
199 ) -> PredictedQuantumGraph:
200 """Make a quantum graph from the pipeline task and internal data
201 repository.
203 Parameters
204 ----------
205 output : `str` or `None`, optional
206 Name of the output chained collection to embed within the quantum
207 graph. Note that this does not actually create this collection.
208 output_run : `str`, optional
209 Name of the `~lsst.daf.butler.CollectionType.RUN` collection for
210 execution outputs. Note that this does not actually create this
211 collection.
212 insert_mocked_inputs : `bool`, optional
213 Whether to automatically insert datasets for all overall inputs to
214 the pipeline graph whose dataset types have not already been
215 registered. If set to `False`, inputs must be provided by imported
216 YAML files or explicit calls to `insert_datasets`, which provides
217 more fine-grained control over the data IDs of the datasets.
218 register_output_dataset_types : `bool`, optional
219 If `True`, register all output dataset types.
221 Returns
222 -------
223 qg : `..quantum_graph.PredictedQuantumGraph`
224 Quantum graph. Datastore records will not be attached, since the
225 test helper does not actually have a datastore.
226 """
227 return (
228 self.make_quantum_graph_builder(
229 insert_mocked_inputs=insert_mocked_inputs,
230 register_output_dataset_types=register_output_dataset_types,
231 output_run=output_run,
232 )
233 .finish(output=output, attach_datastore_records=False)
234 .assemble()
235 )
237 def make_quantum_graph_builder(
238 self,
239 *,
240 output_run: str = "output_run",
241 insert_mocked_inputs: bool = True,
242 register_output_dataset_types: bool = True,
243 ) -> AllDimensionsQuantumGraphBuilder:
244 """Make a quantum graph builder from the pipeline task and internal
245 data repository.
247 Parameters
248 ----------
249 output_run : `str`, optional
250 Name of the `~lsst.daf.butler.CollectionType.RUN` collection for
251 execution outputs. Note that this does not actually create this
252 collection.
253 insert_mocked_inputs : `bool`, optional
254 Whether to automatically insert datasets for all overall inputs to
255 the pipeline graph whose dataset types have not already been
256 registered. If set to `False`, inputs must be provided by imported
257 YAML files or explicit calls to `insert_datasets`, which provides
258 more fine-grained control over the data IDs of the datasets.
259 register_output_dataset_types : `bool`, optional
260 If `True`, register all output dataset types.
262 Returns
263 -------
264 builder : \
265 `..all_dimensions_quantum_graph_builder.AllDimensionsQuantumGraphBuilder`
266 Quantum graph builder. Note that
267 ``attach_datastore_records=False`` must be passed to `build`, since
268 the helper's butler does not have a datastore.
269 """
270 if insert_mocked_inputs:
271 self.pipeline_graph.resolve(self.butler.registry)
272 for _, dataset_type_node in self.pipeline_graph.iter_overall_inputs():
273 assert dataset_type_node is not None, "pipeline graph is resolved."
274 if self.butler.registry.registerDatasetType(dataset_type_node.dataset_type):
275 self.insert_datasets(dataset_type_node.dataset_type, register=False)
276 builder = AllDimensionsQuantumGraphBuilder(
277 self.pipeline_graph,
278 self.butler,
279 input_collections=[self.input_chain],
280 output_run=output_run,
281 )
282 if register_output_dataset_types:
283 self.pipeline_graph.register_dataset_types(self.butler)
284 return builder
286 def insert_datasets(
287 self, dataset_type: DatasetType | str, register: bool = True, *args: Any, **kwargs: Any
288 ) -> list[DatasetRef]:
289 """Insert input datasets into the test repository.
291 Parameters
292 ----------
293 dataset_type : `~lsst.daf.butler.DatasetType` or `str`
294 Dataset type or name. If a name, it must be included in the
295 pipeline graph.
296 register : `bool`, optional
297 Whether to register the dataset type. If `False`, the dataset type
298 must already be registered.
299 *args : `object`
300 Forwarded to `~lsst.daf.butler.query_data_ids`.
301 **kwargs : `object`
302 Forwarded to `~lsst.daf.butler.query_data_ids`.
304 Returns
305 -------
306 refs : `list` [ `lsst.daf.butler.DatasetRef` ]
307 References to the inserted datasets.
309 Notes
310 -----
311 For dataset types with dimensions that are queryable, this queries for
312 all data IDs in the repository (forwarding ``*args`` and ``**kwargs``
313 for e.g. ``where`` strings). For skypix dimensions, this queries for
314 both patches and visit-detector regions (forwarding `*args`` and
315 ``**kwargs`` to both) and uses all overlapping sky pixels. Dataset
316 types with a mix of skypix and queryable dimensions are not supported.
317 """
318 if isinstance(dataset_type, str):
319 self.pipeline_graph.resolve(self.butler.registry)
320 dataset_type = self.pipeline_graph.dataset_types[dataset_type].dataset_type
321 if register:
322 self.butler.registry.registerDatasetType(dataset_type)
323 dimensions = dataset_type.dimensions
324 if dataset_type.dimensions.skypix:
325 if len(dimensions) == 1:
326 (skypix_name,) = dimensions.skypix
327 pixelization = dimensions.universe.skypix_dimensions[skypix_name].pixelization
328 ranges = RangeSet()
329 for patch_record in self.butler.query_dimension_records(
330 "patch", *args, **kwargs, explain=False
331 ):
332 ranges |= pixelization.envelope(patch_record.region)
333 for vdr_record in self.butler.query_dimension_records(
334 "visit_detector_region", *args, **kwargs, explain=False
335 ):
336 ranges |= pixelization.envelope(vdr_record.region)
337 data_ids = []
338 for begin, end in ranges:
339 for index in range(begin, end):
340 data_ids.append(DataCoordinate.from_required_values(dimensions, (index,)))
341 else:
342 raise NotImplementedError(
343 "Can only generate data IDs for queryable dimensions and isolated skypix."
344 )
345 else:
346 data_ids = self.butler.query_data_ids(dimensions, *args, **kwargs, explain=False)
347 return self._insert_datasets_impl(dataset_type, data_ids)
349 @abstractmethod
350 def _insert_datasets_impl(
351 self, dataset_type: DatasetType, data_ids: list[DataCoordinate]
352 ) -> list[DatasetRef]:
353 """Insert datasets after their data IDs have been generated.
355 Parameters
356 ----------
357 dataset_type : `lsst.daf.butler.DatasetType`
358 Type of the datasets.
359 data_ids : `list` [ `lsst.daf.butler.DataCoordinate` ]
360 Data IDs of all datasets.
362 Returns
363 -------
364 refs : `list` [ `lsst.daf.butler.DatasetRef` ]
365 References to the new datasets.
366 """
367 raise NotImplementedError()
369 @abstractmethod
370 def make_single_quantum_executor(
371 self, qg: PredictedQuantumGraph
372 ) -> tuple[SingleQuantumExecutor, LimitedButler]:
373 """Make a single-quantum executor.
375 Parameters
376 ----------
377 qg : `..quantum_graph.PredictedQuantumGraph`
378 Graph whose quanta the executor must be capable of executing.
380 Returns
381 -------
382 executor : `..single_quantum_executor.SingleQuantumExecutor`
383 An executor for a single quantum.
384 butler : `lsst.daf.butler.LimitedButler`
385 The butler that the executor will write to.
386 """
387 raise NotImplementedError()
390class InMemoryRepo(MockRepo):
391 """A test helper that simulates a butler repository for task execution
392 without any disk I/O.
394 Parameters
395 ----------
396 *args : `str` or `lsst.resources.ResourcePath`
397 Butler YAML import files to load into the test repository.
398 registry_config : `lsst.daf.butler.RegistryConfig`, optional
399 Registry configuration for the repository.
400 dimension_config : `lsst.daf.butler.DimensionConfig`, optional
401 Dimension universe configuration for the repository.
402 input_run : `str`, optional
403 Name of a `~lsst.daf.butler.CollectionType.RUN` collection that will be
404 used as an input to quantum graph generation. Input datasets created
405 by the helper are added to this collection.
406 input_chain : `str`, optional
407 Name of a `~lsst.daf.butler.CollectionType.CHAINED` collection that
408 will be the direct input to quantum graph generation. This always
409 includes ``input_run``.
410 use_import_collections_as_input : `bool`, `str`, or \
411 `~collections.abc.Iterable` [ `str`], optional
412 Additional collections from YAML import files to include in
413 ``input_chain``, or `True` to include all such collections (in
414 chain-flattened lexicographical order).
415 data_root : convertible to `lsst.resources.ResourcePath`, optional
416 Root directory to join to each element in ``*args``. Defaults to
417 the `lsst.daf.butler.tests.registry_data` package.
419 Notes
420 -----
421 This helper maintains an `..pipeline_graph.PipelineGraph` and a
422 no-datastore butler backed by an in-memory SQLite database for use in
423 quantum graph generation. It creates a separate in-memory limited butler
424 for execution as needed.
425 """
427 def __init__(
428 self,
429 *args: str | ResourcePath,
430 registry_config: RegistryConfig | None = None,
431 dimension_config: DimensionConfig | None = None,
432 input_run: str = "input_run",
433 input_chain: str = "input_chain",
434 use_import_collections_as_input: bool | str | Iterable[str] = True,
435 data_root: ResourcePathExpression | None = "resource://lsst.daf.butler/tests/registry_data",
436 ):
437 if data_root is not None:
438 data_root = ResourcePath(data_root, forceDirectory=True)
439 args = tuple(data_root.join(arg) for arg in args)
440 butler = create_populated_sqlite_registry(
441 *args, registry_config=registry_config, dimension_config=dimension_config
442 )
443 if use_import_collections_as_input:
444 if use_import_collections_as_input is True:
445 use_import_collections_as_input = sorted(butler.collections.query("*", flatten_chains=True))
446 else:
447 use_import_collections_as_input = ()
448 super().__init__(
449 butler,
450 input_run=input_run,
451 input_chain=input_chain,
452 input_children=list(use_import_collections_as_input),
453 )
455 def _insert_datasets_impl(
456 self, dataset_type: DatasetType, data_ids: list[DataCoordinate]
457 ) -> list[DatasetRef]:
458 return self.butler.registry.insertDatasets(dataset_type, data_ids, run=self.input_run)
460 def make_limited_butler(self) -> InMemoryLimitedButler:
461 """Make a test limited butler for execution.
463 Returns
464 -------
465 limited_butler : `.InMemoryLimitedButler`
466 A limited butler that can be used for task execution.
468 Notes
469 -----
470 This queries the database-only butler used for quantum-graph generation
471 for all datasets in the ``input_chain`` collection, and populates the
472 limited butler with those that have a mock storage class. Other
473 datasets are ignored, so they will appear as though they were present
474 during quantum graph generation but absent during execution.
475 """
476 butler = InMemoryLimitedButler(self.butler.dimensions, self.butler.registry.queryDatasetTypes())
477 for ref in self.butler.query_all_datasets(self.input_chain):
478 if is_mock_name(ref.datasetType.storageClass_name):
479 butler.put(
480 MockDataset(
481 dataset_id=ref.id,
482 dataset_type=ref.datasetType.to_simple(),
483 data_id=dict(ref.dataId.mapping),
484 run=ref.run,
485 ),
486 ref,
487 )
488 return butler
490 def make_single_quantum_executor(
491 self, qg: PredictedQuantumGraph | None = None
492 ) -> tuple[SingleQuantumExecutor, InMemoryLimitedButler]:
493 """Make a single-quantum executor backed by a new limited butler.
495 Parameters
496 ----------
497 qg : `..quantum_graph.PredictedQuantumGraph`
498 Ignored by this implementation.
500 Returns
501 -------
502 executor : `..single_quantum_executor.SingleQuantumExecutor`
503 An executor for a single quantum.
504 butler : `.InMemoryLimitedButler`
505 The butler that the executor will write to.
506 """
507 butler = self.make_limited_butler()
508 return SingleQuantumExecutor(limited_butler_factory=butler.factory), butler
511class DirectButlerRepo(MockRepo):
512 """A test helper for task execution backed by a local direct butler.
514 Parameters
515 ----------
516 butler : `lsst.daf.butler.direct_butler.DirectButler`
517 Butler to write to.
518 *args : `str` or `lsst.resources.ResourcePath`
519 Butler YAML import files to load into the test repository.
520 input_run : `str`, optional
521 Name of a `~lsst.daf.butler.CollectionType.RUN` collection that will be
522 used as an input to quantum graph generation. Input datasets created
523 by the helper are added to this collection.
524 input_chain : `str`, optional
525 Name of a `~lsst.daf.butler.CollectionType.CHAINED` collection that
526 will be the direct input to quantum graph generation. This always
527 includes ``input_run``.
528 use_import_collections_as_input : `bool`, `str`, or \
529 `~collections.abc.Iterable` [ `str`], optional
530 Additional collections from YAML import files to include in
531 ``input_chain``, or `True` to include all such collections (in
532 chain-flattened lexicographical order).
533 data_root : convertible to `lsst.resources.ResourcePath`, optional
534 Root directory to join to each element in ``*args``. Defaults to
535 the `lsst.daf.butler.tests.registry_data` package.
537 Notes
538 -----
539 This helper maintains an `..pipeline_graph.PipelineGraph` and a
540 no-datastore butler backed by an in-memory SQLite database for use in
541 quantum graph generation. It creates a separate in-memory limited butler
542 for execution as needed.
543 """
545 def __init__(
546 self,
547 butler: Butler,
548 *args: str | ResourcePath,
549 input_run: str = "input_run",
550 input_chain: str = "input_chain",
551 use_import_collections_as_input: bool | str | Iterable[str] = True,
552 data_root: ResourcePathExpression | None = "resource://lsst.daf.butler/tests/registry_data",
553 ):
554 if data_root is not None:
555 data_root = ResourcePath(data_root, forceDirectory=True)
556 args = tuple(data_root.join(arg) for arg in args)
557 for arg in args:
558 butler.import_(filename=arg)
559 if use_import_collections_as_input:
560 if use_import_collections_as_input is True:
561 use_import_collections_as_input = sorted(butler.collections.query("*", flatten_chains=True))
562 else:
563 use_import_collections_as_input = ()
564 super().__init__(
565 butler,
566 input_run=input_run,
567 input_chain=input_chain,
568 input_children=list(use_import_collections_as_input),
569 )
571 @classmethod
572 @contextmanager
573 def make_temporary(
574 cls,
575 *args: str | ResourcePath,
576 input_run: str = "input_run",
577 input_chain: str = "input_chain",
578 use_import_collections_as_input: bool | str | Iterable[str] = True,
579 data_root: ResourcePathExpression | None = "resource://lsst.daf.butler/tests/registry_data",
580 **kwargs: Any,
581 ) -> Iterator[tuple[DirectButlerRepo, str]]:
582 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as root:
583 config = Butler.makeRepo(root, **kwargs)
584 butler = Butler.from_config(config, writeable=True)
585 yield (
586 cls(
587 butler,
588 *args,
589 input_run=input_run,
590 input_chain=input_chain,
591 use_import_collections_as_input=use_import_collections_as_input,
592 data_root=data_root,
593 ),
594 root,
595 )
597 def _insert_datasets_impl(
598 self, dataset_type: DatasetType, data_ids: list[DataCoordinate]
599 ) -> list[DatasetRef]:
600 if is_mock_name(dataset_type.storageClass_name):
601 refs: list[DatasetRef] = []
602 for data_id in data_ids:
603 data_id = self.butler.registry.expandDataId(data_id)
604 ref = DatasetRef(dataset_type, data_id, run=self.input_run)
605 self.butler.put(
606 MockDataset(
607 dataset_id=ref.id,
608 dataset_type=ref.datasetType.to_simple(),
609 data_id=dict(ref.dataId.mapping),
610 run=ref.run,
611 ),
612 ref,
613 )
614 refs.append(ref)
615 return refs
616 else:
617 return self.butler.registry.insertDatasets(dataset_type, data_ids, run=self.input_run)
619 def make_single_quantum_executor(
620 self, qg: PredictedQuantumGraph | None = None
621 ) -> tuple[SingleQuantumExecutor, Butler]:
622 """Make a single-quantum executor backed by a new limited butler.
624 Parameters
625 ----------
626 qg : `..quantum_graph.PredictedQuantumGraph`
627 Ignored by this implementation.
629 Returns
630 -------
631 executor : `..single_quantum_executor.SingleQuantumExecutor`
632 An executor for a single quantum.
633 butler : `lsst.daf.butler.Butler`
634 The butler that the executor will write to.
635 """
636 return SingleQuantumExecutor(limited_butler_factory=None, butler=self.butler), self.butler