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

115 statements  

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/>. 

21 

22"""Module defining few methods to generate GraphViz diagrams from pipelines 

23or quantum graphs. 

24""" 

25 

26__all__ = ["graph2dot", "pipeline2dot"] 

27 

28# ------------------------------- 

29# Imports of standard modules -- 

30# ------------------------------- 

31 

32# ----------------------------- 

33# Imports for other modules -- 

34# ----------------------------- 

35from lsst.daf.butler import DimensionUniverse 

36from lsst.pipe.base import Pipeline, iterConnections 

37 

38# ---------------------------------- 

39# Local non-exported definitions -- 

40# ---------------------------------- 

41 

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) 

49 

50 

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) 

57 

58 

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) 

69 

70 

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) 

77 

78 

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) 

85 

86 

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) 

92 

93 

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) 

101 

102 

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) 

108 

109 

110def _makeDSNode(dsRef, allDatasetRefs, file): 

111 """Make new node for dataset if it does not exist. 

112 

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 

123 

124 

125# ------------------------ 

126# Exported definitions -- 

127# ------------------------ 

128 

129 

130def graph2dot(qgraph, file): 

131 """Convert QuantumGraph into GraphViz digraph. 

132 

133 This method is mostly for documentation/presentation purposes. 

134 

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. 

142 

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 

153 

154 print("digraph QuantumGraph {", file=file) 

155 

156 allDatasetRefs = {} 

157 for taskId, taskDef in enumerate(qgraph.taskGraph): 

158 

159 quanta = qgraph.getNodesForTask(taskDef) 

160 for qId, quantumNode in enumerate(quanta): 

161 

162 # node for a task 

163 taskNodeName = "task_{}_{}".format(taskId, qId) 

164 _renderQuantumNode(taskNodeName, taskDef, quantumNode, file) 

165 

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) 

171 

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) 

177 

178 print("}", file=file) 

179 if close: 

180 file.close() 

181 

182 

183def pipeline2dot(pipeline, file): 

184 """Convert Pipeline into GraphViz digraph. 

185 

186 This method is mostly for documentation/presentation purposes. 

187 Unlike other methods this method does not validate graph consistency. 

188 

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. 

196 

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() 

205 

206 def expand_dimensions(dimensions): 

207 """Returns expanded list of dimensions, with special skypix treatment. 

208 

209 Parameters 

210 ---------- 

211 dimensions : `list` [`str`] 

212 

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 

224 

225 # open a file if needed 

226 close = False 

227 if not hasattr(file, "write"): 

228 file = open(file, "w") 

229 close = True 

230 

231 print("digraph Pipeline {", file=file) 

232 

233 allDatasets = set() 

234 if isinstance(pipeline, Pipeline): 

235 pipeline = pipeline.toExpandedPipeline() 

236 for idx, taskDef in enumerate(pipeline): 

237 

238 # node for a task 

239 taskNodeName = "task{}".format(idx) 

240 _renderTaskNode(taskNodeName, taskDef, file, idx) 

241 

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) 

248 

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") 

256 

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) 

263 

264 print("}", file=file) 

265 if close: 

266 file.close()