Coverage for python / lsst / pipe / base / quantum_graph_builder.py: 16%
476 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-31 01:07 -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 attach_datastore_records : `bool`, optional
384 Whether to include datastore records in the graph. Required for
385 `lsst.daf.butler.QuantumBackedButler` execution.
387 Returns
388 -------
389 quantum_graph_skeleton : `QuantumGraphSkeleton`
390 DAG describing processing to be performed.
391 """
392 with self.butler.registry.caching_context():
393 full_skeleton = QuantumGraphSkeleton(self._pipeline_graph.tasks)
394 subgraphs = list(self._pipeline_graph.split_independent())
395 for i, subgraph in enumerate(subgraphs):
396 self.log.info(
397 "Processing pipeline subgraph %d of %d with %d task(s).",
398 i + 1,
399 len(subgraphs),
400 len(subgraph.tasks),
401 )
402 self.log.verbose("Subgraph tasks: [%s]", ", ".join(label for label in subgraph.tasks))
403 subgraph_skeleton = self.process_subgraph(subgraph)
404 full_skeleton.update(subgraph_skeleton)
405 # Loop over tasks to apply skip-existing logic and add missing
406 # prerequisites. The pipeline graph must be topologically sorted,
407 # so a quantum is only processed after any quantum that provides
408 # its inputs has been processed.
409 skipped_quanta: dict[str, list[QuantumKey]] = {}
410 for task_node in self._pipeline_graph.tasks.values():
411 skipped_quanta[task_node.label] = self._resolve_task_quanta(task_node, full_skeleton)
412 # Add any dimension records not handled by the subclass, and
413 # aggregate any that were added directly to data IDs.
414 full_skeleton.attach_dimension_records(self.butler, self._pipeline_graph.get_all_dimensions())
415 # Loop over tasks again to run the adjust hooks.
416 for task_node in self._pipeline_graph.tasks.values():
417 self._adjust_task_quanta(task_node, full_skeleton, skipped_quanta[task_node.label])
418 # Add global init-outputs to the skeleton.
419 for dataset_type in self._global_init_output_types.values():
420 dataset_key = full_skeleton.add_dataset_node(
421 dataset_type.name, self.empty_data_id, is_global_init_output=True
422 )
423 ref = self.empty_dimensions_datasets.outputs_in_the_way.get(dataset_key)
424 if ref is None:
425 ref = DatasetRef(dataset_type, self.empty_data_id, run=self.output_run)
426 full_skeleton.set_dataset_ref(ref, dataset_key)
427 # Remove dataset nodes with no edges that are not global init
428 # outputs, which are generally overall-inputs whose original quanta
429 # end up skipped or with no work to do (we can't remove these along
430 # with the quanta because no quantum knows if its the only
431 # consumer).
432 full_skeleton.remove_orphan_datasets()
433 if attach_datastore_records:
434 self._attach_datastore_records(full_skeleton)
435 return full_skeleton
437 @abstractmethod
438 def process_subgraph(self, subgraph: PipelineGraph) -> QuantumGraphSkeleton:
439 """Build the rough structure for an independent subset of the
440 `.QuantumGraph` and query for relevant existing datasets.
442 Parameters
443 ----------
444 subgraph : `.pipeline_graph.PipelineGraph`
445 Subset of the pipeline graph that should be processed by this call.
446 This is always resolved and topologically sorted. It should not be
447 modified.
449 Returns
450 -------
451 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
452 Class representing an initial quantum graph. See
453 `.quantum_graph_skeleton.QuantumGraphSkeleton` docs for details.
454 After this is returned, the object may be modified in-place in
455 unspecified ways.
457 Notes
458 -----
459 The `.quantum_graph_skeleton.QuantumGraphSkeleton` should associate
460 `lsst.daf.butler.DatasetRef` objects with nodes for existing datasets.
461 In particular:
463 - `.quantum_graph_skeleton.QuantumGraphSkeleton.set_dataset_ref` must
464 be used to associate existing datasets with all overall-input dataset
465 nodes in the skeleton by querying `input_collections`. This includes
466 all standard input nodes and any prerequisite nodes added by the
467 method (prerequisite nodes may also be left out entirely, as the base
468 class can add them later, albeit possibly less efficiently).
469 - `.quantum_graph_skeleton.QuantumGraphSkeleton.set_output_for_skip`
470 must be used to associate existing datasets with output dataset nodes
471 by querying `skip_existing_in`.
472 - `.quantum_graph_skeleton.QuantumGraphSkeleton.add_output_in_the_way`
473 must be used to associated existing outputs with output dataset nodes
474 by querying `output_run` if `output_run_exists` is `True`. Note that
475 the presence of such datasets is not automatically an error, even if
476 `clobber` is `False`, as these may be quanta that will be skipped.
478 `lsst.daf.butler.DatasetRef` objects for existing datasets with empty
479 data IDs in all of the above categories may be found in the
480 `empty_dimensions_datasets` attribute, as these are queried for prior
481 to this call by the base class, but associating them with graph nodes
482 is still this method's responsibility.
484 Dataset types should never be components and should always use the
485 "common" storage class definition in `pipeline_graph.DatasetTypeNode`
486 (which is the data repository definition when the dataset type is
487 registered).
488 """
489 raise NotImplementedError()
491 @final
492 @timeMethod
493 def _resolve_task_quanta(self, task_node: TaskNode, skeleton: QuantumGraphSkeleton) -> list[QuantumKey]:
494 """Process the quanta for one task in a skeleton graph to skip those
495 that have already completed and add missing prerequisite inputs.
497 Parameters
498 ----------
499 task_node : `pipeline_graph.TaskNode`
500 Node for this task in the pipeline graph.
501 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
502 Preliminary quantum graph, to be modified in-place.
504 Returns
505 -------
506 skipped_quanta : `list` [ `.quantum_skeleton_graph.QuantumKey` ]
507 Keys of quanta that were already skipped because their metadata
508 already exists in a ``skip_existing_in`` collections.
510 Notes
511 -----
512 This method modifies ``skeleton`` in-place in several ways:
514 - It associates a `lsst.daf.butler.DatasetRef` with all output datasets
515 and drops input dataset nodes that do not have a
516 `lsst.daf.butler.DatasetRef` already. This ensures producing and
517 consuming tasks start from the same `lsst.daf.butler.DatasetRef`.
518 - It removes quantum nodes that are to be skipped because their outputs
519 already exist in `skip_existing_in`. It also marks their outputs
520 as no longer in the way.
521 - It adds prerequisite dataset nodes and edges that connect them to the
522 quanta that consume them.
523 """
524 # Extract the helper object for the prerequisite inputs of this task,
525 # and tell it to prepare to construct skypix bounds and timespans for
526 # each quantum (these will automatically do nothing if nothing needs
527 # those bounds).
528 task_prerequisite_info = self.prerequisite_info[task_node.label]
529 task_prerequisite_info.update_bounds()
530 # Loop over all quanta for this task, remembering the ones we've
531 # gotten rid of.
532 skipped_quanta = []
533 for quantum_key in skeleton.get_quanta(task_node.label):
534 if self._skip_quantum_if_metadata_exists(task_node, quantum_key, skeleton):
535 skipped_quanta.append(quantum_key)
536 continue
537 quantum_data_id = skeleton[quantum_key]["data_id"]
538 skypix_bounds_builder = task_prerequisite_info.bounds.make_skypix_bounds_builder(quantum_data_id)
539 timespan_builder = task_prerequisite_info.bounds.make_timespan_builder(quantum_data_id)
540 self._update_quantum_for_adjust(
541 quantum_key,
542 skeleton,
543 task_prerequisite_info,
544 skypix_bounds_builder,
545 timespan_builder,
546 )
547 for skipped_quantum in skipped_quanta:
548 skeleton.remove_quantum_node(skipped_quantum, remove_outputs=False)
549 return skipped_quanta
551 @final
552 @timeMethod
553 def _adjust_task_quanta(
554 self, task_node: TaskNode, skeleton: QuantumGraphSkeleton, skipped_quanta: list[QuantumKey]
555 ) -> None:
556 """Process the quanta for one task in a skeleton graph by calling the
557 ``adjust_all_quanta`` and ``adjustQuantum`` hooks.
559 Parameters
560 ----------
561 task_node : `pipeline_graph.TaskNode`
562 Node for this task in the pipeline graph.
563 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
564 Preliminary quantum graph, to be modified in-place.
565 skipped_quanta : `list` [ `.quantum_skeleton_graph.QuantumKey` ]
566 Keys of quanta that were already skipped because their metadata
567 already exists in a ``skip_existing_in`` collections.
569 Notes
570 -----
571 This method modifies ``skeleton`` in-place in several ways:
573 - It adds "inputs", "outputs", and "init_inputs" attributes to the
574 quantum nodes, holding the same `NamedValueMapping` objects needed to
575 construct an actual `Quantum` instances.
576 - It removes quantum nodes whose
577 `~PipelineTaskConnections.adjustQuantum` calls raise `NoWorkFound` or
578 predict no outputs;
579 - It removes the nodes of output datasets that are "adjusted away".
580 - It removes the edges of input datasets that are "adjusted away".
582 The difference between how adjusted inputs and outputs are handled
583 reflects the fact that many quanta can share the same input, but only
584 one produces each output. This can lead to the graph having
585 superfluous isolated nodes after processing is complete, but these
586 should only be removed after all the quanta from all tasks have been
587 processed.
588 """
589 # Give the task a chance to adjust all quanta together. This
590 # operates directly on the skeleton (via a the 'adjuster', which
591 # is just an interface adapter).
592 adjuster = QuantaAdjuster(task_node.label, self._pipeline_graph, skeleton, self.butler)
593 task_node.get_connections().adjust_all_quanta(adjuster)
594 # Loop over all quanta again, remembering those we get rid of in other
595 # ways.
596 no_work_quanta = []
597 for quantum_key in skeleton.get_quanta(task_node.label):
598 adjusted_outputs = self._adapt_quantum_outputs(task_node, quantum_key, skeleton)
599 adjusted_inputs = self._adapt_quantum_inputs(task_node, quantum_key, skeleton)
600 # Give the task's Connections class an opportunity to remove
601 # some inputs, or complain if they are unacceptable. This will
602 # raise if one of the check conditions is not met, which is the
603 # intended behavior.
604 helper = AdjustQuantumHelper(inputs=adjusted_inputs, outputs=adjusted_outputs)
605 quantum_data_id = skeleton[quantum_key]["data_id"]
606 try:
607 helper.adjust_in_place(task_node.get_connections(), task_node.label, quantum_data_id)
608 except NoWorkFound as err:
609 # Do not generate this quantum; it would not produce any
610 # outputs. Remove it and all of the outputs it might have
611 # produced from the skeleton.
612 try:
613 _, connection_name, _ = err.args
614 details = f"not enough datasets for connection {connection_name}."
615 except ValueError:
616 details = str(err)
617 self.log.debug(
618 "No work found for quantum %s of task %s: %s",
619 quantum_key.data_id_values,
620 quantum_key.task_label,
621 details,
622 )
623 no_work_quanta.append(quantum_key)
624 continue
625 if helper.outputs_adjusted:
626 if not any(adjusted_refs for adjusted_refs in helper.outputs.values()):
627 # No outputs also means we don't generate this quantum.
628 self.log.debug(
629 "No outputs predicted for quantum %s of task %s.",
630 quantum_key.data_id_values,
631 quantum_key.task_label,
632 )
633 no_work_quanta.append(quantum_key)
634 continue
635 # Remove output nodes that were not retained by
636 # adjustQuantum.
637 skeleton.remove_dataset_nodes(
638 self._find_removed(skeleton.iter_outputs_of(quantum_key), helper.outputs)
639 )
640 if helper.inputs_adjusted:
641 if not any(bool(adjusted_refs) for adjusted_refs in helper.inputs.values()):
642 raise QuantumGraphBuilderError(
643 f"adjustQuantum implementation for {task_node.label}@{quantum_key.data_id_values} "
644 "returned outputs but no inputs."
645 )
646 # Remove input dataset edges that were not retained by
647 # adjustQuantum. We can't remove the input dataset nodes
648 # because some other quantum might still want them.
649 skeleton.remove_input_edges(
650 quantum_key, self._find_removed(skeleton.iter_inputs_of(quantum_key), helper.inputs)
651 )
652 # Save the adjusted inputs and outputs to the quantum node's
653 # state so we don't have to regenerate those data structures
654 # from the graph.
655 skeleton[quantum_key]["inputs"] = helper.inputs
656 skeleton[quantum_key]["outputs"] = helper.outputs
657 for no_work_quantum in no_work_quanta:
658 skeleton.remove_quantum_node(no_work_quantum, remove_outputs=True)
659 remaining_quanta = skeleton.get_quanta(task_node.label)
660 self._resolve_task_init(task_node, skeleton, bool(skipped_quanta))
661 message_terms = []
662 if no_work_quanta:
663 message_terms.append(f"{len(no_work_quanta)} had no work to do")
664 if skipped_quanta:
665 message_terms.append(f"{len(skipped_quanta)} previously succeeded")
666 if adjuster.n_removed:
667 message_terms.append(f"{adjuster.n_removed} removed by adjust_all_quanta")
668 message_parenthetical = f" ({', '.join(message_terms)})" if message_terms else ""
669 if remaining_quanta:
670 self.log.info(
671 "Generated %s for task %s%s.",
672 _quantum_or_quanta(len(remaining_quanta)),
673 task_node.label,
674 message_parenthetical,
675 )
676 else:
677 self.log.info(
678 "Dropping task %s because no quanta remain%s.", task_node.label, message_parenthetical
679 )
680 skeleton.remove_task(task_node.label)
682 def _skip_quantum_if_metadata_exists(
683 self, task_node: TaskNode, quantum_key: QuantumKey, skeleton: QuantumGraphSkeleton
684 ) -> bool:
685 """Identify and drop quanta that should be skipped because their
686 metadata datasets already exist.
688 Parameters
689 ----------
690 task_node : `pipeline_graph.TaskNode`
691 Node for this task in the pipeline graph.
692 quantum_key : `QuantumKey`
693 Identifier for this quantum in the graph.
694 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
695 Preliminary quantum graph, to be modified in-place.
697 Returns
698 -------
699 skipped : `bool`
700 `True` if the quantum is being skipped and has been removed from
701 the graph, `False` otherwise.
703 Notes
704 -----
705 If the metadata dataset for this quantum exists in the
706 `skip_existing_in` collections, the quantum will be skipped. This
707 causes the quantum node to be removed from the graph. Dataset nodes
708 that were previously the outputs of this quantum will be associated
709 with `lsst.daf.butler.DatasetRef` objects that were found in
710 ``skip_existing_in``, or will be removed if there is no such dataset
711 there. Any output dataset in `output_run` will be removed from the
712 "output in the way" category.
713 """
714 metadata_dataset_key = DatasetKey(
715 task_node.metadata_output.parent_dataset_type_name, quantum_key.data_id_values
716 )
717 if skeleton.get_output_for_skip(metadata_dataset_key):
718 # This quantum's metadata is already present in the the
719 # skip_existing_in collections; we'll skip it. But the presence of
720 # the metadata dataset doesn't guarantee that all of the other
721 # outputs we predicted are present; we have to check.
722 for output_dataset_key in list(skeleton.iter_outputs_of(quantum_key)):
723 # If this dataset was "in the way" (i.e. already in the
724 # output run), it isn't anymore.
725 skeleton.discard_output_in_the_way(output_dataset_key)
726 if (output_ref := skeleton.get_output_for_skip(output_dataset_key)) is not None:
727 # Populate the skeleton graph's node attributes
728 # with the existing DatasetRef, just like a
729 # predicted output of a non-skipped quantum.
730 skeleton.set_dataset_ref(output_ref, output_dataset_key)
731 else:
732 # Remove this dataset from the skeleton graph,
733 # because the quantum that would have produced it
734 # is being skipped and it doesn't already exist.
735 skeleton.remove_dataset_nodes([output_dataset_key])
736 # Removing the quantum node from the graph will happen outside this
737 # function.
738 return True
739 return False
741 @final
742 def _update_quantum_for_adjust(
743 self,
744 quantum_key: QuantumKey,
745 skeleton: QuantumGraphSkeleton,
746 task_prerequisite_info: PrerequisiteInfo,
747 skypix_bounds_builder: SkyPixBoundsBuilder,
748 timespan_builder: TimespanBuilder,
749 ) -> None:
750 """Update the quantum node in the skeleton by finding remaining
751 prerequisite inputs and dropping regular inputs that we now know will
752 not be produced.
754 Parameters
755 ----------
756 quantum_key : `QuantumKey`
757 Identifier for this quantum in the graph.
758 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
759 Preliminary quantum graph, to be modified in-place.
760 task_prerequisite_info : `~prerequisite_helpers.PrerequisiteInfo`
761 Information about the prerequisite inputs to this task.
762 skypix_bounds_builder : `~prerequisite_helpers.SkyPixBoundsBuilder`
763 An object that accumulates the appropriate spatial bounds for a
764 quantum.
765 timespan_builder : `~prerequisite_helpers.TimespanBuilder`
766 An object that accumulates the appropriate timespan for a quantum.
768 Notes
769 -----
770 This first looks for outputs already present in the `output_run` (i.e.
771 "in the way" in the skeleton); if it finds something and `clobber` is
772 `True`, it uses that ref (it's not ideal that both the original dataset
773 and its replacement will have the same UUID, but we don't have space in
774 the quantum graph for two UUIDs, and we need the datastore records of
775 the original there). If `clobber` is `False`, `RuntimeError` is
776 raised. If there is no output already present, a new one with a random
777 UUID is generated. In all cases the dataset node in the skeleton is
778 associated with a `lsst.daf.butler.DatasetRef`.
779 """
780 dataset_key: DatasetKey | PrerequisiteDatasetKey
781 for dataset_key in skeleton.iter_outputs_of(quantum_key):
782 dataset_data_id = skeleton.get_data_id(dataset_key)
783 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
784 if (ref := skeleton.get_output_in_the_way(dataset_key)) is None:
785 ref = DatasetRef(dataset_type_node.dataset_type, dataset_data_id, run=self.output_run)
786 elif not self.clobber:
787 # We intentionally raise here, before running adjustQuantum,
788 # because it'd be weird if we left an old potential output of a
789 # task sitting there in the output collection, just because the
790 # task happened to not actually produce it.
791 raise OutputExistsError(
792 f"Potential output dataset {ref} already exists in the output run "
793 f"{self.output_run}, but clobbering outputs was not expected to be necessary."
794 )
795 skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
796 timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, dataset_data_id)
797 skeleton.set_dataset_ref(ref, dataset_key)
798 quantum_data_id = skeleton.get_data_id(quantum_key)
799 # Process inputs already present in the skeleton - this should include
800 # all regular inputs (including intermediates) and may include some
801 # prerequisites.
802 for dataset_key in list(skeleton.iter_inputs_of(quantum_key)):
803 if (ref := skeleton.get_dataset_ref(dataset_key)) is None:
804 # If the dataset ref hasn't been set either as an existing
805 # input or as an output of an already-processed upstream
806 # quantum, it's not going to be produced; remove it.
807 skeleton.remove_dataset_nodes([dataset_key])
808 continue
809 skypix_bounds_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
810 timespan_builder.handle_dataset(dataset_key.parent_dataset_type_name, ref.dataId)
811 # Query for any prerequisites not handled by process_subgraph. Note
812 # that these were not already in the skeleton graph, so we add them
813 # now.
814 skypix_bounds = skypix_bounds_builder.finish()
815 timespan = timespan_builder.finish()
816 for finder in task_prerequisite_info.finders.values():
817 dataset_keys = []
818 for ref in finder.find(
819 self.butler, self.input_collections, quantum_data_id, skypix_bounds, timespan
820 ):
821 dataset_key = skeleton.add_prerequisite_node(ref)
822 dataset_keys.append(dataset_key)
823 skeleton.add_input_edges(quantum_key, dataset_keys)
825 @final
826 def _adapt_quantum_outputs(
827 self,
828 task_node: TaskNode,
829 quantum_key: QuantumKey,
830 skeleton: QuantumGraphSkeleton,
831 ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
832 """Adapt outputs for a preliminary quantum and put them into the form
833 used by `~lsst.daf.butler.Quantum` and
834 `~PipelineTaskConnections.adjustQuantum`.
836 Parameters
837 ----------
838 task_node : `pipeline_graph.TaskNode`
839 Node for this task in the pipeline graph.
840 quantum_key : `QuantumKey`
841 Identifier for this quantum in the graph.
842 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
843 Preliminary quantum graph, to be modified in-place.
845 Returns
846 -------
847 outputs : `~lsst.daf.butler.NamedKeyDict` [ \
848 `~lsst.daf.butler.DatasetType`, `list` [ \
849 `~lsst.daf.butler.DatasetRef` ] ]
850 All outputs to the task, using the storage class and components
851 defined by the task's own connections.
852 """
853 outputs_by_type: dict[str, list[DatasetRef]] = {}
854 dataset_key: DatasetKey
855 for dataset_key in skeleton.iter_outputs_of(quantum_key):
856 ref = skeleton.get_dataset_ref(dataset_key)
857 assert ref is not None, "Should have been added (or the node removed) in a previous pass."
858 outputs_by_type.setdefault(dataset_key.parent_dataset_type_name, []).append(ref)
859 adapted_outputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
860 for write_edge in task_node.iter_all_outputs():
861 dataset_type_node = self._pipeline_graph.dataset_types[write_edge.parent_dataset_type_name]
862 edge_dataset_type = write_edge.adapt_dataset_type(dataset_type_node.dataset_type)
863 adapted_outputs[edge_dataset_type] = [
864 write_edge.adapt_dataset_ref(ref)
865 for ref in sorted(outputs_by_type.get(write_edge.parent_dataset_type_name, []))
866 ]
867 return adapted_outputs
869 @final
870 def _adapt_quantum_inputs(
871 self,
872 task_node: TaskNode,
873 quantum_key: QuantumKey,
874 skeleton: QuantumGraphSkeleton,
875 ) -> NamedKeyDict[DatasetType, list[DatasetRef]]:
876 """Adapt input datasets for a preliminary quantum into the form used by
877 `~lsst.daf.butler.Quantum` and
878 `~PipelineTaskConnections.adjustQuantum`.
880 Parameters
881 ----------
882 task_node : `pipeline_graph.TaskNode`
883 Node for this task in the pipeline graph.
884 quantum_key : `QuantumKey`
885 Identifier for this quantum in the graph.
886 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
887 Preliminary quantum graph, to be modified in-place.
889 Returns
890 -------
891 inputs : `~lsst.daf.butler.NamedKeyDict` [ \
892 `~lsst.daf.butler.DatasetType`, `list` [ \
893 `~lsst.daf.butler.DatasetRef` ] ]
894 All regular and prerequisite inputs to the task, using the storage
895 class and components defined by the task's own connections.
897 Notes
898 -----
899 This method trims input dataset nodes that are not already associated
900 with a `lsst.daf.butler.DatasetRef`, and queries for prerequisite input
901 nodes that do not exist.
902 """
903 inputs_by_type: dict[str, set[DatasetRef]] = {}
904 dataset_key: DatasetKey | PrerequisiteDatasetKey
905 for dataset_key in list(skeleton.iter_inputs_of(quantum_key)):
906 ref = skeleton.get_dataset_ref(dataset_key)
907 assert ref is not None, "Should have been added (or the node removed) in a previous pass."
908 inputs_by_type.setdefault(dataset_key.parent_dataset_type_name, set()).add(ref)
909 adapted_inputs: NamedKeyDict[DatasetType, list[DatasetRef]] = NamedKeyDict()
910 for read_edge in task_node.iter_all_inputs():
911 dataset_type_node = self._pipeline_graph.dataset_types[read_edge.parent_dataset_type_name]
912 edge_dataset_type = read_edge.adapt_dataset_type(dataset_type_node.dataset_type)
913 if (current_dataset_type := adapted_inputs.keys().get(edge_dataset_type.name)) is None:
914 adapted_inputs[edge_dataset_type] = [
915 read_edge.adapt_dataset_ref(ref)
916 for ref in sorted(inputs_by_type.get(read_edge.parent_dataset_type_name, frozenset()))
917 ]
918 elif current_dataset_type != edge_dataset_type:
919 raise NotImplementedError(
920 f"Task {task_node.label!r} has {edge_dataset_type.name!r} as an input via "
921 "two different connections, with two different storage class overrides. "
922 "This is not yet supported due to limitations in the Quantum data structure."
923 )
924 # If neither the `if` nor the `elif` above match, it means
925 # multiple input connections have exactly the same dataset
926 # type, and hence nothing to do after the first one.
927 return adapted_inputs
929 @final
930 def _resolve_task_init(
931 self, task_node: TaskNode, skeleton: QuantumGraphSkeleton, has_skipped_quanta: bool
932 ) -> None:
933 """Add init-input and init-output dataset nodes and edges for a task to
934 the skeleton.
936 Parameters
937 ----------
938 task_node : `pipeline_graph.TaskNode`
939 Pipeline graph description of the task.
940 skeleton : `QuantumGraphSkeleton`
941 In-progress quantum graph data structure to update in-place.
942 has_skipped_quanta : `bool`
943 Whether any of this task's quanta were skipped because they had
944 already succeeded.
945 """
946 quanta = skeleton.get_quanta(task_node.label)
947 task_init_key = TaskInitKey(task_node.label)
948 if quanta:
949 adapted_inputs: NamedKeyDict[DatasetType, DatasetRef] = NamedKeyDict()
950 # Process init-inputs.
951 input_keys: list[DatasetKey] = []
952 for read_edge in task_node.init.iter_all_inputs():
953 dataset_key = skeleton.add_dataset_node(
954 read_edge.parent_dataset_type_name, self.empty_data_id
955 )
956 skeleton.add_input_edge(task_init_key, dataset_key)
957 if (ref := skeleton.get_dataset_ref(dataset_key)) is None:
958 try:
959 ref = self.empty_dimensions_datasets.inputs[dataset_key]
960 except KeyError:
961 raise InitInputMissingError(
962 f"Overall init-input dataset {read_edge.parent_dataset_type_name!r} "
963 f"needed by task {task_node.label!r} not found in input collection(s) "
964 f"{self.input_collections}."
965 ) from None
966 skeleton.set_dataset_ref(ref, dataset_key)
967 for quantum_key in skeleton.get_quanta(task_node.label):
968 skeleton.add_input_edge(quantum_key, dataset_key)
969 input_keys.append(dataset_key)
970 adapted_ref = read_edge.adapt_dataset_ref(ref)
971 adapted_inputs[adapted_ref.datasetType] = adapted_ref
972 # Save the quantum-adapted init inputs to each quantum, and add
973 # skeleton edges connecting the init inputs to each quantum.
974 for quantum_key in skeleton.get_quanta(task_node.label):
975 skeleton[quantum_key]["init_inputs"] = adapted_inputs
976 # Process init-outputs.
977 adapted_outputs: NamedKeyDict[DatasetType, DatasetRef] = NamedKeyDict()
978 for write_edge in task_node.init.iter_all_outputs():
979 dataset_key = skeleton.add_dataset_node(
980 write_edge.parent_dataset_type_name, self.empty_data_id
981 )
982 if (ref := self.empty_dimensions_datasets.outputs_in_the_way.get(dataset_key)) is None:
983 ref = DatasetRef(
984 self._pipeline_graph.dataset_types[write_edge.parent_dataset_type_name].dataset_type,
985 self.empty_data_id,
986 run=self.output_run,
987 )
988 skeleton.set_dataset_ref(ref, dataset_key)
989 skeleton.add_output_edge(task_init_key, dataset_key)
990 adapted_ref = write_edge.adapt_dataset_ref(ref)
991 adapted_outputs[adapted_ref.datasetType] = adapted_ref
992 skeleton[task_init_key]["inputs"] = adapted_inputs
993 skeleton[task_init_key]["outputs"] = adapted_outputs
994 elif has_skipped_quanta:
995 # No quanta remain for this task, but at least one quantum was
996 # skipped because its outputs were present in the skip_existing_in
997 # collections. This means all init outputs should be present in
998 # the skip_existing_in collections, too, and we need to put those
999 # refs in the graph.
1000 for write_edge in task_node.init.iter_all_outputs():
1001 dataset_key = skeleton.add_dataset_node(
1002 write_edge.parent_dataset_type_name, self.empty_data_id
1003 )
1004 if (ref := self.empty_dimensions_datasets.outputs_for_skip.get(dataset_key)) is None:
1005 raise InitInputMissingError(
1006 f"Init-output dataset {write_edge.parent_dataset_type_name!r} of skipped task "
1007 f"{task_node.label!r} not found in skip-existing-in collection(s) "
1008 f"{self.skip_existing_in}."
1009 ) from None
1010 skeleton.set_dataset_ref(ref, dataset_key)
1011 # If this dataset was "in the way" (i.e. already in the output
1012 # run), it isn't anymore.
1013 skeleton.discard_output_in_the_way(dataset_key)
1014 # No quanta remain in this task, but none were skipped; this means
1015 # they all got pruned because of NoWorkFound conditions. This
1016 # dooms all downstream quanta to the same fate, so we don't bother
1017 # doing anything with the task's init-outputs, since nothing is
1018 # going to consume them.
1020 @final
1021 @timeMethod
1022 def _find_empty_dimension_datasets(self) -> EmptyDimensionsDatasets:
1023 """Query for all dataset types with no dimensions, updating
1024 `empty_dimensions_datasets` in-place.
1026 This includes but is not limited to init inputs and init outputs.
1027 """
1028 inputs: dict[DatasetKey | PrerequisiteDatasetKey, DatasetRef] = {}
1029 outputs_for_skip: dict[DatasetKey, DatasetRef] = {}
1030 outputs_in_the_way: dict[DatasetKey, DatasetRef] = {}
1031 _, dataset_type_nodes = self._pipeline_graph.group_by_dimensions().get(self.universe.empty, ({}, {}))
1032 dataset_types = [node.dataset_type for node in dataset_type_nodes.values()]
1033 dataset_types.extend(self._global_init_output_types.values())
1034 for dataset_type in dataset_types:
1035 key = DatasetKey(dataset_type.name, self.empty_data_id.required_values)
1036 if (
1037 self._pipeline_graph.producer_of(dataset_type.name) is None
1038 and dataset_type.name not in self._global_init_output_types
1039 ):
1040 # Dataset type is an overall input; we always need to try to
1041 # find these.
1042 try:
1043 ref = self.butler.find_dataset(dataset_type.name, collections=self.input_collections)
1044 except MissingDatasetTypeError:
1045 ref = None
1046 if ref is not None:
1047 inputs[key] = ref
1048 elif self.skip_existing_in:
1049 # Dataset type is an intermediate or output; need to find these
1050 # if only they're from previously executed quanta that we might
1051 # skip...
1052 try:
1053 ref = self.butler.find_dataset(dataset_type.name, collections=self.skip_existing_in)
1054 except MissingDatasetTypeError:
1055 ref = None
1056 if ref is not None:
1057 outputs_for_skip[key] = ref
1058 if ref.run == self.output_run:
1059 outputs_in_the_way[key] = ref
1060 if self.output_run_exists and not self.skip_existing_starts_with_output_run:
1061 # ...or if they're in the way and would need to be clobbered
1062 # (and we haven't already found them in the previous block).
1063 try:
1064 ref = self.butler.find_dataset(dataset_type.name, collections=[self.output_run])
1065 except MissingDatasetTypeError:
1066 ref = None
1067 if ref is not None:
1068 outputs_in_the_way[key] = ref
1069 return EmptyDimensionsDatasets(
1070 inputs=inputs, outputs_for_skip=outputs_for_skip, outputs_in_the_way=outputs_in_the_way
1071 )
1073 @final
1074 @timeMethod
1075 def _attach_datastore_records(self, skeleton: QuantumGraphSkeleton) -> None:
1076 """Add datastore records for all overall inputs to a preliminary
1077 quantum graph.
1079 Parameters
1080 ----------
1081 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
1082 Preliminary quantum graph to update in place.
1084 Notes
1085 -----
1086 On return, all quantum nodes in the skeleton graph will have a
1087 "datastore_records" attribute that is a mapping from datastore name
1088 to `lsst.daf.butler.DatastoreRecordData`, as used by
1089 `lsst.daf.butler.Quantum`.
1090 """
1091 self.log.info("Fetching and attaching datastore records for all overall inputs.")
1092 overall_inputs = skeleton.extract_overall_inputs()
1093 exported_records = self.butler._datastore.export_records(overall_inputs.values())
1094 for task_label in self._pipeline_graph.tasks:
1095 if not skeleton.has_task(task_label):
1096 continue
1097 self.log.verbose("Fetching and attaching datastore records for task %s.", task_label)
1098 task_init_key = skeleton.get_task_init_node(task_label)
1099 init_input_ids = {
1100 ref.id
1101 for dataset_key in skeleton.iter_inputs_of(task_init_key)
1102 if (ref := overall_inputs.get(dataset_key)) is not None
1103 }
1104 init_records = {}
1105 if init_input_ids:
1106 for datastore_name, records in exported_records.items():
1107 matching_records = records.subset(init_input_ids)
1108 if matching_records is not None:
1109 init_records[datastore_name] = matching_records
1110 skeleton[task_init_key]["datastore_records"] = init_records
1111 for quantum_key in skeleton.get_quanta(task_label):
1112 quantum_records = {}
1113 input_ids = {
1114 ref.id
1115 for dataset_key in skeleton.iter_inputs_of(quantum_key)
1116 if (ref := overall_inputs.get(dataset_key)) is not None
1117 }
1118 if input_ids:
1119 for datastore_name, records in exported_records.items():
1120 matching_records = records.subset(input_ids)
1121 if matching_records is not None:
1122 quantum_records[datastore_name] = matching_records
1123 skeleton[quantum_key]["datastore_records"] = quantum_records
1125 @final
1126 @timeMethod
1127 def _construct_quantum_graph(
1128 self, skeleton: QuantumGraphSkeleton, metadata: Mapping[str, Any]
1129 ) -> QuantumGraph:
1130 """Construct a `.QuantumGraph` object from the contents of a
1131 fully-processed `.quantum_graph_skeleton.QuantumGraphSkeleton`.
1133 Parameters
1134 ----------
1135 skeleton : `.quantum_graph_skeleton.QuantumGraphSkeleton`
1136 Preliminary quantum graph. Must have "init_inputs", "inputs", and
1137 "outputs" attributes on all quantum nodes, as added by
1138 `_resolve_task_quanta`, as well as a "datastore_records" attribute
1139 as added by `_attach_datastore_records`.
1140 metadata : `~collections.abc.Mapping`
1141 Flexible metadata to add to the graph.
1143 Returns
1144 -------
1145 quantum_graph : `.QuantumGraph`
1146 DAG describing processing to be performed.
1147 """
1148 from .graph import QuantumGraph
1150 self.log.info("Transforming graph skeleton into a QuantumGraph instance.")
1151 quanta: dict[TaskDef, set[Quantum]] = {}
1152 init_inputs: dict[TaskDef, Iterable[DatasetRef]] = {}
1153 init_outputs: dict[TaskDef, Iterable[DatasetRef]] = {}
1154 for task_def in self._pipeline_graph._iter_task_defs():
1155 if not skeleton.has_task(task_def.label):
1156 continue
1157 self.log.verbose("Transforming graph skeleton nodes for task %s.", task_def.label)
1158 task_node = self._pipeline_graph.tasks[task_def.label]
1159 task_init_key = skeleton.get_task_init_node(task_def.label)
1160 task_init_state = skeleton[task_init_key]
1161 init_datastore_records: dict[str, DatastoreRecordData] = task_init_state.get(
1162 "datastore_records", {}
1163 )
1164 init_inputs[task_def] = task_init_state["inputs"].values()
1165 init_outputs[task_def] = task_init_state["outputs"].values()
1166 quanta_for_task: set[Quantum] = set()
1167 for quantum_key in skeleton.get_quanta(task_node.label):
1168 quantum_state = skeleton[quantum_key]
1169 quantum_datastore_records: dict[str, DatastoreRecordData] = quantum_state.get(
1170 "datastore_records", {}
1171 )
1172 quanta_for_task.add(
1173 Quantum(
1174 taskName=task_node.task_class_name,
1175 taskClass=task_node.task_class,
1176 dataId=quantum_state["data_id"],
1177 initInputs=quantum_state["init_inputs"],
1178 inputs=quantum_state["inputs"],
1179 outputs=quantum_state["outputs"],
1180 datastore_records=DatastoreRecordData.merge_mappings(
1181 quantum_datastore_records, init_datastore_records
1182 ),
1183 )
1184 )
1185 quanta[task_def] = quanta_for_task
1187 registry_dataset_types: list[DatasetType] = [
1188 node.dataset_type for node in self._pipeline_graph.dataset_types.values()
1189 ]
1191 all_metadata = self.metadata.to_dict()
1192 all_metadata.update(metadata)
1193 global_init_outputs: list[DatasetRef] = []
1194 for dataset_key in skeleton.global_init_outputs:
1195 ref = skeleton.get_dataset_ref(dataset_key)
1196 assert ref is not None, "Global init input refs should be resolved already."
1197 global_init_outputs.append(ref)
1198 self.log.verbose("Invoking QuantumGraph class constructor.")
1199 result = QuantumGraph(
1200 quanta,
1201 metadata=all_metadata,
1202 universe=self.universe,
1203 initInputs=init_inputs,
1204 initOutputs=init_outputs,
1205 globalInitOutputs=global_init_outputs,
1206 registryDatasetTypes=registry_dataset_types,
1207 )
1208 self.log.info("Graph build complete.")
1209 return result
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 self.log.info("Transforming graph skeleton into PredictedQuantumGraph components.")
1247 components = PredictedQuantumGraphComponents(pipeline_graph=self._pipeline_graph)
1248 components.header.inputs = list(self.input_collections)
1249 components.header.output_run = self.output_run
1250 components.header.output = output
1251 if metadata is not None:
1252 components.header.metadata.update(metadata)
1253 components.dimension_data = DimensionDataAttacher(
1254 records=skeleton.get_dimension_data(),
1255 dimensions=self._pipeline_graph.get_all_dimensions(),
1256 )
1257 components.init_quanta.root = [
1258 PredictedQuantumDatasetsModel.model_construct(
1259 quantum_id=generate_uuidv7(),
1260 task_label="",
1261 outputs={
1262 dataset_key.parent_dataset_type_name: [
1263 PredictedDatasetModel.from_dataset_ref(
1264 cast(DatasetRef, skeleton.get_dataset_ref(dataset_key))
1265 )
1266 ]
1267 for dataset_key in skeleton.global_init_outputs
1268 },
1269 )
1270 ]
1271 for task_node in self._pipeline_graph.tasks.values():
1272 if not skeleton.has_task(task_node.label):
1273 continue
1274 self.log.verbose("Transforming graph skeleton nodes for task %s.", task_node.label)
1275 task_init_key = TaskInitKey(task_node.label)
1276 init_quantum_datasets = PredictedQuantumDatasetsModel.model_construct(
1277 quantum_id=generate_uuidv7(),
1278 task_label=task_node.label,
1279 inputs=self._make_predicted_datasets(
1280 skeleton,
1281 task_node.init.iter_all_inputs(),
1282 skeleton.iter_inputs_of(task_init_key),
1283 ),
1284 outputs=self._make_predicted_datasets(
1285 skeleton,
1286 task_node.init.iter_all_outputs(),
1287 skeleton.iter_outputs_of(task_init_key),
1288 ),
1289 datastore_records={
1290 datastore_name: records.to_simple()
1291 for datastore_name, records in skeleton[task_init_key]
1292 .get("datastore_records", {})
1293 .items()
1294 },
1295 )
1296 components.init_quanta.root.append(init_quantum_datasets)
1297 for quantum_key in skeleton.get_quanta(task_node.label):
1298 quantum_datasets = PredictedQuantumDatasetsModel.model_construct(
1299 quantum_id=generate_uuidv7(),
1300 task_label=task_node.label,
1301 data_coordinate=list(skeleton.get_data_id(quantum_key).full_values),
1302 inputs=self._make_predicted_datasets(
1303 skeleton,
1304 task_node.iter_all_inputs(),
1305 skeleton.iter_inputs_of(quantum_key),
1306 ),
1307 outputs=self._make_predicted_datasets(
1308 skeleton,
1309 task_node.iter_all_outputs(),
1310 skeleton.iter_outputs_of(quantum_key),
1311 ),
1312 datastore_records={
1313 datastore_name: records.to_simple()
1314 for datastore_name, records in skeleton[quantum_key]
1315 .get("datastore_records", {})
1316 .items()
1317 },
1318 )
1319 components.quantum_datasets[quantum_datasets.quantum_id] = quantum_datasets
1320 self.log.verbose("Building the thin summary graph.")
1321 components.set_thin_graph()
1322 components.set_header_counts()
1323 self.log.info("Graph build complete.")
1324 return components
1326 @staticmethod
1327 def _make_predicted_datasets(
1328 skeleton: QuantumGraphSkeleton,
1329 edges: Iterable[Edge],
1330 dataset_keys: Iterable[DatasetKey | PrerequisiteDatasetKey],
1331 ) -> dict[str, list[PredictedDatasetModel]]:
1332 """Make the predicted quantum graph model objects that represent the
1333 datasets from an iterable of pipeline graph edges.
1335 Parameters
1336 ----------
1337 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
1338 Temporary data structure used by the builder to represent the
1339 graph.
1340 edges : `~collections.abc.Iterable` [ `.pipeline_graph.Edge` ]
1341 Pipeline graph edges.
1342 dataset_keys : `~collections.abc.Iterable` [ \
1343 `.quantum_graph_skeleton.DatasetKey` or\
1344 `.quantum_graph_skeleton.PrerequisiteDatasetKey` ]
1345 All nodes in the skeleton that correspond to any of the given
1346 pipeline graph edges.
1348 Returns
1349 -------
1350 predicted_datasets : `dict` [ `str`, \
1351 `list` [ `.quantum_graph.PredictedDatasetModel` ] ]
1352 Mapping of dataset models, keyed by connection name.
1353 """
1354 from .quantum_graph import PredictedDatasetModel
1356 connection_names_by_dataset_type: defaultdict[str, list[str]] = defaultdict(list)
1357 result: dict[str, list[PredictedDatasetModel]] = {}
1358 for edge in edges:
1359 connection_names_by_dataset_type[edge.parent_dataset_type_name].append(edge.connection_name)
1360 result[edge.connection_name] = []
1362 for dataset_key in dataset_keys:
1363 connection_names = connection_names_by_dataset_type.get(dataset_key.parent_dataset_type_name)
1364 if connection_names is None:
1365 # Ignore if this isn't one of the connections we're processing
1366 # (probably an init-input), which would also be predecessor to
1367 # a quantum node, but should be handled separately.
1368 continue
1369 ref = skeleton.get_dataset_ref(dataset_key)
1370 assert ref is not None, "DatasetRefs should have already been added to skeleton."
1371 for connection_name in connection_names:
1372 result[connection_name].append(PredictedDatasetModel.from_dataset_ref(ref))
1373 for refs in result.values():
1374 refs.sort(key=operator.attrgetter("data_coordinate"))
1375 return result
1377 @staticmethod
1378 @final
1379 def _find_removed(
1380 original: Iterable[DatasetKey | PrerequisiteDatasetKey],
1381 adjusted: NamedKeyMapping[DatasetType, Sequence[DatasetRef]],
1382 ) -> set[DatasetKey | PrerequisiteDatasetKey]:
1383 """Identify skeleton-graph dataset nodes that have been removed by
1384 `~PipelineTaskConnections.adjustQuantum`.
1386 Parameters
1387 ----------
1388 original : `~collections.abc.Iterable` [ `DatasetKey` or \
1389 `PrerequisiteDatasetKey` ]
1390 Identifiers for the dataset nodes that were the original neighbors
1391 (inputs or outputs) of a quantum.
1392 adjusted : `~lsst.daf.butler.NamedKeyMapping` [ \
1393 `~lsst.daf.butler.DatasetType`, \
1394 `~collections.abc.Sequence` [ `lsst.daf.butler.DatasetType` ] ]
1395 Adjusted neighbors, in the form used by `lsst.daf.butler.Quantum`.
1397 Returns
1398 -------
1399 removed : `set` [ `DatasetKey` ]
1400 Datasets in ``original`` that have no counterpart in ``adjusted``.
1401 """
1402 result = set(original)
1403 for dataset_type, kept_refs in adjusted.items():
1404 parent_dataset_type_name, _ = DatasetType.splitDatasetTypeName(dataset_type.name)
1405 for kept_ref in kept_refs:
1406 # We don't know if this was a DatasetKey or a
1407 # PrerequisiteDatasetKey; just try both.
1408 result.discard(DatasetKey(parent_dataset_type_name, kept_ref.dataId.required_values))
1409 result.discard(PrerequisiteDatasetKey(parent_dataset_type_name, kept_ref.id.bytes))
1410 return result
1413@dataclasses.dataclass(eq=False, order=False)
1414class EmptyDimensionsDatasets:
1415 """Struct that holds the results of empty-dimensions dataset queries for
1416 `QuantumGraphBuilder`.
1417 """
1419 inputs: Mapping[DatasetKey | PrerequisiteDatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1420 """Overall-input datasets found in `QuantumGraphBuilder.input_collections`.
1422 This may include prerequisite inputs. It does include init-inputs.
1423 It does not include intermediates.
1424 """
1426 outputs_for_skip: Mapping[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1427 """Output datasets found in `QuantumGraphBuilder.skip_existing_in`.
1429 It is unspecified whether this contains init-outputs; there is
1430 no concept of skipping at the init stage, so this is not expected to
1431 matter.
1432 """
1434 outputs_in_the_way: Mapping[DatasetKey, DatasetRef] = dataclasses.field(default_factory=dict)
1435 """Output datasets found in `QuantumGraphBuilder.output_run`.
1437 This includes regular outputs and init-outputs.
1438 """
1441def _quantum_or_quanta(n: int) -> str:
1442 """Correctly pluralize 'quantum' if needed."""
1443 return f"{n} quanta" if n != 1 else "1 quantum"