Coverage for python/lsst/ctrl/mpexec/dotTools.py: 9%
Shortcuts on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of ctrl_mpexec.
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/>.
22"""Module defining few methods to generate GraphViz diagrams from pipelines
23or quantum graphs.
24"""
26__all__ = ["graph2dot", "pipeline2dot"]
28# -------------------------------
29# Imports of standard modules --
30# -------------------------------
32# -----------------------------
33# Imports for other modules --
34# -----------------------------
35from lsst.daf.butler import DimensionUniverse
36from lsst.pipe.base import Pipeline, iterConnections
38# ----------------------------------
39# Local non-exported definitions --
40# ----------------------------------
42# Node styles indexed by node type.
43_STYLES = dict(
44 task=dict(shape="box", style="filled,bold", fillcolor="gray70"),
45 quantum=dict(shape="box", style="filled,bold", fillcolor="gray70"),
46 dsType=dict(shape="box", style="rounded,filled", fillcolor="gray90"),
47 dataset=dict(shape="box", style="rounded,filled", fillcolor="gray90"),
48)
51def _renderNode(file, nodeName, style, labels):
52 """Render GV node"""
53 label = r"\n".join(labels)
54 attrib = dict(_STYLES[style], label=label)
55 attrib = ", ".join([f'{key}="{val}"' for key, val in attrib.items()])
56 print(f'"{nodeName}" [{attrib}];', file=file)
59def _renderTaskNode(nodeName, taskDef, file, idx=None):
60 """Render GV node for a task"""
61 labels = [taskDef.label, taskDef.taskName]
62 if idx is not None:
63 labels.append(f"index: {idx}")
64 if taskDef.connections:
65 # don't print collection of str directly to avoid visually noisy quotes
66 dimensions_str = ", ".join(taskDef.connections.dimensions)
67 labels.append(f"dimensions: {dimensions_str}")
68 _renderNode(file, nodeName, "task", labels)
71def _renderQuantumNode(nodeName, taskDef, quantumNode, file):
72 """Render GV node for a quantum"""
73 labels = [f"{quantumNode.nodeId}", taskDef.label]
74 dataId = quantumNode.quantum.dataId
75 labels.extend(f"{key} = {dataId[key]}" for key in sorted(dataId.keys()))
76 _renderNode(file, nodeName, "quantum", labels)
79def _renderDSTypeNode(name, dimensions, file):
80 """Render GV node for a dataset type"""
81 labels = [name]
82 if dimensions:
83 labels.append("Dimensions: " + ", ".join(dimensions))
84 _renderNode(file, name, "dsType", labels)
87def _renderDSNode(nodeName, dsRef, file):
88 """Render GV node for a dataset"""
89 labels = [dsRef.datasetType.name, f"run: {dsRef.run!r}"]
90 labels.extend(f"{key} = {dsRef.dataId[key]}" for key in sorted(dsRef.dataId.keys()))
91 _renderNode(file, nodeName, "dataset", labels)
94def _renderEdge(fromName, toName, file, **kwargs):
95 """Render GV edge"""
96 if kwargs:
97 attrib = ", ".join([f'{key}="{val}"' for key, val in kwargs.items()])
98 print(f'"{fromName}" -> "{toName}" [{attrib}];', file=file)
99 else:
100 print(f'"{fromName}" -> "{toName}";', file=file)
103def _datasetRefId(dsRef):
104 """Make an identifying string for given ref"""
105 dsId = [dsRef.datasetType.name]
106 dsId.extend(f"{key} = {dsRef.dataId[key]}" for key in sorted(dsRef.dataId.keys()))
107 return ":".join(dsId)
110def _makeDSNode(dsRef, allDatasetRefs, file):
111 """Make new node for dataset if it does not exist.
113 Returns node name.
114 """
115 dsRefId = _datasetRefId(dsRef)
116 nodeName = allDatasetRefs.get(dsRefId)
117 if nodeName is None:
118 idx = len(allDatasetRefs)
119 nodeName = "dsref_{}".format(idx)
120 allDatasetRefs[dsRefId] = nodeName
121 _renderDSNode(nodeName, dsRef, file)
122 return nodeName
125# ------------------------
126# Exported definitions --
127# ------------------------
130def graph2dot(qgraph, file):
131 """Convert QuantumGraph into GraphViz digraph.
133 This method is mostly for documentation/presentation purposes.
135 Parameters
136 ----------
137 qgraph: `pipe.base.QuantumGraph`
138 QuantumGraph instance.
139 file : str or file object
140 File where GraphViz graph (DOT language) is written, can be a file name
141 or file object.
143 Raises
144 ------
145 `OSError` is raised when output file cannot be open.
146 `ImportError` is raised when task class cannot be imported.
147 """
148 # open a file if needed
149 close = False
150 if not hasattr(file, "write"):
151 file = open(file, "w")
152 close = True
154 print("digraph QuantumGraph {", file=file)
156 allDatasetRefs = {}
157 for taskId, taskDef in enumerate(qgraph.taskGraph):
159 quanta = qgraph.getNodesForTask(taskDef)
160 for qId, quantumNode in enumerate(quanta):
162 # node for a task
163 taskNodeName = "task_{}_{}".format(taskId, qId)
164 _renderQuantumNode(taskNodeName, taskDef, quantumNode, file)
166 # quantum inputs
167 for dsRefs in quantumNode.quantum.inputs.values():
168 for dsRef in dsRefs:
169 nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
170 _renderEdge(nodeName, taskNodeName, file)
172 # quantum outputs
173 for dsRefs in quantumNode.quantum.outputs.values():
174 for dsRef in dsRefs:
175 nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
176 _renderEdge(taskNodeName, nodeName, file)
178 print("}", file=file)
179 if close:
180 file.close()
183def pipeline2dot(pipeline, file):
184 """Convert Pipeline into GraphViz digraph.
186 This method is mostly for documentation/presentation purposes.
187 Unlike other methods this method does not validate graph consistency.
189 Parameters
190 ----------
191 pipeline : `pipe.base.Pipeline`
192 Pipeline description.
193 file : str or file object
194 File where GraphViz graph (DOT language) is written, can be a file name
195 or file object.
197 Raises
198 ------
199 `OSError` is raised when output file cannot be open.
200 `ImportError` is raised when task class cannot be imported.
201 `MissingTaskFactoryError` is raised when TaskFactory is needed but not
202 provided.
203 """
204 universe = DimensionUniverse()
206 def expand_dimensions(dimensions):
207 """Returns expanded list of dimensions, with special skypix treatment.
209 Parameters
210 ----------
211 dimensions : `list` [`str`]
213 Returns
214 -------
215 dimensions : `list` [`str`]
216 """
217 dimensions = set(dimensions)
218 skypix_dim = []
219 if "skypix" in dimensions:
220 dimensions.remove("skypix")
221 skypix_dim = ["skypix"]
222 dimensions = universe.extract(dimensions)
223 return list(dimensions.names) + skypix_dim
225 # open a file if needed
226 close = False
227 if not hasattr(file, "write"):
228 file = open(file, "w")
229 close = True
231 print("digraph Pipeline {", file=file)
233 allDatasets = set()
234 if isinstance(pipeline, Pipeline):
235 pipeline = pipeline.toExpandedPipeline()
236 for idx, taskDef in enumerate(pipeline):
238 # node for a task
239 taskNodeName = "task{}".format(idx)
240 _renderTaskNode(taskNodeName, taskDef, file, idx)
242 for attr in iterConnections(taskDef.connections, "inputs"):
243 if attr.name not in allDatasets:
244 dimensions = expand_dimensions(attr.dimensions)
245 _renderDSTypeNode(attr.name, dimensions, file)
246 allDatasets.add(attr.name)
247 _renderEdge(attr.name, taskNodeName, file)
249 for attr in iterConnections(taskDef.connections, "prerequisiteInputs"):
250 if attr.name not in allDatasets:
251 dimensions = expand_dimensions(attr.dimensions)
252 _renderDSTypeNode(attr.name, dimensions, file)
253 allDatasets.add(attr.name)
254 # use dashed line for prerequisite edges to distinguish them
255 _renderEdge(attr.name, taskNodeName, file, style="dashed")
257 for attr in iterConnections(taskDef.connections, "outputs"):
258 if attr.name not in allDatasets:
259 dimensions = expand_dimensions(attr.dimensions)
260 _renderDSTypeNode(attr.name, dimensions, file)
261 allDatasets.add(attr.name)
262 _renderEdge(taskNodeName, attr.name, file)
264 print("}", file=file)
265 if close:
266 file.close()