Coverage for python/lsst/ctrl/bps/bps_utils.py: 28%

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

83 statements  

1# This file is part of ctrl_bps. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://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 <https://www.gnu.org/licenses/>. 

21 

22"""Misc supporting classes and functions for BPS. 

23""" 

24 

25__all__ = [ 

26 "chdir", 

27 "create_job_quantum_graph_filename", 

28 "save_qg_subgraph", 

29 "_create_execution_butler", 

30 "create_count_summary", 

31 "parse_count_summary", 

32 "_dump_pkg_info", 

33 "_dump_env_info", 

34] 

35 

36import contextlib 

37import dataclasses 

38import logging 

39import os 

40import shlex 

41import subprocess 

42from collections import Counter 

43from enum import Enum 

44from pathlib import Path 

45 

46import yaml 

47 

48from lsst.utils.packages import Packages 

49 

50 

51_LOG = logging.getLogger(__name__) 

52 

53 

54class WhenToSaveQuantumGraphs(Enum): 

55 """Values for when to save the job quantum graphs. 

56 """ 

57 QGRAPH = 1 # Must be using single_quantum_clustering algorithm. 

58 TRANSFORM = 2 

59 PREPARE = 3 

60 SUBMIT = 4 

61 NEVER = 5 # Always use full QuantumGraph. 

62 

63 

64@contextlib.contextmanager 

65def chdir(path): 

66 """A chdir function that can be used inside a context. 

67 

68 Parameters 

69 ---------- 

70 path : `str` or `pathlib.Path` 

71 Path to be made current working directory. 

72 """ 

73 cur_dir = os.getcwd() 

74 os.chdir(path) 

75 try: 

76 yield 

77 finally: 

78 os.chdir(cur_dir) 

79 

80 

81def create_job_quantum_graph_filename(config, job, out_prefix=None): 

82 """Create a filename to be used when storing the QuantumGraph 

83 for a job. 

84 

85 Parameters 

86 ---------- 

87 config : `lsst.ctrl.bps.BpsConfig` 

88 BPS configuration. 

89 job : `lsst.ctrl.bps.GenericWorkflowJob` 

90 Job for which the QuantumGraph file is being saved. 

91 out_prefix : `str`, optional 

92 Path prefix for the QuantumGraph filename. If no out_prefix is given, 

93 uses current working directory. 

94 

95 Returns 

96 ------- 

97 full_filename : `str` 

98 The filename for the job's QuantumGraph. 

99 """ 

100 curvals = dataclasses.asdict(job) 

101 if job.tags: 

102 curvals.update(job.tags) 

103 found, subdir = config.search("subDirTemplate", opt={'curvals': curvals}) 

104 if not found: 

105 subdir = "{job.label}" 

106 full_filename = Path("inputs") / subdir / f"quantum_{job.name}.qgraph" 

107 

108 if out_prefix is not None: 

109 full_filename = Path(out_prefix) / full_filename 

110 

111 return str(full_filename) 

112 

113 

114def save_qg_subgraph(qgraph, out_filename, node_ids=None): 

115 """Save subgraph to file. 

116 

117 Parameters 

118 ---------- 

119 qgraph : `lsst.pipe.base.QuantumGraph` 

120 QuantumGraph to save. 

121 out_filename : `str` 

122 Name of the output file. 

123 node_ids : `list` [`lsst.pipe.base.NodeId`] 

124 NodeIds for the subgraph to save to file. 

125 """ 

126 if not os.path.exists(out_filename): 

127 _LOG.debug("Saving QuantumGraph with %d nodes to %s", len(qgraph), out_filename) 

128 if node_ids is None: 

129 qgraph.saveUri(out_filename) 

130 else: 

131 qgraph.subset(qgraph.getQuantumNodeByNodeId(nid) for nid in node_ids).saveUri(out_filename) 

132 else: 

133 _LOG.debug("Skipping saving QuantumGraph to %s because already exists.", out_filename) 

134 

135 

136def _create_execution_butler(config, qgraph_filename, execution_butler_dir, out_prefix): 

