Coverage for python / lsst / pipe / base / pipeline_graph / _pipeline_graph.py: 12%
746 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:13 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29__all__ = ("PipelineGraph", "compare_packages", "log_config_mismatch")
31import gzip
32import itertools
33import json
34import logging
35from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence, Set
36from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypeVar, cast
38import networkx
39import networkx.algorithms.bipartite
40import networkx.algorithms.dag
42from lsst.daf.butler import (
43 Butler,
44 DataCoordinate,
45 DataId,
46 DatasetRef,
47 DatasetType,
48 DimensionGroup,
49 DimensionUniverse,
50 MissingDatasetTypeError,
51)
52from lsst.daf.butler.registry import ConflictingDefinitionError, Registry
53from lsst.resources import ResourcePath, ResourcePathExpression
54from lsst.utils.packages import Packages
56from .._dataset_handle import InMemoryDatasetHandle
57from ..automatic_connection_constants import PACKAGES_INIT_OUTPUT_NAME, PACKAGES_INIT_OUTPUT_STORAGE_CLASS
58from . import expressions
59from ._dataset_types import DatasetTypeNode
60from ._edges import Edge, ReadEdge, WriteEdge
61from ._exceptions import (
62 DuplicateOutputError,
63 EdgesChangedError,
64 InvalidExpressionError,
65 InvalidStepsError,
66 PipelineDataCycleError,
67 PipelineGraphError,
68 PipelineGraphExceptionSafetyError,
69 UnresolvedGraphError,
70)
71from ._mapping_views import DatasetTypeMappingView, TaskMappingView
72from ._nodes import NodeKey, NodeType
73from ._task_subsets import StepDefinitions, TaskSubset
74from ._tasks import TaskImportMode, TaskInitNode, TaskNode, _TaskNodeImportedData
76if TYPE_CHECKING:
77 from ..config import PipelineTaskConfig
78 from ..connections import PipelineTaskConnections
79 from ..pipeline import TaskDef
80 from ..pipelineTask import PipelineTask
83_G = TypeVar("_G", bound=networkx.DiGraph | networkx.MultiDiGraph)
85_LOG = logging.getLogger("lsst.pipe.base.pipeline_graph")
88class PipelineGraph:
89 """A graph representation of fully-configured pipeline.
91 `PipelineGraph` instances are typically constructed by calling
92 `.Pipeline.to_graph`, but in rare cases constructing and then populating an
93 empty one may be preferable.
95 Parameters
96 ----------
97 description : `str`, optional
98 String description for this pipeline.
99 universe : `lsst.daf.butler.DimensionUniverse`, optional
100 Definitions for all butler dimensions. If not provided, some
101 attributes will not be available until `resolve` is called.
102 data_id : `lsst.daf.butler.DataCoordinate` or other data ID, optional
103 Data ID that represents a constraint on all quanta generated by this
104 pipeline. This typically just holds the instrument constraint included
105 in the pipeline definition, if there was one.
106 """
108 ###########################################################################
109 #
110 # Simple Pipeline Graph Inspection Interface:
111 #
112 # - for inspecting graph structure, not modifying it (except to sort and]
113 # resolve);
114 #
115 # - no NodeKey objects, just string dataset type name and task label keys;
116 #
117 # - graph structure is represented as a pair of mappings, with methods to
118 # find neighbors and edges of nodes.
119 #
120 ###########################################################################
122 def __init__(
123 self,
124 *,
125 description: str = "",
126 universe: DimensionUniverse | None = None,
127 data_id: DataId | None = None,
128 ) -> None:
129 self._init_from_args(
130 xgraph=None,
131 sorted_keys=None,
132 task_subsets=None,
133 description=description,
134 universe=universe,
135 data_id=data_id,
136 step_definitions=StepDefinitions(universe),
137 )
139 def __repr__(self) -> str:
140 return f"{type(self).__name__}({self.description!r}, tasks={self.tasks!s})"
142 @property
143 def description(self) -> str:
144 """String description for this pipeline."""
145 return self._description
147 @description.setter
148 def description(self, value: str) -> None:
149 # Docstring in setter.
150 self._description = value
152 @property
153 def universe(self) -> DimensionUniverse:
154 """Definitions for all butler dimensions."""
155 if self._universe is None:
156 raise UnresolvedGraphError("Pipeline graph is not resolved.")
157 return self._universe
159 @property
160 def data_id(self) -> DataCoordinate:
161 """Data ID that represents a constraint on all quanta generated from
162 this pipeline.
164 This is may not be available unless the graph is resolved.
165 """
166 return DataCoordinate.standardize(self._raw_data_id, universe=self.universe)
168 @property
169 def tasks(self) -> TaskMappingView:
170 """A mapping view of the tasks in the graph.
172 This mapping has `str` task label keys and `TaskNode` values. Iteration
173 is topologically and deterministically ordered if and only if `sort`
174 has been called since the last modification to the graph.
175 """
176 return self._tasks
178 @property
179 def dataset_types(self) -> DatasetTypeMappingView:
180 """A mapping view of the dataset types in the graph.
182 This mapping has `str` parent dataset type name keys, but only provides
183 access to its `DatasetTypeNode` values if `resolve` has been called
184 since the last modification involving a task that uses a dataset type.
185 See `DatasetTypeMappingView` for details.
186 """
187 return self._dataset_types
189 @property
190 def task_subsets(self) -> Mapping[str, TaskSubset]:
191 """A mapping of all labeled subsets of tasks.
193 Keys are subset labels, values are sets of task labels. See
194 `TaskSubset` for more information.
196 Use `add_task_subset` to add a new subset. The subsets themselves may
197 be modified in-place.
198 """
199 return self._task_subsets
201 @property
202 def steps(self) -> StepDefinitions:
203 """An ordered iterable of the labels of task subsets that must be
204 executed separately.
206 Steps are intended to be partition of the full pipeline graph - the
207 steps together include all tasks, and each task belongs to exactly one
208 step - but this is only checked when the graph is resolved.
210 The order of the steps is required to be consistent with the graph's
211 topological ordering, and when sorting a graph with steps, that sort
212 order is guaranteed to be consistent with the order of the steps. But
213 there may be other orderings of the steps that are still valid for
214 execution (e.g. two steps could be runnable in parallel).
215 """
216 return self._step_definitions
218 @steps.setter
219 def steps(self, labels: Iterable[str]) -> None:
220 # Docstring on getter.
221 self._step_definitions.assign(labels)
223 def get_task_step(self, task_label: str) -> str:
224 """Return the step that the given task belongs to.
226 Parameters
227 ----------
228 task_label : `str`
229 Label of the task to look up.
231 Returns
232 -------
233 step_label : `str`
234 Step the task belongs to.
236 Raises
237 ------
238 InvalidStepsError
239 Raised if the pipeline has no steps defined.
240 UnresolvedGraphError
241 Raised if the pipeline has been modified since the steps were
242 last verified.
243 """
244 if not self._step_definitions:
245 raise InvalidStepsError("No steps have been defined for this pipeline.")
246 elif self._step_definitions.verified:
247 return self._xgraph.nodes[NodeKey(NodeType.TASK, task_label)]["step"]
248 raise UnresolvedGraphError("Steps have not been verified since the last modification.")
250 @property
251 def is_fully_resolved(self) -> bool:
252 """Whether all of this graph's nodes are resolved and any all have
253 been checked for correctness.
255 A fully-resolved graph is always sorted as well, and a fully-resolved
256 graph with sorted
257 """
258 return (
259 self._universe is not None
260 and self._step_definitions.verified
261 and self._sorted_keys is not None
262 and all(self.dataset_types.is_resolved(k) for k in self.dataset_types)
263 )
265 @property
266 def has_been_sorted(self) -> bool:
267 """Whether this graph's tasks and dataset types have been
268 topologically sorted (with unspecified but deterministic tiebreakers)
269 since the last modification to the graph.
271 If the pipeline graph has step definitions, the sort order is
272 consistent with the step order.
274 This may return `False` if the graph *happens* to be sorted but `sort`
275 was never called,.
276 """
277 return self._sorted_keys is not None
279 def sort(self) -> None:
280 """Sort this graph's nodes topologically with deterministic (but
281 unspecified) tiebreakers.
283 This does nothing if the graph is already known to be sorted.
284 """
285 if self._sorted_keys is None:
286 try:
287 sorted_keys: Sequence[NodeKey] = list(networkx.lexicographical_topological_sort(self._xgraph))
288 except networkx.NetworkXUnfeasible as err: # pragma: no cover
289 # Should't be possible to get here, because we check for cycles
290 # when adding tasks, but we guard against it anyway.
291 cycle = networkx.find_cycle(self._xgraph)
292 raise PipelineDataCycleError(
293 f"Cycle detected while attempting to sort graph: {cycle}."
294 ) from err
295 self._reorder(sorted_keys)
297 def copy(self) -> PipelineGraph:
298 """Return a copy of this graph that copies all mutable state."""
299 xgraph = self._xgraph.copy()
300 step_definitions = self._step_definitions.copy()
301 result = PipelineGraph.__new__(PipelineGraph)
302 result._init_from_args(
303 xgraph,
304 self._sorted_keys,
305 task_subsets={
306 k: TaskSubset(xgraph, v.label, set(v._members), v.description, step_definitions)
307 for k, v in self._task_subsets.items()
308 },
309 description=self._description,
310 universe=self._universe,
311 data_id=self._raw_data_id,
312 step_definitions=step_definitions,
313 )
314 return result
316 def __copy__(self) -> PipelineGraph:
317 # Fully shallow copies are dangerous; we don't want shared mutable
318 # state to lead to broken class invariants.
319 return self.copy()
321 def __deepcopy__(self, memo: dict) -> PipelineGraph:
322 # Genuine deep copies are unnecessary, since we should only ever care
323 # that mutable state is copied.
324 return self.copy()
326 def diff_tasks(self, other: PipelineGraph) -> list[str]:
327 """Compare two pipeline graphs.
329 This only compares graph structure and task classes (including their
330 edges). It does *not* compare full configuration (which is subject to
331 spurious differences due to import-cache state), dataset type
332 resolutions, or sort state.
334 Parameters
335 ----------
336 other : `PipelineGraph`
337 Graph to compare to.
339 Returns
340 -------
341 differences : `list` [ `str` ]
342 List of string messages describing differences between the
343 pipelines. If empty, the graphs have the same tasks and
344 connections.
345 """
346 messages: list[str] = []
347 common_labels: Set[str]
348 if self.tasks.keys() != other.tasks.keys():
349 common_labels = self.tasks.keys() & other.tasks.keys()
350 messages.append(
351 f"Pipelines have different tasks: A & ~B = {list(self.tasks.keys() - common_labels)}, "
352 f"B & ~A = {list(other.tasks.keys() - common_labels)}."
353 )
354 else:
355 common_labels = self.tasks.keys()
356 for label in common_labels:
357 a = self.tasks[label]
358 b = other.tasks[label]
359 if a.task_class != b.task_class:
360 messages.append(
361 f"Task {label!r} has class {a.task_class_name} in A, but {b.task_class_name} in B."
362 )
363 messages.extend(a.diff_edges(b))
364 return messages
366 def producing_edge_of(self, dataset_type_name: str) -> WriteEdge | None:
367 """Return the `WriteEdge` that links the producing task to the named
368 dataset type.
370 Parameters
371 ----------
372 dataset_type_name : `str`
373 Dataset type name. Must not be a component.
375 Returns
376 -------
377 edge : `WriteEdge` or `None`
378 Producing edge or `None` if there isn't one in this graph.
380 Raises
381 ------
382 DuplicateOutputError
383 Raised if there are multiple tasks defined to produce this dataset
384 type. This is only possible if the graph's dataset types are not
385 resolved.
387 Notes
388 -----
389 On resolved graphs, it may be slightly more efficient to use::
391 graph.dataset_types[dataset_type_name].producing_edge
393 but this method works on graphs with unresolved dataset types as well.
394 """
395 producer: str | None = None
396 producing_edge: WriteEdge | None = None
397 for _, _, producing_edge in self._xgraph.in_edges(
398 NodeKey(NodeType.DATASET_TYPE, dataset_type_name), data="instance"
399 ):
400 assert producing_edge is not None, "Should only be None if we never loop."
401 if producer is not None:
402 raise DuplicateOutputError(
403 f"Dataset type {dataset_type_name!r} is produced by both {producing_edge.task_label!r} "
404 f"and {producer!r}."
405 )
406 return producing_edge
408 def consuming_edges_of(self, dataset_type_name: str) -> list[ReadEdge]:
409 """Return the `ReadEdge` objects that link the named dataset type to
410 the tasks that consume it.
412 Parameters
413 ----------
414 dataset_type_name : `str`
415 Dataset type name. Must not be a component.
417 Returns
418 -------
419 edges : `list` [ `ReadEdge` ]
420 Edges that connect this dataset type to the tasks that consume it.
422 Notes
423 -----
424 On resolved graphs, it may be slightly more efficient to use::
426 graph.dataset_types[dataset_type_name].producing_edges
428 but this method works on graphs with unresolved dataset types as well.
429 """
430 return [
431 edge
432 for _, _, edge in self._xgraph.out_edges(
433 NodeKey(NodeType.DATASET_TYPE, dataset_type_name), data="instance"
434 )
435 ]
437 def producer_of(self, dataset_type_name: str) -> TaskNode | TaskInitNode | None:
438 """Return the `TaskNode` or `TaskInitNode` that writes the given
439 dataset type.
441 Parameters
442 ----------
443 dataset_type_name : `str`
444 Dataset type name. Must not be a component.
446 Returns
447 -------
448 edge : `TaskNode`, `TaskInitNode`, or `None`
449 Producing node or `None` if there isn't one in this graph.
451 Raises
452 ------
453 DuplicateOutputError
454 Raised if there are multiple tasks defined to produce this dataset
455 type. This is only possible if the graph's dataset types are not
456 resolved.
457 """
458 if (producing_edge := self.producing_edge_of(dataset_type_name)) is not None:
459 return self._xgraph.nodes[producing_edge.task_key]["instance"]
460 return None
462 def consumers_of(self, dataset_type_name: str) -> list[TaskNode | TaskInitNode]:
463 """Return the `TaskNode` and/or `TaskInitNode` objects that read
464 the given dataset type.
466 Parameters
467 ----------
468 dataset_type_name : `str`
469 Dataset type name. Must not be a component.
471 Returns
472 -------
473 edges : `list` [ `ReadEdge` ]
474 Edges that connect this dataset type to the tasks that consume it.
476 Notes
477 -----
478 On resolved graphs, it may be slightly more efficient to use::
480 graph.dataset_types[dataset_type_name].producing_edges
482 but this method works on graphs with unresolved dataset types as well.
483 """
484 return [
485 self._xgraph.nodes[consuming_edge.task_key]["instance"]
486 for consuming_edge in self.consuming_edges_of(dataset_type_name)
487 ]
489 def inputs_of(self, task_label: str, init: bool = False) -> dict[str, DatasetTypeNode | None]:
490 """Return the dataset types that are inputs to a task.
492 Parameters
493 ----------
494 task_label : `str`
495 Label for the task in the pipeline.
496 init : `bool`, optional
497 If `True`, return init-input dataset types instead of runtime
498 (including prerequisite) inputs.
500 Returns
501 -------
502 inputs : `dict` [ `str`, `DatasetTypeNode` or `None` ]
503 Dictionary parent dataset type name keys and either
504 `DatasetTypeNode` values (if the dataset type has been resolved)
505 or `None` values.
507 Notes
508 -----
509 To get the input edges of a task or task init node (which provide
510 information about storage class overrides nd components) use::
512 graph.tasks[task_label].iter_all_inputs()
514 or
516 graph.tasks[task_label].init.iter_all_inputs()
518 or the various mapping attributes of the `TaskNode` and `TaskInitNode`
519 class.
520 """
521 node: TaskNode | TaskInitNode = self.tasks[task_label] if not init else self.tasks[task_label].init
522 return {
523 edge.parent_dataset_type_name: self._xgraph.nodes[edge.dataset_type_key]["instance"]
524 for edge in node.iter_all_inputs()
525 }
527 def outputs_of(
528 self, task_label: str, init: bool = False, include_automatic_connections: bool = True
529 ) -> dict[str, DatasetTypeNode | None]:
530 """Return the dataset types that are outputs of a task.
532 Parameters
533 ----------
534 task_label : `str`
535 Label for the task in the pipeline.
536 init : `bool`, optional
537 If `True`, return init-output dataset types instead of runtime
538 outputs.
539 include_automatic_connections : `bool`, optional
540 Whether to include automatic connections such as configs, metadata,
541 and logs.
543 Returns
544 -------
545 outputs : `dict` [ `str`, `DatasetTypeNode` or `None` ]
546 Dictionary parent dataset type name keys and either
547 `DatasetTypeNode` values (if the dataset type has been resolved)
548 or `None` values.
550 Notes
551 -----
552 To get the input edges of a task or task init node (which provide
553 information about storage class overrides nd components) use::
555 graph.tasks[task_label].iter_all_outputs()
557 or
559 graph.tasks[task_label].init.iter_all_outputs()
561 or the various mapping attributes of the `TaskNode` and `TaskInitNode`
562 class.
563 """
564 node: TaskNode | TaskInitNode = self.tasks[task_label] if not init else self.tasks[task_label].init
565 iterable = node.iter_all_outputs() if include_automatic_connections else node.outputs.values()
566 return {
567 edge.parent_dataset_type_name: self._xgraph.nodes[edge.dataset_type_key]["instance"]
568 for edge in iterable
569 }
571 def resolve(
572 self,
573 registry: Registry | None = None,
574 dimensions: DimensionUniverse | None = None,
575 dataset_types: Mapping[str, DatasetType] | None = None,
576 visualization_only: bool = False,
577 ) -> None:
578 """Resolve all dimensions and dataset types and check them for
579 consistency.
581 Resolving a graph also causes it to be sorted.
583 Parameters
584 ----------
585 registry : `lsst.daf.butler.Registry`, optional
586 Client for the data repository to resolve against.
587 dimensions : `lsst.daf.butler.DimensionUniverse`, optional
588 Definitions for all dimensions. Takes precedence over
589 ``registry.dimensions`` if both are provided. If neither is
590 provided, defaults to the default dimension universe
591 (``lsst.daf.butler.DimensionUniverse()``).
592 dataset_types : `~collection.abc.Mapping` [ `str`, \
593 `~lsst.daf.butler.DatasetType` ], optional
594 Mapping of dataset types to consider registered. Takes precedence
595 over ``registry.getDatasetType()`` if both are provided.
596 visualization_only : `bool`, optional
597 Resolve the graph as well as possible even when dimensions and
598 storage classes cannot really be determined. This can include
599 using the ``universe.commonSkyPix`` as the assumed dimensions of
600 connections that use the "skypix" placeholder and using "<UNKNOWN>"
601 as a storage class name (which will fail if the storage class
602 itself is ever actually loaded).
604 Notes
605 -----
606 The `universe` attribute is set to ``dimensions`` and used to set all
607 `TaskNode.dimensions` attributes. Dataset type nodes are resolved by
608 first looking for a registry definition, then using the producing
609 task's definition, then looking for consistency between all consuming
610 task definitions.
612 Raises
613 ------
614 ConnectionTypeConsistencyError
615 Raised if a prerequisite input for one task appears as a different
616 kind of connection in any other task.
617 DuplicateOutputError
618 Raised if multiple tasks have the same dataset type as an output.
619 IncompatibleDatasetTypeError
620 Raised if different tasks have different definitions of a dataset
621 type. Different but compatible storage classes are permitted.
622 MissingDatasetTypeError
623 Raised if a dataset type definition is required to exist in the
624 data repository but none was found. This should only occur for
625 dataset types that are not produced by a task in the pipeline and
626 are consumed with different storage classes or as components by
627 tasks in the pipeline.
628 EdgesChangedError
629 Raised if ``check_edges_unchanged=True`` and the edges of a task do
630 change after import and reconfiguration.
631 """
632 get_registered: Callable[[str], DatasetType | None] | None = None
633 if dataset_types is not None:
634 # Ruff seems confused about whether this is used below; it is!
635 get_registered = dataset_types.get
636 elif registry is not None:
638 def get_registered(name: str) -> DatasetType | None:
639 try:
640 return registry.getDatasetType(name)
641 except MissingDatasetTypeError:
642 return None
644 else:
646 def get_registered(name: str) -> None:
647 return None
649 if dimensions is None:
650 if registry is not None:
651 dimensions = registry.dimensions
652 else:
653 dimensions = DimensionUniverse()
655 sort_keys_from_steps = self._resolve_step_flow()
657 node_key: NodeKey
658 updates: dict[NodeKey, TaskNode | DatasetTypeNode] = {}
659 for node_key, node_state in self._xgraph.nodes.items():
660 match node_key.node_type:
661 case NodeType.TASK:
662 task_node: TaskNode = node_state["instance"]
663 new_task_node = task_node._resolved(dimensions)
664 if new_task_node is not task_node:
665 updates[node_key] = new_task_node
666 case NodeType.DATASET_TYPE:
667 dataset_type_node: DatasetTypeNode | None = node_state["instance"]
668 new_dataset_type_node = DatasetTypeNode._from_edges(
669 node_key,
670 self._xgraph,
671 get_registered,
672 dimensions,
673 previous=dataset_type_node,
674 visualization_only=visualization_only,
675 )
676 # Usage of `is`` here is intentional; `_from_edges` returns
677 # `previous=dataset_type_node` if it can determine that it
678 # doesn't need to change.
679 if new_dataset_type_node is not dataset_type_node:
680 updates[node_key] = new_dataset_type_node
681 try:
682 for node_key, node_value in updates.items():
683 self._xgraph.nodes[node_key]["instance"] = node_value
684 except Exception as err: # pragma: no cover
685 # There's no known way to get here, but we want to make it
686 # clear it's a big problem if we do.
687 raise PipelineGraphExceptionSafetyError(
688 "Error during dataset type resolution has left the graph in an inconsistent state."
689 ) from err
691 # If we get an error here, many graph nodes will have been resolved but
692 # the steps will still be marked as unverified and unsorted. That's
693 # still an acceptable state for the pipeline to be in.
694 self._resolve_step_dimensions(dimensions)
695 self._step_definitions._verified = True
697 if sort_keys_from_steps is not None:
698 self._reorder(sort_keys_from_steps)
699 else:
700 self.sort()
701 self._universe = dimensions
703 ###########################################################################
704 #
705 # Graph Modification Interface:
706 #
707 # - methods to add, remove, and replace tasks;
708 #
709 # - methods to add and remove task subsets.
710 #
711 # These are all things that are usually done in a Pipeline before making a
712 # graph at all, but there may be cases where we want to modify the graph
713 # instead. (These are also the methods used to make a graph from a
714 # Pipeline, or make a graph from another graph.)
715 #
716 ###########################################################################
718 def add_task(
719 self,
720 label: str | None,
721 task_class: type[PipelineTask],
722 config: PipelineTaskConfig | None = None,
723 connections: PipelineTaskConnections | None = None,
724 ) -> TaskNode:
725 """Add a new task to the graph.
727 Parameters
728 ----------
729 label : `str` or `None`
730 Label for the task in the pipeline. If `None`, `Task._DefaultName`
731 is used.
732 task_class : `type` [ `PipelineTask` ]
733 Class object for the task.
734 config : `PipelineTaskConfig`, optional
735 Configuration for the task. If not provided, a default-constructed
736 instance of ``task_class.ConfigClass`` is used.
737 connections : `PipelineTaskConnections`, optional
738 Object that describes the dataset types used by the task. If not
739 provided, one will be constructed from the given configuration. If
740 provided, it is assumed that ``config`` has already been validated
741 and frozen.
743 Returns
744 -------
745 node : `TaskNode`
746 The new task node added to the graph.
748 Raises
749 ------
750 ValueError
751 Raised if configuration validation failed when constructing
752 ``connections``.
753 PipelineDataCycleError
754 Raised if the graph is cyclic after this addition.
755 RuntimeError
756 Raised if an unexpected exception (which will be chained) occurred
757 at a stage that may have left the graph in an inconsistent state.
758 Other exceptions should leave the graph unchanged.
760 Notes
761 -----
762 Checks for dataset type consistency and multiple producers do not occur
763 until `resolve` is called, since the resolution depends on both the
764 state of the data repository and all contributing tasks.
766 Adding new tasks removes any existing resolutions of all dataset types
767 it references and marks the graph as unsorted. It is most efficient
768 to add all tasks up front and only then resolve and/or sort the graph.
769 """
770 if label is None:
771 label = task_class._DefaultName
772 if config is None:
773 config = task_class.ConfigClass()
774 _LOG.debug("Adding task %s %s to the pipeline graph", label, task_class)
775 task_node = TaskNode._from_imported_data(
776 key=NodeKey(NodeType.TASK, label),
777 init_key=NodeKey(NodeType.TASK_INIT, label),
778 data=_TaskNodeImportedData.configure(label, task_class, config, connections),
779 universe=self._universe,
780 )
781 self.add_task_nodes([task_node])
782 return task_node
784 def add_task_nodes(self, nodes: Iterable[TaskNode], parent: PipelineGraph | None = None) -> None:
785 """Add one or more existing task nodes to the graph.
787 Parameters
788 ----------
789 nodes : `~collections.abc.Iterable` [ `TaskNode` ]
790 Iterable of task nodes to add. If any tasks have resolved
791 dimensions, they must have the same dimension universe as the rest
792 of the graph.
793 parent : `PipelineGraph`, optional
794 If provided, another `PipelineGraph` from which these nodes were
795 obtained. Any dataset type nodes already present in ``parent``
796 that are referenced by the given tasks will be used in this graph
797 if they are not already present, preserving any dataset type
798 resolutions present in the parent graph. Adding nodes from a
799 parent graph after the graph has its own nodes (e.g. from
800 `add_task`) or nodes from a third graph may result in invalid
801 dataset type resolutions. It is safest to only use this argument
802 when populating an empty graph for the first time.
804 Raises
805 ------
806 PipelineDataCycleError
807 Raised if the graph is cyclic after this addition.
809 Notes
810 -----
811 Checks for dataset type consistency and multiple producers do not occur
812 until `resolve` is called, since the resolution depends on both the
813 state of the data repository and all contributing tasks.
815 Adding new tasks removes any existing resolutions of all dataset types
816 it references (unless ``parent is not None`` and marks the graph as
817 unsorted. It is most efficient to add all tasks up front and only then
818 resolve and/or sort the graph.
819 """
820 node_data: list[tuple[NodeKey, dict[str, Any]]] = []
821 edge_data: list[tuple[NodeKey, NodeKey, str, dict[str, Any]]] = []
822 for task_node in nodes:
823 task_node = task_node._resolved(self._universe)
824 node_data.append(
825 (task_node.key, {"instance": task_node, "bipartite": task_node.key.node_type.bipartite})
826 )
827 node_data.append(
828 (
829 task_node.init.key,
830 {"instance": task_node.init, "bipartite": task_node.init.key.node_type.bipartite},
831 )
832 )
833 # Convert the edge objects attached to the task node to networkx.
834 for read_edge in task_node.init.iter_all_inputs():
835 self._append_graph_data_from_edge(node_data, edge_data, read_edge, parent=parent)
836 for write_edge in task_node.init.iter_all_outputs():
837 self._append_graph_data_from_edge(node_data, edge_data, write_edge, parent=parent)
838 for read_edge in task_node.iter_all_inputs():
839 self._append_graph_data_from_edge(node_data, edge_data, read_edge, parent=parent)
840 for write_edge in task_node.iter_all_outputs():
841 self._append_graph_data_from_edge(node_data, edge_data, write_edge, parent=parent)
842 # Add a special edge (with no Edge instance) that connects the
843 # TaskInitNode to the runtime TaskNode.
844 edge_data.append((task_node.init.key, task_node.key, Edge.INIT_TO_TASK_NAME, {"instance": None}))
845 if not node_data and not edge_data:
846 return
847 # Checks and preparation complete; time to start the actual
848 # modification, during which it's hard to provide strong exception
849 # safety. Start by resetting the sort ordering, if there is one.
850 self._reset()
851 try:
852 self._xgraph.add_nodes_from(node_data)
853 self._xgraph.add_edges_from(edge_data)
854 if not networkx.algorithms.dag.is_directed_acyclic_graph(self._xgraph):
855 cycle = networkx.find_cycle(self._xgraph)
856 raise PipelineDataCycleError(f"Cycle detected while adding tasks: {cycle}.")
857 except Exception:
858 # First try to roll back our changes.
859 try:
860 self._xgraph.remove_edges_from(edge_data)
861 self._xgraph.remove_nodes_from(key for key, _ in node_data)
862 except Exception as err: # pragma: no cover
863 # There's no known way to get here, but we want to make it
864 # clear it's a big problem if we do.
865 raise PipelineGraphExceptionSafetyError(
866 "Error while attempting to revert PipelineGraph modification has left the graph in "
867 "an inconsistent state."
868 ) from err
869 # Successfully rolled back; raise the original exception.
870 raise
872 def reconfigure_tasks(
873 self,
874 *args: tuple[str, PipelineTaskConfig],
875 check_edges_unchanged: bool = False,
876 assume_edges_unchanged: bool = False,
877 **kwargs: PipelineTaskConfig,
878 ) -> None:
879 """Update the configuration for one or more tasks.
881 Parameters
882 ----------
883 *args : `tuple` [ `str`, `.PipelineTaskConfig` ]
884 Positional arguments are each a 2-tuple of task label and new
885 config object. Note that the same arguments may also be passed as
886 ``**kwargs``, which is usually more readable, but task labels in
887 ``*args`` are not required to be valid Python identifiers.
888 check_edges_unchanged : `bool`, optional
889 If `True`, require the edges (connections) of the modified tasks to
890 remain unchanged after the configuration updates, and verify that
891 this is the case.
892 assume_edges_unchanged : `bool`, optional
893 If `True`, the caller declares that the edges (connections) of the
894 modified tasks will remain unchanged after the configuration
895 updates, and that it is unnecessary to check this.
896 **kwargs : `.PipelineTaskConfig`
897 New config objects or overrides to apply to copies of the current
898 config objects, with task labels as the keywords.
900 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(
1701 self,
1702 butler: Butler,
1703 include_packages: bool = True,
1704 *,
1705 include_inputs: bool = True,
1706 include_configs: bool = True,
1707 include_logs: bool = True,
1708 ) -> None:
1709 """Register all dataset types in a data repository.
1711 Parameters
1712 ----------
1713 butler : `~lsst.daf.butler.Butler`
1714 Data repository client.
1715 include_packages : `bool`, optional
1716 Whether to include the special "packages" dataset type that records
1717 software versions (this is not associated with a task and hence is
1718 not considered part of the pipeline graph in other respects, but it
1719 does get written with other provenance datasets).
1720 include_inputs : `bool`, optional
1721 Whether to register overall-input dataset types as well as outputs.
1722 include_configs : `bool`, optional
1723 Whether to register task config dataset types.
1724 include_logs : `bool`, optional
1725 Whether to register task log dataset types.
1726 """
1727 dataset_types = {
1728 node.name: node.dataset_type
1729 for node in self.dataset_types.values()
1730 if include_inputs or self.producer_of(node.name) is not None
1731 }
1732 if include_packages:
1733 dataset_types[self.packages_dataset_type.name] = self.packages_dataset_type
1734 if not include_configs:
1735 for task_node in self.tasks.values():
1736 del dataset_types[task_node.init.config_output.dataset_type_name]
1737 if not include_logs:
1738 for task_node in self.tasks.values():
1739 if task_node.log_output is not None:
1740 del dataset_types[task_node.log_output.dataset_type_name]
1741 for dataset_type in dataset_types.values():
1742 butler.registry.registerDatasetType(dataset_type)
1744 def check_dataset_type_registrations(self, butler: Butler, include_packages: bool = True) -> None:
1745 """Check that dataset type registrations in a data repository match
1746 the definitions in this pipeline graph.
1748 Parameters
1749 ----------
1750 butler : `~lsst.daf.butler.Butler`
1751 Data repository client.
1752 include_packages : `bool`, optional
1753 Whether to include the special "packages" dataset type that records
1754 software versions (this is not associated with a task and hence is
1755 not considered part of the pipeline graph in other respects, but it
1756 does get written with other provenance datasets).
1758 Raises
1759 ------
1760 lsst.daf.butler.MissingDatasetTypeError
1761 Raised if one or more non-optional-input or output dataset types in
1762 the pipeline is not registered at all.
1763 lsst.daf.butler.ConflictingDefinitionError
1764 Raised if the definition in the data repository is not identical
1765 to the definition in the pipeline graph.
1767 Notes
1768 -----
1769 Note that dataset type definitions that are storage-class-conversion
1770 compatible but not identical are not permitted by these checks, because
1771 the expectation is that these differences are handled by `resolve`,
1772 which makes the pipeline graph use the data repository definitions.
1773 This method is intended to check that none of those definitions have
1774 changed.
1775 """
1776 dataset_types = [node.dataset_type for node in self.dataset_types.values()]
1777 if include_packages:
1778 dataset_types.append(self.packages_dataset_type)
1779 missing_dataset_types: list[str] = []
1780 for dataset_type in dataset_types:
1781 try:
1782 expected = butler.registry.getDatasetType(dataset_type.name)
1783 except MissingDatasetTypeError:
1784 expected = None
1785 if expected is None:
1786 # The user probably forgot to register dataset types
1787 # at least once (which should be an error),
1788 # but we could also get here if this is an optional input for
1789 # which no datasets were found in this repo (not an error).
1790 if (
1791 not (
1792 self.producer_of(dataset_type.name) is None
1793 and all(
1794 self.tasks[input_edge.task_label].is_optional(input_edge.connection_name)
1795 for input_edge in self.consuming_edges_of(dataset_type.name)
1796 )
1797 )
1798 or dataset_type.name == PACKAGES_INIT_OUTPUT_NAME
1799 ):
1800 missing_dataset_types.append(dataset_type.name)
1801 elif expected != dataset_type:
1802 raise ConflictingDefinitionError(
1803 f"DatasetType definition in registry has changed since the pipeline graph was resolved: "
1804 f"{dataset_type} (graph) != {expected} (registry)."
1805 )
1806 if missing_dataset_types:
1807 plural = "s" if len(missing_dataset_types) != 1 else ""
1808 raise MissingDatasetTypeError(
1809 f"Missing dataset type definition{plural}: {', '.join(missing_dataset_types)}. "
1810 "Dataset types have to be registered in advance (on the command-line, either via "
1811 "`butler register-dataset-type` or the `--register-dataset-types` option to `pipetask run`."
1812 )
1814 def instantiate_tasks(
1815 self,
1816 get_init_input: Callable[[DatasetType], Any] | None = None,
1817 init_outputs: list[tuple[Any, DatasetType]] | None = None,
1818 labels: Iterable[str] | None = None,
1819 ) -> list[PipelineTask]:
1820 """Instantiate all tasks in the pipeline.
1822 Parameters
1823 ----------
1824 get_init_input : `~collections.abc.Callable`, optional
1825 Callable that accepts a single `~lsst.daf.butler.DatasetType`
1826 parameter and returns the init-input dataset associated with that
1827 dataset type. Must respect the storage class embedded in the type.
1828 This is optional if the pipeline does not have any overall init
1829 inputs. When a full butler is available,
1830 `lsst.daf.butler.Butler.get` can be used directly here.
1831 init_outputs : `list`, optional
1832 A list of ``(obj, dataset type)`` init-output dataset pairs, to be
1833 appended to in-place. Both the object and the dataset type will
1834 correspond to the storage class of the output connection, which
1835 may not be the same as the storage class on the graph's dataset
1836 type node.
1837 labels : `~collections.abc.Iterable` [ `str` ], optional
1838 The labels of tasks to instantiate. If not provided, all tasks in
1839 the graph will be instantiated.
1841 Returns
1842 -------
1843 tasks : `list`
1844 Constructed `PipelineTask` instances.
1845 """
1846 if not self.is_fully_resolved:
1847 raise UnresolvedGraphError("Pipeline graph must be fully resolved before instantiating tasks.")
1848 empty_data_id = DataCoordinate.make_empty(self.universe)
1849 labels = set(labels) if labels is not None else self.tasks.keys()
1850 handles: dict[str, InMemoryDatasetHandle] = {}
1851 tasks: list[PipelineTask] = []
1852 for task_node in self.tasks.values():
1853 if task_node.label not in labels:
1854 continue
1855 task_init_inputs: dict[str, Any] = {}
1856 for read_edge in task_node.init.inputs.values():
1857 if (handle := handles.get(read_edge.dataset_type_name)) is not None:
1858 obj = handle.get(storageClass=read_edge.storage_class_name)
1859 elif (
1860 read_edge.component is not None
1861 and (parent_handle := handles.get(read_edge.parent_dataset_type_name)) is not None
1862 ):
1863 obj = parent_handle.get(
1864 storageClass=read_edge.storage_class_name, component=read_edge.component
1865 )
1866 else:
1867 dataset_type_node = self.dataset_types[read_edge.parent_dataset_type_name]
1868 if get_init_input is None:
1869 raise ValueError(
1870 f"Task {task_node.label!r} requires init-input "
1871 f"{read_edge.dataset_type_name} but no 'get_init_input' callback was provided."
1872 )
1873 obj = get_init_input(read_edge.adapt_dataset_type(dataset_type_node.dataset_type))
1874 n_consumers = len(self.consumers_of(dataset_type_node.name))
1875 if (
1876 n_consumers > 1
1877 and read_edge.component is None
1878 and read_edge.storage_class_name == dataset_type_node.storage_class_name
1879 ):
1880 # Caching what we just got is safe in general only
1881 # if there was no storage class conversion, since
1882 # a->b and a->c does not imply b->c.
1883 handles[read_edge.dataset_type_name] = InMemoryDatasetHandle(
1884 obj,
1885 storageClass=dataset_type_node.storage_class,
1886 dataId=empty_data_id,
1887 copy=True,
1888 )
1889 task_init_inputs[read_edge.connection_name] = obj
1890 task = task_node.task_class(
1891 config=task_node.config, initInputs=task_init_inputs, name=task_node.label
1892 )
1893 tasks.append(task)
1894 for write_edge in task_node.init.outputs.values():
1895 dataset_type_node = self.dataset_types[write_edge.parent_dataset_type_name]
1896 obj = getattr(task, write_edge.connection_name)
1897 # We don't immediately coerce obj to the dataset_type_node
1898 # storage class (which should be the repo storage class, if
1899 # there is one) when appending to `init_outputs` because a
1900 # formatter might be able to do a better job of that later;
1901 # instead we pair it with a dataset type that's consistent with
1902 # the in-memory type. We do coerce when populating `handles`,
1903 # though, because going through the dataset_type_node storage
1904 # class is the conversion path we checked when we resolved the
1905 # pipeline graph.
1906 if init_outputs is not None:
1907 init_outputs.append((obj, write_edge.adapt_dataset_type(dataset_type_node.dataset_type)))
1908 n_consumers = len(self.consumers_of(dataset_type_node.name))
1909 if n_consumers > 0:
1910 handles[dataset_type_node.name] = InMemoryDatasetHandle(
1911 dataset_type_node.storage_class.coerce_type(obj),
1912 dataId=empty_data_id,
1913 storageClass=dataset_type_node.storage_class,
1914 copy=(n_consumers > 1),
1915 )
1916 return tasks
1918 def write_init_outputs(self, butler: Butler) -> None:
1919 """Write the init-output datasets for all tasks in the pipeline graph.
1921 Parameters
1922 ----------
1923 butler : `lsst.daf.butler.Butler`
1924 A full butler data repository client with its default run set
1925 to the collection where datasets should be written.
1927 Notes
1928 -----
1929 Datasets that already exist in the butler's output run collection will
1930 not be written.
1932 This method writes outputs with new random dataset IDs and should
1933 hence only be used when writing init-outputs prior to building a
1934 `QuantumGraph`. Use `QuantumGraph.write_init_outputs` if a quantum
1935 graph has already been built.
1936 """
1937 init_outputs: list[tuple[Any, DatasetType]] = []
1938 self.instantiate_tasks(butler.get, init_outputs)
1939 found_refs: dict[str, DatasetRef] = {}
1940 to_put: list[tuple[Any, DatasetType]] = []
1941 for obj, dataset_type in init_outputs:
1942 if (ref := butler.find_dataset(dataset_type, collections=butler.run)) is not None:
1943 found_refs[dataset_type.name] = ref
1944 else:
1945 to_put.append((obj, dataset_type))
1946 for ref, stored in butler.stored_many(found_refs.values()).items():
1947 if not stored:
1948 raise FileNotFoundError(
1949 f"Init-output dataset {ref.datasetType.name!r} was found in RUN {ref.run!r} "
1950 f"but had not actually been stored (or was stored and later deleted)."
1951 )
1952 for obj, dataset_type in to_put:
1953 butler.put(obj, dataset_type)
1955 def write_configs(self, butler: Butler) -> None:
1956 """Write the config datasets for all tasks in the pipeline graph.
1958 Parameters
1959 ----------
1960 butler : `lsst.daf.butler.Butler`
1961 A full butler data repository client with its default run set
1962 to the collection where datasets should be written.
1964 Notes
1965 -----
1966 Config datasets that already exist in the butler's output run
1967 collection will be checked for consistency.
1969 This method writes outputs with new random dataset IDs and should
1970 hence only be used when writing init-outputs prior to building a
1971 `QuantumGraph`. Use `QuantumGraph.write_configs` if a quantum graph
1972 has already been built.
1974 Raises
1975 ------
1976 lsst.daf.butler.registry.ConflictingDefinitionError
1977 Raised if a config dataset already exists and is not consistent
1978 with the config in the pipeline graph.
1979 """
1980 to_put: list[tuple[PipelineTaskConfig, str]] = []
1981 for task_node in self.tasks.values():
1982 dataset_type_name = task_node.init.config_output.dataset_type_name
1983 if (ref := butler.find_dataset(dataset_type_name, collections=butler.run)) is not None:
1984 old_config = butler.get(ref)
1985 if not task_node.config.compare(old_config, shortcut=False, output=log_config_mismatch):
1986 raise ConflictingDefinitionError(
1987 f"Config does not match existing task config {dataset_type_name!r} in "
1988 "butler; tasks configurations must be consistent within the same run collection"
1989 )
1990 else:
1991 to_put.append((task_node.config, dataset_type_name))
1992 # We do writes at the end to minimize the mess we leave behind when we
1993 # raise an exception.
1994 for config, dataset_type_name in to_put:
1995 butler.put(config, dataset_type_name)
1997 def write_packages(self, butler: Butler) -> None:
1998 """Write the 'packages' dataset for the currently-active software
1999 versions.
2001 Parameters
2002 ----------
2003 butler : `lsst.daf.butler.Butler`
2004 A full butler data repository client with its default run set
2005 to the collection where datasets should be written.
2007 Notes
2008 -----
2009 If the packages dataset already exists, it will be compared to the
2010 versions in the current packages. New packages that weren't present
2011 before are not considered an inconsistency.
2013 This method writes outputs with new random dataset IDs and should
2014 hence only be used when writing init-outputs prior to building a
2015 `QuantumGraph`. Use `QuantumGraph.write_packages` if a quantum graph
2016 has already been built.
2018 Raises
2019 ------
2020 lsst.daf.butler.registry.ConflictingDefinitionError
2021 Raised if the packages dataset already exists and is not consistent
2022 with the current packages.
2023 """
2024 new_packages = Packages.fromSystem()
2025 if (ref := butler.find_dataset(self.packages_dataset_type)) is not None:
2026 packages = butler.get(ref)
2027 if compare_packages(packages, new_packages):
2028 # have to remove existing dataset first; butler has no
2029 # replace option.
2030 butler.pruneDatasets([ref], unstore=True, purge=True)
2031 butler.put(packages, ref)
2032 else:
2033 butler.put(new_packages, self.packages_dataset_type)
2035 def init_output_run(self, butler: Butler) -> None:
2036 """Initialize a new output RUN collection by writing init-output
2037 datasets (including configs and packages).
2039 Parameters
2040 ----------
2041 butler : `lsst.daf.butler.Butler`
2042 A full butler data repository client with its default run set
2043 to the collection where datasets should be written.
2044 """
2045 self.write_configs(butler)
2046 self.write_packages(butler)
2047 self.write_init_outputs(butler)
2049 ###########################################################################
2050 #
2051 # Class- and Package-Private Methods.
2052 #
2053 ###########################################################################
2055 def _iter_task_defs(self) -> Iterator[TaskDef]:
2056 """Iterate over this pipeline as a sequence of `TaskDef` instances.
2058 Notes
2059 -----
2060 This is a package-private method intended to aid in the transition to a
2061 codebase more fully integrated with the `PipelineGraph` class, in which
2062 both `TaskDef` and `PipelineDatasetTypes` are expected to go away, and
2063 much of the functionality on the `Pipeline` class will be moved to
2064 `PipelineGraph` as well.
2066 Raises
2067 ------
2068 TaskNotImportedError
2069 Raised if `TaskNode.is_imported` is `False` for any task.
2070 """
2071 from ..pipeline import TaskDef
2073 for node in self._tasks.values():
2074 yield TaskDef(
2075 config=node.config,
2076 taskClass=node.task_class,
2077 label=node.label,
2078 connections=node.get_connections(),
2079 )
2081 def _init_from_args(
2082 self,
2083 xgraph: networkx.MultiDiGraph | None,
2084 sorted_keys: Sequence[NodeKey] | None,
2085 task_subsets: dict[str, TaskSubset] | None,
2086 description: str,
2087 universe: DimensionUniverse | None,
2088 data_id: DataId | None,
2089 step_definitions: StepDefinitions,
2090 ) -> None:
2091 """Initialize the graph with possibly-nontrivial arguments.
2093 Parameters
2094 ----------
2095 xgraph : `networkx.MultiDiGraph` or `None`
2096 The backing networkx graph, or `None` to create an empty one.
2097 This graph has `NodeKey` instances for nodes and the same structure
2098 as the graph exported by `make_xgraph`, but its nodes and edges
2099 have a single ``instance`` attribute that holds a `TaskNode`,
2100 `TaskInitNode`, `DatasetTypeNode` (or `None`), `ReadEdge`, or
2101 `WriteEdge` instance.
2102 sorted_keys : `~collections.abc.Sequence` [ `NodeKey` ] or `None`
2103 Topologically sorted sequence of node keys, or `None` if the graph
2104 is not sorted.
2105 task_subsets : `dict` [ `str`, `TaskSubset` ]
2106 Labeled subsets of tasks. Values must be constructed with
2107 ``xgraph`` as their parent graph.
2108 description : `str`
2109 String description for this pipeline.
2110 universe : `lsst.daf.butler.DimensionUniverse` or `None`
2111 Definitions of all dimensions.
2112 data_id : `lsst.daf.butler.DataCoordinate` or other data ID mapping.
2113 Data ID that represents a constraint on all quanta generated from
2114 this pipeline.
2115 step_definitions : `StepDefinitions`
2116 Struct holding information about steps.
2118 Notes
2119 -----
2120 Only empty `PipelineGraph` instances should be constructed directly by
2121 users, which sets the signature of ``__init__`` itself, but methods on
2122 `PipelineGraph` and its helper classes need to be able to create them
2123 with state. Those methods can call this after calling ``__new__``
2124 manually, skipping ``__init__``.
2125 """
2126 self._xgraph = xgraph if xgraph is not None else networkx.MultiDiGraph()
2127 self._sorted_keys: Sequence[NodeKey] | None = None
2128 self._task_subsets = task_subsets if task_subsets is not None else {}
2129 self._description = description
2130 self._tasks = TaskMappingView(self._xgraph)
2131 self._dataset_types = DatasetTypeMappingView(self._xgraph)
2132 self._raw_data_id: dict[str, Any]
2133 if isinstance(data_id, DataCoordinate):
2134 if universe is None:
2135 universe = data_id.universe
2136 else:
2137 assert universe is data_id.universe, "data_id.universe and given universe differ"
2138 self._raw_data_id = dict(data_id.required)
2139 elif data_id is None:
2140 self._raw_data_id = {}
2141 else:
2142 self._raw_data_id = dict(data_id)
2143 self._universe = universe
2144 if sorted_keys is not None:
2145 self._reorder(sorted_keys)
2146 self._step_definitions = step_definitions
2148 def _make_bipartite_xgraph_internal(self, init: bool) -> networkx.MultiDiGraph:
2149 """Make a bipartite init-only or runtime-only internal subgraph.
2151 See `make_bipartite_xgraph` for parameters and return values.
2153 Notes
2154 -----
2155 This method returns a view of the `PipelineGraph` object's internal
2156 backing graph, and hence should only be called in methods that copy the
2157 result either explicitly or by running a copying algorithm before
2158 returning it to the user.
2159 """
2160 return self._xgraph.edge_subgraph([edge.key for edge in self.iter_edges(init)])
2162 def _make_task_xgraph_internal(self, init: bool) -> networkx.DiGraph:
2163 """Make a init-only or runtime-only internal task subgraph.
2165 See `make_task_xgraph` for parameters and return values.
2167 Notes
2168 -----
2169 This method returns a view of the `PipelineGraph` object's internal
2170 backing graph, and hence should only be called in methods that copy the
2171 result either explicitly or by running a copying algorithm before
2172 returning it to the user.
2173 """
2174 bipartite_xgraph = self._make_bipartite_xgraph_internal(init=init)
2175 task_keys = [
2176 key
2177 for key, bipartite in bipartite_xgraph.nodes(data="bipartite")
2178 if bipartite == NodeType.TASK.bipartite
2179 ]
2180 return networkx.algorithms.bipartite.projected_graph(networkx.DiGraph(bipartite_xgraph), task_keys)
2182 def _transform_xgraph_state(self, xgraph: _G, skip_edges: bool) -> _G:
2183 """Transform networkx graph attributes in-place from the internal
2184 "instance" attributes to the documented exported attributes.
2186 Parameters
2187 ----------
2188 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
2189 Graph whose state should be transformed.
2190 skip_edges : `bool`
2191 If `True`, do not transform edge state.
2193 Returns
2194 -------
2195 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
2196 The same object passed in, after modification.
2198 Notes
2199 -----
2200 This should be called after making a copy of the internal graph but
2201 before any projection down to just task or dataset type nodes, since
2202 it assumes stateful edges.
2203 """
2204 state: dict[str, Any]
2205 for state in xgraph.nodes.values():
2206 node_value: TaskInitNode | TaskNode | DatasetTypeNode | None = state.pop("instance")
2207 if node_value is not None:
2208 state.update(node_value._to_xgraph_state())
2209 else:
2210 # This is a dataset type node that is not resolved.
2211 state["bipartite"] = NodeType.DATASET_TYPE.bipartite
2212 if not skip_edges:
2213 for _, _, state in xgraph.edges(data=True):
2214 edge: Edge | None = state.pop("instance", None)
2215 if edge is not None:
2216 state.update(edge._to_xgraph_state())
2217 return xgraph
2219 def _replace_task_nodes(
2220 self,
2221 updates: Mapping[str, TaskNode],
2222 check_edges_unchanged: bool,
2223 assume_edges_unchanged: bool,
2224 message_header: str,
2225 ) -> None:
2226 """Replace task nodes and update edges and dataset type nodes
2227 accordingly.
2229 Parameters
2230 ----------
2231 updates : `Mapping` [ `str`, `TaskNode` ]
2232 New task nodes with task label keys. All keys must be task labels
2233 that are already present in the graph.
2234 check_edges_unchanged : `bool`, optional
2235 If `True`, require the edges (connections) of the modified tasks to
2236 remain unchanged after importing and configuring each task, and
2237 verify that this is the case.
2238 assume_edges_unchanged : `bool`, optional
2239 If `True`, the caller declares that the edges (connections) of the
2240 modified tasks will remain unchanged importing and configuring each
2241 task, and that it is unnecessary to check this.
2242 message_header : `str`
2243 Template for `str.format` with a single ``task_label`` placeholder
2244 to use as the first line in `EdgesChangedError` messages that show
2245 the differences between new task edges and old task edges. Should
2246 include the fact that the rest of the message will refer to the old
2247 task as "A" and the new task as "B", and end with a colon.
2249 Raises
2250 ------
2251 ValueError
2252 Raised if ``assume_edges_unchanged`` and ``check_edges_unchanged``
2253 are both `True`, or if a full config is provided for a task after
2254 another full config or an override has already been provided.
2255 EdgesChangedError
2256 Raised if ``check_edges_unchanged=True`` and the edges of a task do
2257 change.
2258 """
2259 deep: dict[str, TaskNode] = {}
2260 shallow: dict[str, TaskNode] = {}
2261 if assume_edges_unchanged:
2262 if check_edges_unchanged:
2263 raise ValueError("Cannot simultaneously assume and check that edges have not changed.")
2264 shallow.update(updates)
2265 else:
2266 for task_label, new_task_node in updates.items():
2267 old_task_node = self.tasks[task_label]
2268 messages = old_task_node.diff_edges(new_task_node)
2269 if messages:
2270 if check_edges_unchanged:
2271 messages.insert(0, message_header.format(task_label=task_label))
2272 raise EdgesChangedError("\n".join(messages))
2273 else:
2274 deep[task_label] = new_task_node
2275 else:
2276 shallow[task_label] = new_task_node
2277 try:
2278 if deep:
2279 removed = self.remove_tasks(deep.keys(), drop_from_subsets=True)
2280 self.add_task_nodes(deep.values())
2281 for replaced_task_node, referencing_subsets in removed:
2282 for subset_label in referencing_subsets:
2283 self._task_subsets[subset_label].add(replaced_task_node.label)
2284 for task_node in shallow.values():
2285 self._xgraph.nodes[task_node.key]["instance"] = task_node
2286 self._xgraph.nodes[task_node.init.key]["instance"] = task_node.init
2287 except PipelineGraphExceptionSafetyError: # pragma: no cover
2288 raise
2289 except Exception as err: # pragma: no cover
2290 # There's no known way to get here, but we want to make it clear
2291 # it's a big problem if we do.
2292 raise PipelineGraphExceptionSafetyError(
2293 "Error while replacing tasks has left the graph in an inconsistent state."
2294 ) from err
2296 def _append_graph_data_from_edge(
2297 self,
2298 node_data: list[tuple[NodeKey, dict[str, Any]]],
2299 edge_data: list[tuple[NodeKey, NodeKey, str, dict[str, Any]]],
2300 edge: Edge,
2301 parent: PipelineGraph | None,
2302 ) -> None:
2303 """Append networkx state dictionaries for an edge and the corresponding
2304 dataset type node.
2306 Parameters
2307 ----------
2308 node_data : `list`
2309 List of node keys and state dictionaries. A node is appended if
2310 one does not already exist for this dataset type.
2311 edge_data : `list`
2312 List of node key pairs, connection names, and state dictionaries
2313 for edges.
2314 edge : `Edge`
2315 New edge being processed.
2316 parent : `PipelineGraph` or `None`
2317 Another pipeline graph whose dataset type nodes should be used
2318 when present.
2319 """
2320 new_dataset_type_node = None
2321 if parent is not None:
2322 new_dataset_type_node = parent._xgraph.nodes[edge.dataset_type_key].get("instance")
2323 if (existing_dataset_type_state := self._xgraph.nodes.get(edge.dataset_type_key)) is not None:
2324 existing_dataset_type_state["instance"] = new_dataset_type_node
2325 else:
2326 node_data.append(
2327 (
2328 edge.dataset_type_key,
2329 {
2330 "instance": new_dataset_type_node,
2331 "bipartite": NodeType.DATASET_TYPE.bipartite,
2332 },
2333 )
2334 )
2335 edge_data.append(
2336 edge.nodes
2337 + (
2338 edge.connection_name,
2339 {"instance": edge},
2340 )
2341 )
2343 def _reorder(self, sorted_keys: Sequence[NodeKey]) -> None:
2344 """Set the order of all views of this graph from the given sorted
2345 sequence of task labels and dataset type names.
2346 """
2347 self._sorted_keys = sorted_keys
2348 self._tasks._reorder(sorted_keys)
2349 self._dataset_types._reorder(sorted_keys)
2351 def _reset(self) -> None:
2352 """Reset all views of this graph following a modification that might
2353 invalidate them.
2354 """
2355 self._sorted_keys = None
2356 self._tasks._reset()
2357 self._dataset_types._reset()
2359 def _resolve_step_flow(self) -> Sequence[NodeKey] | None:
2360 """Check that step definitions are consistent with the ordering of
2361 the graph's nodes and that they partition the task graph.
2363 Returns
2364 -------
2365 sort_keys : `~collections.abc.Sequence` [ `NodeKey` ] or `None`
2366 Sort order for the pipeline graph that is consistent with the
2367 step order, or `None` if there were no steps defined for this
2368 pipeline.
2369 """
2370 if not self._step_definitions:
2371 return None
2372 task_labels_so_far: set[str] = set()
2373 # Inputs we've already seen, and the tasks that wanted them:
2374 inputs_so_far: dict[str, set[str]] = {}
2375 sort_keys: list[NodeKey] = []
2376 keys_already_sorted: set[NodeKey] = set()
2377 for step_label in self.steps:
2378 try:
2379 task_subset = self.task_subsets[step_label]
2380 except KeyError:
2381 raise InvalidStepsError(f"Step {step_label!r} is not a task subset.") from None
2382 if not task_labels_so_far.isdisjoint(task_subset):
2383 raise InvalidStepsError(
2384 f"Step {step_label!r} repeats task(s) {task_labels_so_far & task_subset}."
2385 )
2386 task_labels_so_far.update(task_subset)
2387 # We need a temporary data structure for tracking the inputs of
2388 # just this step to avoid complaining about input-output
2389 # relationships within the step. We'll merge this into the
2390 # cumulative one later.
2391 step_inputs: dict[str, set[str]] = {}
2392 # We'll also gather all task, task-init and dataset type networkx
2393 # graph keys for the step so we can make and sort a step sub-graph.
2394 new_step_keys: set[NodeKey] = set()
2395 for task_label in task_subset:
2396 for input_name in self.inputs_of(task_label):
2397 step_inputs.setdefault(input_name, set()).add(task_label)
2398 # Check that none of the outputs of the tasks in this step were
2399 # expected as inputs of a previous step's tasks.
2400 task_outputs = self.outputs_of(task_label)
2401 if not inputs_so_far.keys().isdisjoint(task_outputs.keys()):
2402 msg: list[str] = []
2403 for input_name in inputs_so_far.keys() & task_outputs.keys():
2404 msg.append(f"{input_name} (used by {inputs_so_far[input_name]})")
2405 raise InvalidStepsError(
2406 f"Task {task_label} in step {step_label!r} produces dataset types "
2407 f"[{', '.join(msg)}], but these are consumed by tasks in earlier "
2408 "steps. Either steps are out of order or the graph is cyclic."
2409 )
2410 task_init_key = NodeKey(NodeType.TASK_INIT, task_label)
2411 task_key = NodeKey(NodeType.TASK, task_label)
2412 new_step_keys.add(task_init_key)
2413 new_step_keys.add(task_key)
2414 new_step_keys.update(self._xgraph.predecessors(task_init_key))
2415 new_step_keys.update(self._xgraph.predecessors(task_key))
2416 new_step_keys.update(self._xgraph.successors(task_init_key))
2417 new_step_keys.update(self._xgraph.successors(task_key))
2418 # Also record the step the task is in as a private xgraph
2419 # attribute so we can look up the step given a task (or,
2420 # indirectly, dataset type); this has to be used with care
2421 # because if the steps haven't been verified it can be wrong.
2422 self._xgraph.nodes[task_init_key]["step"] = step_label
2423 self._xgraph.nodes[task_key]["step"] = step_label
2424 # Drop step input keys that were already either inputs or outputs
2425 # of a previous step, since they'll have already been added to
2426 # sort_keys.
2427 new_step_keys.difference_update(keys_already_sorted)
2428 # Make the step subgraph, sort it, and extend the overall sort_keys
2429 # with result.
2430 step_xgraph = self._xgraph.subgraph(new_step_keys)
2431 sort_keys.extend(networkx.dag.lexicographical_topological_sort(step_xgraph))
2432 keys_already_sorted.update(new_step_keys)
2433 for input_name, consuming_tasks in step_inputs.items():
2434 inputs_so_far.setdefault(input_name, set()).update(consuming_tasks)
2435 if not task_labels_so_far.issuperset(self.tasks.keys()):
2436 # Note that the converse issubset test effectively happens when we
2437 # look up inputs and outputs of each task.
2438 raise InvalidStepsError(f"No step contains task(s) {self.tasks.keys() - task_labels_so_far}.")
2439 return sort_keys
2441 def _resolve_step_dimensions(self, universe: DimensionUniverse) -> None:
2442 """Check that step sharding dimensions are consistent with task and
2443 output dataset dimensions.
2445 Parameters
2446 ----------
2447 universe : `~lsst.daf.butler.DimensionUniverse`
2448 Definitions for all dimensions. Will be attached to the step
2449 definitions by this method.
2450 """
2451 self._step_definitions._universe = universe
2452 for step_label in self.steps:
2453 dimensions = self.steps.get_dimensions(step_label)
2454 for task_label in self.task_subsets[step_label]:
2455 task_node = self.tasks[task_label]
2456 if not _dimensions_compatible(dimensions, task_node.dimensions):
2457 raise InvalidStepsError(
2458 f"Dimensions {task_node.dimensions} of task {task_label!r} are not compatible with "
2459 f"the sharding dimensions {dimensions} of step {step_label!r}."
2460 )
2461 for dataset_type_node in self.outputs_of(task_label).values():
2462 assert dataset_type_node is not None, "dataset types should be resolved first"
2463 if not _dimensions_compatible(dimensions, dataset_type_node.dimensions):
2464 raise InvalidStepsError(
2465 f"Dimensions {dataset_type_node.dimensions} of dataset type "
2466 f"{dataset_type_node.name!r} (produced by task {task_label!r}) are not "
2467 f"compatible with the sharding dimensions {dimensions} of step "
2468 f"{step_label!r}."
2469 )
2471 def _select_expression(self, expr_tree: expressions.Node, task_xgraph: networkx.DiGraph) -> set[NodeKey]:
2472 """Select tasks from a pipeline based on a string expression.
2474 This is the primary implementation method for `select` and
2475 `select_tasks`.
2477 Parameters
2478 ----------
2479 expr_tree : `expressions.Node`
2480 Expression [sub]tree to process (recursively).
2481 task_xgraph : `networkx.DiGraph`
2482 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2484 Returns
2485 -------
2486 selected : `set` [ `NodeKey` ]
2487 Set of `NodeKey` objects for matching tasks (only; no dataset type
2488 or task-init nodes).
2489 """
2490 match expr_tree:
2491 case expressions.IdentifierNode(qualifier=qualifier, label=label):
2492 match self._select_identifier(qualifier, label):
2493 case NodeKey(node_type=NodeType.TASK) as task_key:
2494 return {task_key}
2495 case NodeKey(node_type=NodeType.DATASET_TYPE) as dataset_type_key:
2496 # Since a dataset type can have only one producer, this
2497 # yields 0- (for overall inputs) or 1-element sets.
2498 for producer_key, _ in self._xgraph.in_edges(dataset_type_key):
2499 if producer_key.node_type is NodeType.TASK_INIT:
2500 raise InvalidExpressionError(
2501 f"Init-output dataset type {label!r} cannot be used directly in an "
2502 "expression."
2503 )
2504 return {producer_key}
2505 return set()
2506 case TaskSubset() as task_subset:
2507 return {NodeKey(NodeType.TASK, label) for label in task_subset}
2508 case _: # pragma: no cover
2509 raise AssertionError("Identifier type inconsistent with grammar.")
2510 case expressions.DirectionNode(operator=operator, start=start):
2511 match self._select_identifier(start.qualifier, start.label):
2512 case NodeKey(node_type=NodeType.TASK) as task_key:
2513 if operator.startswith("<"):
2514 return self._select_task_ancestors(
2515 task_key, task_xgraph, inclusive=operator.endswith("=")
2516 )
2517 else:
2518 assert operator.startswith(">"), "Guaranteed by grammar."
2519 return self._select_task_descendants(
2520 task_key, task_xgraph, inclusive=operator.endswith("=")
2521 )
2522 case NodeKey(node_type=NodeType.DATASET_TYPE) as dataset_type_key:
2523 if operator.startswith("<"):
2524 return self._select_dataset_type_ancestors(
2525 dataset_type_key, task_xgraph, inclusive=operator.endswith("=")
2526 )
2527 else:
2528 assert operator.startswith(">"), "Guaranteed by grammar."
2529 return self._select_dataset_type_descendants(
2530 dataset_type_key, task_xgraph, inclusive=operator.endswith("=")
2531 )
2532 case TaskSubset():
2533 raise InvalidExpressionError(
2534 f"Task subset identifier {start!r} cannot be used as the start of an "
2535 "ancestor/descendant search."
2536 )
2537 case _: # pragma: no cover
2538 raise AssertionError("Unexpected parsed identifier result type.")
2539 case expressions.NotNode(operand=operand):
2540 operand_result = self._select_expression(operand, task_xgraph)
2541 return set(task_xgraph.nodes.keys() - operand_result)
2542 case expressions.UnionNode(lhs=lhs, rhs=rhs):
2543 lhs_result = self._select_expression(lhs, task_xgraph)
2544 rhs_result = self._select_expression(rhs, task_xgraph)
2545 return lhs_result.union(rhs_result)
2546 case expressions.IntersectionNode(lhs=lhs, rhs=rhs):
2547 lhs_result = self._select_expression(lhs, task_xgraph)
2548 rhs_result = self._select_expression(rhs, task_xgraph)
2549 return lhs_result.intersection(rhs_result)
2550 case _: # pragma: no cover
2551 raise AssertionError("Expression parse node inconsistent with grammar.")
2553 def _select_task_ancestors(
2554 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2555 ) -> set[NodeKey]:
2556 """Return all task-node ancestors of the given task node, as defined by
2557 the `select` expression language.
2559 Parameters
2560 ----------
2561 start : `NodeKey`
2562 A runtime task node key.
2563 task_xgraph : `networkx.DiGraph`
2564 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2565 inclusive : `bool`
2566 Whether to include the ``start`` node in the results.
2568 Returns
2569 -------
2570 selected : `set` [ `NodeKey` ]
2571 Set of `NodeKey` objects for matching tasks (only; no dataset type
2572 or task-init nodes).
2573 """
2574 result = set(networkx.dag.ancestors(task_xgraph, start))
2575 if inclusive:
2576 result.add(start)
2577 return result
2579 def _select_task_descendants(
2580 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2581 ) -> set[NodeKey]:
2582 """Return all task-node descendants of the given task node, as defined
2583 by the `select` expression language.
2585 Parameters
2586 ----------
2587 start : `NodeKey`
2588 A runtime task node key.
2589 task_xgraph : `networkx.DiGraph`
2590 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2591 inclusive : `bool`
2592 Whether to include the ``start`` node in the results.
2594 Returns
2595 -------
2596 selected : `set` [ `NodeKey` ]
2597 Set of `NodeKey` objects for matching tasks (only; no dataset type
2598 or task-init nodes).
2599 """
2600 result = set(networkx.dag.descendants(task_xgraph, start))
2601 if inclusive:
2602 result.add(start)
2603 return result
2605 def _select_dataset_type_ancestors(
2606 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2607 ) -> set[NodeKey]:
2608 """Return all task-node ancestors of the given dataset type node, as
2609 defined by the `select` expression language.
2611 Parameters
2612 ----------
2613 start : `NodeKey`
2614 A dataset type node key. May not be an init-output.
2615 task_xgraph : `networkx.DiGraph`
2616 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2617 inclusive : `bool`
2618 Whether to include the producer of the ``start`` node in the
2619 results.
2621 Returns
2622 -------
2623 selected : `set` [ `NodeKey` ]
2624 Set of `NodeKey` objects for matching tasks (only; no dataset type
2625 or task-init nodes).
2626 """
2627 result: set[NodeKey] = set()
2628 for producer_key, _ in self._xgraph.in_edges(start):
2629 if producer_key.node_type is NodeType.TASK_INIT:
2630 raise InvalidExpressionError(
2631 f"Init-output dataset type {start.name!r} cannot be used as the "
2632 "starting point for an ancestor ('<' or '<=') search."
2633 )
2634 result.update(networkx.dag.ancestors(task_xgraph, producer_key))
2635 if inclusive:
2636 result.add(producer_key)
2637 return result
2639 def _select_dataset_type_descendants(
2640 self, start: NodeKey, task_xgraph: networkx.DiGraph, inclusive: bool
2641 ) -> set[NodeKey]:
2642 """Return all task-node descendatns of the given dataset type node, as
2643 defined by the `select` expression language.
2645 Parameters
2646 ----------
2647 start : `NodeKey`
2648 A dataset type node key. May not be an init-output if
2649 ``inclusive=True``.
2650 task_xgraph : `networkx.DiGraph`
2651 NetworkX graph of all tasks (runtime nodes only) in the pipeline.
2652 inclusive : `bool`
2653 Whether to include the producer of the ``start`` node in the
2654 results.
2656 Returns
2657 -------
2658 selected : `set` [ `NodeKey` ]
2659 Set of `NodeKey` objects for matching tasks (only; no dataset type
2660 or task-init nodes).
2661 """
2662 result: set[NodeKey] = set()
2663 if inclusive:
2664 for producer_key, _ in self._xgraph.in_edges(start):
2665 if producer_key.node_type is NodeType.TASK_INIT:
2666 raise InvalidExpressionError(
2667 f"Init-output dataset type {start.name!r} cannot be used as the "
2668 "starting point for an includsive descendant ('>=') search."
2669 )
2670 result.add(producer_key)
2671 # We also include tasks that consume a dataset type as an init-input,
2672 # since that can affect their runtime behavior.
2673 consumer_keys: set[NodeKey] = {
2674 (
2675 consumer_key
2676 if consumer_key.node_type is NodeType.TASK
2677 else NodeKey(NodeType.TASK, consumer_key.name)
2678 )
2679 for _, consumer_key in self._xgraph.out_edges(start)
2680 }
2681 for consumer_key in consumer_keys:
2682 result.add(consumer_key)
2683 result.update(networkx.dag.descendants(task_xgraph, consumer_key))
2684 return result
2686 def _select_identifier(
2687 self, qualifier: Literal["T", "D", "S"] | None, label: str
2688 ) -> NodeKey | TaskSubset:
2689 """Return the node key or task subset that corresponds to a `select`
2690 expression identifier.
2692 Parameters
2693 ----------
2694 qualifier : `str` or `None`
2695 Task, dataset type, or task subset qualifier included in the
2696 identifier, if any.
2697 label : `str`
2698 Task label, dataset type name, or task subset label.
2700 Returns
2701 -------
2702 key_or_subset : `NodeKey` or `TaskSubset`
2703 A `NodeKey` for a task or dataset type, or a `TaskSubset` for a
2704 task subset.
2705 """
2706 match qualifier:
2707 case None:
2708 task_key = NodeKey(NodeType.TASK, label)
2709 dataset_type_key = NodeKey(NodeType.DATASET_TYPE, label)
2710 if task_key in self._xgraph.nodes:
2711 if dataset_type_key in self._xgraph.nodes:
2712 raise InvalidExpressionError(
2713 f"{label!r} is both a task label and a dataset type name; "
2714 "prefix with 'T:' or 'D:' (respectively) to specify which."
2715 )
2716 assert label not in self._task_subsets, "Should be prohibited at construction."
2717 return task_key
2718 elif dataset_type_key in self._xgraph.nodes:
2719 if label in self._task_subsets:
2720 raise InvalidExpressionError(
2721 f"{label!r} is both a subset label and a dataset type name; "
2722 "prefix with 'S:' or 'D:' (respectively) to specify which."
2723 )
2724 return dataset_type_key
2725 elif label in self._task_subsets:
2726 return self._task_subsets[label]
2727 else:
2728 raise InvalidExpressionError(
2729 f"{label!r} is not a task label, task subset label, or dataset type name."
2730 )
2731 case "T":
2732 task_key = NodeKey(NodeType.TASK, label)
2733 if task_key not in self._xgraph.nodes:
2734 raise InvalidExpressionError(f"Task with label {label!r} does not exist.")
2735 return task_key
2736 case "D":
2737 dataset_type_key = NodeKey(NodeType.DATASET_TYPE, label)
2738 if dataset_type_key not in self._xgraph.nodes:
2739 raise InvalidExpressionError(f"Dataset type with name {label!r} does not exist.")
2740 return dataset_type_key
2741 case "S":
2742 try:
2743 return self._task_subsets[label]
2744 except KeyError:
2745 raise InvalidExpressionError(f"Task subset with label {label!r} does not exist.")
2746 case _: # pragma: no cover
2747 raise AssertionError("Unexpected identifier qualifier in expression.")
2749 _xgraph: networkx.MultiDiGraph
2750 _sorted_keys: Sequence[NodeKey] | None
2751 _task_subsets: dict[str, TaskSubset]
2752 _step_definitions: StepDefinitions
2753 _description: str
2754 _tasks: TaskMappingView
2755 _dataset_types: DatasetTypeMappingView
2756 _raw_data_id: dict[str, Any]
2757 _universe: DimensionUniverse | None
2760def log_config_mismatch(msg: str) -> None:
2761 """Log messages about configuration mismatch.
2763 Parameters
2764 ----------
2765 msg : `str`
2766 Log message to use.
2767 """
2768 _LOG.fatal("Comparing configuration: %s", msg)
2771def compare_packages(packages: Packages, new_packages: Packages) -> bool:
2772 """Compare two versions of Packages.
2774 Parameters
2775 ----------
2776 packages : `Packages`
2777 Previously recorded package versions. Updated in place to include
2778 any new packages that weren't present before.
2779 new_packages : `Packages`
2780 New set of package versions.
2782 Returns
2783 -------
2784 updated : `bool`
2785 `True` if ``packages`` was updated, `False` if not.
2787 Raises
2788 ------
2789 ConflictingDefinitionError
2790 Raised if versions are inconsistent.
2791 """
2792 diff = new_packages.difference(packages)
2793 if diff:
2794 versions_str = "; ".join(f"{pkg}: {diff[pkg][1]} vs {diff[pkg][0]}" for pkg in diff)
2795 raise ConflictingDefinitionError(f"Package versions mismatch: ({versions_str})")
2796 else:
2797 _LOG.debug("new packages are consistent with old")
2798 # Update the old set of packages in case we have more packages
2799 # that haven't been persisted.
2800 extra = new_packages.extra(packages)
2801 if extra:
2802 _LOG.debug("extra packages: %s", extra)
2803 packages.update(new_packages)
2804 return True
2805 return False
2808def _dimensions_compatible(dimensions: DimensionGroup, object_dimensions: DimensionGroup) -> bool:
2809 if dimensions.issubset(object_dimensions):
2810 # Easy typical case.
2811 return True
2812 # Hard case: if any sharding dimensions that are not in the object
2813 # dimensions are related to something in the object dimensions by a
2814 # many-to-many join table, this is okay too. The main use case here is to
2815 # let {exposure, [detector]} satisfy sharding dimensions of
2816 # {visit, [detector]}.
2817 universe = dimensions.universe
2818 unmatched_sharding = dimensions.required - object_dimensions.required
2819 unmatched_object = object_dimensions.required - dimensions.required
2820 unmatched_union_group = dimensions.union(object_dimensions)
2821 for element_name in unmatched_union_group.elements:
2822 if (
2823 (element := universe[element_name]).defines_relationships
2824 and not element.dimensions.isdisjoint(unmatched_sharding)
2825 and not element.dimensions.isdisjoint(unmatched_object)
2826 ):
2827 return True
2828 return False