Coverage for python / lsst / pipe / base / pipeline_graph / _pipeline_graph.py: 12%
746 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-24 01:12 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29__all__ = ("PipelineGraph", "compare_packages", "log_config_mismatch")
31import gzip
32import itertools
33import json
34import logging
35from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence, Set
36from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypeVar, cast
38import networkx
39import networkx.algorithms.bipartite
40import networkx.algorithms.dag
42from lsst.daf.butler import (
43 Butler,
44 DataCoordinate,
45 DataId,
46 DatasetRef,
47 DatasetType,
48 DimensionGroup,
49 DimensionUniverse,
50 MissingDatasetTypeError,
51)
52from lsst.daf.butler.registry import ConflictingDefinitionError, Registry
53from lsst.resources import ResourcePath, ResourcePathExpression
54from lsst.utils.packages import Packages
56from .._dataset_handle import InMemoryDatasetHandle
57from ..automatic_connection_constants import PACKAGES_INIT_OUTPUT_NAME, PACKAGES_INIT_OUTPUT_STORAGE_CLASS
58from . import expressions
59from ._dataset_types import DatasetTypeNode
60from ._edges import Edge, ReadEdge, WriteEdge
61from ._exceptions import (
62 DuplicateOutputError,
63 EdgesChangedError,
64 InvalidExpressionError,
65 InvalidStepsError,
66 PipelineDataCycleError,
67 PipelineGraphError,
68 PipelineGraphExceptionSafetyError,
69 UnresolvedGraphError,
70)
71from ._mapping_views import DatasetTypeMappingView, TaskMappingView
72from ._nodes import NodeKey, NodeType
73from ._task_subsets import StepDefinitions, TaskSubset
74from ._tasks import TaskImportMode, TaskInitNode, TaskNode, _TaskNodeImportedData
76if TYPE_CHECKING:
77 from ..config import PipelineTaskConfig
78 from ..connections import PipelineTaskConnections
79 from ..pipeline import TaskDef
80 from ..pipelineTask import PipelineTask
83_G = TypeVar("_G", bound=networkx.DiGraph | networkx.MultiDiGraph)
85_LOG = logging.getLogger("lsst.pipe.base.pipeline_graph")
88class PipelineGraph:
89 """A graph representation of fully-configured pipeline.
91 `PipelineGraph` instances are typically constructed by calling
92 `.Pipeline.to_graph`, but in rare cases constructing and then populating an
93 empty one may be preferable.
95 Parameters
96 ----------
97 description : `str`, optional
98 String description for this pipeline.
99 universe : `lsst.daf.butler.DimensionUniverse`, optional
100 Definitions for all butler dimensions. If not provided, some
101 attributes will not be available until `resolve` is called.
102 data_id : `lsst.daf.butler.DataCoordinate` or other data ID, optional
103 Data ID that represents a constraint on all quanta generated by this
104 pipeline. This typically just holds the instrument constraint included
105 in the pipeline definition, if there was one.
106 """
108 ###########################################################################
109 #
110 # Simple Pipeline Graph Inspection Interface:
111 #
112 # - for inspecting graph structure, not modifying it (except to sort and]
113 # resolve);
114 #
115 # - no NodeKey objects, just string dataset type name and task label keys;
116 #
117 # - graph structure is represented as a pair of mappings, with methods to
118 # find neighbors and edges of nodes.
119 #
120 ###########################################################################
122 def __init__(
123 self,
124 *,
125 description: str = "",
126 universe: DimensionUniverse | None = None,
127 data_id: DataId | None = None,
128 ) -> None:
129 self._init_from_args(
130 xgraph=None,
131 sorted_keys=None,
132 task_subsets=None,
133 description=description,
134 universe=universe,
135 data_id=data_id,
136 step_definitions=StepDefinitions(universe),
137 )
139 def __repr__(self) -> str:
140 return f"{type(self).__name__}({self.description!r}, tasks={self.tasks!s})"
142 @property
143 def description(self) -> str:
144 """String description for this pipeline."""
145 return self._description
147 @description.setter
148 def description(self, value: str) -> None:
149 # Docstring in setter.
150 self._description = value
152 @property
153 def universe(self) -> DimensionUniverse:
154 """Definitions for all butler dimensions."""
155 if self._universe is None:
156 raise UnresolvedGraphError("Pipeline graph is not resolved.")
157 return self._universe
159 @property
160 def data_id(self) -> DataCoordinate:
161 """Data ID that represents a constraint on all quanta generated from
162 this pipeline.
164 This is may not be available unless the graph is resolved.
165 """
166 return DataCoordinate.standardize(self._raw_data_id, universe=self.universe)
168 @property
169 def tasks(self) -> TaskMappingView:
170 """A mapping view of the tasks in the graph.
172 This mapping has `str` task label keys and `TaskNode` values. Iteration
173 is topologically and deterministically ordered if and only if `sort`
174 has been called since the last modification to the graph.
175 """
176 return self._tasks
178 @property
179 def dataset_types(self) -> DatasetTypeMappingView:
180 """A mapping view of the dataset types in the graph.
182 This mapping has `str` parent dataset type name keys, but only provides
183 access to its `DatasetTypeNode` values if `resolve` has been called
184 since the last modification involving a task that uses a dataset type.
185 See `DatasetTypeMappingView` for details.
186 """
187 return self._dataset_types
189 @property
190 def task_subsets(self) -> Mapping[str, TaskSubset]:
191 """A mapping of all labeled subsets of tasks.
193 Keys are subset labels, values are sets of task labels. See
194 `TaskSubset` for more information.
196 Use `add_task_subset` to add a new subset. The subsets themselves may
197 be modified in-place.
198 """
199 return self._task_subsets
201 @property
202 def steps(self) -> StepDefinitions:
203 """An ordered iterable of the labels of task subsets that must be
204 executed separately.
206 Steps are intended to be partition of the full pipeline graph - the
207 steps together include all tasks, and each task belongs to exactly one
208 step - but this is only checked when the graph is resolved.
210 The order of the steps is required to be consistent with the graph's
211 topological ordering, and when sorting a graph with steps, that sort
212 order is guaranteed to be consistent with the order of the steps. But
213 there may be other orderings of the steps that are still valid for
214 execution (e.g. two steps could be runnable in parallel).
215 """
216 return self._step_definitions
218 @steps.setter
219 def steps(self, labels: Iterable[str]) -> None:
220 # Docstring on getter.
221 self._step_definitions.assign(labels)
223 def get_task_step(self, task_label: str) -> str:
224 """Return the step that the given task belongs to.
226 Parameters
227 ----------
228 task_label : `str`
229 Label of the task to look up.
231 Returns
232 -------
233 step_label : `str`
234 Step the task belongs to.
236 Raises
237 ------
238 InvalidStepsError
239 Raised if the pipeline has no steps defined.
240 UnresolvedGraphError
241 Raised if the pipeline has been modified since the steps were
242 last verified.
243 """
244 if not self._step_definitions:
245 raise InvalidStepsError("No steps have been defined for this pipeline.")
246 elif self._step_definitions.verified:
247 return self._xgraph.nodes[NodeKey(NodeType.TASK, task_label)]["step"]
248 raise UnresolvedGraphError("Steps have not been verified since the last modification.")
250 @property
251 def is_fully_resolved(self) -> bool:
252 """Whether all of this graph's nodes are resolved and any all have
253 been checked for correctness.
255 A fully-resolved graph is always sorted as well, and a fully-resolved
256 graph with sorted
257 """
258 return (
259 self._universe is not None
260 and self._step_definitions.verified
261 and self._sorted_keys is not None
262 and all(self.dataset_types.is_resolved(k) for k in self.dataset_types)
263 )
265 @property
266 def has_been_sorted(self) -> bool:
267 """Whether this graph's tasks and dataset types have been
268 topologically sorted (with unspecified but deterministic tiebreakers)
269 since the last modification to the graph.
271 If the pipeline graph has step definitions, the sort order is
272 consistent with the step order.
274 This may return `False` if the graph *happens* to be sorted but `sort`
275 was never called,.
276 """
277 return self._sorted_keys is not None
279 def sort(self) -> None:
280 """Sort this graph's nodes topologically with deterministic (but
281 unspecified) tiebreakers.
283 This does nothing if the graph is already known to be sorted.
284 """
285 if self._sorted_keys is None:
286 try:
287 sorted_keys: Sequence[NodeKey] = list(networkx.lexicographical_topological_sort(self._xgraph))
288 except networkx.NetworkXUnfeasible as err: # pragma: no cover
289 # Should't be possible to get here, because we check for cycles
290 # when adding tasks, but we guard against it anyway.
291 cycle = networkx.find_cycle(self._xgraph)
292 raise PipelineDataCycleError(
293 f"Cycle detected while attempting to sort graph: {cycle}."
294 ) from err
295 self._reorder(sorted_keys)
297 def copy(self) -> PipelineGraph:
298 """Return a copy of this graph that copies all mutable state."""
299 xgraph = self._xgraph.copy()
300 step_definitions = self._step_definitions.copy()
301 result = PipelineGraph.__new__(PipelineGraph)
302 result._init_from_args(
303 xgraph,
304 self._sorted_keys,
305 task_subsets={
306 k: TaskSubset(xgraph, v.label, set(v._members), v.description, step_definitions)
307 for k, v in self._task_subsets.items()
308 },
309 description=self._description,
310 universe=self._universe,
311 data_id=self._raw_data_id,
312 step_definitions=step_definitions,
313 )
314 return result
316 def __copy__(self) -> PipelineGraph:
317 # Fully shallow copies are dangerous; we don't want shared mutable
318 # state to lead to broken class invariants.
319 return self.copy()
321 def __deepcopy__(self, memo: dict) -> PipelineGraph:
322 # Genuine deep copies are unnecessary, since we should only ever care
323 # that mutable state is copied.
324 return self.copy()
326 def diff_tasks(self, other: PipelineGraph) -> list[str]:
327 """Compare two pipeline graphs.
329 This only compares graph structure and task classes (including their
330 edges). It does *not* compare full configuration (which is subject to
331 spurious differences due to import-cache state), dataset type
332 resolutions, or sort state.
334 Parameters
335 ----------
336 other : `PipelineGraph`
337 Graph to compare to.
339 Returns
340 -------
341 differences : `list` [ `str` ]
342 List of string messages describing differences between the
343 pipelines. If empty, the graphs have the same tasks and
344 connections.
345 """
346 messages: list[str] = []
347 common_labels: Set[str]
348 if self.tasks.keys() != other.tasks.keys():
349 common_labels = self.tasks.keys() & other.tasks.keys()
350 messages.append(
351 f"Pipelines have different tasks: A & ~B = {list(self.tasks.keys() - common_labels)}, "
352 f"B & ~A = {list(other.tasks.keys() - common_labels)}."
353 )
354 else:
355 common_labels = self.tasks.keys()
356 for label in common_labels:
357 a = self.tasks[label]
358 b = other.tasks[label]
359 if a.task_class != b.task_class:
360 messages.append(
361 f"Task {label!r} has class {a.task_class_name} in A, but {b.task_class_name} in B."
362 )
363 messages.extend(a.diff_edges(b))
364 return messages
366 def producing_edge_of(self, dataset_type_name: str) -> WriteEdge | None:
367 """Return the `WriteEdge` that links the producing task to the named
368 dataset type.
370 Parameters
371 ----------
372 dataset_type_name : `str`
373 Dataset type name. Must not be a component.
375 Returns
376 -------
377 edge : `WriteEdge` or `None`
378 Producing edge or `None` if there isn't one in this graph.
380 Raises
381 ------
382 DuplicateOutputError
383 Raised if there are multiple tasks defined to produce this dataset
384 type. This is only possible if the graph's dataset types are not
385 resolved.
387 Notes
388 -----
389 On resolved graphs, it may be slightly more efficient to use::
391 graph.dataset_types[dataset_type_name].producing_edge
393 but this method works on graphs with unresolved dataset types as well.
394 """
395 producer: str | None = None
396 producing_edge: WriteEdge | None = None
397 for _, _, producing_edge in self._xgraph.in_edges(
398 NodeKey(NodeType.DATASET_TYPE, dataset_type_name), data="instance"
399 ):
400 assert producing_edge is not None, "Should only be None if we never loop."
401 if producer is not None:
402 raise DuplicateOutputError(
403 f"Dataset type {dataset_type_name!r} is produced by both {producing_edge.task_label!r} "
404 f"and {producer!r}."
405 )
406 return producing_edge
408 def consuming_edges_of(self, dataset_type_name: str) -> list[ReadEdge]:
409 """Return the `ReadEdge` objects that link the named dataset type to
410 the tasks that consume it.
412 Parameters
413 ----------
414 dataset_type_name : `str`
415 Dataset type name. Must not be a component.
417 Returns
418 -------
419 edges : `list` [ `ReadEdge` ]
420 Edges that connect this dataset type to the tasks that consume it.
422 Notes
423 -----
424 On resolved graphs, it may be slightly more efficient to use::
426 graph.dataset_types[dataset_type_name].producing_edges
428 but this method works on graphs with unresolved dataset types as well.
429 """
430 return [
431 edge
432 for _, _, edge in self._xgraph.out_edges(
433 NodeKey(NodeType.DATASET_TYPE, dataset_type_name), data="instance"
434 )
435 ]
437 def producer_of(self, dataset_type_name: str) -> TaskNode | TaskInitNode | None:
438 """Return the `TaskNode` or `TaskInitNode` that writes the given
439 dataset type.
441 Parameters
442 ----------
443 dataset_type_name : `str`
444 Dataset type name. Must not be a component.
446 Returns
447 -------
448 edge : `TaskNode`, `TaskInitNode`, or `None`
449 Producing node or `None` if there isn't one in this graph.
451 Raises
452 ------
453 DuplicateOutputError
454 Raised if there are multiple tasks defined to produce this dataset
455 type. This is only possible if the graph's dataset types are not
456 resolved.
457 """
458 if (producing_edge := self.producing_edge_of(dataset_type_name)) is not None:
459 return self._xgraph.nodes[producing_edge.task_key]["instance"]
460 return None
462 def consumers_of(self, dataset_type_name: str) -> list[TaskNode | TaskInitNode]:
463 """Return the `TaskNode` and/or `TaskInitNode` objects that read
464 the given dataset type.
466 Parameters
467 ----------
468 dataset_type_name : `str`
469 Dataset type name. Must not be a component.
471 Returns
472 -------
473 edges : `list` [ `ReadEdge` ]
474 Edges that connect this dataset type to the tasks that consume it.
476 Notes
477 -----
478 On resolved graphs, it may be slightly more efficient to use::
480 graph.dataset_types[dataset_type_name].producing_edges
482 but this method works on graphs with unresolved dataset types as well.
483 """
484 return [
485 self._xgraph.nodes[consuming_edge.task_key]["instance"]
486 for consuming_edge in self.consuming_edges_of(dataset_type_name)
487 ]
489 def inputs_of(self, task_label: str, init: bool = False) -> dict[str, DatasetTypeNode | None]:
490 """Return the dataset types that are inputs to a task.
492 Parameters
493 ----------
494 task_label : `str`
495 Label for the task in the pipeline.
496 init : `bool`, optional
497 If `True`, return init-input dataset types instead of runtime
498 (including prerequisite) inputs.
500 Returns
501 -------
502 inputs : `dict` [ `str`, `DatasetTypeNode` or `None` ]
503 Dictionary parent dataset type name keys and either
504 `DatasetTypeNode` values (if the dataset type has been resolved)
505 or `None` values.
507 Notes
508 -----
509 To get the input edges of a task or task init node (which provide
510 information about storage class overrides nd components) use::
512 graph.tasks[task_label].iter_all_inputs()
514 or
516 graph.tasks[task_label].init.iter_all_inputs()
518 or the various mapping attributes of the `TaskNode` and `TaskInitNode`
519 class.
520 """
521 node: TaskNode | TaskInitNode = self.tasks[task_label] if not init else self.tasks[task_label].init
522 return {
523 edge.parent_dataset_type_name: self._xgraph.nodes[edge.dataset_type_key]["instance"]
524 for edge in node.iter_all_inputs()
525 }
527 def outputs_of(
528 self, task_label: str, init: bool = False, include_automatic_connections: bool = True
529 ) -> dict[str, DatasetTypeNode | None]:
530 """Return the dataset types that are outputs of a task.
532 Parameters
533 ----------
534 task_label : `str`
535 Label for the task in the pipeline.
536 init : `bool`, optional
537 If `True`, return init-output dataset types instead of runtime
538 outputs.
539 include_automatic_connections : `bool`, optional
540 Whether to include automatic connections such as configs, metadata,
541 and logs.
543 Returns
544 -------
545 outputs : `dict` [ `str`, `DatasetTypeNode` or `None` ]
546 Dictionary parent dataset type name keys and either
547 `DatasetTypeNode` values (if the dataset type has been resolved)
548 or `None` values.
550 Notes
551 -----
552 To get the input edges of a task or task init node (which provide
553 information about storage class overrides nd components) use::
555 graph.tasks[task_label].iter_all_outputs()
557 or
559 graph.tasks[task_label].init.iter_all_outputs()
561 or the various mapping attributes of the `TaskNode` and `TaskInitNode`
562 class.
563 """
564 node: TaskNode | TaskInitNode = self.tasks[task_label] if not init else self.tasks[task_label].init
565 iterable = node.iter_all_outputs() if include_automatic_connections else node.outputs.values()
566 return {
567 edge.parent_dataset_type_name: self._xgraph.nodes[edge.dataset_type_key]["instance"]
568 for edge in iterable
569 }
571 def resolve(
572 self,
573 registry: Registry | None = None,
574 dimensions: DimensionUniverse | None = None,
575 dataset_types: Mapping[str, DatasetType] | None = None,
576 visualization_only: bool = False,
577 ) -> None:
578 """Resolve all dimensions and dataset types and check them for
579 consistency.
581 Resolving a graph also causes it to be sorted.
583 Parameters
584 ----------
585 registry : `lsst.daf.butler.Registry`, optional
586 Client for the data repository to resolve against.
587 dimensions : `lsst.daf.butler.DimensionUniverse`, optional
588 Definitions for all dimensions. Takes precedence over
589 ``registry.dimensions`` if both are provided. If neither is
590 provided, defaults to the default dimension universe
591 (``lsst.daf.butler.DimensionUniverse()``).
592 dataset_types : `~collection.abc.Mapping` [ `str`, \
593 `~lsst.daf.butler.DatasetType` ], optional
594 Mapping of dataset types to consider registered. Takes precedence
595 over ``registry.getDatasetType()`` if both are provided.
596 visualization_only : `bool`, optional
597 Resolve the graph as well as possible even when dimensions and
598 storage classes cannot really be determined. This can include
599 using the ``universe.commonSkyPix`` as the assumed dimensions of
600 connections that use the "skypix" placeholder and using "<UNKNOWN>"
601 as a storage class name (which will fail if the storage class
602 itself is ever actually loaded).
604 Notes
605 -----
606 The `universe` attribute is set to ``dimensions`` and used to set all
607 `TaskNode.dimensions` attributes. Dataset type nodes are resolved by
608 first looking for a registry definition, then using the producing
609 task's definition, then looking for consistency between all consuming
610 task definitions.
612 Raises
613 ------
614 ConnectionTypeConsistencyError
615 Raised if a prerequisite input for one task appears as a different
616 kind of connection in any other task.
617 DuplicateOutputError
618 Raised if multiple tasks have the same dataset type as an output.
619 IncompatibleDatasetTypeError
620 Raised if different tasks have different definitions of a dataset
621 type. Different but compatible storage classes are permitted.
622 MissingDatasetTypeError
623 Raised if a dataset type definition is required to exist in the
624 data repository but none was found. This should only occur for
625 dataset types that are not produced by a task in the pipeline and
626 are consumed with different storage classes or as components by
627 tasks in the pipeline.
628 EdgesChangedError
629 Raised if ``check_edges_unchanged=True`` and the edges of a task do
630 change after import and reconfiguration.
631 """
632 get_registered: Callable[[str], DatasetType | None] | None = None
633 if dataset_types is not None:
634 # Ruff seems confused about whether this is used below; it is!
635 get_registered = dataset_types.get
636 elif registry is not None:
638 def get_registered(name: str) -> DatasetType | None:
639 try:
640 return registry.getDatasetType(name)
641 except MissingDatasetTypeError:
642 return None
644 else:
646 def get_registered(name: str) -> None:
647 return None
649 if dimensions is None:
650 if registry is not None:
651 dimensions = registry.dimensions
652 else:
653 dimensions = DimensionUniverse()
655 sort_keys_from_steps = self._resolve_step_flow()
657 node_key: NodeKey
658 updates: dict[NodeKey, TaskNode | DatasetTypeNode] = {}
659 for node_key, node_state in self._xgraph.nodes.items():
660 match node_key.node_type:
661 case NodeType.TASK:
662 task_node: TaskNode = node_state["instance"]
663 new_task_node = task_node._resolved(dimensions)
664 if new_task_node is not task_node:
665 updates[node_key] = new_task_node
666 case NodeType.DATASET_TYPE:
667 dataset_type_node: DatasetTypeNode | None = node_state["instance"]
668 new_dataset_type_node = DatasetTypeNode._from_edges(
669 node_key,
670 self._xgraph,
671 get_registered,
672 dimensions,
673 previous=dataset_type_node,
674 visualization_only=visualization_only,
675 )
676 # Usage of `is`` here is intentional; `_from_edges` returns
677 # `previous=dataset_type_node` if it can determine that it
678 # doesn't need to change.
679 if new_dataset_type_node is not dataset_type_node:
680 updates[node_key] = new_dataset_type_node
681 try:
682 for node_key, node_value in updates.items():
683 self._xgraph.nodes[node_key]["instance"] = node_value
684 except Exception as err: # pragma: no cover
685 # There's no known way to get here, but we want to make it
686 # clear it's a big problem if we do.
687 raise PipelineGraphExceptionSafetyError(
688 "Error during dataset type resolution has left the graph in an inconsistent state."
689 ) from err
691 # If we get an error here, many graph nodes will have been resolved but
692 # the steps will still be marked as unverified and unsorted. That's
693 # still an acceptable state for the pipeline to be in.
694 self._resolve_step_dimensions(dimensions)
695 self._step_definitions._verified = True
697 if sort_keys_from_steps is not None:
698 self._reorder(sort_keys_from_steps)
699 else:
700 self.sort()
701 self._universe = dimensions
703 ###########################################################################
704 #
705 # Graph Modification Interface:
706 #
707 # - methods to add, remove, and replace tasks;
708 #
709 # - methods to add and remove task subsets.
710 #
711 # These are all things that are usually done in a Pipeline before making a
712 # graph at all, but there may be cases where we want to modify the graph
713 # instead. (These are also the methods used to make a graph from a
714 # Pipeline, or make a graph from another graph.)
715 #
716 ###########################################################################
718 def add_task(
719 self,
720 label: str | None,
721 task_class: type[PipelineTask],
722 config: PipelineTaskConfig | None = None,
723 connections: PipelineTaskConnections | None = None,
724 ) -> TaskNode:
725 """Add a new task to the graph.
727 Parameters
728 ----------
729 label : `str` or `None`
730 Label for the task in the pipeline. If `None`, `Task._DefaultName`
731 is used.
732 task_class : `type` [ `PipelineTask` ]
733 Class object for the task.
734 config : `PipelineTaskConfig`, optional
735 Configuration for the task. If not provided, a default-constructed
736 instance of ``task_class.ConfigClass`` is used.
737 connections : `PipelineTaskConnections`, optional
738 Object that describes the dataset types used by the task. If not
739 provided, one will be constructed from the given configuration. If
740 provided, it is assumed that ``config`` has already been validated
741 and frozen.
743 Returns
744 -------
745 node : `TaskNode`
746 The new task node added to the graph.
748 Raises
749 ------
750 ValueError
751 Raised if configuration validation failed when constructing
752 ``connections``.
753 PipelineDataCycleError
754 Raised if the graph is cyclic after this addition.
755 RuntimeError
756 Raised if an unexpected exception (which will be chained) occurred
757 at a stage that may have left the graph in an inconsistent state.
758 Other exceptions should leave the graph unchanged.
760 Notes
761 -----
762 Checks for dataset type consistency and multiple producers do not occur
763 until `resolve` is called, since the resolution depends on both the
764 state of the data repository and all contributing tasks.
766 Adding new tasks removes any existing resolutions of all dataset types
767 it references and marks the graph as unsorted. It is most efficient
768 to add all tasks up front and only then resolve and/or sort the graph.
769 """
770 if label is None:
771 label = task_class._DefaultName
772 if config is None:
773 config = task_class.ConfigClass()
774 _LOG.debug("Adding task %s %s to the pipeline graph", label, task_class)
775 task_node = TaskNode._from_imported_data(
776 key=NodeKey(NodeType.TASK, label),
777 init_key=NodeKey(NodeType.TASK_INIT, label),
778 data=_TaskNodeImportedData.configure(label, task_class, config, connections),
779 universe=self._universe,
780 )
781 self.add_task_nodes([task_node])
782 return task_node
784 def add_task_nodes(self, nodes: Iterable[TaskNode], parent: PipelineGraph | None = None) -> None:
785 """Add one or more existing task nodes to the graph.
787 Parameters
788 ----------
789 nodes : `~collections.abc.Iterable` [ `TaskNode` ]
790 Iterable of task nodes to add. If any tasks have resolved
791 dimensions, they must have the same dimension universe as the rest
792 of the graph.
793 parent : `PipelineGraph`, optional
794 If provided, another `PipelineGraph` from which these nodes were
795 obtained. Any dataset type nodes already present in ``parent``
796 that are referenced by the given tasks will be used in this graph
797 if they are not already present, preserving any dataset type
798 resolutions present in the parent graph. Adding nodes from a
799 parent graph after the graph has its own nodes (e.g. from
800 `add_task`) or nodes from a third graph may result in invalid
801 dataset type resolutions. It is safest to only use this argument
802 when populating an empty graph for the first time.
804 Raises
805 ------
806 PipelineDataCycleError
807 Raised if the graph is cyclic after this addition.
809 Notes
810 -----
811 Checks for dataset type consistency and multiple producers do not occur
812 until `resolve` is called, since the resolution depends on both the
813 state of the data repository and all contributing tasks.
815 Adding new tasks removes any existing resolutions of all dataset types
816 it references (unless ``parent is not None`` and marks the graph as
817 unsorted. It is most efficient to add all tasks up front and only then
818 resolve and/or sort the graph.
819 """
820 node_data: list[tuple[NodeKey, dict[str, Any]]] = []
821 edge_data: list[tuple[NodeKey, NodeKey, str, dict[str, Any]]] = []
822 for task_node in nodes:
823 task_node = task_node._resolved(self._universe)
824 node_data.append(
825 (task_node.key, {"instance": task_node, "bipartite": task_node.key.node_type.bipartite})
826 )
827 node_data.append(
828 (
829 task_node.init.key,
830 {"instance": task_node.init, "bipartite": task_node.init.key.node_type.bipartite},
831 )
832 )
833 # Convert the edge objects attached to the task node to networkx.
834 for read_edge in task_node.init.iter_all_inputs():
835 self._append_graph_data_from_edge(node_data, edge_data, read_edge, parent=parent)
836 for write_edge in task_node.init.iter_all_outputs():
837 self._append_graph_data_from_edge(node_data, edge_data, write_edge, parent=parent)
838 for read_edge in task_node.iter_all_inputs():
839 self._append_graph_data_from_edge(node_data, edge_data, read_edge, parent=parent)
840 for write_edge in task_node.iter_all_outputs():
841 self._append_graph_data_from_edge(node_data, edge_data, write_edge, parent=parent)
842 # Add a special edge (with no Edge instance) that connects the
843 # TaskInitNode to the runtime TaskNode.
844 edge_data.append((task_node.init.key, task_node.key, Edge.INIT_TO_TASK_NAME, {"instance": None}))
845 if not node_data and not edge_data:
846 return
847 # Checks and preparation complete; time to start the actual
848 # modification, during which it's hard to provide strong exception
849 # safety. Start by resetting the sort ordering, if there is one.
850 self._reset()
851 try:
852 self._xgraph.add_nodes_from(node_data)
853 self._xgraph.add_edges_from(edge_data)
854 if not networkx.algorithms.dag.is_directed_acyclic_graph(self._xgraph):
855 cycle = networkx.find_cycle(self._xgraph)
856 raise PipelineDataCycleError(f"Cycle detected while adding tasks: {cycle}.")
857 except Exception:
858 # First try to roll back our changes.
859 try:
860 self._xgraph.remove_edges_from(edge_data)
861 self._xgraph.remove_nodes_from(key for key, _ in node_data)
862 except Exception as err: # pragma: no cover
863 # There's no known way to get here, but we want to make it
864 # clear it's a big problem if we do.
865 raise PipelineGraphExceptionSafetyError(
866 "Error while attempting to revert PipelineGraph modification has left the graph in "
867 "an inconsistent state."
868 ) from err
869 # Successfully rolled back; raise the original exception.
870 raise
872 def reconfigure_tasks(
873 self,
874 *args: tuple[str, PipelineTaskConfig],
875 check_edges_unchanged: bool = False,
876 assume_edges_unchanged: bool = False,
877 **kwargs: PipelineTaskConfig,
878 ) -> None:
879 """Update the configuration for one or more tasks.
881 Parameters
882 ----------
883 *args : `tuple` [ `str`, `.PipelineTaskConfig` ]
884 Positional arguments are each a 2-tuple of task label and new
885 config object. Note that the same arguments may also be passed as
886 ``**kwargs``, which is usually more readable, but task labels in
887 ``*args`` are not required to be valid Python identifiers.
888 check_edges_unchanged : `bool`, optional
889 If `True`, require the edges (connections) of the modified tasks to
890 remain unchanged after the configuration updates, and verify that
891 this is the case.
892 assume_edges_unchanged : `bool`, optional
893 If `True`, the caller declares that the edges (connections) of the
894 modified tasks will remain unchanged after the configuration
895 updates, and that it is unnecessary to check this.
896 **kwargs : `.PipelineTaskConfig`
897 New config objects or overrides to apply to copies of the current
898 config objects, with task labels as the keywords.
900 Returns
901 -------
902 None
904 Raises
905 ------
906 ValueError
907 Raised if ``assume_edges_unchanged`` and ``check_edges_unchanged``
908 are both `True`, or if the same task appears twice.
909 EdgesChangedError
910 Raised if ``check_edges_unchanged=True`` and the edges of a task do
911 change.
913 Notes
914 -----
915 If reconfiguring a task causes its edges to change, any dataset type
916 nodes connected to that task (not just those whose edges have changed!)
917 will be unresolved.
918 """
919 new_configs: dict[str, PipelineTaskConfig] = {}
920 for task_label, config_update in itertools.chain(args, kwargs.items()):
921 if new_configs.setdefault(task_label, config_update) is not config_update:
922 raise ValueError(f"Config for {task_label!r} provided more than once.")
923 updates = {
924 task_label: self.tasks[task_label]._reconfigured(config, rebuild=not assume_edges_unchanged)
925 for task_label, config in new_configs.items()
926 }
927 self._replace_task_nodes(
928 updates,
929 check_edges_unchanged=check_edges_unchanged,
930 assume_edges_unchanged=assume_edges_unchanged,
931 message_header=(
932 "Unexpected change in edges for task {task_label!r} from original config (A) to "
933 "new configs (B):"
934 ),
935 )
937 def remove_tasks(
938 self, labels: Iterable[str], drop_from_subsets: bool = True
939 ) -> list[tuple[TaskNode, set[str]]]:
940 """Remove one or more tasks from the graph.
942 Parameters
943 ----------
944 labels : `~collections.abc.Iterable` [ `str` ]
945 Iterable of the labels of the tasks to remove.
946 drop_from_subsets : `bool`, optional
947 If `True`, drop each removed task from any subset in which it
948 currently appears. If `False`, raise `PipelineGraphError` if any
949 such subsets exist.
951 Returns
952 -------
953 nodes_and_subsets : `list` [ `tuple` [ `TaskNode`, `set` [ `str` ] ] ]
954 List of nodes removed and the labels of task subsets that
955 referenced them.
957 Raises
958 ------
959 PipelineGraphError
960 Raised if ``drop_from_subsets`` is `False` and the task is still
961 part of one or more subsets.
963 Notes
964 -----
965 Removing a task will cause dataset nodes with no other referencing
966 tasks to be removed. Any other dataset type nodes referenced by a
967 removed task will be reset to an "unresolved" state.
968 """
969 task_nodes_and_subsets = []
970 dataset_types: set[NodeKey] = set()
971 nodes_to_remove = set()
972 for label in labels:
973 task_node: TaskNode = self._xgraph.nodes[NodeKey(NodeType.TASK, label)]["instance"]
974 # Find task subsets that reference this task.
975 referencing_subsets = {
976 subset_label
977 for subset_label, task_subset in self.task_subsets.items()
978 if label in task_subset
979 }
980 if not drop_from_subsets and referencing_subsets:
981 raise PipelineGraphError(
982 f"Task {label!r} is still referenced by subset(s) {referencing_subsets}."
983 )
984 task_nodes_and_subsets.append((task_node, referencing_subsets))
985 # Find dataset types referenced by this task.
986 dataset_types.update(self._xgraph.predecessors(task_node.key))
987 dataset_types.update(self._xgraph.successors(task_node.key))
988 dataset_types.update(self._xgraph.predecessors(task_node.init.key))
989 dataset_types.update(self._xgraph.successors(task_node.init.key))
990 # Since there's an edge between the task and its init node, we'll
991 # have added those two nodes here, too, and we don't want that.
992 dataset_types.remove(task_node.init.key)
993 dataset_types.remove(task_node.key)
994 # Mark the task node and its init node for removal from the graph.
995 nodes_to_remove.add(task_node.key)
996 nodes_to_remove.add(task_node.init.key)
997 # Process the referenced datasets to see which ones are orphaned and
998 # need to be removed vs. just unresolved.
999 nodes_to_unresolve = []
1000 for dataset_type_key in dataset_types:
1001 related_tasks = set()
1002 related_tasks.update(self._xgraph.predecessors(dataset_type_key))
1003 related_tasks.update(self._xgraph.successors(dataset_type_key))
1004 related_tasks.difference_update(nodes_to_remove)
1005 if not related_tasks:
1006 nodes_to_remove.add(dataset_type_key)
1007 else:
1008 nodes_to_unresolve.append(dataset_type_key)
1009 # Checks and preparation complete; time to start the actual
1010 # modification, during which it's hard to provide strong exception
1011 # safety. Start by resetting the sort ordering.
1012 self._reset()
1013 try:
1014 for dataset_type_key in nodes_to_unresolve:
1015 self._xgraph.nodes[dataset_type_key]["instance"] = None
1016 for task_node, referencing_subsets in task_nodes_and_subsets:
1017 for subset_label in referencing_subsets:
1018 self._task_subsets[subset_label].remove(task_node.label)
1019 self._xgraph.remove_nodes_from(nodes_to_remove)
1020 except Exception as err: # pragma: no cover
1021 # There's no known way to get here, but we want to make it
1022 # clear it's a big problem if we do.
1023 raise PipelineGraphExceptionSafetyError(
1024 "Error during task removal has left the graph in an inconsistent state."
1025 ) from err
1026 return task_nodes_and_subsets
1028 def add_task_subset(self, subset_label: str, task_labels: Iterable[str], description: str = "") -> None:
1029 """Add a label for a set of tasks that are already in the pipeline.
1031 Parameters
1032 ----------
1033 subset_label : `str`
1034 Label for this set of tasks.
1035 task_labels : `~collections.abc.Iterable` [ `str` ]
1036 Labels of the tasks to include in the set. All must already be
1037 included in the graph.
1038 description : `str`, optional
1039 String description to associate with this label.
1040 """
1041 subset = TaskSubset(self._xgraph, subset_label, set(task_labels), description, self._step_definitions)
1042 self._task_subsets[subset_label] = subset
1044 def remove_task_subset(self, subset_label: str) -> None:
1045 """Remove a labeled set of tasks.
1047 Parameters
1048 ----------
1049 subset_label : `str`
1050 Label for this set of tasks.
1052 Notes
1053 -----
1054 If this subset is a step, it is also removed from the step definitions.
1055 """
1056 try:
1057 self._step_definitions.remove(subset_label)
1058 except KeyError:
1059 pass
1060 del self._task_subsets[subset_label]
1062 ###########################################################################
1063 #
1064 # NetworkX Export Interface:
1065 #
1066 # - methods to export the PipelineGraph's content (or various subsets
1067 # thereof) as NetworkX objects.
1068 #
1069 # These are particularly useful when writing tools to visualize the graph,
1070 # while providing options for which aspects of the graph (tasks, dataset
1071 # types, or both) to include, since all exported graphs have similar
1072 # attributes regardless of their structure.
1073 #
1074 ###########################################################################
1076 def make_xgraph(self) -> networkx.MultiDiGraph:
1077 """Export a networkx representation of the full pipeline graph,
1078 including both init and runtime edges.
1080 Returns
1081 -------
1082 xgraph : `networkx.MultiDiGraph`
1083 Directed acyclic graph with parallel edges.
1085 Notes
1086 -----
1087 The returned graph uses `NodeKey` instances for nodes. Parallel edges
1088 represent the same dataset type appearing in multiple connections for
1089 the same task, and are hence rare. The connection name is used as the
1090 edge key to disambiguate those parallel edges.
1092 Almost all edges connect dataset type nodes to task or task init nodes
1093 or vice versa, but there is also a special edge that connects each task
1094 init node to its runtime node. The existence of these edges makes the
1095 graph not quite bipartite, though its init-only and runtime-only
1096 subgraphs are bipartite.
1098 See `TaskNode`, `TaskInitNode`, `DatasetTypeNode`, `ReadEdge`, and
1099 `WriteEdge` for the descriptive node and edge attributes added.
1100 """
1101 return self._transform_xgraph_state(self._xgraph.copy(), skip_edges=False)
1103 def make_bipartite_xgraph(self, init: bool = False) -> networkx.MultiDiGraph:
1104 """Return a bipartite networkx representation of just the runtime or
1105 init-time pipeline graph.
1107 Parameters
1108 ----------
1109 init : `bool`, optional
1110 If `True` (`False` is default) return the graph of task
1111 initialization nodes and init input/output dataset types, instead
1112 of the graph of runtime task nodes and regular
1113 input/output/prerequisite dataset types.
1115 Returns
1116 -------
1117 xgraph : `networkx.MultiDiGraph`
1118 Directed acyclic graph with parallel edges.
1120 Notes
1121 -----
1122 The returned graph uses `NodeKey` instances for nodes. Parallel edges
1123 represent the same dataset type appearing in multiple connections for
1124 the same task, and are hence rare. The connection name is used as the
1125 edge key to disambiguate those parallel edges.
1127 This graph is bipartite because each dataset type node only has edges
1128 that connect it to a task [init] node, and vice versa.
1130 See `TaskNode`, `TaskInitNode`, `DatasetTypeNode`, `ReadEdge`, and
1131 `WriteEdge` for the descriptive node and edge attributes added.
1132 """
1133 return self._transform_xgraph_state(
1134 self._make_bipartite_xgraph_internal(init).copy(), skip_edges=False
1135 )
1137 def make_task_xgraph(self, init: bool = False) -> networkx.DiGraph:
1138 """Return a networkx representation of just the tasks in the pipeline.
1140 Parameters
1141 ----------
1142 init : `bool`, optional
1143 If `True` (`False` is default) return the graph of task
1144 initialization nodes, instead of the graph of runtime task nodes.
1146 Returns
1147 -------
1148 xgraph : `networkx.DiGraph`
1149 Directed acyclic graph with no parallel edges.
1151 Notes
1152 -----
1153 The returned graph uses `NodeKey` instances for nodes. The dataset
1154 types that link these tasks are not represented at all; edges have no
1155 attributes, and there are no parallel edges.
1157 See `TaskNode` and `TaskInitNode` for the descriptive node and
1158 attributes added.
1159 """
1160 return self._transform_xgraph_state(self._make_task_xgraph_internal(init), skip_edges=True)
1162 def make_dataset_type_xgraph(self, init: bool = False) -> networkx.DiGraph:
1163 """Return a networkx representation of just the dataset types in the
1164 pipeline.
1166 Parameters
1167 ----------
1168 init : `bool`, optional
1169 If `True` (`False` is default) return the graph of init input and
1170 output dataset types, instead of the graph of runtime (input,
1171 output, prerequisite input) dataset types.
1173 Returns
1174 -------
1175 xgraph : `networkx.DiGraph`
1176 Directed acyclic graph with no parallel edges.
1178 Notes
1179 -----
1180 The returned graph uses `NodeKey` instances for nodes. The tasks that
1181 link these tasks are not represented at all; edges have no attributes,
1182 and there are no parallel edges.
1184 See `DatasetTypeNode` for the descriptive node and attributes added.
1185 """
1186 bipartite_xgraph = self._make_bipartite_xgraph_internal(init)
1187 dataset_type_keys = [
1188 key
1189 for key, bipartite in bipartite_xgraph.nodes(data="bipartite")
1190 if bipartite == NodeType.DATASET_TYPE.bipartite
1191 ]
1192 return self._transform_xgraph_state(
1193 networkx.algorithms.bipartite.projected_graph(
1194 networkx.DiGraph(bipartite_xgraph), dataset_type_keys
1195 ),
1196 skip_edges=True,
1197 )
1199 ###########################################################################
1200 #
1201 # Expression-based Selection Interface.
1202 #
1203 ###########################################################################
1205 def select_tasks(self, expression: str) -> set[str]:
1206 """Return the tasks that match an expression.
1208 Parameters
1209 ----------
1210 expression : `str`
1211 String expression to evaluate. See
1212 :ref:`pipeline-graph-subset-expressions`.
1214 Returns
1215 -------
1216 task_labels : `set` [ `str` ]
1217 Set of matching task labels.
1218 """
1219 task_xgraph = self._make_task_xgraph_internal(init=False)
1220 expr_tree = expressions.parse(expression)
1221 matching_task_keys = self._select_expression(expr_tree, task_xgraph)
1222 return {key.name for key in matching_task_keys}
1224 def select(self, expression: str) -> PipelineGraph:
1225 """Return a new pipeline graph with the tasks that match an expression.
1227 Parameters
1228 ----------
1229 expression : `str`
1230 String expression to evaluate. See
1231 :ref:`pipeline-graph-subset-expressions`.
1233 Returns
1234 -------
1235 new_graph : `PipelineGraph`
1236 New pipeline graph with just the matching tasks.
1238 Notes
1239 -----
1240 All resolved dataset type nodes will be preserved.
1242 If `has_been_sorted`, the new graph will be sorted as well.
1244 Task subsets will not be included in the returned graph.
1245 """
1246 selected_tasks = self.select_tasks(expression)
1247 new_pipeline_graph = PipelineGraph(universe=self._universe, data_id=self._raw_data_id)
1248 new_pipeline_graph.add_task_nodes(
1249 [self.tasks[task_label] for task_label in selected_tasks], parent=self
1250 )
1251 if self.has_been_sorted:
1252 new_pipeline_graph.sort()
1253 return new_pipeline_graph
1255 ###########################################################################
1256 #
1257 # Serialization Interface.
1258 #
1259 # Serialization of PipelineGraphs is currently experimental and may not be
1260 # retained in the future. All serialization methods are
1261 # underscore-prefixed to ensure nobody mistakes them for a stable interface
1262 # (let a lone a stable file format).
1263 #
1264 ###########################################################################
1266 @classmethod
1267 def _read_stream(
1268 cls, stream: BinaryIO, import_mode: TaskImportMode = TaskImportMode.REQUIRE_CONSISTENT_EDGES
1269 ) -> PipelineGraph:
1270 """Read a serialized `PipelineGraph` from a file-like object.
1272 Parameters
1273 ----------
1274 stream : `BinaryIO`
1275 File-like object opened for binary reading, containing
1276 gzip-compressed JSON.
1277 import_mode : `TaskImportMode`, optional
1278 Whether to import tasks, and how to reconcile any differences
1279 between the imported task's connections and the those that were
1280 persisted with the graph. Default is to check that they are the
1281 same.
1283 Returns
1284 -------
1285 graph : `PipelineGraph`
1286 Deserialized pipeline graph.
1288 Raises
1289 ------
1290 PipelineGraphReadError
1291 Raised if the serialized `PipelineGraph` is not self-consistent.
1292 EdgesChangedError
1293 Raised if ``import_mode`` is
1294 `TaskImportMode.REQUIRED_CONSISTENT_EDGES` and the edges of a task
1295 did change after import and reconfiguration.
1297 Notes
1298 -----
1299 `PipelineGraph` serialization is currently experimental and may be
1300 removed or significantly changed in the future, with no deprecation
1301 period.
1302 """
1303 from .io import SerializedPipelineGraph
1305 with gzip.open(stream, "rb") as uncompressed_stream:
1306 data = json.load(uncompressed_stream)
1307 serialized_graph = SerializedPipelineGraph.model_validate(data)
1308 return serialized_graph.deserialize(import_mode)
1310 @classmethod
1311 def _read_uri(
1312 cls,
1313 uri: ResourcePathExpression,
1314 import_mode: TaskImportMode = TaskImportMode.REQUIRE_CONSISTENT_EDGES,
1315 ) -> PipelineGraph:
1316 """Read a serialized `PipelineGraph` from a file at a URI.
1318 Parameters
1319 ----------
1320 uri : convertible to `lsst.resources.ResourcePath`
1321 URI to a gzip-compressed JSON file containing a serialized pipeline
1322 graph.
1323 import_mode : `TaskImportMode`, optional
1324 Whether to import tasks, and how to reconcile any differences
1325 between the imported task's connections and the those that were
1326 persisted with the graph. Default is to check that they are the
1327 same.
1329 Returns
1330 -------
1331 graph : `PipelineGraph`
1332 Deserialized pipeline graph.
1334 Raises
1335 ------
1336 PipelineGraphReadError
1337 Raised if the serialized `PipelineGraph` is not self-consistent.
1338 EdgesChangedError
1339 Raised if ``import_mode`` is
1340 `TaskImportMode.REQUIRED_CONSISTENT_EDGES` and the edges of a task
1341 did change after import and reconfiguration.
1343 Notes
1344 -----
1345 `PipelineGraph` serialization is currently experimental and may be
1346 removed or significantly changed in the future, with no deprecation
1347 period.
1348 """
1349 uri = ResourcePath(uri)
1350 with uri.open("rb") as stream:
1351 return cls._read_stream(cast(BinaryIO, stream), import_mode=import_mode)
1353 def _write_stream(self, stream: BinaryIO) -> None:
1354 """Write the pipeline to a file-like object.
1356 Parameters
1357 ----------
1358 stream
1359 File-like object opened for binary writing.
1361 Notes
1362 -----
1363 `PipelineGraph` serialization is currently experimental and may be
1364 removed or significantly changed in the future, with no deprecation
1365 period.
1367 The file format is gzipped JSON, and is intended to be human-readable,
1368 but it should not be considered a stable public interface for outside
1369 code, which should always use `PipelineGraph` methods (or at least the
1370 `io.SerializedPipelineGraph` class) to read these files.
1371 """
1372 from .io import SerializedPipelineGraph
1374 with gzip.open(stream, mode="wb") as compressed_stream:
1375 compressed_stream.write(
1376 SerializedPipelineGraph.serialize(self).model_dump_json(exclude_defaults=True).encode("utf-8")
1377 )
1379 def _write_uri(self, uri: ResourcePathExpression) -> None:
1380 """Write the pipeline to a file given a URI.
1382 Parameters
1383 ----------
1384 uri : convertible to `lsst.resources.ResourcePath`
1385 URI to write to . May have ``.json.gz`` or no extension (which
1386 will cause a ``.json.gz`` extension to be added).
1388 Notes
1389 -----
1390 `PipelineGraph` serialization is currently experimental and may be
1391 removed or significantly changed in the future, with no deprecation
1392 period.
1394 The file format is gzipped JSON, and is intended to be human-readable,
1395 but it should not be considered a stable public interface for outside
1396 code, which should always use `PipelineGraph` methods (or at least the
1397 `io.SerializedPipelineGraph` class) to read these files.
1398 """
1399 uri = ResourcePath(uri)
1400 extension = uri.getExtension()
1401 if not extension:
1402 uri = uri.updatedExtension(".json.gz")
1403 elif extension != ".json.gz":
1404 raise ValueError("Expanded pipeline files should always have a .json.gz extension.")
1405 with uri.open(mode="wb") as stream:
1406 self._write_stream(cast(BinaryIO, stream))
1408 def _import_and_configure(
1409 self, import_mode: TaskImportMode = TaskImportMode.REQUIRE_CONSISTENT_EDGES
1410 ) -> None:
1411 """Import the `PipelineTask` classes referenced by all task nodes and
1412 update those nodes accordingly.
1414 Parameters
1415 ----------
1416 import_mode : `TaskImportMode`, optional
1417 Whether to import tasks, and how to reconcile any differences
1418 between the imported task's connections and the those that were
1419 persisted with the graph. Default is to check that they are the
1420 same. This method does nothing if this is
1421 `TaskImportMode.DO_NOT_IMPORT`.
1423 Raises
1424 ------
1425 EdgesChangedError
1426 Raised if ``import_mode`` is
1427 `TaskImportMode.REQUIRED_CONSISTENT_EDGES` and the edges of a task
1428 did change after import and reconfiguration.
1430 Notes
1431 -----
1432 This method shouldn't need to be called unless the graph was
1433 deserialized without importing and configuring immediately, which is
1434 not the default behavior (but it can greatly speed up deserialization).
1435 If all tasks have already been imported this does nothing.
1437 Importing and configuring a task can change its
1438 `~TaskNode.task_class_name` or `~TaskClass.get_config_str` output,
1439 usually because the software used to read a serialized graph is newer
1440 than the software used to write it (e.g. a new config option has been
1441 added, or the task was moved to a new module with a forwarding alias
1442 left behind). These changes are allowed by
1443 `TaskImportMode.REQUIRE_CONSISTENT_EDGES`.
1445 If importing and configuring a task causes its edges to change, any
1446 dataset type nodes linked to those edges will be reset to the
1447 unresolved state.
1448 """
1449 if import_mode is TaskImportMode.DO_NOT_IMPORT:
1450 return
1451 rebuild = (
1452 import_mode is TaskImportMode.REQUIRE_CONSISTENT_EDGES
1453 or import_mode is TaskImportMode.OVERRIDE_EDGES
1454 )
1455 updates: dict[str, TaskNode] = {}
1456 node_key: NodeKey
1457 for node_key, node_state in self._xgraph.nodes.items():
1458 if node_key.node_type is NodeType.TASK:
1459 task_node: TaskNode = node_state["instance"]
1460 new_task_node = task_node._imported_and_configured(rebuild)
1461 if new_task_node is not task_node:
1462 updates[task_node.label] = new_task_node
1463 self._replace_task_nodes(
1464 updates,
1465 check_edges_unchanged=(import_mode is TaskImportMode.REQUIRE_CONSISTENT_EDGES),
1466 assume_edges_unchanged=(import_mode is TaskImportMode.ASSUME_CONSISTENT_EDGES),
1467 message_header=(
1468 "In task with label {task_label!r}, persisted edges (A)"
1469 "differ from imported and configured edges (B):"
1470 ),
1471 )
1473 ###########################################################################
1474 #
1475 # Advanced PipelineGraph Inspection Interface:
1476 #
1477 # - methods to iterate over all nodes and edges, utilizing NodeKeys;
1478 #
1479 # - methods to find overall inputs and group nodes by their dimensions,
1480 # which are important operations for QuantumGraph generation.
1481 #
1482 ###########################################################################
1484 def iter_edges(self, init: bool = False) -> Iterator[Edge]:
1485 """Iterate over edges in the graph.
1487 Parameters
1488 ----------
1489 init : `bool`, optional
1490 If `True` (`False` is default) iterate over the edges between task
1491 initialization node and init input/output dataset types, instead of
1492 the runtime task nodes and regular input/output/prerequisite
1493 dataset types.
1495 Returns
1496 -------
1497 edges : `~collections.abc.Iterator` [ `Edge` ]
1498 A lazy iterator over `Edge` (`WriteEdge` or `ReadEdge`) instances.
1500 Notes
1501 -----
1502 This method always returns _either_ init edges or runtime edges, never
1503 both. The full (internal) graph that contains both also includes a
1504 special edge that connects each task init node to its runtime node;
1505 that is also never returned by this method, since it is never a part of
1506 the init-only or runtime-only subgraphs.
1507 """
1508 edge: Edge
1509 for _, _, edge in self._xgraph.edges(data="instance"):
1510 if edge is not None and edge.is_init == init:
1511 yield edge
1513 def iter_nodes(
1514 self,
1515 ) -> Iterator[
1516 tuple[Literal[NodeType.TASK_INIT], str, TaskInitNode]
1517 | tuple[Literal[NodeType.TASK], str, TaskInitNode]
1518 | tuple[Literal[NodeType.DATASET_TYPE], str, DatasetTypeNode | None]
1519 ]:
1520 """Iterate over nodes in the graph.
1522 Returns
1523 -------
1524 nodes : `~collections.abc.Iterator` [ `tuple` ]
1525 A lazy iterator over all of the nodes in the graph. Each yielded
1526 element is a tuple of:
1528 - the node type enum value (`NodeType`);
1529 - the string name for the node (task label or parent dataset type
1530 name);
1531 - the node value (`TaskNode`, `TaskInitNode`, `DatasetTypeNode`,
1532 or `None` for dataset type nodes that have not been resolved).
1533 """
1534 key: NodeKey
1535 if self._sorted_keys is not None:
1536 for key in self._sorted_keys:
1537 yield key.node_type, key.name, self._xgraph.nodes[key]["instance"] # type: ignore
1538 else:
1539 for key, node in self._xgraph.nodes(data="instance"):
1540 yield key.node_type, key.name, node # type: ignore
1542 def iter_overall_inputs(self) -> Iterator[tuple[str, DatasetTypeNode | None]]:
1543 """Iterate over all of the dataset types that are consumed but not
1544 produced by the graph.
1546 Returns
1547 -------
1548 dataset_types : `~collections.abc.Iterator` [ `tuple` ]
1549 A lazy iterator over the overall-input dataset types (including
1550 overall init inputs and prerequisites). Each yielded element is a
1551 tuple of:
1553 - the parent dataset type name;
1554 - the resolved `DatasetTypeNode`, or `None` if the dataset type has
1555 not been resolved.
1556 """
1557 for generation in networkx.algorithms.dag.topological_generations(self._xgraph):
1558 key: NodeKey
1559 for key in generation:
1560 # While we expect all tasks to have at least one input and
1561 # hence never appear in the first topological generation, that
1562 # is not true of task init nodes.
1563 if key.node_type is NodeType.DATASET_TYPE:
1564 yield key.name, self._xgraph.nodes[key]["instance"]
1565 return
1567 def group_by_dimensions(
1568 self, prerequisites: bool = False
1569 ) -> dict[DimensionGroup, tuple[dict[str, TaskNode], dict[str, DatasetTypeNode]]]:
1570 """Group this graph's tasks and dataset types by their dimensions.
1572 Parameters
1573 ----------
1574 prerequisites : `bool`, optional
1575 If `True`, include prerequisite dataset types as well as regular
1576 input and output datasets (including intermediates).
1578 Returns
1579 -------
1580 groups : `dict` [ `~lsst.daf.butler.DimensionGroup`, `tuple` ]
1581 A dictionary of groups keyed by `~lsst.daf.butler.DimensionGroup`,
1582 in which each value is a tuple of:
1584 - a `dict` of `TaskNode` instances, keyed by task label
1585 - a `dict` of `DatasetTypeNode` instances, keyed by
1586 dataset type name.
1588 that have those dimensions.
1590 Notes
1591 -----
1592 Init inputs and outputs are always included, but always have empty
1593 dimensions and are hence are all grouped together.
1594 """
1595 result: dict[DimensionGroup, tuple[dict[str, TaskNode], dict[str, DatasetTypeNode]]] = {}
1596 next_new_value: tuple[dict[str, TaskNode], dict[str, DatasetTypeNode]] = ({}, {})
1597 for task_label, task_node in self.tasks.items():
1598 if task_node.dimensions is None:
1599 raise UnresolvedGraphError(f"Task with label {task_label!r} has not been resolved.")
1600 if (group := result.setdefault(task_node.dimensions, next_new_value)) is next_new_value:
1601 next_new_value = ({}, {}) # make new lists for next time
1602 group[0][task_node.label] = task_node
1603 for dataset_type_name, dataset_type_node in self.dataset_types.items():
1604 if dataset_type_node is None:
1605 raise UnresolvedGraphError(f"Dataset type {dataset_type_name!r} has not been resolved.")
1606 if not dataset_type_node.is_prerequisite or prerequisites:
1607 if (
1608 group := result.setdefault(dataset_type_node.dataset_type.dimensions, next_new_value)
1609 ) is next_new_value:
1610 next_new_value = ({}, {}) # make new lists for next time
1611 group[1][dataset_type_node.name] = dataset_type_node
1612 return result
1614 def get_all_dimensions(self, prerequisites: bool = True) -> DimensionGroup:
1615 """Return all dimensions used in this graph's tasks and dataset types.
1617 Parameters
1618 ----------
1619 prerequisites : `bool`, optional
1620 If `False`, do not include the dimensions that are only used by
1621 prerequisite input dataset types.
1623 Returns
1624 -------
1625 dimensions : `~lsst.daf.butler.DimensionGroup`.
1626 All dimensions in this pipeline.
1627 """
1628 return DimensionGroup.union(
1629 *self.group_by_dimensions(prerequisites=prerequisites).keys(),
1630 universe=self.universe,
1631 )
1633 def split_independent(self) -> Iterable[PipelineGraph]:
1634 """Iterate over independent subgraphs that together comprise this
1635 pipeline graph.
1637 Returns
1638 -------
1639 subgraphs : `Iterable` [ `PipelineGraph` ]
1640 An iterable over component subgraphs that could be run
1641 independently (they have only overall inputs in common). May be a
1642 lazy iterator.
1644 Notes
1645 -----
1646 All resolved dataset type nodes will be preserved.
1648 If there is only one component, ``self`` may be returned as the only
1649 element in the iterable.
1651 If `has_been_sorted`, all subgraphs will be sorted as well.
1653 Task subsets will not be included in the returned graphs.
1654 """
1655 # Having an overall input in common isn't enough to make subgraphs
1656 # dependent on each other, so we want to look for connected component
1657 # subgraphs of the task-only projected graph.
1658 bipartite_xgraph = self._make_bipartite_xgraph_internal(init=False)
1659 task_keys = {
1660 key
1661 for key, bipartite in bipartite_xgraph.nodes(data="bipartite")
1662 if bipartite == NodeType.TASK.bipartite
1663 }
1664 task_xgraph = networkx.algorithms.bipartite.projected_graph(
1665 networkx.DiGraph(bipartite_xgraph), task_keys
1666 )
1667 # "Weakly" connected means connected in only one direction, which is
1668 # the only kind of "connected" a DAG can ever be.
1669 for component_task_keys in networkx.algorithms.weakly_connected_components(task_xgraph):
1670 if component_task_keys == task_keys:
1671 yield self
1672 return
1673 else:
1674 component_subgraph = PipelineGraph(universe=self._universe, data_id=self._raw_data_id)
1675 component_subgraph.add_task_nodes(
1676 [self._xgraph.nodes[key]["instance"] for key in component_task_keys], parent=self
1677 )
1678 if self.has_been_sorted:
1679 component_subgraph.sort()
1680 yield component_subgraph
1682 ###########################################################################
1683 #
1684 # Data repository/collection initialization
1685 #
1686 ###########################################################################
1688 @property
1689 def packages_dataset_type(self) -> DatasetType:
1690 """The special "packages" dataset type that records software versions.
1692 This is not associated with a task and hence is
1693 not considered part of the pipeline graph in other respects, but it
1694 does get written with other provenance datasets.
1695 """
1696 if self._universe is None:
1697 raise UnresolvedGraphError(
1698 "PipelineGraph must be resolved in order to get the packages dataset type."
1699 )
1700 return DatasetType(
1701 PACKAGES_INIT_OUTPUT_NAME, self._universe.empty, PACKAGES_INIT_OUTPUT_STORAGE_CLASS
1702 )
1704 def register_dataset_types(
1705 self,
1706 butler: Butler,
1707 include_packages: bool = True,
1708 *,
1709 include_inputs: bool = True,
1710 include_configs: bool = True,
1711 include_logs: bool = True,
1712 ) -> None:
1713 """Register all dataset types in a data repository.
1715 Parameters
1716 ----------
1717 butler : `~lsst.daf.butler.Butler`
1718 Data repository client.
1719 include_packages : `bool`, optional
1720 Whether to include the special "packages" dataset type that records
1721 software versions (this is not associated with a task and hence is
1722 not considered part of the pipeline graph in other respects, but it
1723 does get written with other provenance datasets).
1724 include_inputs : `bool`, optional
1725 Whether to register overall-input dataset types as well as outputs.
1726 include_configs : `bool`, optional
1727 Whether to register task config dataset types.
1728 include_logs : `bool`, optional
1729 Whether to register task log dataset types.
1730 """
1731 dataset_types = {
1732 node.name: node.dataset_type
1733 for node in self.dataset_types.values()
1734 if include_inputs or self.producer_of(node.name) is not None
1735 }
1736 if include_packages:
1737 dataset_types[self.packages_dataset_type.name] = self.packages_dataset_type
1738 if not include_configs:
1739 for task_node in self.tasks.values():
1740 del dataset_types[task_node.init.config_output.dataset_type_name]
1741 if not include_logs:
1742 for task_node in self.tasks.values():
1743 if task_node.log_output is not None:
1744 del dataset_types[task_node.log_output.dataset_type_name]
1745 for dataset_type in dataset_types.values():
1746 butler.registry.registerDatasetType(dataset_type)
1748 def check_dataset_type_registrations(self, butler: Butler, include_packages: bool = True) -> None:
1749 """Check that dataset type registrations in a data repository match
1750 the definitions in this pipeline graph.
1752 Parameters
1753 ----------
1754 butler : `~lsst.daf.butler.Butler`
1755 Data repository client.
1756 include_packages : `bool`, optional
1757 Whether to include the special "packages" dataset type that records
1758 software versions (this is not associated with a task and hence is
1759 not considered part of the pipeline graph in other respects, but it
1760 does get written with other provenance datasets).
1762 Returns
1763 -------
1764 None
1766 Raises
1767 ------
1768 lsst.daf.butler.MissingDatasetTypeError
1769 Raised if one or more non-optional-input or output dataset types in
1770 the pipeline is not registered at all.
1771 lsst.daf.butler.ConflictingDefinitionError
1772 Raised if the definition in the data repository is not identical
1773 to the definition in the pipeline graph.
1775 Notes
1776 -----
1777 Note that dataset type definitions that are storage-class-conversion
1778 compatible but not identical are not permitted by these checks, because
1779 the expectation is that these differences are handled by `resolve`,
1780 which makes the pipeline graph use the data repository definitions.
1781 This method is intended to check that none of those definitions have
1782 changed.
1783 """
1784 dataset_types = [node.dataset_type for node in self.dataset_types.values()]
1785 if include_packages:
1786 dataset_types.append(self.packages_dataset_type)
1787 missing_dataset_types: list[str] = []
1788 for dataset_type in dataset_types:
1789 try:
1790 expected = butler.registry.getDatasetType(dataset_type.name)
1791 except MissingDatasetTypeError:
1792 expected = None
1793 if expected is None:
1794 # The user probably forgot to register dataset types
1795 # at least once (which should be an error),
1796 # but we could also get here if this is an optional input for
1797 # which no datasets were found in this repo (not an error).
1798 if (
1799 not (
1800 self.producer_of(dataset_type.name) is None
1801 and all(
1802 self.tasks[input_edge.task_label].is_optional(input_edge.connection_name)
1803 for input_edge in self.consuming_edges_of(dataset_type.name)
1804 )
1805 )
1806 or dataset_type.name == PACKAGES_INIT_OUTPUT_NAME
1807 ):
1808 missing_dataset_types.append(dataset_type.name)
1809 elif expected != dataset_type:
1810 raise ConflictingDefinitionError(
1811 f"DatasetType definition in registry has changed since the pipeline graph was resolved: "
1812 f"{dataset_type} (graph) != {expected} (registry)."
1813 )
1814 if missing_dataset_types:
1815 plural = "s" if len(missing_dataset_types) != 1 else ""
1816 raise MissingDatasetTypeError(
1817 f"Missing dataset type definition{plural}: {', '.join(missing_dataset_types)}. "
1818 "Dataset types have to be registered in advance (on the command-line, either via "
1819 "`butler register-dataset-type` or the `--register-dataset-types` option to `pipetask run`."
1820 )
1822 def instantiate_tasks(
1823 self,
1824 get_init_input: Callable[[DatasetType], Any] | None = None,
1825 init_outputs: list[tuple[Any, DatasetType]] | None = None,
1826 labels: Iterable[str] | None = None,
1827 ) -> list[PipelineTask]:
1828 """Instantiate all tasks in the pipeline.
1830 Parameters
1831 ----------
1832 get_init_input : `~collections.abc.Callable`, optional
1833 Callable that accepts a single `~lsst.daf.butler.DatasetType`
1834 parameter and returns the init-input dataset associated with that
1835 dataset type. Must respect the storage class embedded in the type.
1836 This is optional if the pipeline does not have any overall init
1837 inputs. When a full butler is available,
1838 `lsst.daf.butler.Butler.get` can be used directly here.
1839 init_outputs : `list`, optional
1840 A list of ``(obj, dataset type)`` init-output dataset pairs, to be
1841 appended to in-place. Both the object and the dataset type will
1842 correspond to the storage class of the output connection, which
1843 may not be the same as the storage class on the graph's dataset
1844 type node.
1845 labels : `~collections.abc.Iterable` [ `str` ], optional
1846 The labels of tasks to instantiate. If not provided, all tasks in
1847 the graph will be instantiated.
1849 Returns
1850 -------
1851 tasks : `list`
1852 Constructed `PipelineTask` instances.
1853 """
1854 if not self.is_fully_resolved:
1855 raise UnresolvedGraphError("Pipeline graph must be fully resolved before instantiating tasks.")
1856 empty_data_id = DataCoordinate.make_empty(self.universe)
1857 labels = set(labels) if labels is not None else self.tasks.keys()
1858 handles: dict[str, InMemoryDatasetHandle] = {}
1859 tasks: list[PipelineTask] = []
1860 for task_node in self.tasks.values():
1861 if task_node.label not in labels:
1862 continue
1863 task_init_inputs: dict[str, Any] = {}
1864 for read_edge in task_node.init.inputs.values():
1865 if (handle := handles.get(read_edge.dataset_type_name)) is not None:
1866 obj = handle.get(storageClass=read_edge.storage_class_name)
1867 elif (
1868 read_edge.component is not None
1869 and (parent_handle := handles.get(read_edge.parent_dataset_type_name)) is not None
1870 ):
1871 obj = parent_handle.get(
1872 storageClass=read_edge.storage_class_name, component=read_edge.component
1873 )
1874 else:
1875 dataset_type_node = self.dataset_types[read_edge.parent_dataset_type_name]
1876 if get_init_input is None:
1877 raise ValueError(
1878 f"Task {task_node.label!r} requires init-input "
1879 f"{read_edge.dataset_type_name} but no 'get_init_input' callback was provided."
1880 )
1881 obj = get_init_input(read_edge.adapt_dataset_type(dataset_type_node.dataset_type))
1882 n_consumers = len(self.consumers_of(dataset_type_node.name))
1883 if (
1884 n_consumers > 1
1885 and read_edge.component is None
1886 and read_edge.storage_class_name == dataset_type_node.storage_class_name
1887 ):
1888 # Caching what we just got is safe in general only
1889 # if there was no storage class conversion, since
1890 # a->b and a->c does not imply b->c.
1891 handles[read_edge.dataset_type_name] = InMemoryDatasetHandle(
1892 obj,
1893 storageClass=dataset_type_node.storage_class,
1894 dataId=empty_data_id,
1895 copy=True,
1896 )
1897 task_init_inputs[read_edge.connection_name] = obj
1898 task = task_node.task_class(
1899 config=task_node.config, initInputs=task_init_inputs, name=task_node.label
1900 )
1901 tasks.append(task)
1902 for write_edge in task_node.init.outputs.values():
1903 dataset_type_node = self.dataset_types[write_edge.parent_dataset_type_name]
1904 obj = getattr(task, write_edge.connection_name)
1905 # We don't immediately coerce obj to the dataset_type_node
1906 # storage class (which should be the repo storage class, if
1907 # there is one) when appending to `init_outputs` because a
1908 # formatter might be able to do a better job of that later;
1909 # instead we pair it with a dataset type that's consistent with
1910 # the in-memory type. We do coerce when populating `handles`,
1911 # though, because going through the dataset_type_node storage
1912 # class is the conversion path we checked when we resolved the
1913 # pipeline graph.
1914 if init_outputs is not None:
1915 init_outputs.append((obj, write_edge.adapt_dataset_type(dataset_type_node.dataset_type)))
1916 n_consumers = len(self.consumers_of(dataset_type_node.name))
1917 if n_consumers > 0:
1918 handles[dataset_type_node.name] = InMemoryDatasetHandle(
1919 dataset_type_node.storage_class.coerce_type(obj),
1920 dataId=empty_data_id,
1921 storageClass=dataset_type_node.storage_class,
1922 copy=(n_consumers > 1),
1923 )
1924 return tasks
1926 def write_init_outputs(self, butler: Butler) -> None:
1927 """Write the init-output datasets for all tasks in the pipeline graph.
1929 Parameters
1930 ----------
1931 butler : `lsst.daf.butler.Butler`
1932 A full butler data repository client with its default run set
1933 to the collection where datasets should be written.
1935 Notes
1936 -----
1937 Datasets that already exist in the butler's output run collection will
1938 not be written.
1940 This method writes outputs with new random dataset IDs and should
1941 hence only be used when writing init-outputs prior to building a
1942 `QuantumGraph`. Use `QuantumGraph.write_init_outputs` if a quantum
1943 graph has already been built.
1944 """
1945 init_outputs: list[tuple[Any, DatasetType]] = []
1946 self.instantiate_tasks(butler.get, init_outputs)
1947 found_refs: dict[str, DatasetRef] = {}
1948 to_put: list[tuple[Any, DatasetType]] = []
1949 for obj, dataset_type in init_outputs:
1950 if (ref := butler.find_dataset(dataset_type, collections=butler.run)) is not None:
1951 found_refs[dataset_type.name] = ref
1952 else:
1953 to_put.append((obj, dataset_type))
1954 for ref, stored in butler.stored_many(found_refs.values()).items():
1955 if not stored:
1956 raise FileNotFoundError(
1957 f"Init-output dataset {ref.datasetType.name!r} was found in RUN {ref.run!r} "
1958 f"but had not actually been stored (or was stored and later deleted)."
1959 )
1960 for obj, dataset_type in to_put:
1961 butler.put(obj, dataset_type)
1963 def write_configs(self, butler: Butler) -> None:
1964 """Write the config datasets for all tasks in the pipeline graph.
1966 Parameters
1967 ----------
1968 butler : `lsst.daf.butler.Butler`
1969 A full butler data repository client with its default run set
1970 to the collection where datasets should be written.
1972 Notes
1973 -----
1974 Config datasets that already exist in the butler's output run
1975 collection will be checked for consistency.
1977 This method writes outputs with new random dataset IDs and should
1978 hence only be used when writing init-outputs prior to building a
1979 `QuantumGraph`. Use `QuantumGraph.write_configs` if a quantum graph
1980 has already been built.
1982 Raises
1983 ------
1984 lsst.daf.butler.registry.ConflictingDefinitionError
1985 Raised if a config dataset already exists and is not consistent
1986 with the config in the pipeline graph.
1987 """
1988 to_put: list[tuple[PipelineTaskConfig, str]] = []
1989 for task_node in self.tasks.values():
1990 dataset_type_name = task_node.init.config_output.dataset_type_name
1991 if (ref := butler.find_dataset(dataset_type_name, collections=butler.run)) is not None:
1992 old_config = butler.get(ref)
1993 if not task_node.config.compare(old_config, shortcut=False, output=log_config_mismatch):
1994 raise ConflictingDefinitionError(
1995 f"Config does not match existing task config {dataset_type_name!r} in "
1996 "butler; tasks configurations must be consistent within the same run collection"
1997 )
1998 else:
1999 to_put.append((task_node.config, dataset_type_name))
2000 # We do writes at the end to minimize the mess we leave behind when we
2001 # raise an exception.
2002 for config, dataset_type_name in to_put:
2003 butler.put(config, dataset_type_name)
2005 def write_packages(self, butler: Butler) -> None:
2006 """Write the 'packages' dataset for the currently-active software
2007 versions.
2009 Parameters
2010 ----------
2011 butler : `lsst.daf.butler.Butler`
2012 A full butler data repository client with its default run set
2013 to the collection where datasets should be written.
2015 Notes
2016 -----
2017 If the packages dataset already exists, it will be compared to the
2018 versions in the current packages. New packages that weren't present
2019 before are not considered an inconsistency.
2021 This method writes outputs with new random dataset IDs and should
2022 hence only be used when writing init-outputs prior to building a
2023 `QuantumGraph`. Use `QuantumGraph.write_packages` if a quantum graph
2024 has already been built.
2026 Raises
2027 ------
2028 lsst.daf.butler.registry.ConflictingDefinitionError
2029 Raised if the packages dataset already exists and is not consistent
2030 with the current packages.
2031 """
2032 new_packages = Packages.fromSystem()
2033 if (ref := butler.find_dataset(self.packages_dataset_type)) is not None:
2034 packages = butler.get(ref)
2035 if compare_packages(packages, new_packages):
2036 # have to remove existing dataset first; butler has no
2037 # replace option.
2038 butler.pruneDatasets([ref], unstore=True, purge=True)
2039 butler.put(packages, ref)
2040 else:
2041 butler.put(new_packages, self.packages_dataset_type)
2043 def init_output_run(self, butler: Butler) -> None:
2044 """Initialize a new output RUN collection by writing init-output
2045 datasets (including configs and packages).
2047 Parameters
2048 ----------
2049 butler : `lsst.daf.butler.Butler`
2050 A full butler data repository client with its default run set
2051 to the collection where datasets should be written.
2052 """
2053 self.write_configs(butler)
2054 self.write_packages(butler)
2055 self.write_init_outputs(butler)
2057 ###########################################################################
2058 #
2059 # Class- and Package-Private Methods.
2060 #
2061 ###########################################################################
2063 def _iter_task_defs(self) -> Iterator[TaskDef]:
2064 """Iterate over this pipeline as a sequence of `TaskDef` instances.
2066 Notes
2067 -----
2068 This is a package-private method intended to aid in the transition to a
2069 codebase more fully integrated with the `PipelineGraph` class, in which
2070 both `TaskDef` and `PipelineDatasetTypes` are expected to go away, and
2071 much of the functionality on the `Pipeline` class will be moved to
2072 `PipelineGraph` as well.
2074 Raises
2075 ------
2076 TaskNotImportedError
2077 Raised if `TaskNode.is_imported` is `False` for any task.
2078 """
2079 from ..pipeline import TaskDef
2081 for node in self._tasks.values():
2082 yield TaskDef(
2083 config=node.config,
2084 taskClass=node.task_class,
2085 label=node.label,
2086 connections=node.get_connections(),
2087 )
2089 def _init_from_args(
2090 self,
2091 xgraph: networkx.MultiDiGraph | None,
2092 sorted_keys: Sequence[NodeKey] | None,
2093 task_subsets: dict[str, TaskSubset] | None,
2094 description: str,
2095 universe: DimensionUniverse | None,
2096 data_id: DataId | None,
2097 step_definitions: StepDefinitions,
2098 ) -> None:
2099 """Initialize the graph with possibly-nontrivial arguments.
2101 Parameters
2102 ----------
2103 xgraph : `networkx.MultiDiGraph` or `None`
2104 The backing networkx graph, or `None` to create an empty one.
2105 This graph has `NodeKey` instances for nodes and the same structure
2106 as the graph exported by `make_xgraph`, but its nodes and edges
2107 have a single ``instance`` attribute that holds a `TaskNode`,
2108 `TaskInitNode`, `DatasetTypeNode` (or `None`), `ReadEdge`, or
2109 `WriteEdge` instance.
2110 sorted_keys : `~collections.abc.Sequence` [ `NodeKey` ] or `None`
2111 Topologically sorted sequence of node keys, or `None` if the graph
2112 is not sorted.
2113 task_subsets : `dict` [ `str`, `TaskSubset` ]
2114 Labeled subsets of tasks. Values must be constructed with
2115 ``xgraph`` as their parent graph.
2116 description : `str`
2117 String description for this pipeline.
2118 universe : `lsst.daf.butler.DimensionUniverse` or `None`
2119 Definitions of all dimensions.
2120 data_id : `lsst.daf.butler.DataCoordinate` or other data ID mapping.
2121 Data ID that represents a constraint on all quanta generated from
2122 this pipeline.
2123 step_definitions : `StepDefinitions`
2124 Struct holding information about steps.
2126 Notes
2127 -----
2128 Only empty `PipelineGraph` instances should be constructed directly by
2129 users, which sets the signature of ``__init__`` itself, but methods on
2130 `PipelineGraph` and its helper classes need to be able to create them
2131 with state. Those methods can call this after calling ``__new__``
2132 manually, skipping ``__init__``.
2133 """
2134 self._xgraph = xgraph if xgraph is not None else networkx.MultiDiGraph()
2135 self._sorted_keys: Sequence[NodeKey] | None = None
2136 self._task_subsets = task_subsets if task_subsets is not None else {}
2137 self._description = description
2138 self._tasks = TaskMappingView(self._xgraph)
2139 self._dataset_types = DatasetTypeMappingView(self._xgraph)
2140 self._raw_data_id: dict[str, Any]
2141 if isinstance(data_id, DataCoordinate):
2142 if universe is None:
2143 universe = data_id.universe
2144 else:
2145 assert universe is data_id.universe, "data_id.universe and given universe differ"
2146 self._raw_data_id = dict(data_id.required)
2147 elif data_id is None:
2148 self._raw_data_id = {}
2149 else:
2150 self._raw_data_id = dict(data_id)
2151 self._universe = universe
2152 if sorted_keys is not None:
2153 self._reorder(sorted_keys)
2154 self._step_definitions = step_definitions
2156 def _make_bipartite_xgraph_internal(self, init: bool) -> networkx.MultiDiGraph:
2157 """Make a bipartite init-only or runtime-only internal subgraph.
2159 See `make_bipartite_xgraph` for parameters and return values.
2161 Notes
2162 -----
2163 This method returns a view of the `PipelineGraph` object's internal
2164 backing graph, and hence should only be called in methods that copy the
2165 result either explicitly or by running a copying algorithm before
2166 returning it to the user.
2167 """
2168 return self._xgraph.edge_subgraph([edge.key for edge in self.iter_edges(init)])
2170 def _make_task_xgraph_internal(self, init: bool) -> networkx.DiGraph:
2171 """Make a init-only or runtime-only internal task subgraph.
2173 See `make_task_xgraph` for parameters and return values.
2175 Notes
2176 -----
2177 This method returns a view of the `PipelineGraph` object's internal
2178 backing graph, and hence should only be called in methods that copy the
2179 result either explicitly or by running a copying algorithm before
2180 returning it to the user.
2181 """
2182 bipartite_xgraph = self._make_bipartite_xgraph_internal(init=init)
2183 task_keys = [
2184 key
2185 for key, bipartite in bipartite_xgraph.nodes(data="bipartite")
2186 if bipartite == NodeType.TASK.bipartite
2187 ]
2188 return networkx.algorithms.bipartite.projected_graph(networkx.DiGraph(bipartite_xgraph), task_keys)
2190 def _transform_xgraph_state(self, xgraph: _G, skip_edges: bool) -> _G:
2191 """Transform networkx graph attributes in-place from the internal
2192 "instance" attributes to the documented exported attributes.
2194 Parameters
2195 ----------
2196 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
2197 Graph whose state should be transformed.
2198 skip_edges : `bool`
2199 If `True`, do not transform edge state.
2201 Returns
2202 -------
2203 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
2204 The same object passed in, after modification.
2206 Notes
2207 -----
2208 This should be called after making a copy of the internal graph but
2209 before any projection down to just task or dataset type nodes, since
2210 it assumes stateful edges.
2211 """
2212 state: dict[str, Any]
2213 for state in xgraph.nodes.values():
2214 node_value: TaskInitNode | TaskNode | DatasetTypeNode | None = state.pop("instance")
2215 if node_value is not None:
2216 state.update(node_value._to_xgraph_state())
2217 else:
2218 # This is a dataset type node that is not resolved.
2219 state["bipartite"] = NodeType.DATASET_TYPE.bipartite
2220 if not skip_edges:
2221 for _, _, state in xgraph.edges(data=True):
2222 edge: Edge | None = state.pop("instance", None)
2223 if edge is not None:
2224 state.update(edge._to_xgraph_state())
2225 return xgraph
2227 def _replace_task_nodes(
2228 self,
2229 updates: Mapping[str, TaskNode],
2230 check_edges_unchanged: bool,
2231 assume_edges_unchanged: bool,
2232 message_header: str,
2233 ) -> None:
2234 """Replace task nodes and update edges and dataset type nodes
2235 accordingly.
2237 Parameters
2238 ----------
2239 updates : `Mapping` [ `str`, `TaskNode` ]
2240 New task nodes with task label keys. All keys must be task labels
2241 that are already present in the graph.
2242 check_edges_unchanged : `bool`, optional
2243 If `True`, require the edges (connections) of the modified tasks to
2244 remain unchanged after importing and configuring each task, and
2245 verify that this is the case.
2246 assume_edges_unchanged : `bool`, optional
2247 If `True`, the caller declares that the edges (connections) of the
2248 modified tasks will remain unchanged importing and configuring each
2249 task, and that it is unnecessary to check this.
2250 message_header : `str`
2251 Template for `str.format` with a single ``task_label`` placeholder
2252 to use as the first line in `EdgesChangedError` messages that show
2253 the differences between new task edges and old task edges. Should
2254 include the fact that the rest of the message will refer to the old
2255 task as "A" and the new task as "B", and end with a colon.
2257 Raises
2258 ------
2259 ValueError
2260 Raised if ``assume_edges_unchanged`` and ``check_edges_unchanged``
2261 are both `True`, or if a full config is provided for a task after
2262 another full config or an override has already been provided.
2263 EdgesChangedError
2264 Raised if ``check_edges_unchanged=True`` and the edges of a task do
2265 change.
2266 """
2267 deep: dict[str, TaskNode] = {}
2268 shallow: dict[str, TaskNode] = {}
2269 if assume_edges_unchanged:
2270 if check_edges_unchanged:
2271 raise ValueError("Cannot simultaneously assume and check that edges have not changed.")
2272 shallow.update(updates)
2273 else:
2274 for task_label, new_task_node in updates.items():
2275 old_task_node = self.tasks[task_label]
2276 messages = old_task_node.diff_edges(new_task_node)
2277 if messages:
2278 if check_edges_unchanged:
2279 messages.insert(0, message_header.format(task_label=task_label))
2280 raise EdgesChangedError("\n".join(messages))
2281 else:
2282 deep[task_label] = new_task_node
2283 else:
2284 shallow[task_label] = new_task_node
2285 try:
2286 if deep:
2287 removed = self.remove_tasks(deep.keys(), drop_from_subsets=True)
2288 self.add_task_nodes(deep.values())
2289 for replaced_task_node, referencing_subsets in removed:
2290 for subset_label in referencing_subsets:
2291 self._task_subsets[subset_label].add(replaced_task_node.label)
2292 for task_node in shallow.values():
2293 self._xgraph.nodes[task_node.key]["instance"] = task_node
2294 self._xgraph.nodes[task_node.init.key]["instance"] = task_node.init
2295 except PipelineGraphExceptionSafetyError: # pragma: no cover
2296 raise
2297 except Exception as err: # pragma: no cover
2298 # There's no known way to get here, but we want to make it clear
2299 # it's a big problem if we do.
2300 raise PipelineGraphExceptionSafetyError(
2301 "Error while replacing tasks has left the graph in an inconsistent state."
2302 ) from err
2304 def _append_graph_data_from_edge(
2305 self,
2306 node_data: list[tuple[NodeKey, dict[str, Any]]],
2307 edge_data: list[tuple[NodeKey, NodeKey, str, dict[str, Any]]],
2308 edge: Edge,
2309 parent: PipelineGraph | None,
2310 ) -> None:
2311 """Append networkx state dictionaries for an edge and the corresponding
2312 dataset type node.
2314 Parameters
2315 ----------
2316 node_data : `list`
2317 List of node keys and state dictionaries. A node is appended if
2318 one does not already exist for this dataset type.
2319 edge_data : `list`
2320 List of node key pairs, connection names, and state dictionaries
2321 for edges.
2322 edge : `Edge`
2323 New edge being processed.
2324 parent : `PipelineGraph` or `None`
2325 Another pipeline graph whose dataset type nodes should be used
2326 when present.
2327 """
2328 new_dataset_type_node = None
2329 if parent is not None:
2330 new_dataset_type_node = parent._xgraph.nodes[edge.dataset_type_key].get("instance")
2331 if (existing_dataset_type_state := self._xgraph.nodes.get(edge.dataset_type_key)) is not None:
2332 existing_dataset_type_state["instance"] = new_dataset_type_node
2333 else:
2334 node_data.append(
2335 (
2336 edge.dataset_type_key,
2337 {
2338 "instance": new_dataset_type_node,
2339 "bipartite": NodeType.DATASET_TYPE.bipartite,
2340 },
2341 )
2342 )
2343 edge_data.append(
2344 edge.nodes
2345 + (
2346 edge.connection_name,
2347 {"instance": edge},
2348 )
2349 )
2351 def _reorder(self, sorted_keys: Sequence[NodeKey]) -> None:
2352 """Set the order of all views of this graph from the given sorted
2353 sequence of task labels and dataset type names.
2354 """
2355 self._sorted_keys = sorted_keys
2356 self._tasks._reorder(sorted_keys)
2357 self._dataset_types._reorder(sorted_keys)
2359 def _reset(self) -> None:
2360 """Reset all views of this graph following a modification that might
2361 invalidate them.
2362 """
2363 self._sorted_keys = None
2364 self._tasks._reset()
2365 self._dataset_types._reset()
2367 def _resolve_step_flow(self) -> Sequence[NodeKey] | None:
2368 """Check that step definitions are consistent with the ordering of
2369 the graph's nodes and that they partition the task graph.
2371 Returns
2372 -------
2373 sort_keys : `~collections.abc.Sequence` [ `NodeKey` ] or `None`
2374 Sort order for the pipeline graph that is consistent with the
2375 step order, or `None` if there were no steps defined for this
2376 pipeline.
2377 """
2378 if not self._step_definitions:
2379 return None
2380 task_labels_so_far: set[str] = set()
2381 # Inputs we've already seen, and the tasks that wanted them:
2382 inputs_so_far: dict[str, set[str]] = {}
2383 sort_keys: list[NodeKey] = []
2384 keys_already_sorted: set[NodeKey] = set()
2385 for step_label in self.steps:
2386 try:
2387 task_subset = self.task_subsets[step_label]
2388 except KeyError:
2389 raise InvalidStepsError(f"Step {step_label!r} is not a task subset.") from None
2390 if not task_labels_so_far.isdisjoint(task_subset):
2391 raise InvalidStepsError(
2392 f"Step {step_label!r} repeats task(s) {task_labels_so_far & task_subset}."
2393 )
2394 task_labels_so_far.update(task_subset)
2395 # We need a temporary data structure for tracking the inputs of
2396 # just this step to avoid complaining about input-output
2397 # relationships within the step. We'll merge this into the
2398 # cumulative one later.
2399 step_inputs: dict[str, set[str]] = {}
2400 # We'll also gather all task, task-init and dataset type networkx
2401 # graph keys for the step so we can make and sort a step sub-graph.
2402 new_step_keys: set[NodeKey] = set()
2403 for task_label in task_subset:
2404 for input_name in self.inputs_of(task_label):
2405 step_inputs.setdefault(input_name, set()).add(task_label)
2406 # Check that none of the outputs of the tasks in this step were
2407 # expected as inputs of a previous step's tasks.
2408 task_outputs = self.outputs_of(task_label)
2409 if not inputs_so_far.keys().isdisjoint(task_outputs.keys()):
2410 msg: list[str] = []
2411 for input_name in inputs_so_far.keys() & task_outputs.keys():
2412 msg.append(f"{input_name} (used by {inputs_so_far[input_name]})")
2413 raise InvalidStepsError(
2414 f"Task {task_label} in step {step_label!r} produces dataset types "
2415 f"[{', '.join(msg)}], but these are consumed by tasks in earlier "
2416 "steps. Either steps are out of order or the graph is cyclic."
2417 )
2418 task_init_key = NodeKey(NodeType.TASK_INIT, task_label)
2419 task_key = NodeKey(NodeType.TASK, task_label)
2420 new_step_keys.add(task_init_key)
2421 new_step_keys.add(task_key)
2422 new_step_keys.update(self._xgraph.predecessors(task_init_key))
2423 new_step_keys.update(self._xgraph.predecessors(task_key))
2424 new_step_keys.update(self._xgraph.successors(task_init_key))
2425 new_step_keys.update(self._xgraph.successors(task_key))
2426 # Also record the step the task is in as a private xgraph
2427 # attribute so we can look up the step given a task (or,
2428 # indirectly, dataset type); this has to be used with care
2429 # because if the steps haven't been verified it can be wrong.
2430 self._xgraph.nodes[task_init_key]["step"] = step_label
2431 self._xgraph.nodes[task_key]["step"] = step_label
2432 # Drop step input keys that were already either inputs or outputs
2433 # of a previous step, since they'll have already been added to
2434 # sort_keys.
2435 new_step_keys.difference_update(keys_already_sorted)
2436 # Make the step subgraph, sort it, and extend the overall sort_keys
2437 # with result.
2438 step_xgraph = self._xgraph.subgraph(new_step_keys)
2439 sort_keys.extend(networkx.dag.lexicographical_topological_sort(step_xgraph))
2440 keys_already_sorted.update(new_step_keys)
2441 for input_name, consuming_tasks in step_inputs.items():
2442 inputs_so_far.setdefault(input_name, set()).update(consuming_tasks)
2443 if not task_labels_so_far.issuperset(self.tasks.keys()):
2444 # Note that the converse issubset test effectively happens when we
2445 # look up inputs and outputs of each task.
2446 raise InvalidStepsError(f"No step contains task(s) {self.tasks.keys() - task_labels_so_far}.")
2447 return sort_keys
2449 def _resolve_step_dimensions(self, universe: DimensionUniverse) -> None:
2450 """Check that step sharding dimensions are consistent with task and
2451 output dataset dimensions.
2453 Parameters
2454 ----------
2455 universe : `~lsst.daf.butler.DimensionUniverse`
2456 Definitions for all dimensions. Will be attached to the step
2457 definitions by this method.
2458 """
2459 self._step_definitions._universe = universe
2460 for step_label in self.steps:
2461 dimensions = self.steps.get_dimensions(step_label)
2462 for task_label in self.task_subsets[step_label]:
2463 task_node = self.tasks[task_label]
2464 if not _dimensions_compatible(dimensions, task_node.dimensions):
2465 raise InvalidStepsError(
2466 f"Dimensions {task_node.dimensions} of task {task_label!r} are not compatible with "
2467 f"the sharding dimensions {dimensions} of step {step_label!r}."
2468 )
2469 for dataset_type_node in self.outputs_of(task_label).values():
2470 assert dataset_type_node is not None, "dataset types should be resolved first"
2471 if not _dimensions_compatible(dimensions, dataset_type_node.dimensions):
2472 raise InvalidStepsError(
2473 f"Dimensions {dataset_type_node.dimensions} of dataset type "
2474 f"{dataset_type_node.name!r} (produced by task {task_label!r}) are not "
2475 f"compatible with the sharding dimensions {dimensions} of step "
2476 f"{step_label!r}."
2477 )
2479 def _select_expression(self, expr_tree: expressions.Node, task_xgraph: networkx.DiGraph) -> set[NodeKey]:
2480 """Select tasks from a pipeline based on a string expression.
2482 This is the primary implementation method for `select` and
2483 `select_tasks`.
2485 Parameters
2486 ----------
2487 expr_tree : `expressions.Node`
2488 Expression [sub]tree to process (recursively).
2489 task_xgraph : `networkx.DiGraph`
2490 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2492 Returns
2493 -------
2494 selected : `set` [ `NodeKey` ]
2495 Set of `NodeKey` objects for matching tasks (only; no dataset type
2496 or task-init nodes).
2497 """
2498 match expr_tree:
2499 case expressions.IdentifierNode(qualifier=qualifier, label=label):
2500 match self._select_identifier(qualifier, label):
2501 case NodeKey(node_type=NodeType.TASK) as task_key:
2502 return {task_key}
2503 case NodeKey(node_type=NodeType.DATASET_TYPE) as dataset_type_key:
2504 # Since a dataset type can have only one producer, this
2505 # yields 0- (for overall inputs) or 1-element sets.
2506 for producer_key, _ in self._xgraph.in_edges(dataset_type_key):
2507 if producer_key.node_type is NodeType.TASK_INIT:
2508 raise InvalidExpressionError(
2509 f"Init-output dataset type {label!r} cannot be used directly in an "
2510 "expression."
2511 )
2512 return {producer_key}
2513 return set()
2514 case TaskSubset() as task_subset:
2515 return {NodeKey(NodeType.TASK, label) for label in task_subset}
2516 case _: # pragma: no cover
2517 raise AssertionError("Identifier type inconsistent with grammar.")
2518 case expressions.DirectionNode(operator=operator, start=start):
2519 match self._select_identifier(start.qualifier, start.label):
2520 case NodeKey(node_type=NodeType.TASK) as task_key:
2521 if operator.startswith("<"):
2522 return self._select_task_ancestors(
2523 task_key, task_xgraph, inclusive=operator.endswith("=")
2524 )
2525 else:
2526 assert operator.startswith(">"), "Guaranteed by grammar."
2527 return self._select_task_descendants(
2528 task_key, task_xgraph, inclusive=operator.endswith("=")
2529 )
2530 case NodeKey(node_type=NodeType.DATASET_TYPE) as dataset_type_key:
2531 if operator.startswith("<"):
2532 return self._select_dataset_type_ancestors(
2533 dataset_type_key, task_xgraph, inclusive=operator.endswith("=")
2534 )
2535 else:
2536 assert operator.startswith(">"), "Guaranteed by grammar."
2537 return self._select_dataset_type_descendants(
2538 dataset_type_key, task_xgraph, inclusive=operator.endswith("=")
2539 )
2540 case TaskSubset():
2541 raise InvalidExpressionError(
2542 f"Task subset identifier {start!r} cannot be used as the start of an "
2543 "ancestor/descendant search."
2544 )
2545 case _: # pragma: no cover
2546 raise AssertionError("Unexpected parsed identifier result type.")
2547 case expressions.NotNode(operand=operand):
2548 operand_result = self._select_expression(operand, task_xgraph)
2549 return set(task_xgraph.nodes.keys() - operand_result)
2550 case expressions.UnionNode(lhs=lhs, rhs=rhs):
2551 lhs_result = self._select_expression(lhs, task_xgraph)
2552 rhs_result = self._select_expression(rhs, task_xgraph)
2553 return lhs_result.union(rhs_result)
2554 case expressions.IntersectionNode(lhs=lhs, rhs=rhs):
2555 lhs_result = self._select_expression(lhs, task_xgraph)
2556 rhs_result = self._select_expression(rhs, task_xgraph)
2557 return lhs_result.intersection(rhs_result)
2558 case _: # pragma: no cover
2559 raise AssertionError("Expression parse node inconsistent with grammar.")
2561 def _select_task_ancestors(
2562 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2563 ) -> set[NodeKey]:
2564 """Return all task-node ancestors of the given task node, as defined by
2565 the `select` expression language.
2567 Parameters
2568 ----------
2569 start : `NodeKey`
2570 A runtime task node key.
2571 task_xgraph : `networkx.DiGraph`
2572 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2573 inclusive : `bool`
2574 Whether to include the ``start`` node in the results.
2576 Returns
2577 -------
2578 selected : `set` [ `NodeKey` ]
2579 Set of `NodeKey` objects for matching tasks (only; no dataset type
2580 or task-init nodes).
2581 """
2582 result = set(networkx.dag.ancestors(task_xgraph, start))
2583 if inclusive:
2584 result.add(start)
2585 return result
2587 def _select_task_descendants(
2588 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2589 ) -> set[NodeKey]:
2590 """Return all task-node descendants of the given task node, as defined
2591 by the `select` expression language.
2593 Parameters
2594 ----------
2595 start : `NodeKey`
2596 A runtime task node key.
2597 task_xgraph : `networkx.DiGraph`
2598 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2599 inclusive : `bool`
2600 Whether to include the ``start`` node in the results.
2602 Returns
2603 -------
2604 selected : `set` [ `NodeKey` ]
2605 Set of `NodeKey` objects for matching tasks (only; no dataset type
2606 or task-init nodes).
2607 """
2608 result = set(networkx.dag.descendants(task_xgraph, start))
2609 if inclusive:
2610 result.add(start)
2611 return result
2613 def _select_dataset_type_ancestors(
2614 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2615 ) -> set[NodeKey]:
2616 """Return all task-node ancestors of the given dataset type node, as
2617 defined by the `select` expression language.
2619 Parameters
2620 ----------
2621 start : `NodeKey`
2622 A dataset type node key. May not be an init-output.
2623 task_xgraph : `networkx.DiGraph`
2624 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2625 inclusive : `bool`
2626 Whether to include the producer of the ``start`` node in the
2627 results.
2629 Returns
2630 -------
2631 selected : `set` [ `NodeKey` ]
2632 Set of `NodeKey` objects for matching tasks (only; no dataset type
2633 or task-init nodes).
2634 """
2635 result: set[NodeKey] = set()
2636 for producer_key, _ in self._xgraph.in_edges(start):
2637 if producer_key.node_type is NodeType.TASK_INIT:
2638 raise InvalidExpressionError(
2639 f"Init-output dataset type {start.name!r} cannot be used as the "
2640 "starting point for an ancestor ('<' or '<=') search."
2641 )
2642 result.update(networkx.dag.ancestors(task_xgraph, producer_key))
2643 if inclusive:
2644 result.add(producer_key)
2645 return result
2647 def _select_dataset_type_descendants(
2648 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2649 ) -> set[NodeKey]:
2650 """Return all task-node descendatns of the given dataset type node, as
2651 defined by the `select` expression language.
2653 Parameters
2654 ----------
2655 start : `NodeKey`
2656 A dataset type node key. May not be an init-output if
2657 ``inclusive=True``.
2658 task_xgraph : `networkx.DiGraph`
2659 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2660 inclusive : `bool`
2661 Whether to include the producer of the ``start`` node in the
2662 results.
2664 Returns
2665 -------
2666 selected : `set` [ `NodeKey` ]
2667 Set of `NodeKey` objects for matching tasks (only; no dataset type
2668 or task-init nodes).
2669 """
2670 result: set[NodeKey] = set()
2671 if inclusive:
2672 for producer_key, _ in self._xgraph.in_edges(start):
2673 if producer_key.node_type is NodeType.TASK_INIT:
2674 raise InvalidExpressionError(
2675 f"Init-output dataset type {start.name!r} cannot be used as the "
2676 "starting point for an includsive descendant ('>=') search."
2677 )
2678 result.add(producer_key)
2679 # We also include tasks that consume a dataset type as an init-input,
2680 # since that can affect their runtime behavior.
2681 consumer_keys: set[NodeKey] = {
2682 (
2683 consumer_key
2684 if consumer_key.node_type is NodeType.TASK
2685 else NodeKey(NodeType.TASK, consumer_key.name)
2686 )
2687 for _, consumer_key in self._xgraph.out_edges(start)
2688 }
2689 for consumer_key in consumer_keys:
2690 result.add(consumer_key)
2691 result.update(networkx.dag.descendants(task_xgraph, consumer_key))
2692 return result
2694 def _select_identifier(
2695 self, qualifier: Literal["T", "D", "S"] | None, label: str
2696 ) -> NodeKey | TaskSubset:
2697 """Return the node key or task subset that corresponds to a `select`
2698 expression identifier.
2700 Parameters
2701 ----------
2702 qualifier : `str` or `None`
2703 Task, dataset type, or task subset qualifier included in the
2704 identifier, if any.
2705 label : `str`
2706 Task label, dataset type name, or task subset label.
2708 Returns
2709 -------
2710 key_or_subset : `NodeKey` or `TaskSubset`
2711 A `NodeKey` for a task or dataset type, or a `TaskSubset` for a
2712 task subset.
2713 """
2714 match qualifier:
2715 case None:
2716 task_key = NodeKey(NodeType.TASK, label)
2717 dataset_type_key = NodeKey(NodeType.DATASET_TYPE, label)
2718 if task_key in self._xgraph.nodes:
2719 if dataset_type_key in self._xgraph.nodes:
2720 raise InvalidExpressionError(
2721 f"{label!r} is both a task label and a dataset type name; "
2722 "prefix with 'T:' or 'D:' (respectively) to specify which."
2723 )
2724 assert label not in self._task_subsets, "Should be prohibited at construction."
2725 return task_key
2726 elif dataset_type_key in self._xgraph.nodes:
2727 if label in self._task_subsets:
2728 raise InvalidExpressionError(
2729 f"{label!r} is both a subset label and a dataset type name; "
2730 "prefix with 'S:' or 'D:' (respectively) to specify which."
2731 )
2732 return dataset_type_key
2733 elif label in self._task_subsets:
2734 return self._task_subsets[label]
2735 else:
2736 raise InvalidExpressionError(
2737 f"{label!r} is not a task label, task subset label, or dataset type name."
2738 )
2739 case "T":
2740 task_key = NodeKey(NodeType.TASK, label)
2741 if task_key not in self._xgraph.nodes:
2742 raise InvalidExpressionError(f"Task with label {label!r} does not exist.")
2743 return task_key
2744 case "D":
2745 dataset_type_key = NodeKey(NodeType.DATASET_TYPE, label)
2746 if dataset_type_key not in self._xgraph.nodes:
2747 raise InvalidExpressionError(f"Dataset type with name {label!r} does not exist.")
2748 return dataset_type_key
2749 case "S":
2750 try:
2751 return self._task_subsets[label]
2752 except KeyError:
2753 raise InvalidExpressionError(f"Task subset with label {label!r} does not exist.")
2754 case _: # pragma: no cover
2755 raise AssertionError("Unexpected identifier qualifier in expression.")
2757 _xgraph: networkx.MultiDiGraph
2758 _sorted_keys: Sequence[NodeKey] | None
2759 _task_subsets: dict[str, TaskSubset]
2760 _step_definitions: StepDefinitions
2761 _description: str
2762 _tasks: TaskMappingView
2763 _dataset_types: DatasetTypeMappingView
2764 _raw_data_id: dict[str, Any]
2765 _universe: DimensionUniverse | None
2768def log_config_mismatch(msg: str) -> None:
2769 """Log messages about configuration mismatch.
2771 Parameters
2772 ----------
2773 msg : `str`
2774 Log message to use.
2775 """
2776 _LOG.fatal("Comparing configuration: %s", msg)
2779def compare_packages(packages: Packages, new_packages: Packages) -> bool:
2780 """Compare two versions of Packages.
2782 Parameters
2783 ----------
2784 packages : `Packages`
2785 Previously recorded package versions. Updated in place to include
2786 any new packages that weren't present before.
2787 new_packages : `Packages`
2788 New set of package versions.
2790 Returns
2791 -------
2792 updated : `bool`
2793 `True` if ``packages`` was updated, `False` if not.
2795 Raises
2796 ------
2797 ConflictingDefinitionError
2798 Raised if versions are inconsistent.
2799 """
2800 diff = new_packages.difference(packages)
2801 if diff:
2802 versions_str = "; ".join(f"{pkg}: {diff[pkg][1]} vs {diff[pkg][0]}" for pkg in diff)
2803 raise ConflictingDefinitionError(f"Package versions mismatch: ({versions_str})")
2804 else:
2805 _LOG.debug("new packages are consistent with old")
2806 # Update the old set of packages in case we have more packages
2807 # that haven't been persisted.
2808 extra = new_packages.extra(packages)
2809 if extra:
2810 _LOG.debug("extra packages: %s", extra)
2811 packages.update(new_packages)
2812 return True
2813 return False
2816def _dimensions_compatible(dimensions: DimensionGroup, object_dimensions: DimensionGroup) -> bool:
2817 if dimensions.issubset(object_dimensions):
2818 # Easy typical case.
2819 return True
2820 # Hard case: if any sharding dimensions that are not in the object
2821 # dimensions are related to something in the object dimensions by a
2822 # many-to-many join table, this is okay too. The main use case here is to
2823 # let {exposure, [detector]} satisfy sharding dimensions of
2824 # {visit, [detector]}.
2825 universe = dimensions.universe
2826 unmatched_sharding = dimensions.required - object_dimensions.required
2827 unmatched_object = object_dimensions.required - dimensions.required
2828 unmatched_union_group = dimensions.union(object_dimensions)
2829 for element_name in unmatched_union_group.elements:
2830 if (
2831 (element := universe[element_name]).defines_relationships
2832 and not element.dimensions.isdisjoint(unmatched_sharding)
2833 and not element.dimensions.isdisjoint(unmatched_object)
2834 ):
2835 return True
2836 return False