Coverage for python / lsst / pipe / base / quantum_graph_builder.py: 16%
476 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-17 16:30 -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"""The base class for the QuantumGraph-generation algorithm and various
29helper classes.
30"""
32from __future__ import annotations
34__all__ = (
35 "EmptyDimensionsDatasets",
36 "OutputExistsError",
37 "PrerequisiteMissingError",
38 "QuantumGraphBuilder",
39 "QuantumGraphBuilderError",
40)
42import dataclasses
43import operator
44from abc import ABC, abstractmethod
45from collections import defaultdict
46from collections.abc import Iterable, Mapping, Sequence
47from typing import TYPE_CHECKING, Any, cast, final
49from lsst.daf.butler import (
50 Butler,
51 CollectionType,
52 DataCoordinate,
53 DatasetRef,
54 DatasetType,
55 DimensionDataAttacher,
56 DimensionUniverse,
57 NamedKeyDict,
58 NamedKeyMapping,
59 Quantum,
60)
61from lsst.daf.butler._rubin import generate_uuidv7
62from lsst.daf.butler.datastore.record_data import DatastoreRecordData
63from lsst.daf.butler.registry import MissingCollectionError, MissingDatasetTypeError
64from lsst.utils.logging import LsstLogAdapter, getLogger
65from lsst.utils.timer import timeMethod
67from . import automatic_connection_constants as acc
68from ._status import NoWorkFound
69from ._task_metadata import TaskMetadata
70from .connections import AdjustQuantumHelper, QuantaAdjuster
71from .pipeline_graph import Edge, PipelineGraph, TaskNode
72from .prerequisite_helpers import PrerequisiteInfo, SkyPixBoundsBuilder, TimespanBuilder
73from .quantum_graph_skeleton import (
74 DatasetKey,
75 PrerequisiteDatasetKey,
76 QuantumGraphSkeleton,
77 QuantumKey,
78 TaskInitKey,
79)
81if TYPE_CHECKING:
82 from .graph import QuantumGraph
83 from .pipeline import TaskDef
84 from .quantum_graph import PredictedDatasetModel, PredictedQuantumGraphComponents
87class QuantumGraphBuilderError(Exception):
88 """Base class for exceptions generated by QuantumGraphBuilder."""
90 pass
93class OutputExistsError(QuantumGraphBuilderError):
94 """Exception generated when output datasets already exist."""
96 pass
99class PrerequisiteMissingError(QuantumGraphBuilderError):
100 """Exception generated when a prerequisite dataset does not exist."""
102 pass
105class InitInputMissingError(QuantumGraphBuilderError):
106 """Exception generated when an init-input dataset does not exist."""
108 pass
111class QuantumGraphBuilder(ABC):
112 """An abstract base class for building `.QuantumGraph` objects from a
113 pipeline.
115 Parameters
116 ----------
117 pipeline_graph : `.pipeline_graph.PipelineGraph`
118 Pipeline to build a `.QuantumGraph` from, as a graph. Will be resolved
119 in-place with the given butler (any existing resolution is ignored).
120 butler : `lsst.daf.butler.Butler`
121 Client for the data repository. Should be read-only.
122 input_collections : `~collections.abc.Sequence` [ `str` ], optional
123 Collections to search for overall-input datasets. If not provided,
124 ``butler.collections`` is used (and must not be empty).
125 output_run : `str`, optional
126 Output `~lsst.daf.butler.CollectionType.RUN` collection. If not
127 provided, ``butler.run`` is used (and must not be `None`).
128 skip_existing_in : `~collections.abc.Sequence` [ `str` ], optional
129 Collections to search for outputs that already exist for the purpose of
130 skipping quanta that have already been run.
131 clobber : `bool`, optional
132 Whether to raise if predicted outputs already exist in ``output_run``
133 (not including those quanta that would be skipped because they've
134 already been run). This never actually clobbers outputs; it just
135 informs the graph generation algorithm whether execution will run with
136 clobbering enabled. This is ignored if ``output_run`` does not exist.
138 Notes
139 -----
140 Constructing a `QuantumGraphBuilder` will run queries for existing datasets
141 with empty data IDs (including but not limited to init inputs and outputs),
142 in addition to resolving the given pipeline graph and testing for existence
143 of the ``output`` run collection.
145 The `build` method splits the pipeline graph into independent subgraphs,
146 then calls the abstract method `process_subgraph` on each, to allow
147 concrete implementations to populate the rough graph structure (the
148 `~.quantum_graph_skeleton.QuantumGraphSkeleton` class), including searching
149 for existing datasets. The `build` method then:
151 - assembles `lsst.daf.butler.Quantum` instances from all data IDs in the
152 skeleton;
153 - looks for existing outputs found in ``skip_existing_in`` to see if any
154 quanta should be skipped;
155 - calls `PipelineTaskConnections.adjustQuantum` on all quanta, adjusting
156 downstream quanta appropriately when preliminary predicted outputs are
157 rejected (pruning nodes that will not have the inputs they need to run);
158 - attaches datastore records and registry dataset types to the graph.
160 In addition to implementing `process_subgraph`, derived classes are
161 generally expected to add new construction keyword-only arguments to
162 control the data IDs of the quantum graph, while forwarding all of the
163 arguments defined in the base class to `super`.
164 """
166 def __init__(
167 self,
168 pipeline_graph: PipelineGraph,
169 butler: Butler,
170 *,
171 input_collections: Sequence[str] | None = None,
172 output_run: str | None = None,
173 skip_existing_in: Sequence[str] = (),
174 clobber: bool = False,
175 ):
176 self.log = getLogger(__name__)
177 self.metadata = TaskMetadata()
178 self._pipeline_graph = pipeline_graph
179 if input_collections is None:
180 input_collections = butler.collections.defaults
181 if not input_collections:
182 raise ValueError("No input collections provided.")
183 self.input_collections = input_collections
184 if output_run is None:
185 output_run = butler.run
186 if not output_run:
187 raise ValueError("No output RUN collection provided.")
188 self.butler = butler.clone(collections=input_collections)
189 self.output_run = output_run
190 self.skip_existing_in = skip_existing_in
191 self.empty_data_id = DataCoordinate.make_empty(butler.dimensions)
192 self.clobber = clobber
193 # See whether the output run already exists.
194 self.output_run_exists = False
195 try:
196 if self.butler.registry.getCollectionType(self.output_run) is not CollectionType.RUN:
197 raise RuntimeError(f"{self.output_run!r} is not a RUN collection.")
198 self.output_run_exists = True
199 except MissingCollectionError:
200 # If the run doesn't exist we never need to clobber. This is not
201 # an error so you can run with clobber=True the first time you
202 # attempt some processing as well as all subsequent times, instead
203 # of forcing the user to make the first attempt different.
204 self.clobber = False
205 # We need to know whether the skip_existing_in collection sequence
206 # starts with the output run collection, as an optimization to avoid
207 # queries later.
208 try:
209 skip_existing_in_flat = self.butler.collections.query(self.skip_existing_in, flatten_chains=True)
210 except MissingCollectionError:
211 skip_existing_in_flat = []
212 if not skip_existing_in_flat:
213 self.skip_existing_in = []
214 if self.skip_existing_in and self.output_run_exists:
215 self.skip_existing_starts_with_output_run = self.output_run == skip_existing_in_flat[0]
216 else:
217 self.skip_existing_starts_with_output_run = False
218 try:
219 packages_storage_class = butler.get_dataset_type(acc.PACKAGES_INIT_OUTPUT_NAME).storageClass_name
220 except MissingDatasetTypeError:
221 packages_storage_class = acc.PACKAGES_INIT_OUTPUT_STORAGE_CLASS
222 self._global_init_output_types = {
223 acc.PACKAGES_INIT_OUTPUT_NAME: DatasetType(
224 acc.PACKAGES_INIT_OUTPUT_NAME,
225 self.universe.empty,
226 packages_storage_class,
227 )
228 }
229 with self.butler.registry.caching_context():
230 self._pipeline_graph.resolve(self.butler.registry)
231 self.empty_dimensions_datasets = self._find_empty_dimension_datasets()
232 self.prerequisite_info = {
233 task_node.label: PrerequisiteInfo(task_node, self._pipeline_graph)
234 for task_node in pipeline_graph.tasks.values()
235 }
237 log: LsstLogAdapter
238 """Logger to use for all quantum-graph generation messages.
240 General and per-task status messages should be logged at `~logging.INFO`
241 level or higher, per-dataset-type status messages should be logged at
242 `~lsst.utils.logging.VERBOSE` or higher, and per-data-ID status messages
243 should be logged at `logging.DEBUG` or higher.
244 """
246 metadata: TaskMetadata
247 """Metadata to store in the QuantumGraph.
249 The `TaskMetadata` class is used here primarily in order to enable
250 resource-usage collection with the `lsst.utils.timer.timeMethod` decorator.
251 """
253 butler: Butler
254 """Client for the data repository.
256 Should be read-only.
257 """
259 input_collections: Sequence[str]
260 """Collections to search for overall-input datasets.
261 """
263 output_run: str
264 """Output `~lsst.daf.butler.CollectionType.RUN` collection.
265 """
267 skip_existing_in: Sequence[str]
268 """Collections to search for outputs that already exist for the purpose
269 of skipping quanta that have already been run.
270 """
272 clobber: bool
273 """Whether to raise if predicted outputs already exist in ``output_run``
275 This never actually clobbers outputs; it just informs the graph generation
276 algorithm whether execution will run with clobbering enabled. This is
277 always `False` if `output_run_exists` is `False`.
278 """
280 empty_data_id: DataCoordinate
281 """An empty data ID in the data repository's dimension universe.
282 """
284 output_run_exists: bool
285 """Whether the output run exists in the data repository already.
286 """
288 skip_existing_starts_with_output_run: bool
289 """Whether the `skip_existing_in` sequence begins with `output_run`.
291 If this is true, any dataset found in `output_run` can be used to
292 short-circuit queries in `skip_existing_in`.
293 """
295 empty_dimensions_datasets: EmptyDimensionsDatasets
296 """Struct holding datasets with empty dimensions that have already been
297 found in the data repository.
298 """
300 prerequisite_info: Mapping[str, PrerequisiteInfo]
301 """Helper objects for finding prerequisite inputs, organized by task label.
303 Subclasses that find prerequisites should remove the
304 covered `~prerequisite_helpers.PrerequisiteFinder` objects from this
305 attribute.
306 """
308 @property
309 def universe(self) -> DimensionUniverse:
310 """Definitions of all data dimensions."""
311 return self.butler.dimensions
313 @final
314 @timeMethod
315 def build(
316 self, metadata: Mapping[str, Any] | None = None, attach_datastore_records: bool = True
317 ) -> QuantumGraph:
318 """Build the quantum graph, returning an old `QuantumGraph` instance.
320 Parameters
321 ----------
322 metadata : `~collections.abc.Mapping`, optional
323 Flexible metadata to add to the quantum graph.
324 attach_datastore_records : `bool`, optional
325 Whether to include datastore records in the graph. Required for
326 `lsst.daf.butler.QuantumBackedButler` execution.
328 Returns
329 -------
330 quantum_graph : `.QuantumGraph`
331 DAG describing processing to be performed.
333 Notes
334 -----
335 External code is expected to construct a `QuantumGraphBuilder` and then
336 call this method exactly once. See class documentation for details on
337 what it does.
338 """
339 skeleton = self._build_skeleton(attach_datastore_records=attach_datastore_records)
340 if metadata is None:
341 metadata = {
342 "input": list(self.input_collections),
343 "output_run": self.output_run,
344 }
345 return self._construct_quantum_graph(skeleton, metadata)
347 def finish(
348 self,
349 output: str | None = None,
350 metadata: Mapping[str, Any] | None = None,
351 attach_datastore_records: bool = True,
352 ) -> PredictedQuantumGraphComponents:
353 """Return quantum graph components that can be used to save or
354 construct a `PredictedQuantumGraph` instance.
356 Parameters
357 ----------
358 output : `str` or `None`, optional
359 Output `~lsst.daf.butler.CollectionType.CHAINED` collection that
360 combines the input and output collections.
361 metadata : `~collections.abc.Mapping`, optional
362 Mapping of JSON-friendly metadata. Collection information, the
363 current user, and the current timestamp are automatically
364 included.
365 attach_datastore_records : `bool`, optional
366 Whether to include datastore records for overall inputs for
367 `~lsst.daf.butler.QuantumBackedButler`.
369 Returns
370 -------
371 components : `.quantum_graph.PredictedQuantumGraphComponents`
372 Components that can be used to construct a graph object and/or save
373 it to disk.
374 """
375 skeleton = self._build_skeleton(attach_datastore_records=attach_datastore_records)
376 return self._construct_components(skeleton, output=output, metadata=metadata)
378 def _build_skeleton(self, attach_datastore_records: bool = True) -> QuantumGraphSkeleton:
379 """Build a complete skeleton for the quantum graph.
381 Parameters
382 ----------
383 metadata : `~collections.abc.Mapping`, optional
384 Flexible metadata to add to the quantum graph.
385 attach_datastore_records : `bool`, optional
386 Whether to include datastore records in the graph. Required for
387 `lsst.daf.butler.QuantumBackedButler` execution.
389 Returns
390 -------
391 quantum_graph_skeleton : `QuantumGraphSkeleton`
392 DAG describing processing to be performed.
393 """
394 with self.butler.registry.caching_context():
395 full_skeleton = QuantumGraphSkeleton(self._pipeline_graph.tasks)
396 subgraphs = list(self._pipeline_graph.split_independent())
397 for i, subgraph in enumerate(subgraphs):
398 self.log.info(
399 "Processing pipeline subgraph %d of %d with %d task(s).",
400 i + 1,
401 len(subgraphs),
402 len(subgraph.tasks),
403 )
404 self.log.verbose("Subgraph tasks: [%s]", ", ".join(label for label in subgraph.tasks))
405 subgraph_skeleton = self.process_subgraph(subgraph)
406 full_skeleton.update(subgraph_skeleton)
407 # Loop over tasks to apply skip-existing logic and add missing
408 # prerequisites. The pipeline graph must be topologically sorted,
409 # so a quantum is only processed after any quantum that provides
410 # its inputs has been processed.
411 skipped_quanta: dict[str, list[QuantumKey]] = {}
412 for task_node in self._pipeline_graph.tasks.values():
413 skipped_quanta[task_node.label] = self._resolve_task_quanta(task_node, full_skeleton)
414 # Add any dimension records not handled by the subclass, and
415 # aggregate any that were added directly to data IDs.
416 full_skeleton.attach_dimension_records(self.butler, self._pipeline_graph.get_all_dimensions())
417 # Loop over tasks again to run the adjust hooks.
418 for task_node in self._pipeline_graph.tasks.values():
419 self._adjust_task_quanta(task_node, full_skeleton, skipped_quanta[task_node.label])
420 # Add global init-outputs to the skeleton.
421 for dataset_type in self._global_init_output_types.values():
422 dataset_key = full_skeleton.add_dataset_node(
423 dataset_type.name, self.empty_data_id, is_global_init_output=True
424 )
425 ref = self.empty_dimensions_datasets.outputs_in_the_way.get(dataset_key)
426 if ref is None:
427 ref = DatasetRef(dataset_type, self.empty_data_id, run=self.output_run)
428 full_skeleton.set_dataset_ref(ref, dataset_key)
429 # Remove dataset nodes with no edges that are not global init
430 # outputs, which are generally overall-inputs whose original quanta
431 # end up skipped or with no work to do (we can't remove these along
432 # with the quanta because no quantum knows if its the only
433 # consumer).
434 full_skeleton.remove_orphan_datasets()
435 if attach_datastore_records:
436 self._attach_datastore_records(full_skeleton)
437 return full_skeleton
439 @abstractmethod
440 def process_subgraph(self, subgraph: PipelineGraph) -> QuantumGraphSkeleton:
441 """Build the rough structure for an independent subset of the
442 `.QuantumGraph` and query for relevant existing datasets.
444 Parameters
445 ----------
446 subgraph : `.pipeline_graph.PipelineGraph`
447 Subset of the pipeline graph that should be processed by this call.
448 This is always resolved and topologically sorted. It should not be
449 modified.
451 Returns
452 -------
453 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
454 Class representing an initial quantum graph. See
455 `.quantum_graph_skeleton.QuantumGraphSkeleton` docs for details.
456 After this is returned, the object may be modified in-place in
457 unspecified ways.
459 Notes
460 -----
461 The `.quantum_graph_skeleton.QuantumGraphSkeleton` should associate
462 `lsst.daf.butler.DatasetRef` objects with nodes for existing datasets.
463 In particular:
465 - `.quantum_graph_skeleton.QuantumGraphSkeleton.set_dataset_ref` must
466 be used to associate existing datasets with all overall-input dataset
467 nodes in the skeleton by querying `input_collections`. This includes
468 all standard input nodes and any prerequisite nodes added by the
469 method (prerequisite nodes may also be left out entirely, as the base
470 class can add them later, albeit possibly less efficiently).
471 - `.quantum_graph_skeleton.QuantumGraphSkeleton.set_output_for_skip`
472 must be used to associate existing datasets with output dataset nodes
473 by querying `skip_existing_in`.
474 - `.quantum_graph_skeleton.QuantumGraphSkeleton.add_output_in_the_way`
475 must be used to associated existing outputs with output dataset nodes
476 by querying `output_run` if `output_run_exists` is `True`. Note that
477 the presence of such datasets is not automatically an error, even if
478 `clobber` is `False`, as these may be quanta that will be skipped.
480 `lsst.daf.butler.DatasetRef` objects for existing datasets with empty
481 data IDs in all of the above categories may be found in the
482 `empty_dimensions_datasets` attribute, as these are queried for prior
483 to this call by the base class, but associating them with graph nodes
484 is still this method's responsibility.
486 Dataset types should never be components and should always use the
487 "common" storage class definition in `pipeline_graph.DatasetTypeNode`
488 (which is the data repository definition when the dataset type is
489 registered).
490 """
491 raise NotImplementedError()
493 @final
494 @timeMethod
495 def _resolve_task_quanta(self, task_node: TaskNode, skeleton: QuantumGraphSkeleton) -> list[QuantumKey]:
496 """Process the quanta for one task in a skeleton graph to skip those
497 that have already completed and add missing prerequisite inputs.
499 Parameters
500 ----------
501 task_node : `pipeline_graph.TaskNode`
502 Node for this task in the pipeline graph.
503 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
504 Preliminary quantum graph, to be modified in-place.
506 Returns
507 -------
508 skipped_quanta : `list` [ `.quantum_skeleton_graph.QuantumKey` ]
509 Keys of quanta that were already skipped because their metadata
510 already exists in a ``skip_existing_in`` collections.
512 Notes
513 -----
514 This method modifies ``skeleton`` in-place in several ways:
516 - It associates a `lsst.daf.butler.DatasetRef` with all output datasets
517 and drops input dataset nodes that do not have a
518 `lsst.daf.butler.DatasetRef` already. This ensures producing and
519 consuming tasks start from the same `lsst.daf.butler.DatasetRef`.
520 - It removes quantum nodes that are to be skipped because their outputs
521 already exist in `skip_existing_in`. It also marks their outputs
522 as no longer in the way.
523 - It adds prerequisite dataset nodes and edges that connect them to the
524 quanta that consume them.
525 """
526 # Extract the helper object for the prerequisite inputs of this task,
527 # and tell it to prepare to construct skypix bounds and timespans for
528 # each quantum (these will automatically do nothing if nothing needs
529 # those bounds).
530 task_prerequisite_info = self.prerequisite_info[task_node.label]
531 task_prerequisite_info.update_bounds()
532 # Loop over all quanta for this task, remembering the ones we've
533 # gotten rid of.
534 skipped_quanta = []
535 for quantum_key in skeleton.get_quanta(task_node.label):
536 if self._skip_quantum_if_metadata_exists(task_node, quantum_key, skeleton):
537 skipped_quanta.append(quantum_key)
538 continue
539 quantum_data_id = skeleton[quantum_key]["data_id"]
540 skypix_bounds_builder = task_prerequisite_info.bounds.make_skypix_bounds_builder(quantum_data_id)
541 timespan_builder = task_prerequisite_info.bounds.make_timespan_builder(quantum_data_id)
542 self._update_quantum_for_adjust(
543 quantum_key,
544 skeleton,
545 task_prerequisite_info,
546 skypix_bounds_builder,
547 timespan_builder,
548 )
549 for skipped_quantum in skipped_quanta:
550 skeleton.remove_quantum_node(skipped_quantum, remove_outputs=False)
551 return skipped_quanta
553 @final
554 @timeMethod
555 def _adjust_task_quanta(
556 self, task_node: TaskNode, skeleton: QuantumGraphSkeleton, skipped_quanta: list[QuantumKey]
557 ) -> None:
558 """Process the quanta for one task in a skeleton graph by calling the
559 ``adjust_all_quanta`` and ``adjustQuantum`` hooks.
561 Parameters
562 ----------
563 task_node : `pipeline_graph.TaskNode`
564 Node for this task in the pipeline graph.
565 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
566 Preliminary quantum graph, to be modified in-place.
567 skipped_quanta : `list` [ `.quantum_skeleton_graph.QuantumKey` ]
568 Keys of quanta that were already skipped because their metadata
569 already exists in a ``skip_existing_in`` collections.
571 Notes
572 -----
573 This method modifies ``skeleton`` in-place in several ways:
575 - It adds "inputs", "outputs", and "init_inputs" attributes to the
576 quantum nodes, holding the same `NamedValueMapping` objects needed to
577 construct an actual `Quantum` instances.
578 - It removes quantum nodes whose
579 `~PipelineTaskConnections.adjustQuantum` calls raise `NoWorkFound` or
580 predict no outputs;
581 - It removes the nodes of output datasets that are "adjusted away".
582 - It removes the edges of input datasets that are "adjusted away".
584 The difference between how adjusted inputs and outputs are handled
585 reflects the fact that many quanta can share the same input, but only
586 one produces each output. This can lead to the graph having
587 superfluous isolated nodes after processing is complete, but these
588 should only be removed after all the quanta from all tasks have been
589 processed.
590 """
591 # Give the task a chance to adjust all quanta together. This
592 # operates directly on the skeleton (via a the 'adjuster', which
593 # is just an interface adapter).
594 adjuster = QuantaAdjuster(task_node.label, self._pipeline_graph, skeleton, self.butler)
595 task_node.get_connections().adjust_all_quanta(adjuster)
596 # Loop over all quanta again, remembering those we get rid of in other
597 # ways.
598 no_work_quanta = []
599 for quantum_key in skeleton.get_quanta(task_node.label):
600 adjusted_outputs = self._adapt_quantum_outputs(task_node, quantum_key, skeleton)
601 adjusted_inputs = self._adapt_quantum_inputs(task_node, quantum_key, skeleton)
602 # Give the task's Connections class an opportunity to remove
603 # some inputs, or complain if they are unacceptable. This will
604 # raise if one of the check conditions is not met, which is the
605 # intended behavior.
606 helper = AdjustQuantumHelper(inputs=adjusted_inputs, outputs=adjusted_outputs)
607 quantum_data_id = skeleton[quantum_key]["data_id"]
608 try:
609 helper.adjust_in_place(task_node.get_connections(), task_node.label, quantum_data_id)
610 except NoWorkFound as err:
611 # Do not generate this quantum; it would not produce any
612 # outputs. Remove it and all of the outputs it might have
613 # produced from the skeleton.
614 try:
615 _, connection_name, _ = err.args
616 details = f"not enough datasets for connection {connection_name}."
617 except ValueError:
618 details = str(err)
619 self.log.debug(
620 "No work found for quantum %s of task %s: %s",
621 quantum_key.data_id_values,
622 quantum_key.task_label,
623 details,
624 )
625 no_work_quanta.append(quantum_key)
626 continue
627 if helper.outputs_adjusted:
628 if not any(adjusted_refs for adjusted_refs in helper.outputs.values()):
629 # No outputs also means we don't generate this quantum.
630 self.log.debug(
631 "No outputs predicted for quantum %s of task %s.",
632 quantum_key.data_id_values,
633 quantum_key.task_label,
634 )
635 no_work_quanta.append(quantum_key)
636 continue
637 # Remove output nodes that were not retained by
638 # adjustQuantum.
639 skeleton.remove_dataset_nodes(
640 self._find_removed(skeleton.iter_outputs_of(quantum_key), helper.outputs)
641 )
642 if helper.inputs_adjusted:
643 if not any(bool(adjusted_refs) for adjusted_refs in helper.inputs.values()):
644 raise QuantumGraphBuilderError(
645 f"adjustQuantum implementation for {task_node.label}@{quantum_key.data_id_values} "
646 "returned outputs but no inputs."
647 )
648 # Remove input dataset edges that were not retained by
649 # adjustQuantum. We can't remove the input dataset nodes
650 # because some other quantum might still want them.
651 skeleton.remove_input_edges(
652 quantum_key, self._find_removed(skeleton.iter_inputs_of(quantum_key), helper.inputs)
653 )
654 # Save the adjusted inputs and outputs to the quantum node's
655 # state so we don't have to regenerate those data structures
656 # from the graph.
657 skeleton[quantum_key]["inputs"] = helper.inputs
658 skeleton[quantum_key]["outputs"] = helper.outputs
659 for no_work_quantum in no_work_quanta:
660 skeleton.remove_quantum_node(no_work_quantum, remove_outputs=True)
661 remaining_quanta = skeleton.get_quanta(task_node.label)
662 self._resolve_task_init(task_node, skeleton, bool(skipped_quanta))
663 message_terms = []
664 if no_work_quanta:
665 message_terms.append(f"{len(no_work_quanta)} had no work to do")
666 if skipped_quanta:
667 message_terms.append(f"{len(skipped_quanta)} previously succeeded")
668 if adjuster.n_removed:
669 message_terms.append(f"{adjuster.n_removed} removed by adjust_all_quanta")
670 message_parenthetical = f" ({', '.join(message_terms)})" if message_terms else ""
671 if remaining_quanta:
672 self.log.info(
673 "Generated %s for task %s%s.",
674 _quantum_or_quanta(len(remaining_quanta)),
675 task_node.label,
676 message_parenthetical,
677 )
678 else:
679 self.log.info(
680 "Dropping task %s because no quanta remain%s.", task_node.label, message_parenthetical
681 )
682 skeleton.remove_task(task_node.label)
684 def _skip_quantum_if_metadata_exists(
685 self, task_node: TaskNode, quantum_key: QuantumKey, skeleton: QuantumGraphSkeleton
686 ) -> bool:
687 """Identify and drop quanta that should be skipped because their
688 metadata datasets already exist.
690 Parameters
691 ----------
692 task_node : `pipeline_graph.TaskNode`
693 Node for this task in the pipeline graph.
694 quantum_key : `QuantumKey`
695 Identifier for this quantum in the graph.
696 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
697 Preliminary quantum graph, to be modified in-place.
699 Returns
700 -------
701 skipped : `bool`
702 `True` if the quantum is being skipped and has been removed from
703 the graph, `False` otherwise.
705 Notes
706 -----
707 If the metadata dataset for this quantum exists in the
708 `skip_existing_in` collections, the quantum will be skipped. This
709 causes the quantum node to be removed from the graph. Dataset nodes
710 that were previously the outputs of this quantum will be associated
711 with `lsst.daf.butler.DatasetRef` objects that were found in
712 ``skip_existing_in``, or will be removed if there is no such dataset
713 there. Any output dataset in `output_run` will be removed from the
714 "output in the way" category.
715 """
716 metadata_dataset_key = DatasetKey(
717 task_node.metadata_output.parent_dataset_type_name, quantum_key.data_id_values
718 )
719 if skeleton.get_output_for_skip(metadata_dataset_key):
720 # This quantum's metadata is already present in the the
721 # skip_existing_in collections; we'll skip it. But the presence of
722 # the metadata dataset doesn't guarantee that all of the other
723 # outputs we predicted are present; we have to check.
724 for output_dataset_key in list(skeleton.iter_outputs_of(quantum_key)):
725 # If this dataset was "in the way" (i.e. already in the
726 # output run), it isn't anymore.
727 skeleton.discard_output_in_the_way(output_dataset_key)
728 if (output_ref := skeleton.get_output_for_skip(output_dataset_key)) is not None:
729 # Populate the skeleton graph's node attributes
730 # with the existing DatasetRef, just like a
731 # predicted output of a non-skipped quantum.
732 skeleton.set_dataset_ref(output_ref, output_dataset_key)
733 else:
734 # Remove this dataset from the skeleton graph,
735 # because the quantum that would have produced it
736 # is being skipped and it doesn't already exist.
737 skeleton.remove_dataset_nodes([output_dataset_key])
738 # Removing the quantum node from the graph will happen outside this
739 # function.
740 return True
741 return False
743 @final
744 def _update_quantum_for_adjust(
745 self,
746 quantum_key: QuantumKey,
747 skeleton: QuantumGraphSkeleton,
748 task_prerequisite_info: PrerequisiteInfo,
749 skypix_bounds_builder: SkyPixBoundsBuilder,
750 timespan_builder: TimespanBuilder,
751 ) -> None:
752 """Update the quantum node in the skeleton by finding remaining
753 prerequisite inputs and dropping regular inputs that we now know will
754 not be produced.
756 Parameters
757 ----------
758 quantum_key : `QuantumKey`
759 Identifier for this quantum in the graph.
760 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
761 Preliminary quantum graph, to be modified in-place.
762 task_prerequisite_info : `~prerequisite_helpers.PrerequisiteInfo`
763 Information about the prerequisite inputs to this task.
764 skypix_bounds_builder : `~prerequisite_helpers.SkyPixBoundsBuilder`
765 An object that accumulates the appropriate spatial bounds for a
766 quantum.
767 timespan_builder : `~prerequisite_helpers.TimespanBuilder`
768 An object that accumulates the appropriate timespan for a quantum.
770 Notes
771 -----
772 This first looks for outputs already present in the `output_run` (i.e.
773 "in the way" in the skeleton); if it finds something and `clobber` is
774 `True`, it uses that ref (it's not ideal that both the original dataset
775 and its replacement will have the same UUID, but we don't have space in
776 the quantum graph for two UUIDs, and we need the datastore records of
777 the original there). If `clobber` is `False`, `RuntimeError` is
778 raised. If there is no output already present, a new one with a random
779 UUID is generated. In all cases the dataset node in the skeleton is
780 associated with a `lsst.daf.butler.DatasetRef`.
781 """
782 dataset_key: DatasetKey | PrerequisiteDatasetKey
783 for dataset_key in skeleton.iter_outputs_of(quantum_key):
784 dataset_data_id = skeleton.get_data_id(dataset_key)
785 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
786 if (ref := skeleton.get_output_in_the_way(dataset_key)) is None:
787 ref = DatasetRef(dataset_type_node.dataset_type, dataset_data_id, run=self.output_run)
788 elif not self.clobber:
789 # We intentionally raise here, before running adjustQuantum,
790 # because it'd be weird if we left an old potential output of a
791 # task sitting there in the output collection, just because the
792 # task happened to not actually produce it.
793 raise OutputExistsError(
794 f"Potential output dataset {ref} already exists in the output run "
795 f"{self.output_run}, but clobbering outputs was not expected to be necessary."
796 )
797 skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
798 timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
799 skeleton.set_dataset_ref(ref, dataset_key)
800 quantum_data_id = skeleton.get_data_id(quantum_key)
801 # Process inputs already present in the skeleton - this should include
802 # all regular inputs (including intermediates) and may include some
803 # prerequisites.
804 for dataset_key in list(skeleton.iter_inputs_of(quantum_key)):
805 if (ref := skeleton.get_dataset_ref(dataset_key)) is None:
806 # If the dataset ref hasn't been set either as an existing
807 # input or as an output of an already-processed upstream
808 # quantum, it's not going to be produced; remove it.
809 skeleton.remove_dataset_nodes([dataset_key])
810 continue
811 skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
812 timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
813 # Query for any prerequisites not handled by process_subgraph. Note
814 # that these were not already in the skeleton graph, so we add them
815 # now.
816 skypix_bounds = skypix_bounds_builder.finish()
817 timespan = timespan_builder.finish()
818 for finder in task_prerequisite_info.finders.values():
819 dataset_keys = []
820 for ref in finder.find(
821 self.butler, self.input_collections, quantum_data_id, skypix_bounds, timespan
822 ):
823 dataset_key = skeleton.add_prerequisite_node(ref)
824 dataset_keys.append(dataset_key)
825 skeleton.add_input_edges(quantum_key, dataset_keys)
827 @final
828 def _adapt_quantum_outputs(
829 self,
830 task_node: TaskNode,
831 quantum_key: QuantumKey,
832 skeleton: QuantumGraphSkeleton,
833 ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
834 """Adapt outputs for a preliminary quantum and put them into the form
835 used by `~lsst.daf.butler.Quantum` and
836 `~PipelineTaskConnections.adjustQuantum`.
838 Parameters
839 ----------
840 task_node : `pipeline_graph.TaskNode`
841 Node for this task in the pipeline graph.
842 quantum_key : `QuantumKey`
843 Identifier for this quantum in the graph.
844 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
845 Preliminary quantum graph, to be modified in-place.
847 Returns
848 -------
849 outputs : `~lsst.daf.butler.NamedKeyDict` [ \
850 `~lsst.daf.butler.DatasetType`, `list` [ \
851 `~lsst.daf.butler.DatasetRef` ] ]
852 All outputs to the task, using the storage class and components
853 defined by the task's own connections.
854 """
855 outputs_by_type: dict[str, list[DatasetRef]] = {}
856 dataset_key: DatasetKey
857 for dataset_key in skeleton.iter_outputs_of(quantum_key):
858 ref = skeleton.get_dataset_ref(dataset_key)
859 assert ref is not None, "Should have been added (or the node removed) in a previous pass."
860 outputs_by_type.setdefault(dataset_key.parent_dataset_type_name, []).append(ref)
861 adapted_outputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
862 for write_edge in task_node.iter_all_outputs():
863 dataset_type_node = self._pipeline_graph.dataset_types[write_edge.parent_dataset_type_name]
864 edge_dataset_type = write_edge.adapt_dataset_type(dataset_type_node.dataset_type)
865 adapted_outputs[edge_dataset_type] = [
866 write_edge.adapt_dataset_ref(ref)
867 for ref in sorted(outputs_by_type.get(write_edge.parent_dataset_type_name, []))
868 ]
869 return adapted_outputs
871 @final
872 def _adapt_quantum_inputs(
873 self,
874 task_node: TaskNode,
875 quantum_key: QuantumKey,
876 skeleton: QuantumGraphSkeleton,
877 ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
878 """Adapt input datasets for a preliminary quantum into the form used by
879 `~lsst.daf.butler.Quantum` and
880 `~PipelineTaskConnections.adjustQuantum`.
882 Parameters
883 ----------
884 task_node : `pipeline_graph.TaskNode`
885 Node for this task in the pipeline graph.
886 quantum_key : `QuantumKey`
887 Identifier for this quantum in the graph.
888 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
889 Preliminary quantum graph, to be modified in-place.
890 skypix_bounds_builder : `~prerequisite_helpers.SkyPixBoundsBuilder`
891 An object that accumulates the appropriate spatial bounds for a
892 quantum.
893 timespan_builder : `~prerequisite_helpers.TimespanBuilder`
894 An object that accumulates the appropriate timespan for a quantum.
896 Returns
897 -------
898 inputs : `~lsst.daf.butler.NamedKeyDict` [ \
899 `~lsst.daf.butler.DatasetType`, `list` [ \
900 `~lsst.daf.butler.DatasetRef` ] ]
901 All regular and prerequisite inputs to the task, using the storage
902 class and components defined by the task's own connections.
904 Notes
905 -----
906 This method trims input dataset nodes that are not already associated
907 with a `lsst.daf.butler.DatasetRef`, and queries for prerequisite input
908 nodes that do not exist.
909 """
910 inputs_by_type: dict[str, set[DatasetRef]] = {}
911 dataset_key: DatasetKey | PrerequisiteDatasetKey
912 for dataset_key in list(skeleton.iter_inputs_of(quantum_key)):
913 ref = skeleton.get_dataset_ref(dataset_key)
914 assert ref is not None, "Should have been added (or the node removed) in a previous pass."
915 inputs_by_type.setdefault(dataset_key.parent_dataset_type_name, set()).add(ref)
916 adapted_inputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
917 for read_edge in task_node.iter_all_inputs():
918 dataset_type_node = self._pipeline_graph.dataset_types[read_edge.parent_dataset_type_name]
919 edge_dataset_type = read_edge.adapt_dataset_type(dataset_type_node.dataset_type)
920 if (current_dataset_type := adapted_inputs.keys().get(edge_dataset_type.name)) is None:
921 adapted_inputs[edge_dataset_type] = [
922 read_edge.adapt_dataset_ref(ref)
923 for ref in sorted(inputs_by_type.get(read_edge.parent_dataset_type_name, frozenset()))
924 ]
925 elif current_dataset_type != edge_dataset_type:
926 raise NotImplementedError(
927 f"Task {task_node.label!r} has {edge_dataset_type.name!r} as an input via "
928 "two different connections, with two different storage class overrides. "
929 "This is not yet supported due to limitations in the Quantum data structure."
930 )
931 # If neither the `if` nor the `elif` above match, it means
932 # multiple input connections have exactly the same dataset
933 # type, and hence nothing to do after the first one.
934 return adapted_inputs
936 @final
937 def _resolve_task_init(
938 self, task_node: TaskNode, skeleton: QuantumGraphSkeleton, has_skipped_quanta: bool
939 ) -> None:
940 """Add init-input and init-output dataset nodes and edges for a task to
941 the skeleton.
943 Parameters
944 ----------
945 task_node : `pipeline_graph.TaskNode`
946 Pipeline graph description of the task.
947 skeleton : `QuantumGraphSkeleton`
948 In-progress quantum graph data structure to update in-place.
949 has_skipped_quanta : `bool`
950 Whether any of this task's quanta were skipped because they had
951 already succeeded.
952 """
953 quanta = skeleton.get_quanta(task_node.label)
954 task_init_key = TaskInitKey(task_node.label)
955 if quanta:
956 adapted_inputs: NamedKeyDict[DatasetType, DatasetRef] = NamedKeyDict()
957 # Process init-inputs.
958 input_keys: list[DatasetKey] = []
959 for read_edge in task_node.init.iter_all_inputs():
960 dataset_key = skeleton.add_dataset_node(
961 read_edge.parent_dataset_type_name, self.empty_data_id
962 )
963 skeleton.add_input_edge(task_init_key, dataset_key)
964 if (ref := skeleton.get_dataset_ref(dataset_key)) is None:
965 try:
966 ref = self.empty_dimensions_datasets.inputs[dataset_key]
967 except KeyError:
968 raise InitInputMissingError(
969 f"Overall init-input dataset {read_edge.parent_dataset_type_name!r} "
970 f"needed by task {task_node.label!r} not found in input collection(s) "
971 f"{self.input_collections}."
972 ) from None
973 skeleton.set_dataset_ref(ref, dataset_key)
974 for quantum_key in skeleton.get_quanta(task_node.label):
975 skeleton.add_input_edge(quantum_key, dataset_key)
976 input_keys.append(dataset_key)
977 adapted_ref = read_edge.adapt_dataset_ref(ref)
978 adapted_inputs[adapted_ref.datasetType] = adapted_ref
979 # Save the quantum-adapted init inputs to each quantum, and add
980 # skeleton edges connecting the init inputs to each quantum.
981 for quantum_key in skeleton.get_quanta(task_node.label):
982 skeleton[quantum_key]["init_inputs"] = adapted_inputs
983 # Process init-outputs.
984 adapted_outputs: NamedKeyDict[DatasetType, DatasetRef] = NamedKeyDict()
985 for write_edge in task_node.init.iter_all_outputs():
986 dataset_key = skeleton.add_dataset_node(
987 write_edge.parent_dataset_type_name, self.empty_data_id
988 )
989 if (ref := self.empty_dimensions_datasets.outputs_in_the_way.get(dataset_key)) is None:
990 ref = DatasetRef(
991 self._pipeline_graph.dataset_types[write_edge.parent_dataset_type_name].dataset_type,
992 self.empty_data_id,
993 run=self.output_run,
994 )
995 skeleton.set_dataset_ref(ref, dataset_key)
996 skeleton.add_output_edge(task_init_key, dataset_key)
997 adapted_ref = write_edge.adapt_dataset_ref(ref)
998 adapted_outputs[adapted_ref.datasetType] = adapted_ref
999 skeleton[task_init_key]["inputs"] = adapted_inputs
1000 skeleton[task_init_key]["outputs"] = adapted_outputs
1001 elif has_skipped_quanta:
1002 # No quanta remain for this task, but at least one quantum was
1003 # skipped because its outputs were present in the skip_existing_in
1004 # collections. This means all init outputs should be present in
1005 # the skip_existing_in collections, too, and we need to put those
1006 # refs in the graph.
1007 for write_edge in task_node.init.iter_all_outputs():
1008 dataset_key = skeleton.add_dataset_node(
1009 write_edge.parent_dataset_type_name, self.empty_data_id
1010 )
1011 if (ref := self.empty_dimensions_datasets.outputs_for_skip.get(dataset_key)) is None:
1012 raise InitInputMissingError(
1013 f"Init-output dataset {write_edge.parent_dataset_type_name!r} of skipped task "
1014 f"{task_node.label!r} not found in skip-existing-in collection(s) "
1015 f"{self.skip_existing_in}."
1016 ) from None
1017 skeleton.set_dataset_ref(ref, dataset_key)
1018 # If this dataset was "in the way" (i.e. already in the output
1019 # run), it isn't anymore.
1020 skeleton.discard_output_in_the_way(dataset_key)
1021 # No quanta remain in this task, but none were skipped; this means
1022 # they all got pruned because of NoWorkFound conditions. This
1023 # dooms all downstream quanta to the same fate, so we don't bother
1024 # doing anything with the task's init-outputs, since nothing is
1025 # going to consume them.
1027 @final
1028 @timeMethod
1029 def _find_empty_dimension_datasets(self) -> EmptyDimensionsDatasets:
1030 """Query for all dataset types with no dimensions, updating
1031 `empty_dimensions_datasets` in-place.
1033 This includes but is not limited to init inputs and init outputs.
1034 """
1035 inputs: dict[DatasetKey | PrerequisiteDatasetKey, DatasetRef] = {}
1036 outputs_for_skip: dict[DatasetKey, DatasetRef] = {}
1037 outputs_in_the_way: dict[DatasetKey, DatasetRef] = {}
1038 _, dataset_type_nodes = self._pipeline_graph.group_by_dimensions().get(self.universe.empty, ({}, {}))
1039 dataset_types = [node.dataset_type for node in dataset_type_nodes.values()]
1040 dataset_types.extend(self._global_init_output_types.values())
1041 for dataset_type in dataset_types:
1042 key = DatasetKey(dataset_type.name, self.empty_data_id.required_values)
1043 if (
1044 self._pipeline_graph.producer_of(dataset_type.name) is None
1045 and dataset_type.name not in self._global_init_output_types
1046 ):
1047 # Dataset type is an overall input; we always need to try to
1048 # find these.
1049 try:
1050 ref = self.butler.find_dataset(dataset_type.name, collections=self.input_collections)
1051 except MissingDatasetTypeError:
1052 ref = None
1053 if ref is not None:
1054 inputs[key] = ref
1055 elif self.skip_existing_in:
1056 # Dataset type is an intermediate or output; need to find these
1057 # if only they're from previously executed quanta that we might
1058 # skip...
1059 try:
1060 ref = self.butler.find_dataset(dataset_type.name, collections=self.skip_existing_in)
1061 except MissingDatasetTypeError:
1062 ref = None
1063 if ref is not None:
1064 outputs_for_skip[key] = ref
1065 if ref.run == self.output_run:
1066 outputs_in_the_way[key] = ref
1067 if self.output_run_exists and not self.skip_existing_starts_with_output_run:
1068 # ...or if they're in the way and would need to be clobbered
1069 # (and we haven't already found them in the previous block).
1070 try:
1071 ref = self.butler.find_dataset(dataset_type.name, collections=[self.output_run])
1072 except MissingDatasetTypeError:
1073 ref = None
1074 if ref is not None:
1075 outputs_in_the_way[key] = ref
1076 return EmptyDimensionsDatasets(
1077 inputs=inputs, outputs_for_skip=outputs_for_skip, outputs_in_the_way=outputs_in_the_way
1078 )
1080 @final
1081 @timeMethod
1082 def _attach_datastore_records(self, skeleton: QuantumGraphSkeleton) -> None:
1083 """Add datastore records for all overall inputs to a preliminary
1084 quantum graph.
1086 Parameters
1087 ----------
1088 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
1089 Preliminary quantum graph to update in place.
1091 Notes
1092 -----
1093 On return, all quantum nodes in the skeleton graph will have a
1094 "datastore_records" attribute that is a mapping from datastore name
1095 to `lsst.daf.butler.DatastoreRecordData`, as used by
1096 `lsst.daf.butler.Quantum`.
1097 """
1098 self.log.info("Fetching and attaching datastore records for all overall inputs.")
1099 overall_inputs = skeleton.extract_overall_inputs()
1100 exported_records = self.butler._datastore.export_records(overall_inputs.values())
1101 for task_label in self._pipeline_graph.tasks:
1102 if not skeleton.has_task(task_label):
1103 continue
1104 self.log.verbose("Fetching and attaching datastore records for task %s.", task_label)
1105 task_init_key = skeleton.get_task_init_node(task_label)
1106 init_input_ids = {
1107 ref.id
1108 for dataset_key in skeleton.iter_inputs_of(task_init_key)
1109 if (ref := overall_inputs.get(dataset_key)) is not None
1110 }
1111 init_records = {}
1112 if init_input_ids:
1113 for datastore_name, records in exported_records.items():
1114 matching_records = records.subset(init_input_ids)
1115 if matching_records is not None:
1116 init_records[datastore_name] = matching_records
1117 skeleton[task_init_key]["datastore_records"] = init_records
1118 for quantum_key in skeleton.get_quanta(task_label):
1119 quantum_records = {}
1120 input_ids = {
1121 ref.id
1122 for dataset_key in skeleton.iter_inputs_of(quantum_key)
1123 if (ref := overall_inputs.get(dataset_key)) is not None
1124 }
1125 if input_ids:
1126 for datastore_name, records in exported_records.items():
1127 matching_records = records.subset(input_ids)
1128 if matching_records is not None:
1129 quantum_records[datastore_name] = matching_records
1130 skeleton[quantum_key]["datastore_records"] = quantum_records
1132 @final
1133 @timeMethod
1134 def _construct_quantum_graph(
1135 self, skeleton: QuantumGraphSkeleton, metadata: Mapping[str, Any]
1136 ) -> QuantumGraph:
1137 """Construct a `.QuantumGraph` object from the contents of a
1138 fully-processed `.quantum_graph_skeleton.QuantumGraphSkeleton`.
1140 Parameters
1141 ----------
1142 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
1143 Preliminary quantum graph. Must have "init_inputs", "inputs", and
1144 "outputs" attributes on all quantum nodes, as added by
1145 `_resolve_task_quanta`, as well as a "datastore_records" attribute
1146 as added by `_attach_datastore_records`.
1147 metadata : `Mapping`
1148 Flexible metadata to add to the graph.
1150 Returns
1151 -------
1152 quantum_graph : `.QuantumGraph`
1153 DAG describing processing to be performed.
1154 """
1155 from .graph import QuantumGraph
1157 self.log.info("Transforming graph skeleton into a QuantumGraph instance.")
1158 quanta: dict[TaskDef, set[Quantum]] = {}
1159 init_inputs: dict[TaskDef, Iterable[DatasetRef]] = {}
1160 init_outputs: dict[TaskDef, Iterable[DatasetRef]] = {}
1161 for task_def in self._pipeline_graph._iter_task_defs():
1162 if not skeleton.has_task(task_def.label):
1163 continue
1164 self.log.verbose("Transforming graph skeleton nodes for task %s.", task_def.label)
1165 task_node = self._pipeline_graph.tasks[task_def.label]
1166 task_init_key = skeleton.get_task_init_node(task_def.label)
1167 task_init_state = skeleton[task_init_key]
1168 init_datastore_records: dict[str, DatastoreRecordData] = task_init_state.get(
1169 "datastore_records", {}
1170 )
1171 init_inputs[task_def] = task_init_state["inputs"].values()
1172 init_outputs[task_def] = task_init_state["outputs"].values()
1173 quanta_for_task: set[Quantum] = set()
1174 for quantum_key in skeleton.get_quanta(task_node.label):
1175 quantum_state = skeleton[quantum_key]
1176 quantum_datastore_records: dict[str, DatastoreRecordData] = quantum_state.get(
1177 "datastore_records", {}
1178 )
1179 quanta_for_task.add(
1180 Quantum(
1181 taskName=task_node.task_class_name,
1182 taskClass=task_node.task_class,
1183 dataId=quantum_state["data_id"],
1184 initInputs=quantum_state["init_inputs"],
1185 inputs=quantum_state["inputs"],
1186 outputs=quantum_state["outputs"],
1187 datastore_records=DatastoreRecordData.merge_mappings(
1188 quantum_datastore_records, init_datastore_records
1189 ),
1190 )
1191 )
1192 quanta[task_def] = quanta_for_task
1194 registry_dataset_types: list[DatasetType] = [
1195 node.dataset_type for node in self._pipeline_graph.dataset_types.values()
1196 ]
1198 all_metadata = self.metadata.to_dict()
1199 all_metadata.update(metadata)
1200 global_init_outputs: list[DatasetRef] = []
1201 for dataset_key in skeleton.global_init_outputs:
1202 ref = skeleton.get_dataset_ref(dataset_key)
1203 assert ref is not None, "Global init input refs should be resolved already."
1204 global_init_outputs.append(ref)
1205 self.log.verbose("Invoking QuantumGraph class constructor.")
1206 result = QuantumGraph(
1207 quanta,
1208 metadata=all_metadata,
1209 universe=self.universe,
1210 initInputs=init_inputs,
1211 initOutputs=init_outputs,
1212 globalInitOutputs=global_init_outputs,
1213 registryDatasetTypes=registry_dataset_types,
1214 )
1215 self.log.info("Graph build complete.")
1216 return result
1218 @final
1219 @timeMethod
1220 def _construct_components(
1221 self,
1222 skeleton: QuantumGraphSkeleton,
1223 output: str | None,
1224 metadata: Mapping[str, Any] | None,
1225 ) -> PredictedQuantumGraphComponents:
1226 """Return quantum graph components from a completed skeleton.
1228 Parameters
1229 ----------
1230 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
1231 Temporary data structure used by the builder to represent the
1232 graph.
1233 output : `str` or `None`, optional
1234 Output `~lsst.daf.butler.CollectionType.CHAINED` collection that
1235 combines the input and output collections.
1236 metadata : `~collections.abc.Mapping`, optional
1237 Mapping of JSON-friendly metadata. Collection information, the
1238 current user, and the current timestamp are automatically
1239 included.
1241 Returns
1242 -------
1243 components : `.quantum_graph.PredictedQuantumGraphComponents`
1244 Components that can be used to construct a graph object and/or save
1245 it to disk.
1246 """
1247 from .quantum_graph import (
1248 PredictedDatasetModel,
1249 PredictedQuantumDatasetsModel,
1250 PredictedQuantumGraphComponents,
1251 )
1253 self.log.info("Transforming graph skeleton into PredictedQuantumGraph components.")
1254 components = PredictedQuantumGraphComponents(pipeline_graph=self._pipeline_graph)
1255 components.header.inputs = list(self.input_collections)
1256 components.header.output_run = self.output_run
1257 components.header.output = output
1258 if metadata is not None:
1259 components.header.metadata.update(metadata)
1260 components.dimension_data = DimensionDataAttacher(
1261 records=skeleton.get_dimension_data(),
1262 dimensions=self._pipeline_graph.get_all_dimensions(),
1263 )
1264 components.init_quanta.root = [
1265 PredictedQuantumDatasetsModel.model_construct(
1266 quantum_id=generate_uuidv7(),
1267 task_label="",
1268 outputs={
1269 dataset_key.parent_dataset_type_name: [
1270 PredictedDatasetModel.from_dataset_ref(
1271 cast(DatasetRef, skeleton.get_dataset_ref(dataset_key))
1272 )
1273 ]
1274 for dataset_key in skeleton.global_init_outputs
1275 },
1276 )
1277 ]
1278 for task_node in self._pipeline_graph.tasks.values():
1279 if not skeleton.has_task(task_node.label):
1280 continue
1281 self.log.verbose("Transforming graph skeleton nodes for task %s.", task_node.label)
1282 task_init_key = TaskInitKey(task_node.label)
1283 init_quantum_datasets = PredictedQuantumDatasetsModel.model_construct(
1284 quantum_id=generate_uuidv7(),
1285 task_label=task_node.label,
1286 inputs=self._make_predicted_datasets(
1287 skeleton,
1288 task_node.init.iter_all_inputs(),
1289 skeleton.iter_inputs_of(task_init_key),
1290 ),
1291 outputs=self._make_predicted_datasets(
1292 skeleton,
1293 task_node.init.iter_all_outputs(),
1294 skeleton.iter_outputs_of(task_init_key),
1295 ),
1296 datastore_records={
1297 datastore_name: records.to_simple()
1298 for datastore_name, records in skeleton[task_init_key]
1299 .get("datastore_records", {})
1300 .items()
1301 },
1302 )
1303 components.init_quanta.root.append(init_quantum_datasets)
1304 for quantum_key in skeleton.get_quanta(task_node.label):
1305 quantum_datasets = PredictedQuantumDatasetsModel.model_construct(
1306 quantum_id=generate_uuidv7(),
1307 task_label=task_node.label,
1308 data_coordinate=list(skeleton.get_data_id(quantum_key).full_values),
1309 inputs=self._make_predicted_datasets(
1310 skeleton,
1311 task_node.iter_all_inputs(),
1312 skeleton.iter_inputs_of(quantum_key),
1313 ),
1314 outputs=self._make_predicted_datasets(
1315 skeleton,
1316 task_node.iter_all_outputs(),
1317 skeleton.iter_outputs_of(quantum_key),
1318 ),
1319 datastore_records={
1320 datastore_name: records.to_simple()
1321 for datastore_name, records in skeleton[quantum_key]
1322 .get("datastore_records", {})
1323 .items()
1324 },
1325 )
1326 components.quantum_datasets[quantum_datasets.quantum_id] = quantum_datasets
1327 self.log.verbose("Building the thin summary graph.")
1328 components.set_thin_graph()
1329 components.set_header_counts()
1330 self.log.info("Graph build complete.")
1331 return components
1333 @staticmethod
1334 def _make_predicted_datasets(
1335 skeleton: QuantumGraphSkeleton,
1336 edges: Iterable[Edge],
1337 dataset_keys: Iterable[DatasetKey | PrerequisiteDatasetKey],
1338 ) -> dict[str, list[PredictedDatasetModel]]:
1339 """Make the predicted quantum graph model objects that represent the
1340 datasets from an iterable of pipeline graph edges.
1342 Parameters
1343 ----------
1344 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
1345 Temporary data structure used by the builder to represent the
1346 graph.
1347 edges : `~collections.abc.Iterable` [ `.pipeline_graph.Edge` ]
1348 Pipeline graph edges.
1349 dataset_keys : `~collections.abc.Iterable` [ \
1350 `.quantum_graph_skeleton.DatasetKey` or\
1351 `.quantum_graph_skeleton.PrerequisiteDatasetKey` ]
1352 All nodes in the skeleton that correspond to any of the given
1353 pipeline graph edges.
1355 Returns
1356 -------
1357 predicted_datasets : `dict` [ `str`, \
1358 `list` [ `.quantum_graph.PredictedDatasetModel` ] ]
1359 Mapping of dataset models, keyed by connection name.
1360 """
1361 from .quantum_graph import PredictedDatasetModel
1363 connection_names_by_dataset_type: defaultdict[str, list[str]] = defaultdict(list)
1364 result: dict[str, list[PredictedDatasetModel]] = {}
1365 for edge in edges:
1366 connection_names_by_dataset_type[edge.parent_dataset_type_name].append(edge.connection_name)
1367 result[edge.connection_name] = []
1369 for dataset_key in dataset_keys:
1370 connection_names = connection_names_by_dataset_type.get(dataset_key.parent_dataset_type_name)
1371 if connection_names is None:
1372 # Ignore if this isn't one of the connections we're processing
1373 # (probably an init-input), which would also be predecessor to
1374 # a quantum node, but should be handled separately.
1375 continue
1376 ref = skeleton.get_dataset_ref(dataset_key)
1377 assert ref is not None, "DatasetRefs should have already been added to skeleton."
1378 for connection_name in connection_names:
1379 result[connection_name].append(PredictedDatasetModel.from_dataset_ref(ref))
1380 for refs in result.values():
1381 refs.sort(key=operator.attrgetter("data_coordinate"))
1382 return result
1384 @staticmethod
1385 @final
1386 def _find_removed(
1387 original: Iterable[DatasetKey | PrerequisiteDatasetKey],
1388 adjusted: NamedKeyMapping[DatasetType, Sequence[DatasetRef]],
1389 ) -> set[DatasetKey | PrerequisiteDatasetKey]:
1390 """Identify skeleton-graph dataset nodes that have been removed by
1391 `~PipelineTaskConnections.adjustQuantum`.
1393 Parameters
1394 ----------
1395 original : `~collections.abc.Iterable` [ `DatasetKey` or \
1396 `PrerequisiteDatasetKey` ]
1397 Identifiers for the dataset nodes that were the original neighbors
1398 (inputs or outputs) of a quantum.
1399 adjusted : `~lsst.daf.butler.NamedKeyMapping` [ \
1400 `~lsst.daf.butler.DatasetType`, \
1401 `~collections.abc.Sequence` [ `lsst.daf.butler.DatasetType` ] ]
1402 Adjusted neighbors, in the form used by `lsst.daf.butler.Quantum`.
1404 Returns
1405 -------
1406 removed : `set` [ `DatasetKey` ]
1407 Datasets in ``original`` that have no counterpart in ``adjusted``.
1408 """
1409 result = set(original)
1410 for dataset_type, kept_refs in adjusted.items():
1411 parent_dataset_type_name, _ = DatasetType.splitDatasetTypeName(dataset_type.name)
1412 for kept_ref in kept_refs:
1413 # We don't know if this was a DatasetKey or a
1414 # PrerequisiteDatasetKey; just try both.
1415 result.discard(DatasetKey(parent_dataset_type_name, kept_ref.dataId.required_values))
1416 result.discard(PrerequisiteDatasetKey(parent_dataset_type_name, kept_ref.id.bytes))
1417 return result
1420@dataclasses.dataclass(eq=False, order=False)
1421class EmptyDimensionsDatasets:
1422 """Struct that holds the results of empty-dimensions dataset queries for
1423 `QuantumGraphBuilder`.
1424 """
1426 inputs: Mapping[DatasetKey | PrerequisiteDatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1427 """Overall-input datasets found in `QuantumGraphBuilder.input_collections`.
1429 This may include prerequisite inputs. It does include init-inputs.
1430 It does not include intermediates.
1431 """
1433 outputs_for_skip: Mapping[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1434 """Output datasets found in `QuantumGraphBuilder.skip_existing_in`.
1436 It is unspecified whether this contains init-outputs; there is
1437 no concept of skipping at the init stage, so this is not expected to
1438 matter.
1439 """
1441 outputs_in_the_way: Mapping[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1442 """Output datasets found in `QuantumGraphBuilder.output_run`.
1444 This includes regular outputs and init-outputs.
1445 """
1448def _quantum_or_quanta(n: int) -> str:
1449 """Correctly pluralize 'quantum' if needed."""
1450 return f"{n} quanta" if n != 1 else "1 quantum"