137 """Create the execution butler for use by the compute jobs. 

138 

139 Parameters 

140 ---------- 

141 config : `lsst.ctrl.bps.BpsConfig` 

142 BPS configuration. 

143 qgraph_filename : `str` 

144 Run QuantumGraph filename. 

145 execution_butler_dir : `str` 

146 Directory in which to create the execution butler. 

147 out_prefix : `str` or None 

148 Prefix for output filename to contain both stdout and stderr. 

149 

150 Raises 

151 ------ 

152 CalledProcessError 

153 Raised if command to create execution butler exits with non-zero 

154 exit code. 

155 """ 

156 _, command = config.search(".executionButler.createCommand", 

157 opt={"curvals": {"executionButlerDir": execution_butler_dir, 

158 "qgraphFile": qgraph_filename}, 

159 "replaceVars": True}) 

160 out_filename = "execution_butler_creation.out" 

161 if out_prefix is not None: 

162 out_filename = os.path.join(out_prefix, out_filename) 

163 

164 # When creating the execution Butler, handle separately errors related 

165 # to creating the log file and errors directly related to creating 

166 # the execution Butler itself. 

167 opening = "cannot create the execution Butler" 

168 try: 

169 with open(out_filename, "w", encoding="utf-8") as fh: 

170 print(command, file=fh) 

171 print("\n", file=fh) # Note: want a blank line 

172 subprocess.run(shlex.split(command), shell=False, check=True, stdout=fh, stderr=subprocess.STDOUT) 

173 except OSError as exc: 

174 raise type(exc)(f"{opening}: {exc.strerror}") from None 

175 except subprocess.SubprocessError as exc: 

176 raise RuntimeError(f"{opening}, see '{out_filename}' for details") from exc 

177 

178 

179def create_count_summary(counts): 

180 """Create summary from count mapping. 

181 

182 Parameters 

183 ---------- 

184 counts : `collections.Counter` or `dict` [`str`, `int`] 

185 Mapping of counts to keys. 

186 

187 Returns 

188 ------- 

189 summary : `str` 

190 Semi-colon delimited string of key:count pairs. 

191 (e.g. "key1:cnt1;key2;cnt2") Parsable by 

192 parse_count_summary(). 

193 """ 

194 summary = "" 

195 if isinstance(counts, dict): 

196 summary = ";".join([f"{key}:{counts[key]}" for key in counts]) 

197 return summary 

198 

199 

200def parse_count_summary(summary): 

201 """Parse summary into count mapping. 

202 

203 Parameters 

204 ---------- 

205 summary : `str` 

206 Semi-colon delimited string of key:count pairs. 

207 

208 Returns 

209 ------- 

210 counts : `collections.Counter` 

211 Mapping representation of given summary for easier 

212 individual count lookup. 

213 """ 

214 counts = Counter() 

215 for part in summary.split(";"): 

216 label, count = part.split(":") 

217 counts[label] = count 

218 return counts 

219 

220 

221def _dump_pkg_info(filename): 

222 """Save information about versions of packages in use for future reference. 

223 

224 Parameters 

225 ---------- 

226 filename : `str` 

227 The name of the file where to save the information about the versions 

228 of the packages. 

229 """ 

230 file = Path(filename) 

231 if file.suffix.lower() not in {".yaml", ".yml"}: 

232 file = file.with_suffix(f"{file.suffix}.yaml") 

233 packages = Packages.fromSystem() 

234 packages.write(str(file)) 

235 

236 

237def _dump_env_info(filename): 

238 """Save information about runtime environment for future reference. 

239 

240 Parameters 

241 ---------- 

242 filename : `str` 

243 The name of the file where to save the information about the runtime 

244 environment. 

245 """ 

246 file = Path(filename) 

247 if file.suffix.lower() not in {".yaml", ".yml"}: 

248 file = file.with_suffix(f"{file.suffix}.yaml") 

249 with open(file, "w", encoding="utf-8") as fh: 

250 yaml.dump(dict(os.environ), fh)