Coverage for python / lsst / pipe / base / graph / graph.py: 15%
586 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:14 -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__ = ("IncompatibleGraphError", "QuantumGraph")
31import datetime
32import getpass
33import io
34import json
35import lzma
36import os
37import struct
38import sys
39import time
40import uuid
41from collections import defaultdict, deque
42from collections.abc import Generator, Iterable, Iterator, Mapping, MutableMapping
43from itertools import chain
44from types import MappingProxyType
45from typing import Any, BinaryIO, TypeVar
47import networkx as nx
48from networkx.drawing.nx_agraph import write_dot
50import lsst.utils.logging
51from lsst.daf.butler import (
52 Config,
53 DatasetId,
54 DatasetRef,
55 DatasetType,
56 DimensionRecordsAccumulator,
57 DimensionUniverse,
58 LimitedButler,
59 Quantum,
60 QuantumBackedButler,
61)
62from lsst.daf.butler._rubin import generate_uuidv7
63from lsst.daf.butler.datastore.record_data import DatastoreRecordData
64from lsst.daf.butler.persistence_context import PersistenceContextVars
65from lsst.daf.butler.registry import ConflictingDefinitionError
66from lsst.resources import ResourcePath, ResourcePathExpression
67from lsst.utils.introspection import get_full_type_name
68from lsst.utils.packages import Packages
70from ..config import PipelineTaskConfig
71from ..connections import iterConnections
72from ..pipeline import TaskDef
73from ..pipeline_graph import PipelineGraph, compare_packages, log_config_mismatch
74from ._implDetails import DatasetTypeName, _DatasetTracker
75from ._loadHelpers import LoadHelper
76from ._versionDeserializers import DESERIALIZER_MAP
77from .graphSummary import QgraphSummary, QgraphTaskSummary
78from .quantumNode import BuildId, QuantumNode
80_T = TypeVar("_T", bound="QuantumGraph")
81_LOG = lsst.utils.logging.getLogger(__name__)
83# modify this constant any time the on disk representation of the save file
84# changes, and update the load helpers to behave properly for each version.
85SAVE_VERSION = 3
87# Strings used to describe the format for the preamble bytes in a file save
88# The base is a big endian encoded unsigned short that is used to hold the
89# file format version. This allows reading version bytes and determine which
90# loading code should be used for the rest of the file
91STRUCT_FMT_BASE = ">H"
92#
93# Version 1
94# This marks a big endian encoded format with an unsigned short, an unsigned
95# long long, and an unsigned long long in the byte stream
96# Version 2
97# A big endian encoded format with an unsigned long long byte stream used to
98# indicate the total length of the entire header.
99STRUCT_FMT_STRING = {1: ">QQ", 2: ">Q"}
101# magic bytes that help determine this is a graph save
102MAGIC_BYTES = b"qgraph4\xf6\xe8\xa9"
105class IncompatibleGraphError(Exception):
106 """Exception class to indicate that a lookup by NodeId is impossible due
107 to incompatibilities.
108 """
110 pass
113class QuantumGraph:
114 """QuantumGraph is a directed acyclic graph of `QuantumNode` objects.
116 This data structure represents a concrete workflow generated from a
117 `Pipeline`.
119 Parameters
120 ----------
121 quanta : `~collections.abc.Mapping` [ `TaskDef`, \
122 `set` [ `~lsst.daf.butler.Quantum` ] ]
123 This maps tasks (and their configs) to the sets of data they are to
124 process.
125 metadata : Optional `~collections.abc.Mapping` of `str` to primitives
126 This is an optional parameter of extra data to carry with the graph.
127 Entries in this mapping should be able to be serialized in JSON.
128 universe : `~lsst.daf.butler.DimensionUniverse`, optional
129 The dimensions in which quanta can be defined. Need only be provided if
130 no quanta have data IDs.
131 initInputs : `~collections.abc.Mapping`, optional
132 Maps tasks to their InitInput dataset refs. Dataset refs can be either
133 resolved or non-resolved. Presently the same dataset refs are included
134 in each `~lsst.daf.butler.Quantum` for the same task.
135 initOutputs : `~collections.abc.Mapping`, optional
136 Maps tasks to their InitOutput dataset refs. Dataset refs can be either
137 resolved or non-resolved. For intermediate resolved refs their dataset
138 ID must match ``initInputs`` and Quantum ``initInputs``.
139 globalInitOutputs : iterable [ `~lsst.daf.butler.DatasetRef` ], optional
140 Dataset refs for some global objects produced by pipeline. These
141 objects include task configurations and package versions. Typically
142 they have an empty DataId, but there is no real restriction on what
143 can appear here.
144 registryDatasetTypes : iterable [ `~lsst.daf.butler.DatasetType` ], \
145 optional
146 Dataset types which are used by this graph, their definitions must
147 match registry. If registry does not define dataset type yet, then
148 it should match one that will be created later.
150 Raises
151 ------
152 ValueError
153 Raised if the graph is pruned such that some tasks no longer have nodes
154 associated with them.
155 """
157 def __init__(
158 self,
159 quanta: Mapping[TaskDef, set[Quantum]],
160 metadata: Mapping[str, Any] | None = None,
161 universe: DimensionUniverse | None = None,
162 initInputs: Mapping[TaskDef, Iterable[DatasetRef]] | None = None,
163 initOutputs: Mapping[TaskDef, Iterable[DatasetRef]] | None = None,
164 globalInitOutputs: Iterable[DatasetRef] | None = None,
165 registryDatasetTypes: Iterable[DatasetType] | None = None,
166 ):
167 self._buildGraphs(
168 quanta,
169 metadata=metadata,
170 universe=universe,
171 initInputs=initInputs,
172 initOutputs=initOutputs,
173 globalInitOutputs=globalInitOutputs,
174 registryDatasetTypes=registryDatasetTypes,
175 )
177 def _buildGraphs(
178 self,
179 quanta: Mapping[TaskDef, set[Quantum]],
180 *,
181 _quantumToNodeId: Mapping[Quantum, uuid.UUID] | None = None,
182 _buildId: BuildId | None = None,
183 metadata: Mapping[str, Any] | None = None,
184 universe: DimensionUniverse | None = None,
185 initInputs: Mapping[TaskDef, Iterable[DatasetRef]] | None = None,
186 initOutputs: Mapping[TaskDef, Iterable[DatasetRef]] | None = None,
187 globalInitOutputs: Iterable[DatasetRef] | None = None,
188 registryDatasetTypes: Iterable[DatasetType] | None = None,
189 ) -> None:
190 """Build the graph that is used to store the relation between tasks,
191 and the graph that holds the relations between quanta
192 """
193 # Save packages to metadata
194 self._metadata = dict(metadata) if metadata is not None else {}
195 self._metadata.setdefault("packages", Packages.fromSystem())
196 self._metadata.setdefault("user", getpass.getuser())
197 self._metadata.setdefault("time", f"{datetime.datetime.now()}")
198 self._metadata.setdefault("full_command", " ".join(sys.argv))
200 self._buildId = _buildId if _buildId is not None else BuildId(f"{time.time()}-{os.getpid()}")
201 # Data structure used to identify relations between
202 # DatasetTypeName -> TaskDef.
203 self._datasetDict = _DatasetTracker(createInverse=True)
205 # Temporary graph that will have dataset UUIDs (as raw bytes) and
206 # QuantumNode objects as nodes; will be collapsed down to just quanta
207 # later.
208 bipartite_graph = nx.DiGraph()
210 self._nodeIdMap: dict[uuid.UUID, QuantumNode] = {}
211 self._taskToQuantumNode: MutableMapping[TaskDef, set[QuantumNode]] = defaultdict(set)
212 for taskDef, quantumSet in quanta.items():
213 connections = taskDef.connections
215 # For each type of connection in the task, add a key to the
216 # `_DatasetTracker` for the connections name, with a value of
217 # the TaskDef in the appropriate field
218 for inpt in iterConnections(connections, ("inputs", "prerequisiteInputs", "initInputs")):
219 # Have to handle components in inputs.
220 dataset_name, _, _ = inpt.name.partition(".")
221 self._datasetDict.addConsumer(DatasetTypeName(dataset_name), taskDef)
223 for output in iterConnections(connections, ("outputs",)):
224 # Have to handle possible components in outputs.
225 dataset_name, _, _ = output.name.partition(".")
226 self._datasetDict.addProducer(DatasetTypeName(dataset_name), taskDef)
228 # For each `Quantum` in the set of all `Quantum` for this task,
229 # add a key to the `_DatasetTracker` that is a `DatasetRef` for one
230 # of the individual datasets inside the `Quantum`, with a value of
231 # a newly created QuantumNode to the appropriate input/output
232 # field.
233 for quantum in quantumSet:
234 if quantum.dataId is not None:
235 if universe is None:
236 universe = quantum.dataId.universe
237 elif universe != quantum.dataId.universe:
238 raise RuntimeError(
239 "Mismatched dimension universes in QuantumGraph construction: "
240 f"{universe} != {quantum.dataId.universe}. "
241 )
243 if _quantumToNodeId:
244 if (nodeId := _quantumToNodeId.get(quantum)) is None:
245 raise ValueError(
246 "If _quantuMToNodeNumber is not None, all quanta must have an "
247 "associated value in the mapping"
248 )
249 else:
250 nodeId = generate_uuidv7()
252 inits = quantum.initInputs.values()
253 inputs = quantum.inputs.values()
254 value = QuantumNode(quantum, taskDef, nodeId)
255 self._taskToQuantumNode[taskDef].add(value)
256 self._nodeIdMap[nodeId] = value
258 bipartite_graph.add_node(value, bipartite=0)
259 for dsRef in chain(inits, inputs):
260 # unfortunately, `Quantum` allows inits to be individual
261 # `DatasetRef`s or an Iterable of such, so there must
262 # be an instance check here
263 if isinstance(dsRef, Iterable):
264 for sub in dsRef:
265 bipartite_graph.add_node(sub.id.bytes, bipartite=1)
266 bipartite_graph.add_edge(sub.id.bytes, value)
267 else:
268 assert isinstance(dsRef, DatasetRef)
269 if dsRef.isComponent():
270 dsRef = dsRef.makeCompositeRef()
271 bipartite_graph.add_node(dsRef.id.bytes, bipartite=1)
272 bipartite_graph.add_edge(dsRef.id.bytes, value)
273 for dsRef in chain.from_iterable(quantum.outputs.values()):
274 bipartite_graph.add_node(dsRef.id.bytes, bipartite=1)
275 bipartite_graph.add_edge(value, dsRef.id.bytes)
277 # Dimension universe
278 if universe is None:
279 raise RuntimeError(
280 "Dimension universe or at least one quantum with a data ID "
281 "must be provided when constructing a QuantumGraph."
282 )
283 self._universe = universe
285 # Make graph of quanta relations, by projecting out the dataset nodes
286 # in the bipartite_graph, leaving just the quanta.
287 self._connectedQuanta = nx.algorithms.bipartite.projected_graph(
288 bipartite_graph, self._nodeIdMap.values()
289 )
290 self._count = len(self._connectedQuanta)
292 # Graph of task relations, used in various methods
293 self._taskGraph = self._datasetDict.makeNetworkXGraph()
295 # convert default dict into a regular to prevent accidental key
296 # insertion
297 self._taskToQuantumNode = dict(self._taskToQuantumNode.items())
299 self._initInputRefs: dict[str, list[DatasetRef]] = {}
300 self._initOutputRefs: dict[str, list[DatasetRef]] = {}
301 self._globalInitOutputRefs: list[DatasetRef] = []
302 self._registryDatasetTypes: list[DatasetType] = []
303 if initInputs is not None:
304 self._initInputRefs = {taskDef.label: list(refs) for taskDef, refs in initInputs.items()}
305 if initOutputs is not None:
306 self._initOutputRefs = {taskDef.label: list(refs) for taskDef, refs in initOutputs.items()}
307 if globalInitOutputs is not None:
308 self._globalInitOutputRefs = list(globalInitOutputs)
309 if registryDatasetTypes is not None:
310 self._registryDatasetTypes = list(registryDatasetTypes)
312 # PipelineGraph is current constructed on first use.
313 # TODO DM-40442: use PipelineGraph instead of TaskDef
314 # collections.
315 self._pipeline_graph: PipelineGraph | None = None
317 @property
318 def pipeline_graph(self) -> PipelineGraph:
319 """A graph representation of the tasks and dataset types in the quantum
320 graph.
321 """
322 if self._pipeline_graph is None:
323 # Construct into a temporary for strong exception safety.
324 pipeline_graph = PipelineGraph()
325 for task_def in self._taskToQuantumNode.keys():
326 pipeline_graph.add_task(
327 task_def.label, task_def.taskClass, task_def.config, connections=task_def.connections
328 )
329 dataset_types = {dataset_type.name: dataset_type for dataset_type in self._registryDatasetTypes}
330 pipeline_graph.resolve(dimensions=self._universe, dataset_types=dataset_types)
331 self._pipeline_graph = pipeline_graph
332 return self._pipeline_graph
334 def get_task_quanta(self, label: str) -> Mapping[uuid.UUID, Quantum]:
335 """Return the quanta associated with the given task label.
337 Parameters
338 ----------
339 label : `str`
340 Task label.
342 Returns
343 -------
344 quanta : `~collections.abc.Mapping` [ uuid.UUID, `Quantum` ]
345 Mapping from quantum ID to quantum. Empty if ``label`` does not
346 correspond to a task in this graph.
347 """
348 task_def = self.findTaskDefByLabel(label)
349 if not task_def:
350 return {}
351 return {node.nodeId: node.quantum for node in self.getNodesForTask(task_def)}
353 @property
354 def taskGraph(self) -> nx.DiGraph:
355 """A graph representing the relations between the tasks inside
356 the quantum graph (`networkx.DiGraph`).
357 """
358 return self._taskGraph
360 @property
361 def graph(self) -> nx.DiGraph:
362 """A graph representing the relations between all the `QuantumNode`
363 objects (`networkx.DiGraph`).
365 The graph should usually be iterated over, or passed to methods of this
366 class, but sometimes direct access to the ``networkx`` object may be
367 helpful.
368 """
369 return self._connectedQuanta
371 @property
372 def inputQuanta(self) -> Iterable[QuantumNode]:
373 """The nodes that are inputs to the graph (iterable [`QuantumNode`]).
375 These are the nodes that do not depend on any other nodes in the
376 graph.
377 """
378 return (q for q, n in self._connectedQuanta.in_degree if n == 0)
380 @property
381 def outputQuanta(self) -> Iterable[QuantumNode]:
382 """The nodes that are outputs of the graph (iterable [`QuantumNode`]).
384 These are the nodes that have no nodes that depend on them in the
385 graph.
386 """
387 return [q for q, n in self._connectedQuanta.out_degree if n == 0]
389 @property
390 def allDatasetTypes(self) -> tuple[DatasetTypeName, ...]:
391 """All the data set type names that are present in the graph
392 (`tuple` [`str`]).
394 These types do not include global init-outputs.
395 """
396 return tuple(self._datasetDict.keys())
398 @property
399 def isConnected(self) -> bool:
400 """Whether all of the nodes in the graph are connected, ignoring
401 directionality of connections (`bool`).
402 """
403 return nx.is_weakly_connected(self._connectedQuanta)
405 def getQuantumNodeByNodeId(self, nodeId: uuid.UUID) -> QuantumNode:
406 """Lookup a `QuantumNode` from an id associated with the node.
408 Parameters
409 ----------
410 nodeId : `NodeId`
411 The number associated with a node.
413 Returns
414 -------
415 node : `QuantumNode`
416 The node corresponding with input number.
418 Raises
419 ------
420 KeyError
421 Raised if the requested nodeId is not in the graph.
422 """
423 return self._nodeIdMap[nodeId]
425 def getQuantaForTask(self, taskDef: TaskDef) -> frozenset[Quantum]:
426 """Return all the `~lsst.daf.butler.Quantum` associated with a
427 `TaskDef`.
429 Parameters
430 ----------
431 taskDef : `TaskDef`
432 The `TaskDef` for which `~lsst.daf.butler.Quantum` are to be
433 queried.
435 Returns
436 -------
437 quanta : `frozenset` of `~lsst.daf.butler.Quantum`
438 The `set` of `~lsst.daf.butler.Quantum` that is associated with the
439 specified `TaskDef`.
440 """
441 return frozenset(node.quantum for node in self._taskToQuantumNode.get(taskDef, ()))
443 def getNumberOfQuantaForTask(self, taskDef: TaskDef) -> int:
444 """Return the number of `~lsst.daf.butler.Quantum` associated with
445 a `TaskDef`.
447 Parameters
448 ----------
449 taskDef : `TaskDef`
450 The `TaskDef` for which `~lsst.daf.butler.Quantum` are to be
451 queried.
453 Returns
454 -------
455 count : `int`
456 The number of `~lsst.daf.butler.Quantum` that are associated with
457 the specified `TaskDef`.
458 """
459 return len(self._taskToQuantumNode.get(taskDef, ()))
461 def getNodesForTask(self, taskDef: TaskDef) -> frozenset[QuantumNode]:
462 r"""Return all the `QuantumNode`\s associated with a `TaskDef`.
464 Parameters
465 ----------
466 taskDef : `TaskDef`
467 The `TaskDef` for which `~lsst.daf.butler.Quantum` are to be
468 queried.
470 Returns
471 -------
472 nodes : `frozenset` [ `QuantumNode` ]
473 A `frozenset` of `QuantumNode` that is associated with the
474 specified `TaskDef`.
475 """
476 return frozenset(self._taskToQuantumNode[taskDef])
478 def findTasksWithInput(self, datasetTypeName: DatasetTypeName) -> Iterable[TaskDef]:
479 """Find all tasks that have the specified dataset type name as an
480 input.
482 Parameters
483 ----------
484 datasetTypeName : `str`
485 A string representing the name of a dataset type to be queried,
486 can also accept a `DatasetTypeName` which is a `~typing.NewType` of
487 `str` for type safety in static type checking.
489 Returns
490 -------
491 tasks : iterable of `TaskDef`
492 `TaskDef` objects that have the specified `DatasetTypeName` as an
493 input, list will be empty if no tasks use specified
494 `DatasetTypeName` as an input.
496 Raises
497 ------
498 KeyError
499 Raised if the `DatasetTypeName` is not part of the `QuantumGraph`.
500 """
501 return (c for c in self._datasetDict.getConsumers(datasetTypeName))
503 def findTaskWithOutput(self, datasetTypeName: DatasetTypeName) -> TaskDef | None:
504 """Find all tasks that have the specified dataset type name as an
505 output.
507 Parameters
508 ----------
509 datasetTypeName : `str`
510 A string representing the name of a dataset type to be queried,
511 can also accept a `DatasetTypeName` which is a `~typing.NewType` of
512 `str` for type safety in static type checking.
514 Returns
515 -------
516 result : `TaskDef` or `None`
517 `TaskDef` that outputs `DatasetTypeName` as an output or `None` if
518 none of the tasks produce this `DatasetTypeName`.
520 Raises
521 ------
522 KeyError
523 Raised if the `DatasetTypeName` is not part of the `QuantumGraph`.
524 """
525 return self._datasetDict.getProducer(datasetTypeName)
527 def tasksWithDSType(self, datasetTypeName: DatasetTypeName) -> Iterable[TaskDef]:
528 """Find all tasks that are associated with the specified dataset type
529 name.
531 Parameters
532 ----------
533 datasetTypeName : `str`
534 A string representing the name of a dataset type to be queried,
535 can also accept a `DatasetTypeName` which is a `~typing.NewType` of
536 `str` for type safety in static type checking.
538 Returns
539 -------
540 result : iterable of `TaskDef`
541 `TaskDef` objects that are associated with the specified
542 `DatasetTypeName`.
544 Raises
545 ------
546 KeyError
547 Raised if the `DatasetTypeName` is not part of the `QuantumGraph`.
548 """
549 return self._datasetDict.getAll(datasetTypeName)
551 def findTaskDefByName(self, taskName: str) -> list[TaskDef]:
552 """Determine which `TaskDef` objects in this graph are associated
553 with a `str` representing a task name (looks at the ``taskName``
554 property of `TaskDef` objects).
556 Returns a list of `TaskDef` objects as a `PipelineTask` may appear
557 multiple times in a graph with different labels.
559 Parameters
560 ----------
561 taskName : `str`
562 Name of a task to search for.
564 Returns
565 -------
566 result : `list` of `TaskDef`
567 List of the `TaskDef` objects that have the name specified.
568 Multiple values are returned in the case that a task is used
569 multiple times with different labels.
570 """
571 results = []
572 for task in self._taskToQuantumNode:
573 split = task.taskName.split(".")
574 if split[-1] == taskName:
575 results.append(task)
576 return results
578 def findTaskDefByLabel(self, label: str) -> TaskDef | None:
579 """Determine which `TaskDef` objects in this graph are associated
580 with a `str` representing a tasks label.
582 Parameters
583 ----------
584 label : `str`
585 Name of a task to search for.
587 Returns
588 -------
589 result : `TaskDef`
590 `TaskDef` objects that has the specified label.
591 """
592 for task in self._taskToQuantumNode:
593 if label == task.label:
594 return task
595 return None
597 def findQuantaWithDSType(self, datasetTypeName: DatasetTypeName) -> set[Quantum]:
598 r"""Return all the `~lsst.daf.butler.Quantum` that contain a specified
599 `DatasetTypeName`.
601 Parameters
602 ----------
603 datasetTypeName : `str`
604 The name of the dataset type to search for as a string,
605 can also accept a `DatasetTypeName` which is a `~typing.NewType` of
606 `str` for type safety in static type checking.
608 Returns
609 -------
610 result : `set` of `QuantumNode` objects
611 A `set` of `QuantumNode`\s that contain specified
612 `DatasetTypeName`.
614 Raises
615 ------
616 KeyError
617 Raised if the `DatasetTypeName` is not part of the `QuantumGraph`.
618 """
619 tasks = self._datasetDict.getAll(datasetTypeName)
620 result: set[Quantum] = set()
621 result = result.union(quantum for task in tasks for quantum in self.getQuantaForTask(task))
622 return result
624 def checkQuantumInGraph(self, quantum: Quantum) -> bool:
625 """Check if specified quantum appears in the graph as part of a node.
627 Parameters
628 ----------
629 quantum : `lsst.daf.butler.Quantum`
630 The quantum to search for.
632 Returns
633 -------
634 in_graph : `bool`
635 The result of searching for the quantum.
636 """
637 return any(quantum == node.quantum for node in self)
639 def writeDotGraph(self, output: str | io.BufferedIOBase) -> None:
640 """Write out the graph as a dot graph.
642 Parameters
643 ----------
644 output : `str` or `io.BufferedIOBase`
645 Either a filesystem path to write to, or a file handle object.
646 """
647 write_dot(self._connectedQuanta, output)
649 def subset(self: _T, nodes: QuantumNode | Iterable[QuantumNode]) -> _T:
650 """Create a new graph object that contains the subset of the nodes
651 specified as input. Node number is preserved.
653 Parameters
654 ----------
655 nodes : `QuantumNode` or iterable of `QuantumNode`
656 Nodes from which to create subset.
658 Returns
659 -------
660 graph : instance of graph type
661 An instance of the type from which the subset was created.
662 """
663 if not isinstance(nodes, Iterable):
664 nodes = (nodes,)
665 quantumSubgraph = self._connectedQuanta.subgraph(nodes).nodes
666 quantumMap = defaultdict(set)
668 dataset_type_names: set[str] = set()
669 node: QuantumNode
670 for node in quantumSubgraph:
671 quantumMap[node.taskDef].add(node.quantum)
672 dataset_type_names.update(
673 dstype.name
674 for dstype in chain(
675 node.quantum.inputs.keys(), node.quantum.outputs.keys(), node.quantum.initInputs.keys()
676 )
677 )
679 # May need to trim dataset types from registryDatasetTypes.
680 for taskDef in quantumMap:
681 if refs := self.initOutputRefs(taskDef):
682 dataset_type_names.update(ref.datasetType.name for ref in refs)
683 dataset_type_names.update(ref.datasetType.name for ref in self._globalInitOutputRefs)
684 registryDatasetTypes = [
685 dstype for dstype in self._registryDatasetTypes if dstype.name in dataset_type_names
686 ]
688 # convert to standard dict to prevent accidental key insertion
689 quantumDict: dict[TaskDef, set[Quantum]] = dict(quantumMap.items())
690 # Create an empty graph, and then populate it with custom mapping
691 newInst = type(self)({}, universe=self._universe)
692 # TODO: Do we need to copy initInputs/initOutputs?
693 newInst._buildGraphs(
694 quantumDict,
695 _quantumToNodeId={n.quantum: n.nodeId for n in nodes},
696 _buildId=self._buildId,
697 metadata=self._metadata,
698 universe=self._universe,
699 globalInitOutputs=self._globalInitOutputRefs,
700 registryDatasetTypes=registryDatasetTypes,
701 )
702 return newInst
704 def subsetToConnected(self: _T) -> tuple[_T, ...]:
705 """Generate a list of subgraphs where each is connected.
707 Returns
708 -------
709 result : `list` of `QuantumGraph`
710 A list of graphs that are each connected.
711 """
712 return tuple(
713 self.subset(connectedSet)
714 for connectedSet in nx.weakly_connected_components(self._connectedQuanta)
715 )
717 def determineInputsToQuantumNode(self, node: QuantumNode) -> set[QuantumNode]:
718 """Return a set of `QuantumNode` that are direct inputs to a specified
719 node.
721 Parameters
722 ----------
723 node : `QuantumNode`
724 The node of the graph for which inputs are to be determined.
726 Returns
727 -------
728 inputs : `set` of `QuantumNode`
729 All the nodes that are direct inputs to specified node.
730 """
731 return set(self._connectedQuanta.predecessors(node))
733 def determineOutputsOfQuantumNode(self, node: QuantumNode) -> set[QuantumNode]:
734 """Return a set of `QuantumNode` that are direct outputs of a specified
735 node.
737 Parameters
738 ----------
739 node : `QuantumNode`
740 The node of the graph for which outputs are to be determined.
742 Returns
743 -------
744 outputs : `set` of `QuantumNode`
745 All the nodes that are direct outputs to specified node.
746 """
747 return set(self._connectedQuanta.successors(node))
749 def determineConnectionsOfQuantumNode(self: _T, node: QuantumNode) -> _T:
750 """Return a graph of `QuantumNode` that are direct inputs and outputs
751 of a specified node.
753 Parameters
754 ----------
755 node : `QuantumNode`
756 The node of the graph for which connected nodes are to be
757 determined.
759 Returns
760 -------
761 graph : graph of `QuantumNode`
762 All the nodes that are directly connected to specified node.
763 """
764 nodes = self.determineInputsToQuantumNode(node).union(self.determineOutputsOfQuantumNode(node))
765 nodes.add(node)
766 return self.subset(nodes)
768 def determineAncestorsOfQuantumNode(self: _T, node: QuantumNode) -> _T:
769 """Return a graph of the specified node and all the ancestor nodes
770 directly reachable by walking edges.
772 Parameters
773 ----------
774 node : `QuantumNode`
775 The node for which all ancestors are to be determined.
777 Returns
778 -------
779 ancestors : graph of `QuantumNode`
780 Graph of node and all of its ancestors.
781 """
782 predecessorNodes = nx.ancestors(self._connectedQuanta, node)
783 predecessorNodes.add(node)
784 return self.subset(predecessorNodes)
786 def findCycle(self) -> list[tuple[QuantumNode, QuantumNode]]:
787 """Check a graph for the presense of cycles and returns the edges of
788 any cycles found, or an empty list if there is no cycle.
790 Returns
791 -------
792 result : `list` of `tuple` of [ `QuantumNode`, `QuantumNode` ]
793 A list of any graph edges that form a cycle, or an empty list if
794 there is no cycle. Empty list to so support if graph.find_cycle()
795 syntax as an empty list is falsy.
796 """
797 try:
798 return nx.find_cycle(self._connectedQuanta)
799 except nx.NetworkXNoCycle:
800 return []
802 def saveUri(self, uri: ResourcePathExpression) -> None:
803 """Save `QuantumGraph` to the specified URI.
805 Parameters
806 ----------
807 uri : convertible to `~lsst.resources.ResourcePath`
808 URI to where the graph should be saved.
809 """
810 path = ResourcePath(uri)
811 match path.getExtension():
812 case ".qgraph":
813 buffer = self._buildSaveObject()
814 path.write(buffer) # type: ignore # Ignore because bytearray is safe to use in place of bytes
815 case ".qg":
816 from ..quantum_graph import PredictedQuantumGraphComponents
818 pqg = PredictedQuantumGraphComponents.from_old_quantum_graph(self)
819 pqg.write(path)
820 case ext:
821 raise TypeError(f"Can currently only save a graph in .qgraph or .qg format, not {ext!r}.")
823 @property
824 def metadata(self) -> MappingProxyType[str, Any]:
825 """Extra data carried with the graph (mapping [`str`] or `None`).
827 The mapping is a dynamic view of this object's metadata. Values should
828 be able to be serialized in JSON.
829 """
830 return MappingProxyType(self._metadata)
832 def get_init_input_refs(self, task_label: str) -> list[DatasetRef]:
833 """Return the DatasetRefs for the given task's init inputs.
835 Parameters
836 ----------
837 task_label : `str`
838 Label of the task.
840 Returns
841 -------
842 refs : `list` [ `lsst.daf.butler.DatasetRef` ]
843 Dataset references. Guaranteed to be a new list, not internal
844 state.
845 """
846 return list(self._initInputRefs.get(task_label, ()))
848 def get_init_output_refs(self, task_label: str) -> list[DatasetRef]:
849 """Return the DatasetRefs for the given task's init outputs.
851 Parameters
852 ----------
853 task_label : `str`
854 Label of the task.
856 Returns
857 -------
858 refs : `list` [ `lsst.daf.butler.DatasetRef` ]
859 Dataset references. Guaranteed to be a new list, not internal
860 state.
861 """
862 return list(self._initOutputRefs.get(task_label, ()))
864 def initInputRefs(self, taskDef: TaskDef) -> list[DatasetRef] | None:
865 """Return DatasetRefs for a given task InitInputs.
867 Parameters
868 ----------
869 taskDef : `TaskDef`
870 Task definition structure.
872 Returns
873 -------
874 refs : `list` [ `~lsst.daf.butler.DatasetRef` ] or `None`
875 DatasetRef for the task InitInput, can be `None`. This can return
876 either resolved or non-resolved reference.
877 """
878 return self._initInputRefs.get(taskDef.label)
880 def initOutputRefs(self, taskDef: TaskDef) -> list[DatasetRef] | None:
881 """Return DatasetRefs for a given task InitOutputs.
883 Parameters
884 ----------
885 taskDef : `TaskDef`
886 Task definition structure.
888 Returns
889 -------
890 refs : `list` [ `~lsst.daf.butler.DatasetRef` ] or `None`
891 DatasetRefs for the task InitOutput, can be `None`. This can return
892 either resolved or non-resolved reference. Resolved reference will
893 match Quantum's initInputs if this is an intermediate dataset type.
894 """
895 return self._initOutputRefs.get(taskDef.label)
897 def globalInitOutputRefs(self) -> list[DatasetRef]:
898 """Return DatasetRefs for global InitOutputs.
900 Returns
901 -------
902 refs : `list` [ `~lsst.daf.butler.DatasetRef` ]
903 DatasetRefs for global InitOutputs.
904 """
905 return self._globalInitOutputRefs
907 def registryDatasetTypes(self) -> list[DatasetType]:
908 """Return dataset types used by this graph, their definitions match
909 dataset types from registry.
911 Returns
912 -------
913 refs : `list` [ `~lsst.daf.butler.DatasetType` ]
914 Dataset types for this graph.
915 """
916 return self._registryDatasetTypes
918 @classmethod
919 def loadUri(
920 cls,
921 uri: ResourcePathExpression,
922 universe: DimensionUniverse | None = None,
923 nodes: Iterable[uuid.UUID | str] | None = None,
924 graphID: BuildId | None = None,
925 minimumVersion: int = 3,
926 ) -> QuantumGraph:
927 """Read `QuantumGraph` from a URI.
929 Parameters
930 ----------
931 uri : convertible to `~lsst.resources.ResourcePath`
932 URI from where to load the graph.
933 universe : `~lsst.daf.butler.DimensionUniverse`, optional
934 If `None` it is loaded from the `QuantumGraph`
935 saved structure. If supplied, the
936 `~lsst.daf.butler.DimensionUniverse` from the loaded `QuantumGraph`
937 will be validated against the supplied argument for compatibility.
938 nodes : iterable of [ `uuid.UUID` | `str` ] or `None`
939 UUIDs that correspond to nodes in the graph. If specified, only
940 these nodes will be loaded. Defaults to None, in which case all
941 nodes will be loaded.
942 graphID : `str` or `None`
943 If specified this ID is verified against the loaded graph prior to
944 loading any Nodes. This defaults to None in which case no
945 validation is done.
946 minimumVersion : `int`
947 Minimum version of a save file to load. Set to -1 to load all
948 versions. Older versions may need to be loaded, and re-saved
949 to upgrade them to the latest format before they can be used in
950 production.
952 Returns
953 -------
954 graph : `QuantumGraph`
955 Resulting QuantumGraph instance.
957 Raises
958 ------
959 TypeError
960 Raised if file contains instance of a type other than
961 `QuantumGraph`.
962 ValueError
963 Raised if one or more of the nodes requested is not in the
964 `QuantumGraph` or if graphID parameter does not match the graph
965 being loaded or if the supplied uri does not point at a valid
966 `QuantumGraph` save file.
967 RuntimeError
968 Raise if Supplied `~lsst.daf.butler.DimensionUniverse` is not
969 compatible with the `~lsst.daf.butler.DimensionUniverse` saved in
970 the graph.
971 """
972 uri = ResourcePath(uri)
973 match uri.getExtension():
974 case ".qgraph":
975 with LoadHelper(uri, minimumVersion, fullRead=(nodes is None)) as loader:
976 qgraph = loader.load(universe, nodes, graphID)
977 case ".qg":
978 from ..quantum_graph import PredictedQuantumGraphReader
980 with PredictedQuantumGraphReader.open(uri, page_size=100000) as qgr:
981 quantum_ids = (
982 [uuid.UUID(q) if not isinstance(q, uuid.UUID) else q for q in nodes]
983 if nodes is not None
984 else None
985 )
986 qgr.read_execution_quanta(quantum_ids)
987 qgraph = qgr.finish().to_old_quantum_graph()
988 case _:
989 raise ValueError(f"Only know how to handle files saved as `.qgraph`, not {uri}")
990 if not isinstance(qgraph, QuantumGraph):
991 raise TypeError(f"QuantumGraph file {uri} contains unexpected object type: {type(qgraph)}")
992 return qgraph
994 @classmethod
995 def readHeader(cls, uri: ResourcePathExpression, minimumVersion: int = 3) -> str | None:
996 """Read the header of a `QuantumGraph` pointed to by the uri parameter
997 and return it as a string.
999 Parameters
1000 ----------
1001 uri : convertible to `~lsst.resources.ResourcePath`
1002 The location of the `QuantumGraph` to load. If the argument is a
1003 string, it must correspond to a valid
1004 `~lsst.resources.ResourcePath` path.
1005 minimumVersion : `int`
1006 Minimum version of a save file to load. Set to -1 to load all
1007 versions. Older versions may need to be loaded, and re-saved
1008 to upgrade them to the latest format before they can be used in
1009 production.
1011 Returns
1012 -------
1013 header : `str` or `None`
1014 The header associated with the specified `QuantumGraph` it there is
1015 one, else `None`.
1017 Raises
1018 ------
1019 ValueError
1020 Raised if the extension of the file specified by uri is not a
1021 `QuantumGraph` extension.
1022 """
1023 uri = ResourcePath(uri)
1024 if uri.getExtension() in {".qgraph"}:
1025 return LoadHelper(uri, minimumVersion).readHeader()
1026 else:
1027 raise ValueError("Only know how to handle files saved as `.qgraph`")
1029 def buildAndPrintHeader(self) -> None:
1030 """Create a header that would be used in a save of this object and
1031 prints it out to standard out.
1032 """
1033 _, header = self._buildSaveObject(returnHeader=True)
1034 print(json.dumps(header))
1036 def save(self, file: BinaryIO) -> None:
1037 """Save QuantumGraph to a file.
1039 Parameters
1040 ----------
1041 file : `io.BufferedIOBase`
1042 File to write data open in binary mode.
1043 """
1044 buffer = self._buildSaveObject()
1045 file.write(buffer) # type: ignore # Ignore because bytearray is safe to use in place of bytes
1047 def _buildSaveObject(self, returnHeader: bool = False) -> bytearray | tuple[bytearray, dict]:
1048 thing = PersistenceContextVars()
1049 result = thing.run(self._buildSaveObjectImpl, returnHeader)
1050 return result
1052 def _buildSaveObjectImpl(self, returnHeader: bool = False) -> bytearray | tuple[bytearray, dict]:
1053 # make some containers
1054 jsonData: deque[bytes] = deque()
1055 # node map is a list because json does not accept mapping keys that
1056 # are not strings, so we store a list of key, value pairs that will
1057 # be converted to a mapping on load
1058 nodeMap = []
1059 taskDefMap = {}
1060 headerData: dict[str, Any] = {}
1062 # Store the QuantumGraph BuildId, this will allow validating BuildIds
1063 # at load time, prior to loading any QuantumNodes. Name chosen for
1064 # unlikely conflicts.
1065 headerData["GraphBuildID"] = self.graphID
1066 headerData["Metadata"] = self._metadata
1068 # Store the universe this graph was created with
1069 universeConfig = self._universe.dimensionConfig
1070 headerData["universe"] = universeConfig.toDict()
1072 # counter for the number of bytes processed thus far
1073 count = 0
1074 # serialize out the task Defs recording the start and end bytes of each
1075 # taskDef
1076 inverseLookup = self._datasetDict.inverse
1077 taskDef: TaskDef
1078 # sort by task label to ensure serialization happens in the same order
1079 for taskDef in self.taskGraph:
1080 # compressing has very little impact on saving or load time, but
1081 # a large impact on on disk size, so it is worth doing
1082 taskDescription: dict[str, Any] = {}
1083 # save the fully qualified name.
1084 taskDescription["taskName"] = get_full_type_name(taskDef.taskClass)
1085 # save the config as a text stream that will be un-persisted on the
1086 # other end
1087 stream = io.StringIO()
1088 taskDef.config.saveToStream(stream)
1089 taskDescription["config"] = stream.getvalue()
1090 taskDescription["label"] = taskDef.label
1091 if (refs := self._initInputRefs.get(taskDef.label)) is not None:
1092 taskDescription["initInputRefs"] = [ref.to_json() for ref in refs]
1093 if (refs := self._initOutputRefs.get(taskDef.label)) is not None:
1094 taskDescription["initOutputRefs"] = [ref.to_json() for ref in refs]
1096 inputs = []
1097 outputs = []
1099 # Determine the connection between all of tasks and save that in
1100 # the header as a list of connections and edges in each task
1101 # this will help in un-persisting, and possibly in a "quick view"
1102 # method that does not require everything to be un-persisted
1103 #
1104 # Typing returns can't be parameter dependent
1105 for connection in inverseLookup[taskDef]: # type: ignore
1106 consumers = self._datasetDict.getConsumers(connection)
1107 producer = self._datasetDict.getProducer(connection)
1108 if taskDef in consumers:
1109 # This checks if the task consumes the connection directly
1110 # from the datastore or it is produced by another task
1111 producerLabel = producer.label if producer is not None else "datastore"
1112 inputs.append((producerLabel, connection))
1113 elif taskDef not in consumers and producer is taskDef:
1114 # If there are no consumers for this tasks produced
1115 # connection, the output will be said to be the datastore
1116 # in which case the for loop will be a zero length loop
1117 if not consumers:
1118 outputs.append(("datastore", connection))
1119 for td in consumers:
1120 outputs.append((td.label, connection))
1122 # dump to json string, and encode that string to bytes and then
1123 # conpress those bytes
1124 dump = lzma.compress(json.dumps(taskDescription).encode(), preset=2)
1125 # record the sizing and relation information
1126 taskDefMap[taskDef.label] = {
1127 "bytes": (count, count + len(dump)),
1128 "inputs": inputs,
1129 "outputs": outputs,
1130 }
1131 count += len(dump)
1132 jsonData.append(dump)
1134 headerData["TaskDefs"] = taskDefMap
1136 # serialize the nodes, recording the start and end bytes of each node
1137 dimAccumulator = DimensionRecordsAccumulator()
1138 for node in self:
1139 # compressing has very little impact on saving or load time, but
1140 # a large impact on on disk size, so it is worth doing
1141 simpleNode = node.to_simple(accumulator=dimAccumulator)
1143 dump = lzma.compress(simpleNode.model_dump_json().encode(), preset=2)
1144 jsonData.append(dump)
1145 nodeMap.append(
1146 (
1147 str(node.nodeId),
1148 {
1149 "bytes": (count, count + len(dump)),
1150 "inputs": [str(n.nodeId) for n in self.determineInputsToQuantumNode(node)],
1151 "outputs": [str(n.nodeId) for n in self.determineOutputsOfQuantumNode(node)],
1152 },
1153 )
1154 )
1155 count += len(dump)
1157 headerData["DimensionRecords"] = {
1158 key: value.model_dump()
1159 for key, value in dimAccumulator.makeSerializedDimensionRecordMapping().items()
1160 }
1162 # need to serialize this as a series of key,value tuples because of
1163 # a limitation on how json cant do anything but strings as keys
1164 headerData["Nodes"] = nodeMap
1166 if self._globalInitOutputRefs:
1167 headerData["GlobalInitOutputRefs"] = [ref.to_json() for ref in self._globalInitOutputRefs]
1169 if self._registryDatasetTypes:
1170 headerData["RegistryDatasetTypes"] = [dstype.to_json() for dstype in self._registryDatasetTypes]
1172 # dump the headerData to json
1173 header_encode = lzma.compress(json.dumps(headerData).encode())
1175 # record the sizes as 2 unsigned long long numbers for a total of 16
1176 # bytes
1177 save_bytes = struct.pack(STRUCT_FMT_BASE, SAVE_VERSION)
1179 fmt_string = DESERIALIZER_MAP[SAVE_VERSION].FMT_STRING()
1180 map_lengths = struct.pack(fmt_string, len(header_encode))
1182 # write each component of the save out in a deterministic order
1183 buffer = bytearray()
1184 buffer.extend(MAGIC_BYTES)
1185 buffer.extend(save_bytes)
1186 buffer.extend(map_lengths)
1187 buffer.extend(header_encode)
1188 # Iterate over the length of jsonData, and for each element pop the
1189 # leftmost element off the deque and write it out. This is to save
1190 # memory, as the memory is added to the buffer object, it is removed
1191 # from from the container.
1192 #
1193 # Only this section needs to worry about memory pressure because
1194 # everything else written to the buffer prior to this data is
1195 # only on the order of kilobytes to low numbers of megabytes.
1196 while jsonData:
1197 buffer.extend(jsonData.popleft())
1198 if returnHeader:
1199 return buffer, headerData
1200 else:
1201 return buffer
1203 @classmethod
1204 def load(
1205 cls,
1206 file: BinaryIO,
1207 universe: DimensionUniverse | None = None,
1208 nodes: Iterable[uuid.UUID] | None = None,
1209 graphID: BuildId | None = None,
1210 minimumVersion: int = 3,
1211 ) -> QuantumGraph:
1212 """Read `QuantumGraph` from a file that was made by `save`.
1214 Parameters
1215 ----------
1216 file : `io.IO` of bytes
1217 File with data open in binary mode.
1218 universe : `~lsst.daf.butler.DimensionUniverse`, optional
1219 If `None` it is loaded from the `QuantumGraph`
1220 saved structure. If supplied, the
1221 `~lsst.daf.butler.DimensionUniverse` from the loaded `QuantumGraph`
1222 will be validated against the supplied argument for compatibility.
1223 nodes : iterable of `uuid.UUID` or `None`
1224 UUIDs that correspond to nodes in the graph. If specified, only
1225 these nodes will be loaded. Defaults to None, in which case all
1226 nodes will be loaded.
1227 graphID : `str` or `None`
1228 If specified this ID is verified against the loaded graph prior to
1229 loading any Nodes. This defaults to None in which case no
1230 validation is done.
1231 minimumVersion : `int`
1232 Minimum version of a save file to load. Set to -1 to load all
1233 versions. Older versions may need to be loaded, and re-saved
1234 to upgrade them to the latest format before they can be used in
1235 production.
1237 Returns
1238 -------
1239 graph : `QuantumGraph`
1240 Resulting QuantumGraph instance.
1242 Raises
1243 ------
1244 TypeError
1245 Raised if data contains instance of a type other than
1246 `QuantumGraph`.
1247 ValueError
1248 Raised if one or more of the nodes requested is not in the
1249 `QuantumGraph` or if graphID parameter does not match the graph
1250 being loaded or if the supplied uri does not point at a valid
1251 `QuantumGraph` save file.
1252 """
1253 with LoadHelper(file, minimumVersion, fullRead=(nodes is None)) as loader:
1254 qgraph = loader.load(universe, nodes, graphID)
1255 if not isinstance(qgraph, QuantumGraph):
1256 raise TypeError(f"QuantumGraph file contains unexpected object type: {type(qgraph)}")
1257 return qgraph
1259 def iterTaskGraph(self) -> Generator[TaskDef, None, None]:
1260 """Iterate over the `taskGraph` attribute in topological order.
1262 Yields
1263 ------
1264 taskDef : `TaskDef`
1265 `TaskDef` objects in topological order.
1266 """
1267 yield from nx.topological_sort(self.taskGraph)
1269 def updateRun(self, run: str, *, metadata_key: str | None = None, update_graph_id: bool = False) -> None:
1270 """Change output run and dataset ID for each output dataset.
1272 Parameters
1273 ----------
1274 run : `str`
1275 New output run name.
1276 metadata_key : `str` or `None`
1277 Specifies matadata key corresponding to output run name to update
1278 with new run name. If `None` or if metadata is missing it is not
1279 updated. If metadata is present but key is missing, it will be
1280 added.
1281 update_graph_id : `bool`, optional
1282 If `True` then also update graph ID with a new unique value.
1283 """
1284 dataset_id_map: dict[DatasetId, DatasetId] = {}
1286 def _update_output_refs(
1287 refs: Iterable[DatasetRef], run: str, dataset_id_map: MutableMapping[DatasetId, DatasetId]
1288 ) -> Iterator[DatasetRef]:
1289 """Update a collection of `~lsst.daf.butler.DatasetRef` with new
1290 run and dataset IDs.
1291 """
1292 for ref in refs:
1293 new_ref = ref.replace(run=run)
1294 dataset_id_map[ref.id] = new_ref.id
1295 yield new_ref
1297 def _update_intermediate_refs(
1298 refs: Iterable[DatasetRef], run: str, dataset_id_map: Mapping[DatasetId, DatasetId]
1299 ) -> Iterator[DatasetRef]:
1300 """Update intermediate references with new run and IDs. Only the
1301 references that appear in ``dataset_id_map`` are updated, others
1302 are returned unchanged.
1303 """
1304 for ref in refs:
1305 if dataset_id := dataset_id_map.get(ref.id):
1306 ref = ref.replace(run=run, id=dataset_id)
1307 yield ref
1309 # Replace quantum output refs first.
1310 for node in self._connectedQuanta:
1311 quantum = node.quantum
1312 outputs = {
1313 dataset_type: tuple(_update_output_refs(refs, run, dataset_id_map))
1314 for dataset_type, refs in quantum.outputs.items()
1315 }
1316 updated_quantum = Quantum(
1317 taskName=quantum.taskName,
1318 dataId=quantum.dataId,
1319 initInputs=quantum.initInputs,
1320 inputs=quantum.inputs,
1321 outputs=outputs,
1322 datastore_records=quantum.datastore_records,
1323 )
1324 node._replace_quantum(updated_quantum)
1326 self._initOutputRefs = {
1327 task_def: list(_update_output_refs(refs, run, dataset_id_map))
1328 for task_def, refs in self._initOutputRefs.items()
1329 }
1330 self._globalInitOutputRefs = list(
1331 _update_output_refs(self._globalInitOutputRefs, run, dataset_id_map)
1332 )
1334 # Update all intermediates from their matching outputs.
1335 for node in self._connectedQuanta:
1336 quantum = node.quantum
1337 inputs = {
1338 dataset_type: tuple(_update_intermediate_refs(refs, run, dataset_id_map))
1339 for dataset_type, refs in quantum.inputs.items()
1340 }
1341 initInputs = list(_update_intermediate_refs(quantum.initInputs.values(), run, dataset_id_map))
1343 updated_quantum = Quantum(
1344 taskName=quantum.taskName,
1345 dataId=quantum.dataId,
1346 initInputs=initInputs,
1347 inputs=inputs,
1348 outputs=quantum.outputs,
1349 datastore_records=quantum.datastore_records,
1350 )
1351 node._replace_quantum(updated_quantum)
1353 self._initInputRefs = {
1354 task_def: list(_update_intermediate_refs(refs, run, dataset_id_map))
1355 for task_def, refs in self._initInputRefs.items()
1356 }
1358 if update_graph_id:
1359 self._buildId = BuildId(f"{time.time()}-{os.getpid()}")
1361 # Update run if given.
1362 if metadata_key is not None:
1363 self._metadata[metadata_key] = run
1365 @property
1366 def graphID(self) -> BuildId:
1367 """The ID generated by the graph at construction time (`str`)."""
1368 return self._buildId
1370 @property
1371 def universe(self) -> DimensionUniverse:
1372 """Dimension universe associated with this graph
1373 (`~lsst.daf.butler.DimensionUniverse`).
1374 """
1375 return self._universe
1377 def __iter__(self) -> Generator[QuantumNode, None, None]:
1378 yield from nx.topological_sort(self._connectedQuanta)
1380 def __len__(self) -> int:
1381 return self._count
1383 def __contains__(self, node: QuantumNode) -> bool:
1384 return self._connectedQuanta.has_node(node)
1386 def __getstate__(self) -> dict:
1387 """Store a compact form of the graph as a list of graph nodes, and a
1388 tuple of task labels and task configs. The full graph can be
1389 reconstructed with this information, and it preserves the ordering of
1390 the graph nodes.
1391 """
1392 universe: DimensionUniverse | None = None
1393 for node in self:
1394 dId = node.quantum.dataId
1395 if dId is None:
1396 continue
1397 universe = dId.universe
1398 return {"reduced": self._buildSaveObject(), "graphId": self._buildId, "universe": universe}
1400 def __setstate__(self, state: dict) -> None:
1401 """Reconstructs the state of the graph from the information persisted
1402 in getstate.
1403 """
1404 buffer = io.BytesIO(state["reduced"])
1405 with LoadHelper(buffer, minimumVersion=3) as loader:
1406 qgraph = loader.load(state["universe"], graphID=state["graphId"])
1408 self._metadata = qgraph._metadata
1409 self._buildId = qgraph._buildId
1410 self._datasetDict = qgraph._datasetDict
1411 self._nodeIdMap = qgraph._nodeIdMap
1412 self._count = len(qgraph)
1413 self._taskToQuantumNode = qgraph._taskToQuantumNode
1414 self._taskGraph = qgraph._taskGraph
1415 self._connectedQuanta = qgraph._connectedQuanta
1416 self._initInputRefs = qgraph._initInputRefs
1417 self._initOutputRefs = qgraph._initOutputRefs
1419 def __eq__(self, other: object) -> bool:
1420 if not isinstance(other, QuantumGraph):
1421 return False
1422 if len(self) != len(other):
1423 return False
1424 for node in self:
1425 if node not in other:
1426 return False
1427 if self.determineInputsToQuantumNode(node) != other.determineInputsToQuantumNode(node):
1428 return False
1429 if self.determineOutputsOfQuantumNode(node) != other.determineOutputsOfQuantumNode(node):
1430 return False
1431 if set(self.allDatasetTypes) != set(other.allDatasetTypes):
1432 return False
1433 return set(self.taskGraph) == set(other.taskGraph)
1435 def getSummary(self) -> QgraphSummary:
1436 """Create summary of graph.
1438 Returns
1439 -------
1440 summary : `QgraphSummary`
1441 Summary of QuantumGraph.
1442 """
1443 inCollection = self.metadata.get("input", None)
1444 if isinstance(inCollection, str):
1445 inCollection = [inCollection]
1446 summary = QgraphSummary(
1447 graphID=self.graphID,
1448 cmdLine=self.metadata.get("full_command", None),
1449 creationUTC=self.metadata.get("time", None),
1450 inputCollection=inCollection,
1451 outputCollection=self.metadata.get("output", None),
1452 outputRun=self.metadata.get("output_run", None),
1453 )
1454 for q in self:
1455 qts = summary.qgraphTaskSummaries.setdefault(
1456 q.taskDef.label, QgraphTaskSummary(taskLabel=q.taskDef.label)
1457 )
1458 qts.numQuanta += 1
1460 for k in q.quantum.inputs.keys():
1461 qts.numInputs[k.name] += 1
1463 for k in q.quantum.outputs.keys():
1464 qts.numOutputs[k.name] += 1
1466 return summary
1468 def make_init_qbb(
1469 self,
1470 butler_config: Config | ResourcePathExpression,
1471 *,
1472 config_search_paths: Iterable[str] | None = None,
1473 ) -> QuantumBackedButler:
1474 """Construct an quantum-backed butler suitable for reading and writing
1475 init input and init output datasets, respectively.
1477 This requires the full graph to have been loaded.
1479 Parameters
1480 ----------
1481 butler_config : `~lsst.daf.butler.Config` or \
1482 `~lsst.resources.ResourcePathExpression`
1483 A butler repository root, configuration filename, or configuration
1484 instance.
1485 config_search_paths : `~collections.abc.Iterable` [ `str` ], optional
1486 Additional search paths for butler configuration.
1488 Returns
1489 -------
1490 qbb : `~lsst.daf.butler.QuantumBackedButler`
1491 A limited butler that can ``get`` init-input datasets and ``put``
1492 init-output datasets.
1493 """
1494 universe = self.universe
1495 # Collect all init input/output dataset IDs.
1496 predicted_inputs: set[DatasetId] = set()
1497 predicted_outputs: set[DatasetId] = set()
1498 pipeline_graph = self.pipeline_graph
1499 for task_label in pipeline_graph.tasks:
1500 predicted_inputs.update(ref.id for ref in self.get_init_input_refs(task_label))
1501 predicted_outputs.update(ref.id for ref in self.get_init_output_refs(task_label))
1502 predicted_outputs.update(ref.id for ref in self.globalInitOutputRefs())
1503 # remove intermediates from inputs
1504 predicted_inputs -= predicted_outputs
1505 # Very inefficient way to extract datastore records from quantum graph,
1506 # we have to scan all quanta and look at their datastore records.
1507 datastore_records: dict[str, DatastoreRecordData] = {}
1508 for quantum_node in self:
1509 for store_name, records in quantum_node.quantum.datastore_records.items():
1510 subset = records.subset(predicted_inputs)
1511 if subset is not None:
1512 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset)
1514 dataset_types = {dstype.name: dstype for dstype in self.registryDatasetTypes()}
1515 # Make butler from everything.
1516 return QuantumBackedButler.from_predicted(
1517 config=butler_config,
1518 predicted_inputs=predicted_inputs,
1519 predicted_outputs=predicted_outputs,
1520 dimensions=universe,
1521 datastore_records=datastore_records,
1522 search_paths=list(config_search_paths) if config_search_paths is not None else None,
1523 dataset_types=dataset_types,
1524 )
1526 def write_init_outputs(self, butler: LimitedButler, skip_existing: bool = True) -> None:
1527 """Write the init-output datasets for all tasks in the quantum graph.
1529 Parameters
1530 ----------
1531 butler : `lsst.daf.butler.LimitedButler`
1532 A limited butler data repository client.
1533 skip_existing : `bool`, optional
1534 If `True` (default) ignore init-outputs that already exist. If
1535 `False`, raise.
1537 Raises
1538 ------
1539 lsst.daf.butler.registry.ConflictingDefinitionError
1540 Raised if an init-output dataset already exists and
1541 ``skip_existing=False``.
1542 """
1543 # Extract init-input and init-output refs from the QG.
1544 input_refs: dict[str, DatasetRef] = {}
1545 output_refs: dict[str, DatasetRef] = {}
1546 for task_node in self.pipeline_graph.tasks.values():
1547 input_refs.update(
1548 {ref.datasetType.name: ref for ref in self.get_init_input_refs(task_node.label)}
1549 )
1550 output_refs.update(
1551 {
1552 ref.datasetType.name: ref
1553 for ref in self.get_init_output_refs(task_node.label)
1554 if ref.datasetType.name != task_node.init.config_output.dataset_type_name
1555 }
1556 )
1557 for ref, is_stored in butler.stored_many(output_refs.values()).items():
1558 if is_stored:
1559 if not skip_existing:
1560 raise ConflictingDefinitionError(f"Init-output dataset {ref} already exists.")
1561 # We'll `put` whatever's left in output_refs at the end.
1562 del output_refs[ref.datasetType.name]
1563 # Instantiate tasks, reading overall init-inputs and gathering
1564 # init-output in-memory objects.
1565 init_outputs: list[tuple[Any, DatasetType]] = []
1566 self.pipeline_graph.instantiate_tasks(
1567 get_init_input=lambda dataset_type: butler.get(
1568 input_refs[dataset_type.name].overrideStorageClass(dataset_type.storageClass)
1569 ),
1570 init_outputs=init_outputs,
1571 )
1572 # Write init-outputs that weren't already present.
1573 for obj, dataset_type in init_outputs:
1574 if new_ref := output_refs.get(dataset_type.name):
1575 assert new_ref.datasetType.storageClass_name == dataset_type.storageClass_name, (
1576 "QG init refs should use task connection storage classes."
1577 )
1578 butler.put(obj, new_ref)
1580 def write_configs(self, butler: LimitedButler, compare_existing: bool = True) -> None:
1581 """Write the config datasets for all tasks in the quantum graph.
1583 Parameters
1584 ----------
1585 butler : `lsst.daf.butler.LimitedButler`
1586 A limited butler data repository client.
1587 compare_existing : `bool`, optional
1588 If `True` check configs that already exist for consistency. If
1589 `False`, always raise if configs already exist.
1591 Raises
1592 ------
1593 lsst.daf.butler.registry.ConflictingDefinitionError
1594 Raised if an config dataset already exists and
1595 ``compare_existing=False``, or if the existing config is not
1596 consistent with the config in the quantum graph.
1597 """
1598 to_put: list[tuple[PipelineTaskConfig, DatasetRef]] = []
1599 for task_node in self.pipeline_graph.tasks.values():
1600 dataset_type_name = task_node.init.config_output.dataset_type_name
1601 (ref,) = [ # noqa: UP027
1602 ref
1603 for ref in self.get_init_output_refs(task_node.label)
1604 if ref.datasetType.name == dataset_type_name
1605 ]
1606 try:
1607 old_config = butler.get(ref)
1608 except (LookupError, FileNotFoundError):
1609 old_config = None
1610 if old_config is not None:
1611 if not compare_existing:
1612 raise ConflictingDefinitionError(f"Config dataset {ref} already exists.")
1613 if not task_node.config.compare(old_config, shortcut=False, output=log_config_mismatch):
1614 raise ConflictingDefinitionError(
1615 f"Config does not match existing task config {dataset_type_name!r} in "
1616 "butler; tasks configurations must be consistent within the same run collection."
1617 )
1618 else:
1619 to_put.append((task_node.config, ref))
1620 # We do writes at the end to minimize the mess we leave behind when we
1621 # raise an exception.
1622 for config, ref in to_put:
1623 butler.put(config, ref)
1625 def write_packages(self, butler: LimitedButler, compare_existing: bool = True) -> None:
1626 """Write the 'packages' dataset for the currently-active software
1627 versions.
1629 Parameters
1630 ----------
1631 butler : `lsst.daf.butler.LimitedButler`
1632 A limited butler data repository client.
1633 compare_existing : `bool`, optional
1634 If `True` check packages that already exist for consistency. If
1635 `False`, always raise if the packages dataset already exists.
1637 Raises
1638 ------
1639 lsst.daf.butler.registry.ConflictingDefinitionError
1640 Raised if the packages dataset already exists and is not consistent
1641 with the current packages.
1642 """
1643 new_packages = Packages.fromSystem()
1644 (ref,) = self.globalInitOutputRefs()
1645 try:
1646 packages = butler.get(ref)
1647 except (LookupError, FileNotFoundError):
1648 packages = None
1649 if packages is not None:
1650 if not compare_existing:
1651 raise ConflictingDefinitionError(f"Packages dataset {ref} already exists.")
1652 if compare_packages(packages, new_packages):
1653 # have to remove existing dataset first; butler has no
1654 # replace option.
1655 butler.pruneDatasets([ref], unstore=True, purge=True)
1656 butler.put(packages, ref)
1657 else:
1658 butler.put(new_packages, ref)
1660 def init_output_run(self, butler: LimitedButler, existing: bool = True) -> None:
1661 """Initialize a new output RUN collection by writing init-output
1662 datasets (including configs and packages).
1664 Parameters
1665 ----------
1666 butler : `lsst.daf.butler.LimitedButler`
1667 A limited butler data repository client.
1668 existing : `bool`, optional
1669 If `True` check or ignore outputs that already exist. If
1670 `False`, always raise if an output dataset already exists.
1672 Raises
1673 ------
1674 lsst.daf.butler.registry.ConflictingDefinitionError
1675 Raised if there are existing init output datasets, and either
1676 ``existing=False`` or their contents are not compatible with this
1677 graph.
1678 """
1679 self.write_configs(butler, compare_existing=existing)
1680 self.write_packages(butler, compare_existing=existing)
1681 self.write_init_outputs(butler, skip_existing=existing)
1683 def get_refs(
1684 self,
1685 *,
1686 include_init_inputs: bool = False,
1687 include_inputs: bool = False,
1688 include_intermediates: bool | None = None,
1689 include_init_outputs: bool = False,
1690 include_outputs: bool = False,
1691 conform_outputs: bool = True,
1692 ) -> tuple[set[DatasetRef], dict[str, DatastoreRecordData]]:
1693 """Get the requested dataset refs from the graph.
1695 Parameters
1696 ----------
1697 include_init_inputs : `bool`, optional
1698 Include init inputs.
1699 include_inputs : `bool`, optional
1700 Include inputs.
1701 include_intermediates : `bool` or `None`, optional
1702 If `None`, no special handling for intermediates is performed.
1703 If `True` intermediates are calculated even if other flags
1704 do not request datasets. If `False` intermediates will be removed
1705 from any results.
1706 include_init_outputs : `bool`, optional
1707 Include init outpus.
1708 include_outputs : `bool`, optional
1709 Include outputs.
1710 conform_outputs : `bool`, optional
1711 Whether any outputs found should have their dataset types conformed
1712 with the registry dataset types.
1714 Returns
1715 -------
1716 refs : `set` [ `lsst.daf.butler.DatasetRef` ]
1717 The requested dataset refs found in the graph.
1718 datastore_records : `dict` [ `str`, \
1719 `lsst.daf.butler.datastore.record_data.DatastoreRecordData` ]
1720 Any datastore records found.
1722 Notes
1723 -----
1724 Conforming and requesting inputs and outputs can result in the same
1725 dataset appearing in the results twice with differing storage classes.
1726 """
1727 datastore_records: dict[str, DatastoreRecordData] = {}
1728 init_input_refs: set[DatasetRef] = set()
1729 init_output_refs: set[DatasetRef] = set(self.globalInitOutputRefs())
1731 if include_intermediates is True:
1732 # Need to enable inputs and outputs even if not explicitly
1733 # requested.
1734 request_include_init_inputs = True
1735 request_include_inputs = True
1736 request_include_init_outputs = True
1737 request_include_outputs = True
1738 else:
1739 request_include_init_inputs = include_init_inputs
1740 request_include_inputs = include_inputs
1741 request_include_init_outputs = include_init_outputs
1742 request_include_outputs = include_outputs
1744 if request_include_init_inputs or request_include_init_outputs:
1745 for task_def in self.iterTaskGraph():
1746 if request_include_init_inputs:
1747 if in_refs := self.initInputRefs(task_def):
1748 init_input_refs.update(in_refs)
1749 if request_include_init_outputs:
1750 if out_refs := self.initOutputRefs(task_def):
1751 init_output_refs.update(out_refs)
1753 input_refs: set[DatasetRef] = set()
1754 output_refs: set[DatasetRef] = set()
1756 for qnode in self:
1757 if request_include_inputs:
1758 for other_refs in qnode.quantum.inputs.values():
1759 input_refs.update(other_refs)
1760 # Inputs can come with datastore records.
1761 for store_name, records in qnode.quantum.datastore_records.items():
1762 datastore_records.setdefault(store_name, DatastoreRecordData()).update(records)
1763 if request_include_outputs:
1764 for other_refs in qnode.quantum.outputs.values():
1765 output_refs.update(other_refs)
1767 # Intermediates are the intersection of inputs and outputs. Must do
1768 # this analysis before conforming since dataset type changes will
1769 # change set membership.
1770 inter_msg = ""
1771 intermediates = set()
1772 if include_intermediates is not None:
1773 intermediates = (input_refs | init_input_refs) & (output_refs | init_output_refs)
1775 if include_intermediates is False:
1776 # Remove intermediates from results.
1777 init_input_refs -= intermediates
1778 input_refs -= intermediates
1779 init_output_refs -= intermediates
1780 output_refs -= intermediates
1781 inter_msg = f"; Intermediates removed: {len(intermediates)}"
1782 intermediates = set()
1783 elif include_intermediates is True:
1784 # Do not mention intermediates if all the input/output flags
1785 # would have resulted in them anyhow.
1786 if (
1787 (request_include_init_inputs is not include_init_inputs)
1788 or (request_include_inputs is not include_inputs)
1789 or (request_include_init_outputs is not include_init_outputs)
1790 or (request_include_outputs is not include_outputs)
1791 ):
1792 inter_msg = f"; including intermediates: {len(intermediates)}"
1794 # Assign intermediates to the relevant category.
1795 if not include_init_inputs:
1796 init_input_refs &= intermediates
1797 if not include_inputs:
1798 input_refs &= intermediates
1799 if not include_init_outputs:
1800 init_output_refs &= intermediates
1801 if not include_outputs:
1802 output_refs &= intermediates
1804 # Conforming can result in an input ref and an output ref appearing
1805 # in the returned results that are identical apart from storage class.
1806 if conform_outputs:
1807 # Get data repository definitions from the QuantumGraph; these can
1808 # have different storage classes than those in the quanta.
1809 dataset_types = {dstype.name: dstype for dstype in self.registryDatasetTypes()}
1811 def _update_ref(ref: DatasetRef) -> DatasetRef:
1812 internal_dataset_type = dataset_types.get(ref.datasetType.name, ref.datasetType)
1813 if internal_dataset_type.storageClass_name != ref.datasetType.storageClass_name:
1814 ref = ref.replace(storage_class=internal_dataset_type.storageClass_name)
1815 return ref
1817 # Convert output_refs to the data repository storage classes, too.
1818 output_refs = {_update_ref(ref) for ref in output_refs}
1819 init_output_refs = {_update_ref(ref) for ref in init_output_refs}
1821 _LOG.verbose(
1822 "Found the following datasets. InitInputs: %d; Inputs: %d; InitOutputs: %s; Outputs: %d%s",
1823 len(init_input_refs),
1824 len(input_refs),
1825 len(init_output_refs),
1826 len(output_refs),
1827 inter_msg,
1828 )
1830 refs = input_refs | init_input_refs | init_output_refs | output_refs
1831 return refs, datastore_records