Coverage for python / lsst / pipe / base / quantum_graph_builder.py: 17%
465 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:14 -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 overall_inputs = skeleton.extract_overall_inputs()
1099 exported_records = self.butler._datastore.export_records(overall_inputs.values())
1100 for task_label in self._pipeline_graph.tasks:
1101 if not skeleton.has_task(task_label):
1102 continue
1103 task_init_key = skeleton.get_task_init_node(task_label)
1104 init_input_ids = {
1105 ref.id
1106 for dataset_key in skeleton.iter_inputs_of(task_init_key)
1107 if (ref := overall_inputs.get(dataset_key)) is not None
1108 }
1109 init_records = {}
1110 if init_input_ids:
1111 for datastore_name, records in exported_records.items():
1112 matching_records = records.subset(init_input_ids)
1113 if matching_records is not None:
1114 init_records[datastore_name] = matching_records
1115 skeleton[task_init_key]["datastore_records"] = init_records
1116 for quantum_key in skeleton.get_quanta(task_label):
1117 quantum_records = {}
1118 input_ids = {
1119 ref.id
1120 for dataset_key in skeleton.iter_inputs_of(quantum_key)
1121 if (ref := overall_inputs.get(dataset_key)) is not None
1122 }
1123 if input_ids:
1124 for datastore_name, records in exported_records.items():
1125 matching_records = records.subset(input_ids)
1126 if matching_records is not None:
1127 quantum_records[datastore_name] = matching_records
1128 skeleton[quantum_key]["datastore_records"] = quantum_records
1130 @final
1131 @timeMethod
1132 def _construct_quantum_graph(
1133 self, skeleton: QuantumGraphSkeleton, metadata: Mapping[str, Any]
1134 ) -> QuantumGraph:
1135 """Construct a `.QuantumGraph` object from the contents of a
1136 fully-processed `.quantum_graph_skeleton.QuantumGraphSkeleton`.
1138 Parameters
1139 ----------
1140 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
1141 Preliminary quantum graph. Must have "init_inputs", "inputs", and
1142 "outputs" attributes on all quantum nodes, as added by
1143 `_resolve_task_quanta`, as well as a "datastore_records" attribute
1144 as added by `_attach_datastore_records`.
1145 metadata : `Mapping`
1146 Flexible metadata to add to the graph.
1148 Returns
1149 -------
1150 quantum_graph : `.QuantumGraph`
1151 DAG describing processing to be performed.
1152 """
1153 from .graph import QuantumGraph
1155 quanta: dict[TaskDef, set[Quantum]] = {}
1156 init_inputs: dict[TaskDef, Iterable[DatasetRef]] = {}
1157 init_outputs: dict[TaskDef, Iterable[DatasetRef]] = {}
1158 for task_def in self._pipeline_graph._iter_task_defs():
1159 if not skeleton.has_task(task_def.label):
1160 continue
1161 task_node = self._pipeline_graph.tasks[task_def.label]
1162 task_init_key = skeleton.get_task_init_node(task_def.label)
1163 task_init_state = skeleton[task_init_key]
1164 init_datastore_records: dict[str, DatastoreRecordData] = task_init_state.get(
1165 "datastore_records", {}
1166 )
1167 init_inputs[task_def] = task_init_state["inputs"].values()
1168 init_outputs[task_def] = task_init_state["outputs"].values()
1169 quanta_for_task: set[Quantum] = set()
1170 for quantum_key in skeleton.get_quanta(task_node.label):
1171 quantum_state = skeleton[quantum_key]
1172 quantum_datastore_records: dict[str, DatastoreRecordData] = quantum_state.get(
1173 "datastore_records", {}
1174 )
1175 quanta_for_task.add(
1176 Quantum(
1177 taskName=task_node.task_class_name,
1178 taskClass=task_node.task_class,
1179 dataId=quantum_state["data_id"],
1180 initInputs=quantum_state["init_inputs"],
1181 inputs=quantum_state["inputs"],
1182 outputs=quantum_state["outputs"],
1183 datastore_records=DatastoreRecordData.merge_mappings(
1184 quantum_datastore_records, init_datastore_records
1185 ),
1186 )
1187 )
1188 quanta[task_def] = quanta_for_task
1190 registry_dataset_types: list[DatasetType] = [
1191 node.dataset_type for node in self._pipeline_graph.dataset_types.values()
1192 ]
1194 all_metadata = self.metadata.to_dict()
1195 all_metadata.update(metadata)
1196 global_init_outputs: list[DatasetRef] = []
1197 for dataset_key in skeleton.global_init_outputs:
1198 ref = skeleton.get_dataset_ref(dataset_key)
1199 assert ref is not None, "Global init input refs should be resolved already."
1200 global_init_outputs.append(ref)
1201 return QuantumGraph(
1202 quanta,
1203 metadata=all_metadata,
1204 universe=self.universe,
1205 initInputs=init_inputs,
1206 initOutputs=init_outputs,
1207 globalInitOutputs=global_init_outputs,
1208 registryDatasetTypes=registry_dataset_types,
1209 )
1211 @final
1212 @timeMethod
1213 def _construct_components(
1214 self,
1215 skeleton: QuantumGraphSkeleton,
1216 output: str | None,
1217 metadata: Mapping[str, Any] | None,
1218 ) -> PredictedQuantumGraphComponents:
1219 """Return quantum graph components from a completed skeleton.
1221 Parameters
1222 ----------
1223 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
1224 Temporary data structure used by the builder to represent the
1225 graph.
1226 output : `str` or `None`, optional
1227 Output `~lsst.daf.butler.CollectionType.CHAINED` collection that
1228 combines the input and output collections.
1229 metadata : `~collections.abc.Mapping`, optional
1230 Mapping of JSON-friendly metadata. Collection information, the
1231 current user, and the current timestamp are automatically
1232 included.
1234 Returns
1235 -------
1236 components : `.quantum_graph.PredictedQuantumGraphComponents`
1237 Components that can be used to construct a graph object and/or save
1238 it to disk.
1239 """
1240 from .quantum_graph import (
1241 PredictedDatasetModel,
1242 PredictedQuantumDatasetsModel,
1243 PredictedQuantumGraphComponents,
1244 )
1246 components = PredictedQuantumGraphComponents(pipeline_graph=self._pipeline_graph)
1247 components.header.inputs = list(self.input_collections)
1248 components.header.output_run = self.output_run
1249 components.header.output = output
1250 if metadata is not None:
1251 components.header.metadata.update(metadata)
1252 components.dimension_data = DimensionDataAttacher(
1253 records=skeleton.get_dimension_data(),
1254 dimensions=self._pipeline_graph.get_all_dimensions(),
1255 )
1256 components.init_quanta.root = [
1257 PredictedQuantumDatasetsModel.model_construct(
1258 quantum_id=generate_uuidv7(),
1259 task_label="",
1260 outputs={
1261 dataset_key.parent_dataset_type_name: [
1262 PredictedDatasetModel.from_dataset_ref(
1263 cast(DatasetRef, skeleton.get_dataset_ref(dataset_key))
1264 )
1265 ]
1266 for dataset_key in skeleton.global_init_outputs
1267 },
1268 )
1269 ]
1270 for task_node in self._pipeline_graph.tasks.values():
1271 if not skeleton.has_task(task_node.label):
1272 continue
1273 task_init_key = TaskInitKey(task_node.label)
1274 init_quantum_datasets = PredictedQuantumDatasetsModel.model_construct(
1275 quantum_id=generate_uuidv7(),
1276 task_label=task_node.label,
1277 inputs=self._make_predicted_datasets(
1278 skeleton,
1279 task_node.init.iter_all_inputs(),
1280 skeleton.iter_inputs_of(task_init_key),
1281 ),
1282 outputs=self._make_predicted_datasets(
1283 skeleton,
1284 task_node.init.iter_all_outputs(),
1285 skeleton.iter_outputs_of(task_init_key),
1286 ),
1287 datastore_records={
1288 datastore_name: records.to_simple()
1289 for datastore_name, records in skeleton[task_init_key]
1290 .get("datastore_records", {})
1291 .items()
1292 },
1293 )
1294 components.init_quanta.root.append(init_quantum_datasets)
1295 for quantum_key in skeleton.get_quanta(task_node.label):
1296 quantum_datasets = PredictedQuantumDatasetsModel.model_construct(
1297 quantum_id=generate_uuidv7(),
1298 task_label=task_node.label,
1299 data_coordinate=list(skeleton.get_data_id(quantum_key).full_values),
1300 inputs=self._make_predicted_datasets(
1301 skeleton,
1302 task_node.iter_all_inputs(),
1303 skeleton.iter_inputs_of(quantum_key),
1304 ),
1305 outputs=self._make_predicted_datasets(
1306 skeleton,
1307 task_node.iter_all_outputs(),
1308 skeleton.iter_outputs_of(quantum_key),
1309 ),
1310 datastore_records={
1311 datastore_name: records.to_simple()
1312 for datastore_name, records in skeleton[quantum_key]
1313 .get("datastore_records", {})
1314 .items()
1315 },
1316 )
1317 components.quantum_datasets[quantum_datasets.quantum_id] = quantum_datasets
1318 components.set_thin_graph()
1319 components.set_header_counts()
1320 return components
1322 @staticmethod
1323 def _make_predicted_datasets(
1324 skeleton: QuantumGraphSkeleton,
1325 edges: Iterable[Edge],
1326 dataset_keys: Iterable[DatasetKey | PrerequisiteDatasetKey],
1327 ) -> dict[str, list[PredictedDatasetModel]]:
1328 """Make the predicted quantum graph model objects that represent the
1329 datasets from an iterable of pipeline graph edges.
1331 Parameters
1332 ----------
1333 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
1334 Temporary data structure used by the builder to represent the
1335 graph.
1336 edges : `~collections.abc.Iterable` [ `.pipeline_graph.Edge` ]
1337 Pipeline graph edges.
1338 dataset_keys : `~collections.abc.Iterable` [ \
1339 `.quantum_graph_skeleton.DatasetKey` or\
1340 `.quantum_graph_skeleton.PrerequisiteDatasetKey` ]
1341 All nodes in the skeleton that correspond to any of the given
1342 pipeline graph edges.
1344 Returns
1345 -------
1346 predicted_datasets : `dict` [ `str`, \
1347 `list` [ `.quantum_graph.PredictedDatasetModel` ] ]
1348 Mapping of dataset models, keyed by connection name.
1349 """
1350 from .quantum_graph import PredictedDatasetModel
1352 connection_names_by_dataset_type: defaultdict[str, list[str]] = defaultdict(list)
1353 result: dict[str, list[PredictedDatasetModel]] = {}
1354 for edge in edges:
1355 connection_names_by_dataset_type[edge.parent_dataset_type_name].append(edge.connection_name)
1356 result[edge.connection_name] = []
1358 for dataset_key in dataset_keys:
1359 connection_names = connection_names_by_dataset_type.get(dataset_key.parent_dataset_type_name)
1360 if connection_names is None:
1361 # Ignore if this isn't one of the connections we're processing
1362 # (probably an init-input), which would also be predecessor to
1363 # a quantum node, but should be handled separately.
1364 continue
1365 ref = skeleton.get_dataset_ref(dataset_key)
1366 assert ref is not None, "DatasetRefs should have already been added to skeleton."
1367 for connection_name in connection_names:
1368 result[connection_name].append(PredictedDatasetModel.from_dataset_ref(ref))
1369 for refs in result.values():
1370 refs.sort(key=operator.attrgetter("data_coordinate"))
1371 return result
1373 @staticmethod
1374 @final
1375 def _find_removed(
1376 original: Iterable[DatasetKey | PrerequisiteDatasetKey],
1377 adjusted: NamedKeyMapping[DatasetType, Sequence[DatasetRef]],
1378 ) -> set[DatasetKey | PrerequisiteDatasetKey]:
1379 """Identify skeleton-graph dataset nodes that have been removed by
1380 `~PipelineTaskConnections.adjustQuantum`.
1382 Parameters
1383 ----------
1384 original : `~collections.abc.Iterable` [ `DatasetKey` or \
1385 `PrerequisiteDatasetKey` ]
1386 Identifiers for the dataset nodes that were the original neighbors
1387 (inputs or outputs) of a quantum.
1388 adjusted : `~lsst.daf.butler.NamedKeyMapping` [ \
1389 `~lsst.daf.butler.DatasetType`, \
1390 `~collections.abc.Sequence` [ `lsst.daf.butler.DatasetType` ] ]
1391 Adjusted neighbors, in the form used by `lsst.daf.butler.Quantum`.
1393 Returns
1394 -------
1395 removed : `set` [ `DatasetKey` ]
1396 Datasets in ``original`` that have no counterpart in ``adjusted``.
1397 """
1398 result = set(original)
1399 for dataset_type, kept_refs in adjusted.items():
1400 parent_dataset_type_name, _ = DatasetType.splitDatasetTypeName(dataset_type.name)
1401 for kept_ref in kept_refs:
1402 # We don't know if this was a DatasetKey or a
1403 # PrerequisiteDatasetKey; just try both.
1404 result.discard(DatasetKey(parent_dataset_type_name, kept_ref.dataId.required_values))
1405 result.discard(PrerequisiteDatasetKey(parent_dataset_type_name, kept_ref.id.bytes))
1406 return result
1409@dataclasses.dataclass(eq=False, order=False)
1410class EmptyDimensionsDatasets:
1411 """Struct that holds the results of empty-dimensions dataset queries for
1412 `QuantumGraphBuilder`.
1413 """
1415 inputs: Mapping[DatasetKey | PrerequisiteDatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1416 """Overall-input datasets found in `QuantumGraphBuilder.input_collections`.
1418 This may include prerequisite inputs. It does include init-inputs.
1419 It does not include intermediates.
1420 """
1422 outputs_for_skip: Mapping[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1423 """Output datasets found in `QuantumGraphBuilder.skip_existing_in`.
1425 It is unspecified whether this contains init-outputs; there is
1426 no concept of skipping at the init stage, so this is not expected to
1427 matter.
1428 """
1430 outputs_in_the_way: Mapping[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1431 """Output datasets found in `QuantumGraphBuilder.output_run`.
1433 This includes regular outputs and init-outputs.
1434 """
1437def _quantum_or_quanta(n: int) -> str:
1438 """Correctly pluralize 'quantum' if needed."""
1439 return f"{n} quanta" if n != 1 else "1 quantum"