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