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

83 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-05-03 02:57 -0700

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <https://www.gnu.org/licenses/>. 

27 

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

29""" 

30 

31__all__ = [ 

32 "chdir", 

33 "create_job_quantum_graph_filename", 

34 "save_qg_subgraph", 

35 "_create_execution_butler", 

36 "create_count_summary", 

37 "parse_count_summary", 

38 "_dump_pkg_info", 

39 "_dump_env_info", 

40] 

41 

42import contextlib 

43import dataclasses 

44import logging 

45import os 

46import shlex 

47import subprocess 

48from collections import Counter 

49from enum import Enum 

50from pathlib import Path 

51 

52import yaml 

53from lsst.utils.packages import Packages 

54 

55_LOG = logging.getLogger(__name__) 

56 

57 

58class WhenToSaveQuantumGraphs(Enum): 

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

60 

61 QGRAPH = 1 # Must be using single_quantum_clustering algorithm. 

62 TRANSFORM = 2 

63 PREPARE = 3 

64 SUBMIT = 4 

65 NEVER = 5 # Always use full QuantumGraph. 

66 

67 

68@contextlib.contextmanager 

69def chdir(path): 

70 """Change working directory. 

71 

72 A chdir function that can be used inside a context. 

73 

74 Parameters 

75 ---------- 

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

77 Path to be made current working directory. 

78 """ 

79 cur_dir = os.getcwd() 

80 os.chdir(path) 

81 try: 

82 yield 

83 finally: 

84 os.chdir(cur_dir) 

85 

86 

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

88 """Create a filename to be used when storing the QuantumGraph for a job. 

89 

90 Parameters 

91 ---------- 

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

93 BPS configuration. 

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

95 Job for which the QuantumGraph file is being saved. 

96 out_prefix : `str`, optional 

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

98 uses current working directory. 

99 

100 Returns 

101 ------- 

102 full_filename : `str` 

103 The filename for the job's QuantumGraph. 

104 """ 

105 curvals = dataclasses.asdict(job) 

106 if job.tags: 

107 curvals.update(job.tags) 

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

109 if not found: 

110 subdir = "{job.label}" 

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

112 

113 if out_prefix is not None: 

114 full_filename = Path(out_prefix) / full_filename 

115 

116 return str(full_filename) 

117 

118 

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

120 """Save subgraph to file. 

121 

122 Parameters 

123 ---------- 

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

125 QuantumGraph to save. 

126 out_filename : `str` 

127 Name of the output file. 

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

129 NodeIds for the subgraph to save to file. 

130 """ 

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

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

133 if node_ids is None: 

134 qgraph.saveUri(out_filename) 

135 else: 

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

137 else: 

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

139 

140 

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

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

143 

144 Parameters 

145 ---------- 

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

147 BPS configuration. 

148 qgraph_filename : `str` 

149 Run QuantumGraph filename. 

150 execution_butler_dir : `str` 

151 Directory in which to create the execution butler. 

152 out_prefix : `str` or None 

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

154 

155 Raises 

156 ------ 

157 CalledProcessError 

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

159 exit code. 

160 """ 

161 _, command = config.search( 

162 ".executionButler.createCommand", 

163 opt={ 

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

165 "replaceVars": True, 

166 }, 

167 ) 

168 out_filename = "execution_butler_creation.out" 

169 if out_prefix is not None: 

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

171 

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

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

174 # the execution Butler itself. 

175 opening = "cannot create the execution Butler" 

176 try: 

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

178 print(command, file=fh) 

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

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

181 except OSError as exc: 

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

183 except subprocess.SubprocessError as exc: 

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

185 

186 

187def create_count_summary(counts): 

188 """Create summary from count mapping. 

189 

190 Parameters 

191 ---------- 

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

193 Mapping of counts to keys. 

194 

195 Returns 

196 ------- 

197 summary : `str` 

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

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

200 parse_count_summary(). 

201 """ 

202 summary = "" 

203 if isinstance(counts, dict): 

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

205 return summary 

206 

207 

208def parse_count_summary(summary): 

209 """Parse summary into count mapping. 

210 

211 Parameters 

212 ---------- 

213 summary : `str` 

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

215 

216 Returns 

217 ------- 

218 counts : `collections.Counter` 

219 Mapping representation of given summary for easier 

220 individual count lookup. 

221 """ 

222 counts = Counter() 

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

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

225 counts[label] = count 

226 return counts 

227 

228 

229def _dump_pkg_info(filename): 

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

231 

232 Parameters 

233 ---------- 

234 filename : `str` 

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

236 of the packages. 

237 """ 

238 file = Path(filename) 

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

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

241 packages = Packages.fromSystem() 

242 packages.write(str(file)) 

243 

244 

245def _dump_env_info(filename): 

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

247 

248 Parameters 

249 ---------- 

250 filename : `str` 

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

252 environment. 

253 """ 

254 file = Path(filename) 

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

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

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

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