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

83 statements  

« prev     ^ index     » next       coverage.py v7.3.0, created at 2023-08-22 09:44 +0000

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 

47from lsst.utils.packages import Packages 

48 

49_LOG = logging.getLogger(__name__) 

50 

51 

52class WhenToSaveQuantumGraphs(Enum): 

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

54 

55 QGRAPH = 1 # Must be using single_quantum_clustering algorithm. 

56 TRANSFORM = 2 

57 PREPARE = 3 

58 SUBMIT = 4 

59 NEVER = 5 # Always use full QuantumGraph. 

60 

61 

62@contextlib.contextmanager 

63def chdir(path): 

64 """Change working directory. 

65 

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 for a job. 

83 

84 Parameters 

85 ---------- 

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

87 BPS configuration. 

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

89 Job for which the QuantumGraph file is being saved. 

90 out_prefix : `str`, optional 

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

92 uses current working directory. 

93 

94 Returns 

95 ------- 

96 full_filename : `str` 

97 The filename for the job's QuantumGraph. 

98 """ 

99 curvals = dataclasses.asdict(job) 

100 if job.tags: 

101 curvals.update(job.tags) 

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

103 if not found: 

104 subdir = "{job.label}" 

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

106 

107 if out_prefix is not None: 

108 full_filename = Path(out_prefix) / full_filename 

109 

110 return str(full_filename) 

111 

112 

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

114 """Save subgraph to file. 

115 

116 Parameters 

117 ---------- 

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

119 QuantumGraph to save. 

120 out_filename : `str` 

121 Name of the output file. 

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

123 NodeIds for the subgraph to save to file. 

124 """ 

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

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

127 if node_ids is None: 

128 qgraph.saveUri(out_filename) 

129 else: 

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

131 else: 

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

133 

134 

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

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

137 

138 Parameters 

139 ---------- 

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

141 BPS configuration. 

142 qgraph_filename : `str` 

143 Run QuantumGraph filename. 

144 execution_butler_dir : `str` 

145 Directory in which to create the execution butler. 

146 out_prefix : `str` or None 

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

148 

149 Raises 

150 ------ 

151 CalledProcessError 

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

153 exit code. 

154 """ 

155 _, command = config.search( 

156 ".executionButler.createCommand", 

157 opt={ 

158 "curvals": {"executionButlerDir": execution_butler_dir, "qgraphFile": qgraph_filename}, 

159 "replaceVars": True, 

160 }, 

161 ) 

162 out_filename = "execution_butler_creation.out" 

163 if out_prefix is not None: 

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

165 

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

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

168 # the execution Butler itself. 

169 opening = "cannot create the execution Butler" 

170 try: 

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

172 print(command, file=fh) 

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

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

175 except OSError as exc: 

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

177 except subprocess.SubprocessError as exc: 

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

179 

180 

181def create_count_summary(counts): 

182 """Create summary from count mapping. 

183 

184 Parameters 

185 ---------- 

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

187 Mapping of counts to keys. 

188 

189 Returns 

190 ------- 

191 summary : `str` 

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

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

194 parse_count_summary(). 

195 """ 

196 summary = "" 

197 if isinstance(counts, dict): 

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

199 return summary 

200 

201 

202def parse_count_summary(summary): 

203 """Parse summary into count mapping. 

204 

205 Parameters 

206 ---------- 

207 summary : `str` 

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

209 

210 Returns 

211 ------- 

212 counts : `collections.Counter` 

213 Mapping representation of given summary for easier 

214 individual count lookup. 

215 """ 

216 counts = Counter() 

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

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

219 counts[label] = count 

220 return counts 

221 

222 

223def _dump_pkg_info(filename): 

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

225 

226 Parameters 

227 ---------- 

228 filename : `str` 

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

230 of the packages. 

231 """ 

232 file = Path(filename) 

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

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

235 packages = Packages.fromSystem() 

236 packages.write(str(file)) 

237 

238 

239def _dump_env_info(filename): 

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

241 

242 Parameters 

243 ---------- 

244 filename : `str` 

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

246 environment. 

247 """ 

248 file = Path(filename) 

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

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

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

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