lsst.pipe.base  18.1.0-5-gbd1decb+2
graph.py
Go to the documentation of this file.
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 program is free software: you can redistribute it and/or modify
10 # it under the terms of the GNU General Public License as published by
11 # the Free Software Foundation, either version 3 of the License, or
12 # (at your option) any later version.
13 #
14 # This program is distributed in the hope that it will be useful,
15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17 # GNU General Public License for more details.
18 #
19 # You should have received a copy of the GNU General Public License
20 # along with this program. If not, see <http://www.gnu.org/licenses/>.
21 from __future__ import annotations
22 
23 """Module defining quantum graph classes and related methods.
24 
25 There could be different representations of the quantum graph depending
26 on the client needs. Presently this module contains graph implementation
27 which is based on requirements of command-line environment. In the future
28 we could add other implementations and methods to convert between those
29 representations.
30 """
31 
32 # "exported" names
33 __all__ = ["QuantumGraph", "QuantumGraphTaskNodes", "QuantumIterData"]
34 
35 # -------------------------------
36 # Imports of standard modules --
37 # -------------------------------
38 from itertools import chain
39 from dataclasses import dataclass
40 from typing import List, FrozenSet, Mapping
41 
42 # -----------------------------
43 # Imports for other modules --
44 # -----------------------------
45 from .pipeline import Pipeline, TaskDef
46 from .pipeTools import orderPipeline
47 from lsst.daf.butler import DataId, Quantum, DatasetRef, DatasetType
48 from lsst.daf.butler.core.utils import NamedKeyDict
49 
50 # ----------------------------------
51 # Local non-exported definitions --
52 # ----------------------------------
53 
54 # ------------------------
55 # Exported definitions --
56 # ------------------------
57 
58 
59 @dataclass
61  """Helper class for iterating over quanta in a graph.
62 
63  The `QuantumGraph.traverse` method needs to return topologically ordered
64  Quanta together with their dependencies. This class is used as a value
65  for the iterator, it contains enumerated Quantum and its dependencies.
66  """
67 
68  __slots__ = ["index", "quantum", "taskDef", "dependencies"]
69 
70  index: int
71  """Index of this Quantum, a unique but arbitrary integer."""
72 
73  quantum: Quantum
74  """Quantum corresponding to a graph node."""
75 
76  taskDef: TaskDef
77  """Task class to be run on this quantum, and corresponding label and
78  config.
79  """
80 
81  dependencies: FrozenSet(int)
82  """Possibly empty set of indices of dependencies for this Quantum.
83  Dependencies include other nodes in the graph; they do not reflect data
84  already in butler (there are no graph nodes for those).
85  """
86 
87 
88 @dataclass
90  """QuantumGraphTaskNodes represents a bunch of nodes in an quantum graph
91  corresponding to a single task.
92 
93  The node in quantum graph is represented by the `PipelineTask` and a
94  single `~lsst.daf.butler.Quantum` instance. One possible representation
95  of the graph is just a list of nodes without edges (edges can be deduced
96  from nodes' quantum inputs and outputs if needed). That representation can
97  be reduced to the list of PipelineTasks (or their corresponding TaskDefs)
98  and the corresponding list of Quanta. This class is used in this reduced
99  representation for a single task, and full `QuantumGraph` is a sequence of
100  tinstances of this class for one or more tasks.
101 
102  Different frameworks may use different graph representation, this
103  representation was based mostly on requirements of command-line
104  executor which does not need explicit edges information.
105  """
106 
107  taskDef: TaskDef
108  """Task defintion for this set of nodes."""
109 
110  quanta: List[Quantum]
111  """List of quanta corresponding to the task."""
112 
113  initInputs: Mapping[DatasetType, DatasetRef]
114  """Datasets that must be loaded or created to construct this task."""
115 
116  initOutputs: Mapping[DatasetType, DatasetRef]
117  """Datasets that may be written after constructing this task."""
118 
119 
120 class QuantumGraph(list):
121  """QuantumGraph is a sequence of `QuantumGraphTaskNodes` objects.
122 
123  Typically the order of the tasks in the list will be the same as the
124  order of tasks in a pipeline (obviously depends on the code which
125  constructs graph).
126 
127  Parameters
128  ----------
129  iterable : iterable of `QuantumGraphTaskNodes`, optional
130  Initial sequence of per-task nodes.
131  """
132  def __init__(self, iterable=None):
133  list.__init__(self, iterable or [])
134  self.initInputs = NamedKeyDict()
135  self.initIntermediates = NamedKeyDict()
136  self.initOutputs = NamedKeyDict()
137 
138  initInputs: NamedKeyDict
139  """Datasets that must be provided to construct one or more Tasks in this
140  graph, and must be obtained from the data repository.
141 
142  This is disjoint with both `initIntermediates` and `initOutputs`.
143  """
144 
145  initIntermediates: NamedKeyDict
146  """Datasets that must be provided to construct one or more Tasks in this
147  graph, but are also produced after constructing a Task in this graph.
148 
149  This is disjoint with both `initInputs` and `initOutputs`.
150  """
151 
152  initOutputs: NamedKeyDict
153  """Datasets that are produced after constructing a Task in this graph,
154  and are not used to construct any other Task in this graph.
155 
156  This is disjoint from both `initInputs` and `initIntermediates`.
157  """
158 
159  def quanta(self):
160  """Iterator over quanta in a graph.
161 
162  Quanta are returned in unspecified order.
163 
164  Yields
165  ------
166  taskDef : `TaskDef`
167  Task definition for a Quantum.
168  quantum : `~lsst.daf.butler.Quantum`
169  Single quantum.
170  """
171  for taskNodes in self:
172  taskDef = taskNodes.taskDef
173  for quantum in taskNodes.quanta:
174  yield taskDef, quantum
175 
176  def countQuanta(self):
177  """Return total count of quanta in a graph.
178 
179  Returns
180  -------
181  count : `int`
182  Number of quanta in a graph.
183  """
184  return sum(len(taskNodes.quanta) for taskNodes in self)
185 
186  def traverse(self):
187  """Return topologically ordered Quanta and their dependencies.
188 
189  This method iterates over all Quanta in topological order, enumerating
190  them during iteration. Returned `QuantumIterData` object contains
191  Quantum instance, its ``index`` and the ``index`` of all its
192  prerequsites (Quanta that produce inputs for this Quantum):
193  - the ``index`` values are generated by an iteration of a
194  QuantumGraph, and are not intrinsic to the QuantumGraph
195  - during iteration, each ID will appear in index before it ever
196  appears in dependencies.
197 
198  Yields
199  ------
200  quantumData : `QuantumIterData`
201  """
202 
203  def orderedTaskNodes(graph):
204  """Return topologically ordered task nodes.
205 
206  Yields
207  ------
208  nodes : `QuantumGraphTaskNodes`
209  """
210  # Tasks in a graph are probably topologically sorted already but there
211  # is no guarantee for that. Just re-construct Pipeline and order tasks
212  # in a pipeline using existing method.
213  nodesMap = {id(item.taskDef): item for item in graph}
214  pipeline = orderPipeline(Pipeline(item.taskDef for item in graph))
215  for taskDef in pipeline:
216  yield nodesMap[id(taskDef)]
217 
218  index = 0
219  outputs = {} # maps (DatasetType.name, DataId) to its producing quantum index
220  for nodes in orderedTaskNodes(self):
221  for quantum in nodes.quanta:
222 
223  # Find quantum dependencies (must be in `outputs` already)
224  prereq = []
225  for dataRef in chain.from_iterable(quantum.predictedInputs.values()):
226  # if data exists in butler then `id` is not None
227  if dataRef.id is None:
228  # Get the base name if this is a component
229  name, component = dataRef.datasetType.nameAndComponent()
230  key = (name, DataId(dataRef.dataId))
231  try:
232  prereq.append(outputs[key])
233  except KeyError:
234  # The Quantum that makes our inputs is not in the graph,
235  # this could happen if we run on a "split graph" which is
236  # usually just one quantum. Check for number of Quanta
237  # in a graph and ignore error if it's just one.
238  # TODO: This code has to be removed or replaced with
239  # something more generic
240  if not (len(self) == 1 and len(self[0].quanta) == 1):
241  raise
242 
243  # Update `outputs` with this quantum outputs
244  for dataRef in chain.from_iterable(quantum.outputs.values()):
245  key = (dataRef.datasetType.name, DataId(dataRef.dataId))
246  outputs[key] = index
247 
248  yield QuantumIterData(index=index, quantum=quantum, taskDef=nodes.taskDef,
249  dependencies=frozenset(prereq))
250  index += 1
def __init__(self, iterable=None)
Definition: graph.py:132
def orderPipeline(pipeline, taskFactory=None)
Definition: pipeTools.py:136