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