Coverage for python/lsst/pipe/base/pipeline_graph/_pipeline_graph.py: 12%
739 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-11 01:30 -0700
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
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 Raises
901 ------
902 ValueError
903 Raised if ``assume_edges_unchanged`` and ``check_edges_unchanged``
904 are both `True`, or if the same task appears twice.
905 EdgesChangedError
906 Raised if ``check_edges_unchanged=True`` and the edges of a task do
907 change.
909 Notes
910 -----
911 If reconfiguring a task causes its edges to change, any dataset type
912 nodes connected to that task (not just those whose edges have changed!)
913 will be unresolved.
914 """
915 new_configs: dict[str, PipelineTaskConfig] = {}
916 for task_label, config_update in itertools.chain(args, kwargs.items()):
917 if new_configs.setdefault(task_label, config_update) is not config_update:
918 raise ValueError(f"Config for {task_label!r} provided more than once.")
919 updates = {
920 task_label: self.tasks[task_label]._reconfigured(config, rebuild=not assume_edges_unchanged)
921 for task_label, config in new_configs.items()
922 }
923 self._replace_task_nodes(
924 updates,
925 check_edges_unchanged=check_edges_unchanged,
926 assume_edges_unchanged=assume_edges_unchanged,
927 message_header=(
928 "Unexpected change in edges for task {task_label!r} from original config (A) to "
929 "new configs (B):"
930 ),
931 )
933 def remove_tasks(
934 self, labels: Iterable[str], drop_from_subsets: bool = True
935 ) -> list[tuple[TaskNode, set[str]]]:
936 """Remove one or more tasks from the graph.
938 Parameters
939 ----------
940 labels : `~collections.abc.Iterable` [ `str` ]
941 Iterable of the labels of the tasks to remove.
942 drop_from_subsets : `bool`, optional
943 If `True`, drop each removed task from any subset in which it
944 currently appears. If `False`, raise `PipelineGraphError` if any
945 such subsets exist.
947 Returns
948 -------
949 nodes_and_subsets : `list` [ `tuple` [ `TaskNode`, `set` [ `str` ] ] ]
950 List of nodes removed and the labels of task subsets that
951 referenced them.
953 Raises
954 ------
955 PipelineGraphError
956 Raised if ``drop_from_subsets`` is `False` and the task is still
957 part of one or more subsets.
959 Notes
960 -----
961 Removing a task will cause dataset nodes with no other referencing
962 tasks to be removed. Any other dataset type nodes referenced by a
963 removed task will be reset to an "unresolved" state.
964 """
965 task_nodes_and_subsets = []
966 dataset_types: set[NodeKey] = set()
967 nodes_to_remove = set()
968 for label in labels:
969 task_node: TaskNode = self._xgraph.nodes[NodeKey(NodeType.TASK, label)]["instance"]
970 # Find task subsets that reference this task.
971 referencing_subsets = {
972 subset_label
973 for subset_label, task_subset in self.task_subsets.items()
974 if label in task_subset
975 }
976 if not drop_from_subsets and referencing_subsets:
977 raise PipelineGraphError(
978 f"Task {label!r} is still referenced by subset(s) {referencing_subsets}."
979 )
980 task_nodes_and_subsets.append((task_node, referencing_subsets))
981 # Find dataset types referenced by this task.
982 dataset_types.update(self._xgraph.predecessors(task_node.key))
983 dataset_types.update(self._xgraph.successors(task_node.key))
984 dataset_types.update(self._xgraph.predecessors(task_node.init.key))
985 dataset_types.update(self._xgraph.successors(task_node.init.key))
986 # Since there's an edge between the task and its init node, we'll
987 # have added those two nodes here, too, and we don't want that.
988 dataset_types.remove(task_node.init.key)
989 dataset_types.remove(task_node.key)
990 # Mark the task node and its init node for removal from the graph.
991 nodes_to_remove.add(task_node.key)
992 nodes_to_remove.add(task_node.init.key)
993 # Process the referenced datasets to see which ones are orphaned and
994 # need to be removed vs. just unresolved.
995 nodes_to_unresolve = []
996 for dataset_type_key in dataset_types:
997 related_tasks = set()
998 related_tasks.update(self._xgraph.predecessors(dataset_type_key))
999 related_tasks.update(self._xgraph.successors(dataset_type_key))
1000 related_tasks.difference_update(nodes_to_remove)
1001 if not related_tasks:
1002 nodes_to_remove.add(dataset_type_key)
1003 else:
1004 nodes_to_unresolve.append(dataset_type_key)
1005 # Checks and preparation complete; time to start the actual
1006 # modification, during which it's hard to provide strong exception
1007 # safety. Start by resetting the sort ordering.
1008 self._reset()
1009 try:
1010 for dataset_type_key in nodes_to_unresolve:
1011 self._xgraph.nodes[dataset_type_key]["instance"] = None
1012 for task_node, referencing_subsets in task_nodes_and_subsets:
1013 for subset_label in referencing_subsets:
1014 self._task_subsets[subset_label].remove(task_node.label)
1015 self._xgraph.remove_nodes_from(nodes_to_remove)
1016 except Exception as err: # pragma: no cover
1017 # There's no known way to get here, but we want to make it
1018 # clear it's a big problem if we do.
1019 raise PipelineGraphExceptionSafetyError(
1020 "Error during task removal has left the graph in an inconsistent state."
1021 ) from err
1022 return task_nodes_and_subsets
1024 def add_task_subset(self, subset_label: str, task_labels: Iterable[str], description: str = "") -> None:
1025 """Add a label for a set of tasks that are already in the pipeline.
1027 Parameters
1028 ----------
1029 subset_label : `str`
1030 Label for this set of tasks.
1031 task_labels : `~collections.abc.Iterable` [ `str` ]
1032 Labels of the tasks to include in the set. All must already be
1033 included in the graph.
1034 description : `str`, optional
1035 String description to associate with this label.
1036 """
1037 subset = TaskSubset(self._xgraph, subset_label, set(task_labels), description, self._step_definitions)
1038 self._task_subsets[subset_label] = subset
1040 def remove_task_subset(self, subset_label: str) -> None:
1041 """Remove a labeled set of tasks.
1043 Parameters
1044 ----------
1045 subset_label : `str`
1046 Label for this set of tasks.
1048 Notes
1049 -----
1050 If this subset is a step, it is also removed from the step definitions.
1051 """
1052 try:
1053 self._step_definitions.remove(subset_label)
1054 except KeyError:
1055 pass
1056 del self._task_subsets[subset_label]
1058 ###########################################################################
1059 #
1060 # NetworkX Export Interface:
1061 #
1062 # - methods to export the PipelineGraph's content (or various subsets
1063 # thereof) as NetworkX objects.
1064 #
1065 # These are particularly useful when writing tools to visualize the graph,
1066 # while providing options for which aspects of the graph (tasks, dataset
1067 # types, or both) to include, since all exported graphs have similar
1068 # attributes regardless of their structure.
1069 #
1070 ###########################################################################
1072 def make_xgraph(self) -> networkx.MultiDiGraph:
1073 """Export a networkx representation of the full pipeline graph,
1074 including both init and runtime edges.
1076 Returns
1077 -------
1078 xgraph : `networkx.MultiDiGraph`
1079 Directed acyclic graph with parallel edges.
1081 Notes
1082 -----
1083 The returned graph uses `NodeKey` instances for nodes. Parallel edges
1084 represent the same dataset type appearing in multiple connections for
1085 the same task, and are hence rare. The connection name is used as the
1086 edge key to disambiguate those parallel edges.
1088 Almost all edges connect dataset type nodes to task or task init nodes
1089 or vice versa, but there is also a special edge that connects each task
1090 init node to its runtime node. The existence of these edges makes the
1091 graph not quite bipartite, though its init-only and runtime-only
1092 subgraphs are bipartite.
1094 See `TaskNode`, `TaskInitNode`, `DatasetTypeNode`, `ReadEdge`, and
1095 `WriteEdge` for the descriptive node and edge attributes added.
1096 """
1097 return self._transform_xgraph_state(self._xgraph.copy(), skip_edges=False)
1099 def make_bipartite_xgraph(self, init: bool = False) -> networkx.MultiDiGraph:
1100 """Return a bipartite networkx representation of just the runtime or
1101 init-time pipeline graph.
1103 Parameters
1104 ----------
1105 init : `bool`, optional
1106 If `True` (`False` is default) return the graph of task
1107 initialization nodes and init input/output dataset types, instead
1108 of the graph of runtime task nodes and regular
1109 input/output/prerequisite dataset types.
1111 Returns
1112 -------
1113 xgraph : `networkx.MultiDiGraph`
1114 Directed acyclic graph with parallel edges.
1116 Notes
1117 -----
1118 The returned graph uses `NodeKey` instances for nodes. Parallel edges
1119 represent the same dataset type appearing in multiple connections for
1120 the same task, and are hence rare. The connection name is used as the
1121 edge key to disambiguate those parallel edges.
1123 This graph is bipartite because each dataset type node only has edges
1124 that connect it to a task [init] node, and vice versa.
1126 See `TaskNode`, `TaskInitNode`, `DatasetTypeNode`, `ReadEdge`, and
1127 `WriteEdge` for the descriptive node and edge attributes added.
1128 """
1129 return self._transform_xgraph_state(
1130 self._make_bipartite_xgraph_internal(init).copy(), skip_edges=False
1131 )
1133 def make_task_xgraph(self, init: bool = False) -> networkx.DiGraph:
1134 """Return a networkx representation of just the tasks in the pipeline.
1136 Parameters
1137 ----------
1138 init : `bool`, optional
1139 If `True` (`False` is default) return the graph of task
1140 initialization nodes, instead of the graph of runtime task nodes.
1142 Returns
1143 -------
1144 xgraph : `networkx.DiGraph`
1145 Directed acyclic graph with no parallel edges.
1147 Notes
1148 -----
1149 The returned graph uses `NodeKey` instances for nodes. The dataset
1150 types that link these tasks are not represented at all; edges have no
1151 attributes, and there are no parallel edges.
1153 See `TaskNode` and `TaskInitNode` for the descriptive node and
1154 attributes added.
1155 """
1156 return self._transform_xgraph_state(self._make_task_xgraph_internal(init), skip_edges=True)
1158 def make_dataset_type_xgraph(self, init: bool = False) -> networkx.DiGraph:
1159 """Return a networkx representation of just the dataset types in the
1160 pipeline.
1162 Parameters
1163 ----------
1164 init : `bool`, optional
1165 If `True` (`False` is default) return the graph of init input and
1166 output dataset types, instead of the graph of runtime (input,
1167 output, prerequisite input) dataset types.
1169 Returns
1170 -------
1171 xgraph : `networkx.DiGraph`
1172 Directed acyclic graph with no parallel edges.
1174 Notes
1175 -----
1176 The returned graph uses `NodeKey` instances for nodes. The tasks that
1177 link these tasks are not represented at all; edges have no attributes,
1178 and there are no parallel edges.
1180 See `DatasetTypeNode` for the descriptive node and attributes added.
1181 """
1182 bipartite_xgraph = self._make_bipartite_xgraph_internal(init)
1183 dataset_type_keys = [
1184 key
1185 for key, bipartite in bipartite_xgraph.nodes(data="bipartite")
1186 if bipartite == NodeType.DATASET_TYPE.bipartite
1187 ]
1188 return self._transform_xgraph_state(
1189 networkx.algorithms.bipartite.projected_graph(
1190 networkx.DiGraph(bipartite_xgraph), dataset_type_keys
1191 ),
1192 skip_edges=True,
1193 )
1195 ###########################################################################
1196 #
1197 # Expression-based Selection Interface.
1198 #
1199 ###########################################################################
1201 def select_tasks(self, expression: str) -> set[str]:
1202 """Return the tasks that match an expression.
1204 Parameters
1205 ----------
1206 expression : `str`
1207 String expression to evaluate. See
1208 :ref:`pipeline-graph-subset-expressions`.
1210 Returns
1211 -------
1212 task_labels : `set` [ `str` ]
1213 Set of matching task labels.
1214 """
1215 task_xgraph = self._make_task_xgraph_internal(init=False)
1216 expr_tree = expressions.parse(expression)
1217 matching_task_keys = self._select_expression(expr_tree, task_xgraph)
1218 return {key.name for key in matching_task_keys}
1220 def select(self, expression: str) -> PipelineGraph:
1221 """Return a new pipeline graph with the tasks that match an expression.
1223 Parameters
1224 ----------
1225 expression : `str`
1226 String expression to evaluate. See
1227 :ref:`pipeline-graph-subset-expressions`.
1229 Returns
1230 -------
1231 new_graph : `PipelineGraph`
1232 New pipeline graph with just the matching tasks.
1234 Notes
1235 -----
1236 All resolved dataset type nodes will be preserved.
1238 If `has_been_sorted`, the new graph will be sorted as well.
1240 Task subsets will not be included in the returned graph.
1241 """
1242 selected_tasks = self.select_tasks(expression)
1243 new_pipeline_graph = PipelineGraph(universe=self._universe, data_id=self._raw_data_id)
1244 new_pipeline_graph.add_task_nodes(
1245 [self.tasks[task_label] for task_label in selected_tasks], parent=self
1246 )
1247 if self.has_been_sorted:
1248 new_pipeline_graph.sort()
1249 return new_pipeline_graph
1251 ###########################################################################
1252 #
1253 # Serialization Interface.
1254 #
1255 # Serialization of PipelineGraphs is currently experimental and may not be
1256 # retained in the future. All serialization methods are
1257 # underscore-prefixed to ensure nobody mistakes them for a stable interface
1258 # (let a lone a stable file format).
1259 #
1260 ###########################################################################
1262 @classmethod
1263 def _read_stream(
1264 cls, stream: BinaryIO, import_mode: TaskImportMode = TaskImportMode.REQUIRE_CONSISTENT_EDGES
1265 ) -> PipelineGraph:
1266 """Read a serialized `PipelineGraph` from a file-like object.
1268 Parameters
1269 ----------
1270 stream : `BinaryIO`
1271 File-like object opened for binary reading, containing
1272 gzip-compressed JSON.
1273 import_mode : `TaskImportMode`, optional
1274 Whether to import tasks, and how to reconcile any differences
1275 between the imported task's connections and the those that were
1276 persisted with the graph. Default is to check that they are the
1277 same.
1279 Returns
1280 -------
1281 graph : `PipelineGraph`
1282 Deserialized pipeline graph.
1284 Raises
1285 ------
1286 PipelineGraphReadError
1287 Raised if the serialized `PipelineGraph` is not self-consistent.
1288 EdgesChangedError
1289 Raised if ``import_mode`` is
1290 `TaskImportMode.REQUIRED_CONSISTENT_EDGES` and the edges of a task
1291 did change after import and reconfiguration.
1293 Notes
1294 -----
1295 `PipelineGraph` serialization is currently experimental and may be
1296 removed or significantly changed in the future, with no deprecation
1297 period.
1298 """
1299 from .io import SerializedPipelineGraph
1301 with gzip.open(stream, "rb") as uncompressed_stream:
1302 data = json.load(uncompressed_stream)
1303 serialized_graph = SerializedPipelineGraph.model_validate(data)
1304 return serialized_graph.deserialize(import_mode)
1306 @classmethod
1307 def _read_uri(
1308 cls,
1309 uri: ResourcePathExpression,
1310 import_mode: TaskImportMode = TaskImportMode.REQUIRE_CONSISTENT_EDGES,
1311 ) -> PipelineGraph:
1312 """Read a serialized `PipelineGraph` from a file at a URI.
1314 Parameters
1315 ----------
1316 uri : convertible to `lsst.resources.ResourcePath`
1317 URI to a gzip-compressed JSON file containing a serialized pipeline
1318 graph.
1319 import_mode : `TaskImportMode`, optional
1320 Whether to import tasks, and how to reconcile any differences
1321 between the imported task's connections and the those that were
1322 persisted with the graph. Default is to check that they are the
1323 same.
1325 Returns
1326 -------
1327 graph : `PipelineGraph`
1328 Deserialized pipeline graph.
1330 Raises
1331 ------
1332 PipelineGraphReadError
1333 Raised if the serialized `PipelineGraph` is not self-consistent.
1334 EdgesChangedError
1335 Raised if ``import_mode`` is
1336 `TaskImportMode.REQUIRED_CONSISTENT_EDGES` and the edges of a task
1337 did change after import and reconfiguration.
1339 Notes
1340 -----
1341 `PipelineGraph` serialization is currently experimental and may be
1342 removed or significantly changed in the future, with no deprecation
1343 period.
1344 """
1345 uri = ResourcePath(uri)
1346 with uri.open("rb") as stream:
1347 return cls._read_stream(cast(BinaryIO, stream), import_mode=import_mode)
1349 def _write_stream(self, stream: BinaryIO) -> None:
1350 """Write the pipeline to a file-like object.
1352 Parameters
1353 ----------
1354 stream
1355 File-like object opened for binary writing.
1357 Notes
1358 -----
1359 `PipelineGraph` serialization is currently experimental and may be
1360 removed or significantly changed in the future, with no deprecation
1361 period.
1363 The file format is gzipped JSON, and is intended to be human-readable,
1364 but it should not be considered a stable public interface for outside
1365 code, which should always use `PipelineGraph` methods (or at least the
1366 `io.SerializedPipelineGraph` class) to read these files.
1367 """
1368 from .io import SerializedPipelineGraph
1370 with gzip.open(stream, mode="wb") as compressed_stream:
1371 compressed_stream.write(
1372 SerializedPipelineGraph.serialize(self).model_dump_json(exclude_defaults=True).encode("utf-8")
1373 )
1375 def _write_uri(self, uri: ResourcePathExpression) -> None:
1376 """Write the pipeline to a file given a URI.
1378 Parameters
1379 ----------
1380 uri : convertible to `lsst.resources.ResourcePath`
1381 URI to write to . May have ``.json.gz`` or no extension (which
1382 will cause a ``.json.gz`` extension to be added).
1384 Notes
1385 -----
1386 `PipelineGraph` serialization is currently experimental and may be
1387 removed or significantly changed in the future, with no deprecation
1388 period.
1390 The file format is gzipped JSON, and is intended to be human-readable,
1391 but it should not be considered a stable public interface for outside
1392 code, which should always use `PipelineGraph` methods (or at least the
1393 `io.SerializedPipelineGraph` class) to read these files.
1394 """
1395 uri = ResourcePath(uri)
1396 extension = uri.getExtension()
1397 if not extension:
1398 uri = uri.updatedExtension(".json.gz")
1399 elif extension != ".json.gz":
1400 raise ValueError("Expanded pipeline files should always have a .json.gz extension.")
1401 with uri.open(mode="wb") as stream:
1402 self._write_stream(cast(BinaryIO, stream))
1404 def _import_and_configure(
1405 self, import_mode: TaskImportMode = TaskImportMode.REQUIRE_CONSISTENT_EDGES
1406 ) -> None:
1407 """Import the `PipelineTask` classes referenced by all task nodes and
1408 update those nodes accordingly.
1410 Parameters
1411 ----------
1412 import_mode : `TaskImportMode`, optional
1413 Whether to import tasks, and how to reconcile any differences
1414 between the imported task's connections and the those that were
1415 persisted with the graph. Default is to check that they are the
1416 same. This method does nothing if this is
1417 `TaskImportMode.DO_NOT_IMPORT`.
1419 Raises
1420 ------
1421 EdgesChangedError
1422 Raised if ``import_mode`` is
1423 `TaskImportMode.REQUIRED_CONSISTENT_EDGES` and the edges of a task
1424 did change after import and reconfiguration.
1426 Notes
1427 -----
1428 This method shouldn't need to be called unless the graph was
1429 deserialized without importing and configuring immediately, which is
1430 not the default behavior (but it can greatly speed up deserialization).
1431 If all tasks have already been imported this does nothing.
1433 Importing and configuring a task can change its
1434 `~TaskNode.task_class_name` or `~TaskClass.get_config_str` output,
1435 usually because the software used to read a serialized graph is newer
1436 than the software used to write it (e.g. a new config option has been
1437 added, or the task was moved to a new module with a forwarding alias
1438 left behind). These changes are allowed by
1439 `TaskImportMode.REQUIRE_CONSISTENT_EDGES`.
1441 If importing and configuring a task causes its edges to change, any
1442 dataset type nodes linked to those edges will be reset to the
1443 unresolved state.
1444 """
1445 if import_mode is TaskImportMode.DO_NOT_IMPORT:
1446 return
1447 rebuild = (
1448 import_mode is TaskImportMode.REQUIRE_CONSISTENT_EDGES
1449 or import_mode is TaskImportMode.OVERRIDE_EDGES
1450 )
1451 updates: dict[str, TaskNode] = {}
1452 node_key: NodeKey
1453 for node_key, node_state in self._xgraph.nodes.items():
1454 if node_key.node_type is NodeType.TASK:
1455 task_node: TaskNode = node_state["instance"]
1456 new_task_node = task_node._imported_and_configured(rebuild)
1457 if new_task_node is not task_node:
1458 updates[task_node.label] = new_task_node
1459 self._replace_task_nodes(
1460 updates,
1461 check_edges_unchanged=(import_mode is TaskImportMode.REQUIRE_CONSISTENT_EDGES),
1462 assume_edges_unchanged=(import_mode is TaskImportMode.ASSUME_CONSISTENT_EDGES),
1463 message_header=(
1464 "In task with label {task_label!r}, persisted edges (A)"
1465 "differ from imported and configured edges (B):"
1466 ),
1467 )
1469 ###########################################################################
1470 #
1471 # Advanced PipelineGraph Inspection Interface:
1472 #
1473 # - methods to iterate over all nodes and edges, utilizing NodeKeys;
1474 #
1475 # - methods to find overall inputs and group nodes by their dimensions,
1476 # which are important operations for QuantumGraph generation.
1477 #
1478 ###########################################################################
1480 def iter_edges(self, init: bool = False) -> Iterator[Edge]:
1481 """Iterate over edges in the graph.
1483 Parameters
1484 ----------
1485 init : `bool`, optional
1486 If `True` (`False` is default) iterate over the edges between task
1487 initialization node and init input/output dataset types, instead of
1488 the runtime task nodes and regular input/output/prerequisite
1489 dataset types.
1491 Returns
1492 -------
1493 edges : `~collections.abc.Iterator` [ `Edge` ]
1494 A lazy iterator over `Edge` (`WriteEdge` or `ReadEdge`) instances.
1496 Notes
1497 -----
1498 This method always returns _either_ init edges or runtime edges, never
1499 both. The full (internal) graph that contains both also includes a
1500 special edge that connects each task init node to its runtime node;
1501 that is also never returned by this method, since it is never a part of
1502 the init-only or runtime-only subgraphs.
1503 """
1504 edge: Edge
1505 for _, _, edge in self._xgraph.edges(data="instance"):
1506 if edge is not None and edge.is_init == init:
1507 yield edge
1509 def iter_nodes(
1510 self,
1511 ) -> Iterator[
1512 tuple[Literal[NodeType.TASK_INIT], str, TaskInitNode]
1513 | tuple[Literal[NodeType.TASK], str, TaskInitNode]
1514 | tuple[Literal[NodeType.DATASET_TYPE], str, DatasetTypeNode | None]
1515 ]:
1516 """Iterate over nodes in the graph.
1518 Returns
1519 -------
1520 nodes : `~collections.abc.Iterator` [ `tuple` ]
1521 A lazy iterator over all of the nodes in the graph. Each yielded
1522 element is a tuple of:
1524 - the node type enum value (`NodeType`);
1525 - the string name for the node (task label or parent dataset type
1526 name);
1527 - the node value (`TaskNode`, `TaskInitNode`, `DatasetTypeNode`,
1528 or `None` for dataset type nodes that have not been resolved).
1529 """
1530 key: NodeKey
1531 if self._sorted_keys is not None:
1532 for key in self._sorted_keys:
1533 yield key.node_type, key.name, self._xgraph.nodes[key]["instance"] # type: ignore
1534 else:
1535 for key, node in self._xgraph.nodes(data="instance"):
1536 yield key.node_type, key.name, node # type: ignore
1538 def iter_overall_inputs(self) -> Iterator[tuple[str, DatasetTypeNode | None]]:
1539 """Iterate over all of the dataset types that are consumed but not
1540 produced by the graph.
1542 Returns
1543 -------
1544 dataset_types : `~collections.abc.Iterator` [ `tuple` ]
1545 A lazy iterator over the overall-input dataset types (including
1546 overall init inputs and prerequisites). Each yielded element is a
1547 tuple of:
1549 - the parent dataset type name;
1550 - the resolved `DatasetTypeNode`, or `None` if the dataset type has
1551 not been resolved.
1552 """
1553 for generation in networkx.algorithms.dag.topological_generations(self._xgraph):
1554 key: NodeKey
1555 for key in generation:
1556 # While we expect all tasks to have at least one input and
1557 # hence never appear in the first topological generation, that
1558 # is not true of task init nodes.
1559 if key.node_type is NodeType.DATASET_TYPE:
1560 yield key.name, self._xgraph.nodes[key]["instance"]
1561 return
1563 def group_by_dimensions(
1564 self, prerequisites: bool = False
1565 ) -> dict[DimensionGroup, tuple[dict[str, TaskNode], dict[str, DatasetTypeNode]]]:
1566 """Group this graph's tasks and dataset types by their dimensions.
1568 Parameters
1569 ----------
1570 prerequisites : `bool`, optional
1571 If `True`, include prerequisite dataset types as well as regular
1572 input and output datasets (including intermediates).
1574 Returns
1575 -------
1576 groups : `dict` [ `~lsst.daf.butler.DimensionGroup`, `tuple` ]
1577 A dictionary of groups keyed by `~lsst.daf.butler.DimensionGroup`,
1578 in which each value is a tuple of:
1580 - a `dict` of `TaskNode` instances, keyed by task label
1581 - a `dict` of `DatasetTypeNode` instances, keyed by
1582 dataset type name.
1584 that have those dimensions.
1586 Notes
1587 -----
1588 Init inputs and outputs are always included, but always have empty
1589 dimensions and are hence are all grouped together.
1590 """
1591 result: dict[DimensionGroup, tuple[dict[str, TaskNode], dict[str, DatasetTypeNode]]] = {}
1592 next_new_value: tuple[dict[str, TaskNode], dict[str, DatasetTypeNode]] = ({}, {})
1593 for task_label, task_node in self.tasks.items():
1594 if task_node.dimensions is None:
1595 raise UnresolvedGraphError(f"Task with label {task_label!r} has not been resolved.")
1596 if (group := result.setdefault(task_node.dimensions, next_new_value)) is next_new_value:
1597 next_new_value = ({}, {}) # make new lists for next time
1598 group[0][task_node.label] = task_node
1599 for dataset_type_name, dataset_type_node in self.dataset_types.items():
1600 if dataset_type_node is None:
1601 raise UnresolvedGraphError(f"Dataset type {dataset_type_name!r} has not been resolved.")
1602 if not dataset_type_node.is_prerequisite or prerequisites:
1603 if (
1604 group := result.setdefault(dataset_type_node.dataset_type.dimensions, next_new_value)
1605 ) is next_new_value:
1606 next_new_value = ({}, {}) # make new lists for next time
1607 group[1][dataset_type_node.name] = dataset_type_node
1608 return result
1610 def get_all_dimensions(self, prerequisites: bool = True) -> DimensionGroup:
1611 """Return all dimensions used in this graph's tasks and dataset types.
1613 Parameters
1614 ----------
1615 prerequisites : `bool`, optional
1616 If `False`, do not include the dimensions that are only used by
1617 prerequisite input dataset types.
1619 Returns
1620 -------
1621 dimensions : `~lsst.daf.butler.DimensionGroup`.
1622 All dimensions in this pipeline.
1623 """
1624 return DimensionGroup.union(
1625 *self.group_by_dimensions(prerequisites=prerequisites).keys(),
1626 universe=self.universe,
1627 )
1629 def split_independent(self) -> Iterable[PipelineGraph]:
1630 """Iterate over independent subgraphs that together comprise this
1631 pipeline graph.
1633 Returns
1634 -------
1635 subgraphs : `Iterable` [ `PipelineGraph` ]
1636 An iterable over component subgraphs that could be run
1637 independently (they have only overall inputs in common). May be a
1638 lazy iterator.
1640 Notes
1641 -----
1642 All resolved dataset type nodes will be preserved.
1644 If there is only one component, ``self`` may be returned as the only
1645 element in the iterable.
1647 If `has_been_sorted`, all subgraphs will be sorted as well.
1649 Task subsets will not be included in the returned graphs.
1650 """
1651 # Having an overall input in common isn't enough to make subgraphs
1652 # dependent on each other, so we want to look for connected component
1653 # subgraphs of the task-only projected graph.
1654 bipartite_xgraph = self._make_bipartite_xgraph_internal(init=False)
1655 task_keys = {
1656 key
1657 for key, bipartite in bipartite_xgraph.nodes(data="bipartite")
1658 if bipartite == NodeType.TASK.bipartite
1659 }
1660 task_xgraph = networkx.algorithms.bipartite.projected_graph(
1661 networkx.DiGraph(bipartite_xgraph), task_keys
1662 )
1663 # "Weakly" connected means connected in only one direction, which is
1664 # the only kind of "connected" a DAG can ever be.
1665 for component_task_keys in networkx.algorithms.weakly_connected_components(task_xgraph):
1666 if component_task_keys == task_keys:
1667 yield self
1668 return
1669 else:
1670 component_subgraph = PipelineGraph(universe=self._universe, data_id=self._raw_data_id)
1671 component_subgraph.add_task_nodes(
1672 [self._xgraph.nodes[key]["instance"] for key in component_task_keys], parent=self
1673 )
1674 if self.has_been_sorted:
1675 component_subgraph.sort()
1676 yield component_subgraph
1678 ###########################################################################
1679 #
1680 # Data repository/collection initialization
1681 #
1682 ###########################################################################
1684 @property
1685 def packages_dataset_type(self) -> DatasetType:
1686 """The special "packages" dataset type that records software versions.
1688 This is not associated with a task and hence is
1689 not considered part of the pipeline graph in other respects, but it
1690 does get written with other provenance datasets.
1691 """
1692 if self._universe is None:
1693 raise UnresolvedGraphError(
1694 "PipelineGraph must be resolved in order to get the packages dataset type."
1695 )
1696 return DatasetType(
1697 PACKAGES_INIT_OUTPUT_NAME, self._universe.empty, PACKAGES_INIT_OUTPUT_STORAGE_CLASS
1698 )
1700 def register_dataset_types(self, butler: Butler, include_packages: bool = True) -> None:
1701 """Register all dataset types in a data repository.
1703 Parameters
1704 ----------
1705 butler : `~lsst.daf.butler.Butler`
1706 Data repository client.
1707 include_packages : `bool`, optional
1708 Whether to include the special "packages" dataset type that records
1709 software versions (this is not associated with a task and hence is
1710 not considered part of the pipeline graph in other respects, but it
1711 does get written with other provenance datasets).
1712 """
1713 dataset_types = [node.dataset_type for node in self.dataset_types.values()]
1714 if include_packages:
1715 dataset_types.append(self.packages_dataset_type)
1716 for dataset_type in dataset_types:
1717 butler.registry.registerDatasetType(dataset_type)
1719 def check_dataset_type_registrations(self, butler: Butler, include_packages: bool = True) -> None:
1720 """Check that dataset type registrations in a data repository match
1721 the definitions in this pipeline graph.
1723 Parameters
1724 ----------
1725 butler : `~lsst.daf.butler.Butler`
1726 Data repository client.
1727 include_packages : `bool`, optional
1728 Whether to include the special "packages" dataset type that records
1729 software versions (this is not associated with a task and hence is
1730 not considered part of the pipeline graph in other respects, but it
1731 does get written with other provenance datasets).
1733 Raises
1734 ------
1735 lsst.daf.butler.MissingDatasetTypeError
1736 Raised if one or more non-optional-input or output dataset types in
1737 the pipeline is not registered at all.
1738 lsst.daf.butler.ConflictingDefinitionError
1739 Raised if the definition in the data repository is not identical
1740 to the definition in the pipeline graph.
1742 Notes
1743 -----
1744 Note that dataset type definitions that are storage-class-conversion
1745 compatible but not identical are not permitted by these checks, because
1746 the expectation is that these differences are handled by `resolve`,
1747 which makes the pipeline graph use the data repository definitions.
1748 This method is intended to check that none of those definitions have
1749 changed.
1750 """
1751 dataset_types = [node.dataset_type for node in self.dataset_types.values()]
1752 if include_packages:
1753 dataset_types.append(self.packages_dataset_type)
1754 missing_dataset_types: list[str] = []
1755 for dataset_type in dataset_types:
1756 try:
1757 expected = butler.registry.getDatasetType(dataset_type.name)
1758 except MissingDatasetTypeError:
1759 expected = None
1760 if expected is None:
1761 # The user probably forgot to register dataset types
1762 # at least once (which should be an error),
1763 # but we could also get here if this is an optional input for
1764 # which no datasets were found in this repo (not an error).
1765 if (
1766 not (
1767 self.producer_of(dataset_type.name) is None
1768 and all(
1769 self.tasks[input_edge.task_label].is_optional(input_edge.connection_name)
1770 for input_edge in self.consuming_edges_of(dataset_type.name)
1771 )
1772 )
1773 or dataset_type.name == PACKAGES_INIT_OUTPUT_NAME
1774 ):
1775 missing_dataset_types.append(dataset_type.name)
1776 elif expected != dataset_type:
1777 raise ConflictingDefinitionError(
1778 f"DatasetType definition in registry has changed since the pipeline graph was resolved: "
1779 f"{dataset_type} (graph) != {expected} (registry)."
1780 )
1781 if missing_dataset_types:
1782 plural = "s" if len(missing_dataset_types) != 1 else ""
1783 raise MissingDatasetTypeError(
1784 f"Missing dataset type definition{plural}: {', '.join(missing_dataset_types)}. "
1785 "Dataset types have to be registered in advance (on the command-line, either via "
1786 "`butler register-dataset-type` or the `--register-dataset-types` option to `pipetask run`."
1787 )
1789 def instantiate_tasks(
1790 self,
1791 get_init_input: Callable[[DatasetType], Any] | None = None,
1792 init_outputs: list[tuple[Any, DatasetType]] | None = None,
1793 labels: Iterable[str] | None = None,
1794 ) -> list[PipelineTask]:
1795 """Instantiate all tasks in the pipeline.
1797 Parameters
1798 ----------
1799 get_init_input : `~collections.abc.Callable`, optional
1800 Callable that accepts a single `~lsst.daf.butler.DatasetType`
1801 parameter and returns the init-input dataset associated with that
1802 dataset type. Must respect the storage class embedded in the type.
1803 This is optional if the pipeline does not have any overall init
1804 inputs. When a full butler is available,
1805 `lsst.daf.butler.Butler.get` can be used directly here.
1806 init_outputs : `list`, optional
1807 A list of ``(obj, dataset type)`` init-output dataset pairs, to be
1808 appended to in-place. Both the object and the dataset type will
1809 correspond to the storage class of the output connection, which
1810 may not be the same as the storage class on the graph's dataset
1811 type node.
1812 labels : `~collections.abc.Iterable` [ `str` ], optional
1813 The labels of tasks to instantiate. If not provided, all tasks in
1814 the graph will be instantiated.
1816 Returns
1817 -------
1818 tasks : `list`
1819 Constructed `PipelineTask` instances.
1820 """
1821 if not self.is_fully_resolved:
1822 raise UnresolvedGraphError("Pipeline graph must be fully resolved before instantiating tasks.")
1823 empty_data_id = DataCoordinate.make_empty(self.universe)
1824 labels = set(labels) if labels is not None else self.tasks.keys()
1825 handles: dict[str, InMemoryDatasetHandle] = {}
1826 tasks: list[PipelineTask] = []
1827 for task_node in self.tasks.values():
1828 if task_node.label not in labels:
1829 continue
1830 task_init_inputs: dict[str, Any] = {}
1831 for read_edge in task_node.init.inputs.values():
1832 if (handle := handles.get(read_edge.dataset_type_name)) is not None:
1833 obj = handle.get(storageClass=read_edge.storage_class_name)
1834 elif (
1835 read_edge.component is not None
1836 and (parent_handle := handles.get(read_edge.parent_dataset_type_name)) is not None
1837 ):
1838 obj = parent_handle.get(
1839 storageClass=read_edge.storage_class_name, component=read_edge.component
1840 )
1841 else:
1842 dataset_type_node = self.dataset_types[read_edge.parent_dataset_type_name]
1843 if get_init_input is None:
1844 raise ValueError(
1845 f"Task {task_node.label!r} requires init-input "
1846 f"{read_edge.dataset_type_name} but no 'get_init_input' callback was provided."
1847 )
1848 obj = get_init_input(read_edge.adapt_dataset_type(dataset_type_node.dataset_type))
1849 n_consumers = len(self.consumers_of(dataset_type_node.name))
1850 if (
1851 n_consumers > 1
1852 and read_edge.component is None
1853 and read_edge.storage_class_name == dataset_type_node.storage_class_name
1854 ):
1855 # Caching what we just got is safe in general only
1856 # if there was no storage class conversion, since
1857 # a->b and a->c does not imply b->c.
1858 handles[read_edge.dataset_type_name] = InMemoryDatasetHandle(
1859 obj,
1860 storageClass=dataset_type_node.storage_class,
1861 dataId=empty_data_id,
1862 copy=True,
1863 )
1864 task_init_inputs[read_edge.connection_name] = obj
1865 task = task_node.task_class(
1866 config=task_node.config, initInputs=task_init_inputs, name=task_node.label
1867 )
1868 tasks.append(task)
1869 for write_edge in task_node.init.outputs.values():
1870 dataset_type_node = self.dataset_types[write_edge.parent_dataset_type_name]
1871 obj = getattr(task, write_edge.connection_name)
1872 # We don't immediately coerce obj to the dataset_type_node
1873 # storage class (which should be the repo storage class, if
1874 # there is one) when appending to `init_outputs` because a
1875 # formatter might be able to do a better job of that later;
1876 # instead we pair it with a dataset type that's consistent with
1877 # the in-memory type. We do coerce when populating `handles`,
1878 # though, because going through the dataset_type_node storage
1879 # class is the conversion path we checked when we resolved the
1880 # pipeline graph.
1881 if init_outputs is not None:
1882 init_outputs.append((obj, write_edge.adapt_dataset_type(dataset_type_node.dataset_type)))
1883 n_consumers = len(self.consumers_of(dataset_type_node.name))
1884 if n_consumers > 0:
1885 handles[dataset_type_node.name] = InMemoryDatasetHandle(
1886 dataset_type_node.storage_class.coerce_type(obj),
1887 dataId=empty_data_id,
1888 storageClass=dataset_type_node.storage_class,
1889 copy=(n_consumers > 1),
1890 )
1891 return tasks
1893 def write_init_outputs(self, butler: Butler) -> None:
1894 """Write the init-output datasets for all tasks in the pipeline graph.
1896 Parameters
1897 ----------
1898 butler : `lsst.daf.butler.Butler`
1899 A full butler data repository client with its default run set
1900 to the collection where datasets should be written.
1902 Notes
1903 -----
1904 Datasets that already exist in the butler's output run collection will
1905 not be written.
1907 This method writes outputs with new random dataset IDs and should
1908 hence only be used when writing init-outputs prior to building a
1909 `QuantumGraph`. Use `QuantumGraph.write_init_outputs` if a quantum
1910 graph has already been built.
1911 """
1912 init_outputs: list[tuple[Any, DatasetType]] = []
1913 self.instantiate_tasks(butler.get, init_outputs)
1914 found_refs: dict[str, DatasetRef] = {}
1915 to_put: list[tuple[Any, DatasetType]] = []
1916 for obj, dataset_type in init_outputs:
1917 if (ref := butler.find_dataset(dataset_type, collections=butler.run)) is not None:
1918 found_refs[dataset_type.name] = ref
1919 else:
1920 to_put.append((obj, dataset_type))
1921 for ref, stored in butler.stored_many(found_refs.values()).items():
1922 if not stored:
1923 raise FileNotFoundError(
1924 f"Init-output dataset {ref.datasetType.name!r} was found in RUN {ref.run!r} "
1925 f"but had not actually been stored (or was stored and later deleted)."
1926 )
1927 for obj, dataset_type in to_put:
1928 butler.put(obj, dataset_type)
1930 def write_configs(self, butler: Butler) -> None:
1931 """Write the config datasets for all tasks in the pipeline graph.
1933 Parameters
1934 ----------
1935 butler : `lsst.daf.butler.Butler`
1936 A full butler data repository client with its default run set
1937 to the collection where datasets should be written.
1939 Notes
1940 -----
1941 Config datasets that already exist in the butler's output run
1942 collection will be checked for consistency.
1944 This method writes outputs with new random dataset IDs and should
1945 hence only be used when writing init-outputs prior to building a
1946 `QuantumGraph`. Use `QuantumGraph.write_configs` if a quantum graph
1947 has already been built.
1949 Raises
1950 ------
1951 lsst.daf.butler.registry.ConflictingDefinitionError
1952 Raised if a config dataset already exists and is not consistent
1953 with the config in the pipeline graph.
1954 """
1955 to_put: list[tuple[PipelineTaskConfig, str]] = []
1956 for task_node in self.tasks.values():
1957 dataset_type_name = task_node.init.config_output.dataset_type_name
1958 if (ref := butler.find_dataset(dataset_type_name, collections=butler.run)) is not None:
1959 old_config = butler.get(ref)
1960 if not task_node.config.compare(old_config, shortcut=False, output=log_config_mismatch):
1961 raise ConflictingDefinitionError(
1962 f"Config does not match existing task config {dataset_type_name!r} in "
1963 "butler; tasks configurations must be consistent within the same run collection"
1964 )
1965 else:
1966 to_put.append((task_node.config, dataset_type_name))
1967 # We do writes at the end to minimize the mess we leave behind when we
1968 # raise an exception.
1969 for config, dataset_type_name in to_put:
1970 butler.put(config, dataset_type_name)
1972 def write_packages(self, butler: Butler) -> None:
1973 """Write the 'packages' dataset for the currently-active software
1974 versions.
1976 Parameters
1977 ----------
1978 butler : `lsst.daf.butler.Butler`
1979 A full butler data repository client with its default run set
1980 to the collection where datasets should be written.
1982 Notes
1983 -----
1984 If the packages dataset already exists, it will be compared to the
1985 versions in the current packages. New packages that weren't present
1986 before are not considered an inconsistency.
1988 This method writes outputs with new random dataset IDs and should
1989 hence only be used when writing init-outputs prior to building a
1990 `QuantumGraph`. Use `QuantumGraph.write_packages` if a quantum graph
1991 has already been built.
1993 Raises
1994 ------
1995 lsst.daf.butler.registry.ConflictingDefinitionError
1996 Raised if the packages dataset already exists and is not consistent
1997 with the current packages.
1998 """
1999 new_packages = Packages.fromSystem()
2000 if (ref := butler.find_dataset(self.packages_dataset_type)) is not None:
2001 packages = butler.get(ref)
2002 if compare_packages(packages, new_packages):
2003 # have to remove existing dataset first; butler has no
2004 # replace option.
2005 butler.pruneDatasets([ref], unstore=True, purge=True)
2006 butler.put(packages, ref)
2007 else:
2008 butler.put(new_packages, self.packages_dataset_type)
2010 def init_output_run(self, butler: Butler) -> None:
2011 """Initialize a new output RUN collection by writing init-output
2012 datasets (including configs and packages).
2014 Parameters
2015 ----------
2016 butler : `lsst.daf.butler.Butler`
2017 A full butler data repository client with its default run set
2018 to the collection where datasets should be written.
2019 """
2020 self.write_configs(butler)
2021 self.write_packages(butler)
2022 self.write_init_outputs(butler)
2024 ###########################################################################
2025 #
2026 # Class- and Package-Private Methods.
2027 #
2028 ###########################################################################
2030 def _iter_task_defs(self) -> Iterator[TaskDef]:
2031 """Iterate over this pipeline as a sequence of `TaskDef` instances.
2033 Notes
2034 -----
2035 This is a package-private method intended to aid in the transition to a
2036 codebase more fully integrated with the `PipelineGraph` class, in which
2037 both `TaskDef` and `PipelineDatasetTypes` are expected to go away, and
2038 much of the functionality on the `Pipeline` class will be moved to
2039 `PipelineGraph` as well.
2041 Raises
2042 ------
2043 TaskNotImportedError
2044 Raised if `TaskNode.is_imported` is `False` for any task.
2045 """
2046 from ..pipeline import TaskDef
2048 for node in self._tasks.values():
2049 yield TaskDef(
2050 config=node.config,
2051 taskClass=node.task_class,
2052 label=node.label,
2053 connections=node.get_connections(),
2054 )
2056 def _init_from_args(
2057 self,
2058 xgraph: networkx.MultiDiGraph | None,
2059 sorted_keys: Sequence[NodeKey] | None,
2060 task_subsets: dict[str, TaskSubset] | None,
2061 description: str,
2062 universe: DimensionUniverse | None,
2063 data_id: DataId | None,
2064 step_definitions: StepDefinitions,
2065 ) -> None:
2066 """Initialize the graph with possibly-nontrivial arguments.
2068 Parameters
2069 ----------
2070 xgraph : `networkx.MultiDiGraph` or `None`
2071 The backing networkx graph, or `None` to create an empty one.
2072 This graph has `NodeKey` instances for nodes and the same structure
2073 as the graph exported by `make_xgraph`, but its nodes and edges
2074 have a single ``instance`` attribute that holds a `TaskNode`,
2075 `TaskInitNode`, `DatasetTypeNode` (or `None`), `ReadEdge`, or
2076 `WriteEdge` instance.
2077 sorted_keys : `~collections.abc.Sequence` [ `NodeKey` ] or `None`
2078 Topologically sorted sequence of node keys, or `None` if the graph
2079 is not sorted.
2080 task_subsets : `dict` [ `str`, `TaskSubset` ]
2081 Labeled subsets of tasks. Values must be constructed with
2082 ``xgraph`` as their parent graph.
2083 description : `str`
2084 String description for this pipeline.
2085 universe : `lsst.daf.butler.DimensionUniverse` or `None`
2086 Definitions of all dimensions.
2087 data_id : `lsst.daf.butler.DataCoordinate` or other data ID mapping.
2088 Data ID that represents a constraint on all quanta generated from
2089 this pipeline.
2090 step_definitions : `StepDefinitions`
2091 Struct holding information about steps.
2093 Notes
2094 -----
2095 Only empty `PipelineGraph` instances should be constructed directly by
2096 users, which sets the signature of ``__init__`` itself, but methods on
2097 `PipelineGraph` and its helper classes need to be able to create them
2098 with state. Those methods can call this after calling ``__new__``
2099 manually, skipping ``__init__``.
2100 """
2101 self._xgraph = xgraph if xgraph is not None else networkx.MultiDiGraph()
2102 self._sorted_keys: Sequence[NodeKey] | None = None
2103 self._task_subsets = task_subsets if task_subsets is not None else {}
2104 self._description = description
2105 self._tasks = TaskMappingView(self._xgraph)
2106 self._dataset_types = DatasetTypeMappingView(self._xgraph)
2107 self._raw_data_id: dict[str, Any]
2108 if isinstance(data_id, DataCoordinate):
2109 if universe is None:
2110 universe = data_id.universe
2111 else:
2112 assert universe is data_id.universe, "data_id.universe and given universe differ"
2113 self._raw_data_id = dict(data_id.required)
2114 elif data_id is None:
2115 self._raw_data_id = {}
2116 else:
2117 self._raw_data_id = dict(data_id)
2118 self._universe = universe
2119 if sorted_keys is not None:
2120 self._reorder(sorted_keys)
2121 self._step_definitions = step_definitions
2123 def _make_bipartite_xgraph_internal(self, init: bool) -> networkx.MultiDiGraph:
2124 """Make a bipartite init-only or runtime-only internal subgraph.
2126 See `make_bipartite_xgraph` for parameters and return values.
2128 Notes
2129 -----
2130 This method returns a view of the `PipelineGraph` object's internal
2131 backing graph, and hence should only be called in methods that copy the
2132 result either explicitly or by running a copying algorithm before
2133 returning it to the user.
2134 """
2135 return self._xgraph.edge_subgraph([edge.key for edge in self.iter_edges(init)])
2137 def _make_task_xgraph_internal(self, init: bool) -> networkx.DiGraph:
2138 """Make a init-only or runtime-only internal task subgraph.
2140 See `make_task_xgraph` for parameters and return values.
2142 Notes
2143 -----
2144 This method returns a view of the `PipelineGraph` object's internal
2145 backing graph, and hence should only be called in methods that copy the
2146 result either explicitly or by running a copying algorithm before
2147 returning it to the user.
2148 """
2149 bipartite_xgraph = self._make_bipartite_xgraph_internal(init=init)
2150 task_keys = [
2151 key
2152 for key, bipartite in bipartite_xgraph.nodes(data="bipartite")
2153 if bipartite == NodeType.TASK.bipartite
2154 ]
2155 return networkx.algorithms.bipartite.projected_graph(networkx.DiGraph(bipartite_xgraph), task_keys)
2157 def _transform_xgraph_state(self, xgraph: _G, skip_edges: bool) -> _G:
2158 """Transform networkx graph attributes in-place from the internal
2159 "instance" attributes to the documented exported attributes.
2161 Parameters
2162 ----------
2163 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
2164 Graph whose state should be transformed.
2165 skip_edges : `bool`
2166 If `True`, do not transform edge state.
2168 Returns
2169 -------
2170 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
2171 The same object passed in, after modification.
2173 Notes
2174 -----
2175 This should be called after making a copy of the internal graph but
2176 before any projection down to just task or dataset type nodes, since
2177 it assumes stateful edges.
2178 """
2179 state: dict[str, Any]
2180 for state in xgraph.nodes.values():
2181 node_value: TaskInitNode | TaskNode | DatasetTypeNode | None = state.pop("instance")
2182 if node_value is not None:
2183 state.update(node_value._to_xgraph_state())
2184 else:
2185 # This is a dataset type node that is not resolved.
2186 state["bipartite"] = NodeType.DATASET_TYPE.bipartite
2187 if not skip_edges:
2188 for _, _, state in xgraph.edges(data=True):
2189 edge: Edge | None = state.pop("instance", None)
2190 if edge is not None:
2191 state.update(edge._to_xgraph_state())
2192 return xgraph
2194 def _replace_task_nodes(
2195 self,
2196 updates: Mapping[str, TaskNode],
2197 check_edges_unchanged: bool,
2198 assume_edges_unchanged: bool,
2199 message_header: str,
2200 ) -> None:
2201 """Replace task nodes and update edges and dataset type nodes
2202 accordingly.
2204 Parameters
2205 ----------
2206 updates : `Mapping` [ `str`, `TaskNode` ]
2207 New task nodes with task label keys. All keys must be task labels
2208 that are already present in the graph.
2209 check_edges_unchanged : `bool`, optional
2210 If `True`, require the edges (connections) of the modified tasks to
2211 remain unchanged after importing and configuring each task, and
2212 verify that this is the case.
2213 assume_edges_unchanged : `bool`, optional
2214 If `True`, the caller declares that the edges (connections) of the
2215 modified tasks will remain unchanged importing and configuring each
2216 task, and that it is unnecessary to check this.
2217 message_header : `str`
2218 Template for `str.format` with a single ``task_label`` placeholder
2219 to use as the first line in `EdgesChangedError` messages that show
2220 the differences between new task edges and old task edges. Should
2221 include the fact that the rest of the message will refer to the old
2222 task as "A" and the new task as "B", and end with a colon.
2224 Raises
2225 ------
2226 ValueError
2227 Raised if ``assume_edges_unchanged`` and ``check_edges_unchanged``
2228 are both `True`, or if a full config is provided for a task after
2229 another full config or an override has already been provided.
2230 EdgesChangedError
2231 Raised if ``check_edges_unchanged=True`` and the edges of a task do
2232 change.
2233 """
2234 deep: dict[str, TaskNode] = {}
2235 shallow: dict[str, TaskNode] = {}
2236 if assume_edges_unchanged:
2237 if check_edges_unchanged:
2238 raise ValueError("Cannot simultaneously assume and check that edges have not changed.")
2239 shallow.update(updates)
2240 else:
2241 for task_label, new_task_node in updates.items():
2242 old_task_node = self.tasks[task_label]
2243 messages = old_task_node.diff_edges(new_task_node)
2244 if messages:
2245 if check_edges_unchanged:
2246 messages.insert(0, message_header.format(task_label=task_label))
2247 raise EdgesChangedError("\n".join(messages))
2248 else:
2249 deep[task_label] = new_task_node
2250 else:
2251 shallow[task_label] = new_task_node
2252 try:
2253 if deep:
2254 removed = self.remove_tasks(deep.keys(), drop_from_subsets=True)
2255 self.add_task_nodes(deep.values())
2256 for replaced_task_node, referencing_subsets in removed:
2257 for subset_label in referencing_subsets:
2258 self._task_subsets[subset_label].add(replaced_task_node.label)
2259 for task_node in shallow.values():
2260 self._xgraph.nodes[task_node.key]["instance"] = task_node
2261 self._xgraph.nodes[task_node.init.key]["instance"] = task_node.init
2262 except PipelineGraphExceptionSafetyError: # pragma: no cover
2263 raise
2264 except Exception as err: # pragma: no cover
2265 # There's no known way to get here, but we want to make it clear
2266 # it's a big problem if we do.
2267 raise PipelineGraphExceptionSafetyError(
2268 "Error while replacing tasks has left the graph in an inconsistent state."
2269 ) from err
2271 def _append_graph_data_from_edge(
2272 self,
2273 node_data: list[tuple[NodeKey, dict[str, Any]]],
2274 edge_data: list[tuple[NodeKey, NodeKey, str, dict[str, Any]]],
2275 edge: Edge,
2276 parent: PipelineGraph | None,
2277 ) -> None:
2278 """Append networkx state dictionaries for an edge and the corresponding
2279 dataset type node.
2281 Parameters
2282 ----------
2283 node_data : `list`
2284 List of node keys and state dictionaries. A node is appended if
2285 one does not already exist for this dataset type.
2286 edge_data : `list`
2287 List of node key pairs, connection names, and state dictionaries
2288 for edges.
2289 edge : `Edge`
2290 New edge being processed.
2291 parent : `PipelineGraph` or `None`
2292 Another pipeline graph whose dataset type nodes should be used
2293 when present.
2294 """
2295 new_dataset_type_node = None
2296 if parent is not None:
2297 new_dataset_type_node = parent._xgraph.nodes[edge.dataset_type_key].get("instance")
2298 if (existing_dataset_type_state := self._xgraph.nodes.get(edge.dataset_type_key)) is not None:
2299 existing_dataset_type_state["instance"] = new_dataset_type_node
2300 else:
2301 node_data.append(
2302 (
2303 edge.dataset_type_key,
2304 {
2305 "instance": new_dataset_type_node,
2306 "bipartite": NodeType.DATASET_TYPE.bipartite,
2307 },
2308 )
2309 )
2310 edge_data.append(
2311 edge.nodes
2312 + (
2313 edge.connection_name,
2314 {"instance": edge},
2315 )
2316 )
2318 def _reorder(self, sorted_keys: Sequence[NodeKey]) -> None:
2319 """Set the order of all views of this graph from the given sorted
2320 sequence of task labels and dataset type names.
2321 """
2322 self._sorted_keys = sorted_keys
2323 self._tasks._reorder(sorted_keys)
2324 self._dataset_types._reorder(sorted_keys)
2326 def _reset(self) -> None:
2327 """Reset all views of this graph following a modification that might
2328 invalidate them.
2329 """
2330 self._sorted_keys = None
2331 self._tasks._reset()
2332 self._dataset_types._reset()
2334 def _resolve_step_flow(self) -> Sequence[NodeKey] | None:
2335 """Check that step definitions are consistent with the ordering of
2336 the graph's nodes and that they partition the task graph.
2338 Returns
2339 -------
2340 sort_keys : `~collections.abc.Sequence` [ `NodeKey` ] or `None`
2341 Sort order for the pipeline graph that is consistent with the
2342 step order, or `None` if there were no steps defined for this
2343 pipeline.
2344 """
2345 if not self._step_definitions:
2346 return None
2347 task_labels_so_far: set[str] = set()
2348 # Inputs we've already seen, and the tasks that wanted them:
2349 inputs_so_far: dict[str, set[str]] = {}
2350 sort_keys: list[NodeKey] = []
2351 keys_already_sorted: set[NodeKey] = set()
2352 for step_label in self.steps:
2353 try:
2354 task_subset = self.task_subsets[step_label]
2355 except KeyError:
2356 raise InvalidStepsError(f"Step {step_label!r} is not a task subset.") from None
2357 if not task_labels_so_far.isdisjoint(task_subset):
2358 raise InvalidStepsError(
2359 f"Step {step_label!r} repeats task(s) {task_labels_so_far & task_subset}."
2360 )
2361 task_labels_so_far.update(task_subset)
2362 # We need a temporary data structure for tracking the inputs of
2363 # just this step to avoid complaining about input-output
2364 # relationships within the step. We'll merge this into the
2365 # cumulative one later.
2366 step_inputs: dict[str, set[str]] = {}
2367 # We'll also gather all task, task-init and dataset type networkx
2368 # graph keys for the step so we can make and sort a step sub-graph.
2369 new_step_keys: set[NodeKey] = set()
2370 for task_label in task_subset:
2371 for input_name in self.inputs_of(task_label):
2372 step_inputs.setdefault(input_name, set()).add(task_label)
2373 # Check that none of the outputs of the tasks in this step were
2374 # expected as inputs of a previous step's tasks.
2375 task_outputs = self.outputs_of(task_label)
2376 if not inputs_so_far.keys().isdisjoint(task_outputs.keys()):
2377 msg: list[str] = []
2378 for input_name in inputs_so_far.keys() & task_outputs.keys():
2379 msg.append(f"{input_name} (used by {inputs_so_far[input_name]})")
2380 raise InvalidStepsError(
2381 f"Task {task_label} in step {step_label!r} produces dataset types "
2382 f"[{', '.join(msg)}], but these are consumed by tasks in earlier "
2383 "steps. Either steps are out of order or the graph is cyclic."
2384 )
2385 task_init_key = NodeKey(NodeType.TASK_INIT, task_label)
2386 task_key = NodeKey(NodeType.TASK, task_label)
2387 new_step_keys.add(task_init_key)
2388 new_step_keys.add(task_key)
2389 new_step_keys.update(self._xgraph.predecessors(task_init_key))
2390 new_step_keys.update(self._xgraph.predecessors(task_key))
2391 new_step_keys.update(self._xgraph.successors(task_init_key))
2392 new_step_keys.update(self._xgraph.successors(task_key))
2393 # Also record the step the task is in as a private xgraph
2394 # attribute so we can look up the step given a task (or,
2395 # indirectly, dataset type); this has to be used with care
2396 # because if the steps haven't been verified it can be wrong.
2397 self._xgraph.nodes[task_init_key]["step"] = step_label
2398 self._xgraph.nodes[task_key]["step"] = step_label
2399 # Drop step input keys that were already either inputs or outputs
2400 # of a previous step, since they'll have already been added to
2401 # sort_keys.
2402 new_step_keys.difference_update(keys_already_sorted)
2403 # Make the step subgraph, sort it, and extend the overall sort_keys
2404 # with result.
2405 step_xgraph = self._xgraph.subgraph(new_step_keys)
2406 sort_keys.extend(networkx.dag.lexicographical_topological_sort(step_xgraph))
2407 keys_already_sorted.update(new_step_keys)
2408 for input_name, consuming_tasks in step_inputs.items():
2409 inputs_so_far.setdefault(input_name, set()).update(consuming_tasks)
2410 if not task_labels_so_far.issuperset(self.tasks.keys()):
2411 # Note that the converse issubset test effectively happens when we
2412 # look up inputs and outputs of each task.
2413 raise InvalidStepsError(f"No step contains task(s) {self.tasks.keys() - task_labels_so_far}.")
2414 return sort_keys
2416 def _resolve_step_dimensions(self, universe: DimensionUniverse) -> None:
2417 """Check that step sharding dimensions are consistent with task and
2418 output dataset dimensions.
2420 Parameters
2421 ----------
2422 universe : `~lsst.daf.butler.DimensionUniverse`
2423 Definitions for all dimensions. Will be attached to the step
2424 definitions by this method.
2425 """
2426 self._step_definitions._universe = universe
2427 for step_label in self.steps:
2428 dimensions = self.steps.get_dimensions(step_label)
2429 for task_label in self.task_subsets[step_label]:
2430 task_node = self.tasks[task_label]
2431 if not _dimensions_compatible(dimensions, task_node.dimensions):
2432 raise InvalidStepsError(
2433 f"Dimensions {task_node.dimensions} of task {task_label!r} are not compatible with "
2434 f"the sharding dimensions {dimensions} of step {step_label!r}."
2435 )
2436 for dataset_type_node in self.outputs_of(task_label).values():
2437 assert dataset_type_node is not None, "dataset types should be resolved first"
2438 if not _dimensions_compatible(dimensions, dataset_type_node.dimensions):
2439 raise InvalidStepsError(
2440 f"Dimensions {dataset_type_node.dimensions} of dataset type "
2441 f"{dataset_type_node.name!r} (produced by task {task_label!r}) are not "
2442 f"compatible with the sharding dimensions {dimensions} of step "
2443 f"{step_label!r}."
2444 )
2446 def _select_expression(self, expr_tree: expressions.Node, task_xgraph: networkx.DiGraph) -> set[NodeKey]:
2447 """Select tasks from a pipeline based on a string expression.
2449 This is the primary implementation method for `select` and
2450 `select_tasks`.
2452 Parameters
2453 ----------
2454 expr_tree : `expressions.Node`
2455 Expression [sub]tree to process (recursively).
2456 task_xgraph : `networkx.DiGraph`
2457 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2459 Returns
2460 -------
2461 selected : `set` [ `NodeKey` ]
2462 Set of `NodeKey` objects for matching tasks (only; no dataset type
2463 or task-init nodes).
2464 """
2465 match expr_tree:
2466 case expressions.IdentifierNode(qualifier=qualifier, label=label):
2467 match self._select_identifier(qualifier, label):
2468 case NodeKey(node_type=NodeType.TASK) as task_key:
2469 return {task_key}
2470 case NodeKey(node_type=NodeType.DATASET_TYPE) as dataset_type_key:
2471 # Since a dataset type can have only one producer, this
2472 # yields 0- (for overall inputs) or 1-element sets.
2473 for producer_key, _ in self._xgraph.in_edges(dataset_type_key):
2474 if producer_key.node_type is NodeType.TASK_INIT:
2475 raise InvalidExpressionError(
2476 f"Init-output dataset type {label!r} cannot be used directly in an "
2477 "expression."
2478 )
2479 return {producer_key}
2480 return set()
2481 case TaskSubset() as task_subset:
2482 return {NodeKey(NodeType.TASK, label) for label in task_subset}
2483 case _: # pragma: no cover
2484 raise AssertionError("Identifier type inconsistent with grammar.")
2485 case expressions.DirectionNode(operator=operator, start=start):
2486 match self._select_identifier(start.qualifier, start.label):
2487 case NodeKey(node_type=NodeType.TASK) as task_key:
2488 if operator.startswith("<"):
2489 return self._select_task_ancestors(
2490 task_key, task_xgraph, inclusive=operator.endswith("=")
2491 )
2492 else:
2493 assert operator.startswith(">"), "Guaranteed by grammar."
2494 return self._select_task_descendants(
2495 task_key, task_xgraph, inclusive=operator.endswith("=")
2496 )
2497 case NodeKey(node_type=NodeType.DATASET_TYPE) as dataset_type_key:
2498 if operator.startswith("<"):
2499 return self._select_dataset_type_ancestors(
2500 dataset_type_key, task_xgraph, inclusive=operator.endswith("=")
2501 )
2502 else:
2503 assert operator.startswith(">"), "Guaranteed by grammar."
2504 return self._select_dataset_type_descendants(
2505 dataset_type_key, task_xgraph, inclusive=operator.endswith("=")
2506 )
2507 case TaskSubset():
2508 raise InvalidExpressionError(
2509 f"Task subset identifier {start!r} cannot be used as the start of an "
2510 "ancestor/descendant search."
2511 )
2512 case _: # pragma: no cover
2513 raise AssertionError("Unexpected parsed identifier result type.")
2514 case expressions.NotNode(operand=operand):
2515 operand_result = self._select_expression(operand, task_xgraph)
2516 return set(task_xgraph.nodes.keys() - operand_result)
2517 case expressions.UnionNode(lhs=lhs, rhs=rhs):
2518 lhs_result = self._select_expression(lhs, task_xgraph)
2519 rhs_result = self._select_expression(rhs, task_xgraph)
2520 return lhs_result.union(rhs_result)
2521 case expressions.IntersectionNode(lhs=lhs, rhs=rhs):
2522 lhs_result = self._select_expression(lhs, task_xgraph)
2523 rhs_result = self._select_expression(rhs, task_xgraph)
2524 return lhs_result.intersection(rhs_result)
2525 case _: # pragma: no cover
2526 raise AssertionError("Expression parse node inconsistent with grammar.")
2528 def _select_task_ancestors(
2529 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2530 ) -> set[NodeKey]:
2531 """Return all task-node ancestors of the given task node, as defined by
2532 the `select` expression language.
2534 Parameters
2535 ----------
2536 start : `NodeKey`
2537 A runtime task node key.
2538 task_xgraph : `networkx.DiGraph`
2539 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2540 inclusive : `bool`
2541 Whether to include the ``start`` node in the results.
2543 Returns
2544 -------
2545 selected : `set` [ `NodeKey` ]
2546 Set of `NodeKey` objects for matching tasks (only; no dataset type
2547 or task-init nodes).
2548 """
2549 result = set(networkx.dag.ancestors(task_xgraph, start))
2550 if inclusive:
2551 result.add(start)
2552 return result
2554 def _select_task_descendants(
2555 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2556 ) -> set[NodeKey]:
2557 """Return all task-node descendants of the given task node, as defined
2558 by the `select` expression language.
2560 Parameters
2561 ----------
2562 start : `NodeKey`
2563 A runtime task node key.
2564 task_xgraph : `networkx.DiGraph`
2565 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2566 inclusive : `bool`
2567 Whether to include the ``start`` node in the results.
2569 Returns
2570 -------
2571 selected : `set` [ `NodeKey` ]
2572 Set of `NodeKey` objects for matching tasks (only; no dataset type
2573 or task-init nodes).
2574 """
2575 result = set(networkx.dag.descendants(task_xgraph, start))
2576 if inclusive:
2577 result.add(start)
2578 return result
2580 def _select_dataset_type_ancestors(
2581 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2582 ) -> set[NodeKey]:
2583 """Return all task-node ancestors of the given dataset type node, as
2584 defined by the `select` expression language.
2586 Parameters
2587 ----------
2588 start : `NodeKey`
2589 A dataset type node key. May not be an init-output.
2590 task_xgraph : `networkx.DiGraph`
2591 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2592 inclusive : `bool`
2593 Whether to include the producer of the ``start`` node in the
2594 results.
2596 Returns
2597 -------
2598 selected : `set` [ `NodeKey` ]
2599 Set of `NodeKey` objects for matching tasks (only; no dataset type
2600 or task-init nodes).
2601 """
2602 result: set[NodeKey] = set()
2603 for producer_key, _ in self._xgraph.in_edges(start):
2604 if producer_key.node_type is NodeType.TASK_INIT:
2605 raise InvalidExpressionError(
2606 f"Init-output dataset type {start.name!r} cannot be used as the "
2607 "starting point for an ancestor ('<' or '<=') search."
2608 )
2609 result.update(networkx.dag.ancestors(task_xgraph, producer_key))
2610 if inclusive:
2611 result.add(producer_key)
2612 return result
2614 def _select_dataset_type_descendants(
2615 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2616 ) -> set[NodeKey]:
2617 """Return all task-node descendatns of the given dataset type node, as
2618 defined by the `select` expression language.
2620 Parameters
2621 ----------
2622 start : `NodeKey`
2623 A dataset type node key. May not be an init-output if
2624 ``inclusive=True``.
2625 task_xgraph : `networkx.DiGraph`
2626 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2627 inclusive : `bool`
2628 Whether to include the producer of the ``start`` node in the
2629 results.
2631 Returns
2632 -------
2633 selected : `set` [ `NodeKey` ]
2634 Set of `NodeKey` objects for matching tasks (only; no dataset type
2635 or task-init nodes).
2636 """
2637 result: set[NodeKey] = set()
2638 if inclusive:
2639 for producer_key, _ in self._xgraph.in_edges(start):
2640 if producer_key.node_type is NodeType.TASK_INIT:
2641 raise InvalidExpressionError(
2642 f"Init-output dataset type {start.name!r} cannot be used as the "
2643 "starting point for an includsive descendant ('>=') search."
2644 )
2645 result.add(producer_key)
2646 # We also include tasks that consume a dataset type as an init-input,
2647 # since that can affect their runtime behavior.
2648 consumer_keys: set[NodeKey] = {
2649 (
2650 consumer_key
2651 if consumer_key.node_type is NodeType.TASK
2652 else NodeKey(NodeType.TASK, consumer_key.name)
2653 )
2654 for _, consumer_key in self._xgraph.out_edges(start)
2655 }
2656 for consumer_key in consumer_keys:
2657 result.add(consumer_key)
2658 result.update(networkx.dag.descendants(task_xgraph, consumer_key))
2659 return result
2661 def _select_identifier(
2662 self, qualifier: Literal["T", "D", "S"] | None, label: str
2663 ) -> NodeKey | TaskSubset:
2664 """Return the node key or task subset that corresponds to a `select`
2665 expression identifier.
2667 Parameters
2668 ----------
2669 qualifier : `str` or `None`
2670 Task, dataset type, or task subset qualifier included in the
2671 identifier, if any.
2672 label : `str`
2673 Task label, dataset type name, or task subset label.
2675 Returns
2676 -------
2677 key_or_subset : `NodeKey` or `TaskSubset`
2678 A `NodeKey` for a task or dataset type, or a `TaskSubset` for a
2679 task subset.
2680 """
2681 match qualifier:
2682 case None:
2683 task_key = NodeKey(NodeType.TASK, label)
2684 dataset_type_key = NodeKey(NodeType.DATASET_TYPE, label)
2685 if task_key in self._xgraph.nodes:
2686 if dataset_type_key in self._xgraph.nodes:
2687 raise InvalidExpressionError(
2688 f"{label!r} is both a task label and a dataset type name; "
2689 "prefix with 'T:' or 'D:' (respectively) to specify which."
2690 )
2691 assert label not in self._task_subsets, "Should be prohibited at construction."
2692 return task_key
2693 elif dataset_type_key in self._xgraph.nodes:
2694 if label in self._task_subsets:
2695 raise InvalidExpressionError(
2696 f"{label!r} is both a subset label and a dataset type name; "
2697 "prefix with 'S:' or 'D:' (respectively) to specify which."
2698 )
2699 return dataset_type_key
2700 elif label in self._task_subsets:
2701 return self._task_subsets[label]
2702 else:
2703 raise InvalidExpressionError(
2704 f"{label!r} is not a task label, task subset label, or dataset type name."
2705 )
2706 case "T":
2707 task_key = NodeKey(NodeType.TASK, label)
2708 if task_key not in self._xgraph.nodes:
2709 raise InvalidExpressionError(f"Task with label {label!r} does not exist.")
2710 return task_key
2711 case "D":
2712 dataset_type_key = NodeKey(NodeType.DATASET_TYPE, label)
2713 if dataset_type_key not in self._xgraph.nodes:
2714 raise InvalidExpressionError(f"Dataset type with name {label!r} does not exist.")
2715 return dataset_type_key
2716 case "S":
2717 try:
2718 return self._task_subsets[label]
2719 except KeyError:
2720 raise InvalidExpressionError(f"Task subset with label {label!r} does not exist.")
2721 case _: # pragma: no cover
2722 raise AssertionError("Unexpected identifier qualifier in expression.")
2724 _xgraph: networkx.MultiDiGraph
2725 _sorted_keys: Sequence[NodeKey] | None
2726 _task_subsets: dict[str, TaskSubset]
2727 _step_definitions: StepDefinitions
2728 _description: str
2729 _tasks: TaskMappingView
2730 _dataset_types: DatasetTypeMappingView
2731 _raw_data_id: dict[str, Any]
2732 _universe: DimensionUniverse | None
2735def log_config_mismatch(msg: str) -> None:
2736 """Log messages about configuration mismatch.
2738 Parameters
2739 ----------
2740 msg : `str`
2741 Log message to use.
2742 """
2743 _LOG.fatal("Comparing configuration: %s", msg)
2746def compare_packages(packages: Packages, new_packages: Packages) -> bool:
2747 """Compare two versions of Packages.
2749 Parameters
2750 ----------
2751 packages : `Packages`
2752 Previously recorded package versions. Updated in place to include
2753 any new packages that weren't present before.
2754 new_packages : `Packages`
2755 New set of package versions.
2757 Returns
2758 -------
2759 updated : `bool`
2760 `True` if ``packages`` was updated, `False` if not.
2762 Raises
2763 ------
2764 ConflictingDefinitionError
2765 Raised if versions are inconsistent.
2766 """
2767 diff = new_packages.difference(packages)
2768 if diff:
2769 versions_str = "; ".join(f"{pkg}: {diff[pkg][1]} vs {diff[pkg][0]}" for pkg in diff)
2770 raise ConflictingDefinitionError(f"Package versions mismatch: ({versions_str})")
2771 else:
2772 _LOG.debug("new packages are consistent with old")
2773 # Update the old set of packages in case we have more packages
2774 # that haven't been persisted.
2775 extra = new_packages.extra(packages)
2776 if extra:
2777 _LOG.debug("extra packages: %s", extra)
2778 packages.update(new_packages)
2779 return True
2780 return False
2783def _dimensions_compatible(dimensions: DimensionGroup, object_dimensions: DimensionGroup) -> bool:
2784 if dimensions.issubset(object_dimensions):
2785 # Easy typical case.
2786 return True
2787 # Hard case: if any sharding dimensions that are not in the object
2788 # dimensions are related to something in the object dimensions by a
2789 # many-to-many join table, this is okay too. The main use case here is to
2790 # let {exposure, [detector]} satisfy sharding dimensions of
2791 # {visit, [detector]}.
2792 universe = dimensions.universe
2793 unmatched_sharding = dimensions.required - object_dimensions.required
2794 unmatched_object = object_dimensions.required - dimensions.required
2795 unmatched_union_group = dimensions.union(object_dimensions)
2796 for element_name in unmatched_union_group.elements:
2797 if (
2798 (element := universe[element_name]).defines_relationships
2799 and not element.dimensions.isdisjoint(unmatched_sharding)
2800 and not element.dimensions.isdisjoint(unmatched_object)
2801 ):
2802 return True
2803 return False