Coverage for python/lsst/ctrl/mpexec/dotTools.py : 7%
Hot-keys 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 iterConnections, Pipeline
38# ----------------------------------
39# Local non-exported definitions --
40# ----------------------------------
43def _renderTaskNode(nodeName, taskDef, file, idx=None):
44 """Render GV node for a task"""
45 label = [taskDef.taskName.rpartition('.')[-1]]
46 if idx is not None:
47 label += ["index: {}".format(idx)]
48 if taskDef.label:
49 label += ["label: {}".format(taskDef.label)]
50 label = r'\n'.join(label)
51 attrib = dict(shape="box",
52 style="filled,bold",
53 fillcolor="gray70",
54 label=label)
55 attrib = ['{}="{}"'.format(key, val) for key, val in attrib.items()]
56 print("{} [{}];".format(nodeName, ", ".join(attrib)), file=file)
59def _renderDSTypeNode(name, dimensions, file):
60 """Render GV node for a dataset type"""
61 label = [name]
62 if dimensions:
63 label += ["Dimensions: " + ", ".join(dimensions)]
64 label = r'\n'.join(label)
65 attrib = dict(shape="box",
66 style="rounded,filled",
67 fillcolor="gray90",
68 label=label)
69 attrib = ['{}="{}"'.format(key, val) for key, val in attrib.items()]
70 print("{} [{}];".format(name, ", ".join(attrib)), file=file)
73def _renderDSNode(nodeName, dsRef, file):
74 """Render GV node for a dataset"""
75 label = [dsRef.datasetType.name]
76 for key in sorted(dsRef.dataId.keys()):
77 label += [str(key) + "=" + str(dsRef.dataId[key])]
78 label = r'\n'.join(label)
79 attrib = dict(shape="box",
80 style="rounded,filled",
81 fillcolor="gray90",
82 label=label)
83 attrib = ['{}="{}"'.format(key, val) for key, val in attrib.items()]
84 print("{} [{}];".format(nodeName, ", ".join(attrib)), file=file)
87def _datasetRefId(dsRef):
88 """Make an idetifying string for given ref"""
89 idStr = str(dsRef.datasetType.name)
90 for key in sorted(dsRef.dataId.keys()):
91 idStr += ":" + str(key) + "=" + str(dsRef.dataId[key])
92 return idStr
95def _makeDSNode(dsRef, allDatasetRefs, file):
96 """Make new node for dataset if it does not exist.
98 Returns node name.
99 """
100 dsRefId = _datasetRefId(dsRef)
101 nodeName = allDatasetRefs.get(dsRefId)
102 if nodeName is None:
103 idx = len(allDatasetRefs)
104 nodeName = "dsref_{}".format(idx)
105 allDatasetRefs[dsRefId] = nodeName
106 _renderDSNode(nodeName, dsRef, file)
107 return nodeName
109# ------------------------
110# Exported definitions --
111# ------------------------
114def graph2dot(qgraph, file):
115 """Convert QuantumGraph into GraphViz digraph.
117 This method is mostly for documentation/presentation purposes.
119 Parameters
120 ----------
121 qgraph: `pipe.base.QuantumGraph`
122 QuantumGraph instance.
123 file : str or file object
124 File where GraphViz graph (DOT language) is written, can be a file name
125 or file object.
127 Raises
128 ------
129 `OSError` is raised when output file cannot be open.
130 `ImportError` is raised when task class cannot be imported.
131 """
132 # open a file if needed
133 close = False
134 if not hasattr(file, "write"):
135 file = open(file, "w")
136 close = True
138 print("digraph QuantumGraph {", file=file)
140 allDatasetRefs = {}
141 for taskId, nodes in enumerate(qgraph):
143 taskDef = nodes.taskDef
145 for qId, quantum in enumerate(nodes.quanta):
147 # node for a task
148 taskNodeName = "task_{}_{}".format(taskId, qId)
149 _renderTaskNode(taskNodeName, taskDef, file)
151 # quantum inputs
152 for dsRefs in quantum.predictedInputs.values():
153 for dsRef in dsRefs:
154 nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
155 print("{} -> {};".format(nodeName, taskNodeName), file=file)
157 # quantum outputs
158 for dsRefs in quantum.outputs.values():
159 for dsRef in dsRefs:
160 nodeName = _makeDSNode(dsRef, allDatasetRefs, file)
161 print("{} -> {};".format(taskNodeName, nodeName), file=file)
163 print("}", file=file)
164 if close:
165 file.close()
168def pipeline2dot(pipeline, file):
169 """Convert Pipeline into GraphViz digraph.
171 This method is mostly for documentation/presentation purposes.
172 Unlike other methods this method does not validate graph consistency.
174 Parameters
175 ----------
176 pipeline : `pipe.base.Pipeline`
177 Pipeline description.
178 file : str or file object
179 File where GraphViz graph (DOT language) is written, can be a file name
180 or file object.
182 Raises
183 ------
184 `OSError` is raised when output file cannot be open.
185 `ImportError` is raised when task class cannot be imported.
186 `MissingTaskFactoryError` is raised when TaskFactory is needed but not
187 provided.
188 """
189 universe = DimensionUniverse()
191 def expand_dimensions(dimensions):
192 """Returns expanded list of dimensions, with special skypix treatment.
194 Parameters
195 ----------
196 dimensions : `list` [`str`]
198 Returns
199 -------
200 dimensions : `list` [`str`]
201 """
202 dimensions = set(dimensions)
203 skypix_dim = []
204 if "skypix" in dimensions:
205 dimensions.remove("skypix")
206 skypix_dim = ["skypix"]
207 dimensions = universe.extract(dimensions)
208 return list(dimensions.names) + skypix_dim
210 # open a file if needed
211 close = False
212 if not hasattr(file, "write"):
213 file = open(file, "w")
214 close = True
216 print("digraph Pipeline {", file=file)
218 allDatasets = set()
219 if isinstance(pipeline, Pipeline):
220 pipeline = pipeline.toExpandedPipeline()
221 for idx, taskDef in enumerate(pipeline):
223 # node for a task
224 taskNodeName = "task{}".format(idx)
225 _renderTaskNode(taskNodeName, taskDef, file, idx)
227 for attr in iterConnections(taskDef.connections, 'inputs'):
228 if attr.name not in allDatasets:
229 dimensions = expand_dimensions(attr.dimensions)
230 _renderDSTypeNode(attr.name, dimensions, file)
231 allDatasets.add(attr.name)
232 print("{} -> {};".format(attr.name, taskNodeName), file=file)
234 for attr in iterConnections(taskDef.connections, 'prerequisiteInputs'):
235 if attr.name not in allDatasets:
236 dimensions = expand_dimensions(attr.dimensions)
237 _renderDSTypeNode(attr.name, dimensions, file)
238 allDatasets.add(attr.name)
239 # use dashed line for prerequisite edges to distinguish them
240 print("{} -> {} [style = dashed];".format(attr.name, taskNodeName), file=file)
242 for attr in iterConnections(taskDef.connections, 'outputs'):
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 print("{} -> {};".format(taskNodeName, attr.name), file=file)
249 print("}", file=file)
250 if close:
251 file.close()