Coverage for python/lsst/pipe/base/tests/simpleQGraph.py: 23%
178 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-19 01:44 -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/>.
28"""Bunch of common classes and methods for use in unit tests."""
30from __future__ import annotations
32__all__ = ["AddTask", "AddTaskConfig", "AddTaskFactoryMock"]
34import logging
35from collections.abc import Iterable, Mapping, MutableMapping
36from typing import TYPE_CHECKING, Any, cast
38import numpy
40import lsst.daf.butler.tests as butlerTests
41import lsst.pex.config as pexConfig
42from lsst.daf.butler import Butler, Config, DataId, DatasetRef, DatasetType, Formatter, LimitedButler
43from lsst.daf.butler.logging import ButlerLogRecords
44from lsst.resources import ResourcePath
45from lsst.utils import doImportType
46from lsst.utils.introspection import get_full_type_name
48from .. import connectionTypes as cT
49from .._instrument import Instrument
50from ..all_dimensions_quantum_graph_builder import AllDimensionsQuantumGraphBuilder
51from ..all_dimensions_quantum_graph_builder import DatasetQueryConstraintVariant as DSQVariant
52from ..automatic_connection_constants import (
53 CONFIG_INIT_OUTPUT_CONNECTION_NAME,
54 LOG_OUTPUT_CONNECTION_NAME,
55 METADATA_OUTPUT_CONNECTION_NAME,
56 PACKAGES_INIT_OUTPUT_NAME,
57 PACKAGES_INIT_OUTPUT_STORAGE_CLASS,
58)
59from ..config import PipelineTaskConfig
60from ..connections import PipelineTaskConnections
61from ..graph import QuantumGraph
62from ..pipeline import Pipeline
63from ..pipeline_graph import PipelineGraph, TaskNode
64from ..pipelineTask import PipelineTask
65from ..struct import Struct
66from ..task import _TASK_FULL_METADATA_TYPE
67from ..taskFactory import TaskFactory
69if TYPE_CHECKING:
70 from lsst.daf.butler import Registry
72_LOG = logging.getLogger(__name__)
75class SimpleInstrument(Instrument):
76 """Simple instrument class suitable for testing.
78 Parameters
79 ----------
80 *args : `~typing.Any`
81 Ignore parameters.
82 **kwargs : `~typing.Any`
83 Unused keyword arguments.
84 """
86 def __init__(self, *args: Any, **kwargs: Any):
87 pass
89 @staticmethod
90 def getName() -> str:
91 return "INSTRU"
93 def getRawFormatter(self, dataId: DataId) -> type[Formatter]:
94 return Formatter
96 def register(self, registry: Registry, *, update: bool = False) -> None:
97 pass
100class AddTaskConnections(
101 PipelineTaskConnections,
102 dimensions=("instrument", "detector"),
103 defaultTemplates={"in_tmpl": "_in", "out_tmpl": "_out"},
104):
105 """Connections for AddTask, has one input and two outputs,
106 plus one init output.
107 """
109 input = cT.Input(
110 name="add_dataset{in_tmpl}",
111 dimensions=["instrument", "detector"],
112 storageClass="NumpyArray",
113 doc="Input dataset type for this task",
114 )
115 output = cT.Output(
116 name="add_dataset{out_tmpl}",
117 dimensions=["instrument", "detector"],
118 storageClass="NumpyArray",
119 doc="Output dataset type for this task",
120 )
121 output2 = cT.Output(
122 name="add2_dataset{out_tmpl}",
123 dimensions=["instrument", "detector"],
124 storageClass="NumpyArray",
125 doc="Output dataset type for this task",
126 )
127 initout = cT.InitOutput(
128 name="add_init_output{out_tmpl}",
129 storageClass="NumpyArray",
130 doc="Init Output dataset type for this task",
131 )
134class AddTaskConfig(PipelineTaskConfig, pipelineConnections=AddTaskConnections):
135 """Config for AddTask."""
137 addend = pexConfig.Field[int](doc="amount to add", default=3)
140class AddTask(PipelineTask):
141 """Trivial PipelineTask for testing, has some extras useful for specific
142 unit tests.
143 """
145 ConfigClass = AddTaskConfig
146 _DefaultName = "add_task"
148 initout = numpy.array([999])
149 """InitOutputs for this task"""
151 taskFactory: AddTaskFactoryMock | None = None
152 """Factory that makes instances"""
154 def run(self, input: int) -> Struct:
155 if self.taskFactory:
156 # do some bookkeeping
157 if self.taskFactory.stopAt == self.taskFactory.countExec:
158 raise RuntimeError("pretend something bad happened")
159 self.taskFactory.countExec += 1
161 self.config = cast(AddTaskConfig, self.config)
162 self.metadata.add("add", self.config.addend)
163 output = input + self.config.addend
164 output2 = output + self.config.addend
165 _LOG.info("input = %s, output = %s, output2 = %s", input, output, output2)
166 return Struct(output=output, output2=output2)
169class AddTaskFactoryMock(TaskFactory):
170 """Special task factory that instantiates AddTask.
172 It also defines some bookkeeping variables used by AddTask to report
173 progress to unit tests.
175 Parameters
176 ----------
177 stopAt : `int`, optional
178 Number of times to call `run` before stopping.
179 """
181 def __init__(self, stopAt: int = -1):
182 self.countExec = 0 # incremented by AddTask
183 self.stopAt = stopAt # AddTask raises exception at this call to run()
185 def makeTask(
186 self,
187 task_node: TaskNode,
188 /,
189 butler: LimitedButler,
190 initInputRefs: Iterable[DatasetRef] | None,
191 ) -> PipelineTask:
192 task = task_node.task_class(config=task_node.config, initInputs=None, name=task_node.label)
193 task.taskFactory = self # type: ignore
194 return task
197def registerDatasetTypes(registry: Registry, pipeline: Pipeline | PipelineGraph) -> None:
198 """Register all dataset types used by tasks in a registry.
200 Copied and modified from `PreExecInit.initializeDatasetTypes`.
202 Parameters
203 ----------
204 registry : `~lsst.daf.butler.Registry`
205 Registry instance.
206 pipeline : `.Pipeline`, or `.pipeline_graph.PipelineGraph`
207 The pipeline whose dataset types should be registered, as a `.Pipeline`
208 instance or `.PipelineGraph` instance.
209 """
210 match pipeline:
211 case PipelineGraph() as pipeline_graph:
212 pass
213 case Pipeline():
214 pipeline_graph = pipeline.to_graph()
215 case _:
216 raise TypeError(f"Unexpected pipeline argument value: {pipeline}.")
217 pipeline_graph.resolve(registry)
218 dataset_types = [node.dataset_type for node in pipeline_graph.dataset_types.values()]
219 dataset_types.append(
220 DatasetType(
221 PACKAGES_INIT_OUTPUT_NAME,
222 {},
223 storageClass=PACKAGES_INIT_OUTPUT_STORAGE_CLASS,
224 universe=registry.dimensions,
225 )
226 )
227 for dataset_type in dataset_types:
228 _LOG.info("Registering %s with registry", dataset_type)
229 registry.registerDatasetType(dataset_type)
232def makeSimplePipeline(nQuanta: int, instrument: str | None = None) -> Pipeline:
233 """Make a simple Pipeline for tests.
235 This is called by `makeSimpleQGraph()` if no pipeline is passed to that
236 function. It can also be used to customize the pipeline used by
237 `makeSimpleQGraph()` function by calling this first and passing the result
238 to it.
240 Parameters
241 ----------
242 nQuanta : `int`
243 The number of quanta to add to the pipeline.
244 instrument : `str` or `None`, optional
245 The importable name of an instrument to be added to the pipeline or
246 if no instrument should be added then an empty string or `None`, by
247 default `None`.
249 Returns
250 -------
251 pipeline : `~lsst.pipe.base.Pipeline`
252 The created pipeline object.
253 """
254 pipeline = Pipeline("test pipeline")
255 # make a bunch of tasks that execute in well defined order (via data
256 # dependencies)
257 for lvl in range(nQuanta):
258 pipeline.addTask(AddTask, f"task{lvl}")
259 pipeline.addConfigOverride(f"task{lvl}", "connections.in_tmpl", lvl)
260 pipeline.addConfigOverride(f"task{lvl}", "connections.out_tmpl", lvl + 1)
261 if instrument:
262 pipeline.addInstrument(instrument)
263 return pipeline
266def makeSimpleButler(
267 root: str, run: str = "test", inMemory: bool = True, config: Config | str | None = None
268) -> Butler:
269 """Create new data butler instance.
271 Parameters
272 ----------
273 root : `str`
274 Path or URI to the root location of the new repository.
275 run : `str`, optional
276 Run collection name.
277 inMemory : `bool`, optional
278 If true make in-memory repository.
279 config : `~lsst.daf.butler.Config`, optional
280 Configuration to use for new Butler, if `None` then default
281 configuration is used. If ``inMemory`` is `True` then configuration
282 is updated to use SQLite registry and file datastore in ``root``.
284 Returns
285 -------
286 butler : `~lsst.daf.butler.Butler`
287 Data butler instance.
288 """
289 root_path = ResourcePath(root, forceDirectory=True)
290 if not root_path.isLocal:
291 raise ValueError(f"Only works with local root not {root_path}")
292 butler_config = Config()
293 if config:
294 butler_config.update(Config(config))
295 if not inMemory:
296 butler_config["registry", "db"] = f"sqlite:///{root_path.ospath}/gen3.sqlite"
297 butler_config["datastore", "cls"] = "lsst.daf.butler.datastores.fileDatastore.FileDatastore"
298 repo = butlerTests.makeTestRepo(str(root_path), {}, config=butler_config)
299 butler = Butler.from_config(butler=repo, run=run)
300 return butler
303def populateButler(
304 pipeline: Pipeline | PipelineGraph,
305 butler: Butler,
306 datasetTypes: dict[str | None, list[str]] | None = None,
307 instrument: str | None = None,
308) -> None:
309 """Populate data butler with data needed for test.
311 Initializes data butler with a bunch of items:
312 - registers dataset types which are defined by pipeline
313 - create dimensions data for (instrument, detector)
314 - adds datasets based on ``datasetTypes`` dictionary, if dictionary is
315 missing then a single dataset with type "add_dataset0" is added.
317 All datasets added to butler have ``dataId={instrument=instrument,
318 detector=0}`` where ``instrument`` is extracted from pipeline, "INSTR" is
319 used if pipeline is missing instrument definition. Type of the dataset is
320 guessed from dataset type name (assumes that pipeline is made of `AddTask`
321 tasks).
323 Parameters
324 ----------
325 pipeline : `.Pipeline` or `.pipeline_graph.PipelineGraph`
326 Pipeline or pipeline graph instance.
327 butler : `~lsst.daf.butler.Butler`
328 Data butler instance.
329 datasetTypes : `dict` [ `str`, `list` ], optional
330 Dictionary whose keys are collection names and values are lists of
331 dataset type names. By default a single dataset of type "add_dataset0"
332 is added to a ``butler.run`` collection.
333 instrument : `str`, optional
334 Fully-qualified name of the instrumemnt class (as it appears in a
335 pipeline). This is needed to propagate the instrument from the
336 original pipeline if a pipeline graph is passed instead.
337 """
338 # Add dataset types to registry
339 match pipeline:
340 case PipelineGraph() as pipeline_graph:
341 pass
342 case Pipeline():
343 pipeline_graph = pipeline.to_graph()
344 if instrument is None:
345 instrument = pipeline.getInstrument()
346 case _:
347 raise TypeError(f"Unexpected pipeline object: {pipeline!r}.")
349 registerDatasetTypes(butler.registry, pipeline_graph)
351 if instrument is not None:
352 instrument_class = doImportType(instrument)
353 instrumentName = cast(Instrument, instrument_class).getName()
354 instrumentClass = get_full_type_name(instrument_class)
355 else:
356 instrumentName = "INSTR"
357 instrumentClass = None
359 # Add all needed dimensions to registry
360 butler.registry.insertDimensionData("instrument", dict(name=instrumentName, class_name=instrumentClass))
361 butler.registry.insertDimensionData("detector", dict(instrument=instrumentName, id=0, full_name="det0"))
363 # Add inputs to butler.
364 if not datasetTypes:
365 datasetTypes = {None: ["add_dataset0"]}
366 for run, dsTypes in datasetTypes.items():
367 if run is not None:
368 butler.registry.registerRun(run)
369 for dsType in dsTypes:
370 if dsType == PACKAGES_INIT_OUTPUT_NAME:
371 # Version is intentionally inconsistent.
372 # Dict is convertible to Packages if Packages is installed.
373 data: Any = {"python": "9.9.99"}
374 butler.put(data, dsType, run=run)
375 else:
376 if dsType.endswith(CONFIG_INIT_OUTPUT_CONNECTION_NAME):
377 # find a config from matching task name or make a new one
378 taskLabel, _, _ = dsType.rpartition("_")
379 task_node = pipeline_graph.tasks.get(taskLabel)
380 if task_node is not None:
381 data = task_node.config
382 else:
383 data = AddTaskConfig()
384 elif dsType.endswith(METADATA_OUTPUT_CONNECTION_NAME):
385 data = _TASK_FULL_METADATA_TYPE()
386 elif dsType.endswith(LOG_OUTPUT_CONNECTION_NAME):
387 data = ButlerLogRecords.from_records([])
388 else:
389 data = numpy.array([0.0, 1.0, 2.0, 5.0])
390 butler.put(data, dsType, run=run, instrument=instrumentName, detector=0)
393def makeSimpleQGraph(
394 nQuanta: int = 5,
395 pipeline: Pipeline | PipelineGraph | None = None,
396 butler: Butler | None = None,
397 root: str | None = None,
398 callPopulateButler: bool = True,
399 run: str = "test",
400 instrument: str | None = None,
401 skipExistingIn: Any = None,
402 inMemory: bool = True,
403 userQuery: str = "",
404 datasetTypes: dict[str | None, list[str]] | None = None,
405 datasetQueryConstraint: DSQVariant = DSQVariant.ALL,
406 makeDatastoreRecords: bool = False,
407 bind: Mapping[str, Any] | None = None,
408 metadata: MutableMapping[str, Any] | None = None,
409) -> tuple[Butler, QuantumGraph]:
410 """Make simple `QuantumGraph` for tests.
412 Makes simple one-task pipeline with AddTask, sets up in-memory registry
413 and butler, fills them with minimal data, and generates QuantumGraph with
414 all of that.
416 Parameters
417 ----------
418 nQuanta : `int`
419 Number of quanta in a graph, only used if ``pipeline`` is None.
420 pipeline : `.Pipeline` or `.pipeline_graph.PipelineGraph`
421 Pipeline or pipeline graph to build the graph from. If `None`, a
422 pipeline is made with `AddTask` and default `AddTaskConfig`.
423 butler : `~lsst.daf.butler.Butler`, optional
424 Data butler instance, if None then new data butler is created by
425 calling `makeSimpleButler`.
426 root : `str`
427 Path or URI to the root location of the new repository. Only used if
428 ``butler`` is None.
429 callPopulateButler : `bool`, optional
430 If True insert datasets into the butler prior to building a graph.
431 If False butler argument must not be None, and must be pre-populated.
432 Defaults to True.
433 run : `str`, optional
434 Name of the RUN collection to add to butler, only used if ``butler``
435 is None.
436 instrument : `str` or `None`, optional
437 The importable name of an instrument to be added to the pipeline or
438 if no instrument should be added then an empty string or `None`, by
439 default `None`. Only used if ``pipeline`` is `None`.
440 skipExistingIn : `~typing.Any`
441 Expressions representing the collections to search for existing
442 output datasets that should be skipped. See
443 :ref:`daf_butler_ordered_collection_searches`.
444 inMemory : `bool`, optional
445 If true make in-memory repository, only used if ``butler`` is `None`.
446 userQuery : `str`, optional
447 The user query to pass to ``makeGraph``, by default an empty string.
448 datasetTypes : `dict` [ `str`, `list` ], optional
449 Dictionary whose keys are collection names and values are lists of
450 dataset type names. By default a single dataset of type "add_dataset0"
451 is added to a ``butler.run`` collection.
452 datasetQueryConstraint : `DatasetQueryConstraintVariant`
453 The query constraint variant that should be used to constrain the
454 query based on dataset existence, defaults to
455 `DatasetQueryConstraintVariant.ALL`.
456 makeDatastoreRecords : `bool`, optional
457 If `True` then add datastore records to generated quanta.
458 bind : `~collections.abc.Mapping`, optional
459 Mapping containing literal values that should be injected into the
460 ``userQuery`` expression, keyed by the identifiers they replace.
461 metadata : `~collections.abc.Mapping`, optional
462 Optional graph metadata.
464 Returns
465 -------
466 butler : `~lsst.daf.butler.Butler`
467 Butler instance.
468 qgraph : `~lsst.pipe.base.QuantumGraph`
469 Quantum graph instance.
470 """
471 match pipeline:
472 case PipelineGraph() as pipeline_graph:
473 pass
474 case Pipeline():
475 pipeline_graph = pipeline.to_graph()
476 if instrument is None:
477 instrument = pipeline.getInstrument()
478 case None:
479 pipeline_graph = makeSimplePipeline(nQuanta=nQuanta, instrument=instrument).to_graph()
480 case _:
481 raise TypeError(f"Unexpected pipeline object: {pipeline!r}.")
483 if butler is None:
484 if root is None:
485 raise ValueError("Must provide `root` when `butler` is None")
486 if callPopulateButler is False:
487 raise ValueError("populateButler can only be False when butler is supplied as an argument")
488 butler = makeSimpleButler(root, run=run, inMemory=inMemory)
490 if callPopulateButler:
491 populateButler(pipeline_graph, butler, datasetTypes=datasetTypes, instrument=instrument)
493 # Make the graph
494 _LOG.debug(
495 "Instantiating QuantumGraphBuilder, "
496 "skip_existing_in=%s, input_collections=%r, output_run=%r, where=%r, bind=%s.",
497 skipExistingIn,
498 butler.collections.defaults,
499 run,
500 userQuery,
501 bind,
502 )
503 if not run:
504 assert butler.run is not None, "Butler must have run defined"
505 run = butler.run
506 builder = AllDimensionsQuantumGraphBuilder(
507 pipeline_graph,
508 butler,
509 skip_existing_in=skipExistingIn if skipExistingIn is not None else [],
510 input_collections=butler.collections.defaults,
511 output_run=run,
512 where=userQuery,
513 bind=bind,
514 dataset_query_constraint=datasetQueryConstraint,
515 )
516 _LOG.debug("Calling QuantumGraphBuilder.build.")
517 if not metadata:
518 metadata = {}
519 metadata["output_run"] = run
521 qgraph = builder.build(metadata=metadata, attach_datastore_records=makeDatastoreRecords)
523 return butler, qgraph