Coverage for python/lsst/ctrl/bps/pre_transform.py: 15%

96 statements  

« prev     ^ index     » next       coverage.py v7.3.3, created at 2023-12-20 17:34 +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 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"""Driver to execute steps outside of BPS that need to be done first 

29including running QuantumGraph generation and reading the QuantumGraph 

30into memory. 

31""" 

32 

33import logging 

34import os 

35import shlex 

36import shutil 

37import subprocess 

38from pathlib import Path 

39 

40from lsst.ctrl.bps.bps_utils import _create_execution_butler 

41from lsst.pipe.base.graph import QuantumGraph 

42from lsst.utils import doImport 

43from lsst.utils.logging import VERBOSE 

44from lsst.utils.timer import time_this, timeMethod 

45 

46_LOG = logging.getLogger(__name__) 

47 

48 

49@timeMethod(logger=_LOG, logLevel=VERBOSE) 

50def acquire_quantum_graph(config, out_prefix=""): 

51 """Read a quantum graph from a file or create one from scratch. 

52 

53 Parameters 

54 ---------- 

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

56 Configuration values for BPS. In particular, looking for qgraphFile. 

57 out_prefix : `str`, optional 

58 Output path for the QuantumGraph and stdout/stderr from generating 

59 the QuantumGraph. Default value is empty string. 

60 

61 Returns 

62 ------- 

63 qgraph_filename : `str` 

64 Name of file containing QuantumGraph that was read into qgraph. 

65 qgraph : `lsst.pipe.base.graph.QuantumGraph` 

66 A QuantumGraph read in from pre-generated file or one that is the 

67 result of running code that generates it. 

68 execution_butler_dir : `str` or None 

69 The directory containing the execution butler if user-provided or 

70 created during this submission step. 

71 """ 

72 # consistently name execution butler directory 

73 _, execution_butler_dir = config.search("executionButlerTemplate") 

74 if not execution_butler_dir.startswith("/"): 

75 execution_butler_dir = os.path.join(config["submitPath"], execution_butler_dir) 

76 _, when_create = config.search(".executionButler.whenCreate") 

77 

78 # Check to see if user provided pre-generated QuantumGraph. 

79 found, input_qgraph_filename = config.search("qgraphFile") 

80 if found and input_qgraph_filename: 

81 if out_prefix is not None: 

82 # Save a copy of the QuantumGraph file in out_prefix. 

83 _LOG.info("Copying quantum graph from '%s'", input_qgraph_filename) 

84 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed copying quantum graph"): 

85 qgraph_filename = os.path.join(out_prefix, os.path.basename(input_qgraph_filename)) 

86 shutil.copy2(input_qgraph_filename, qgraph_filename) 

87 else: 

88 # Use QuantumGraph file in original given location. 

89 qgraph_filename = input_qgraph_filename 

90 

91 # Update the output run in the user provided quantum graph. 

92 if "finalJob" in config: 

93 update_quantum_graph(config, qgraph_filename, out_prefix) 

94 

95 # Copy Execution Butler if user provided (shouldn't provide execution 

96 # butler if not providing QuantumGraph) 

97 if when_create.upper() == "USER_PROVIDED": 

98 found, user_exec_butler_dir = config.search(".executionButler.executionButlerDir") 

99 if not found: 

100 raise KeyError("Missing .executionButler.executionButlerDir for when_create == USER_PROVIDED") 

101 

102 # Save a copy of the execution butler file in out_prefix. 

103 _LOG.info("Copying execution butler to '%s'", user_exec_butler_dir) 

104 with time_this( 

105 log=_LOG, level=logging.INFO, prefix=None, msg="Completed copying execution butler" 

106 ): 

107 shutil.copytree(user_exec_butler_dir, execution_butler_dir) 

108 else: 

109 if when_create.upper() == "USER_PROVIDED": 

110 raise KeyError("Missing qgraphFile to go with provided executionButlerDir") 

111 

112 # Run command to create the QuantumGraph. 

113 _LOG.info("Creating quantum graph") 

114 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed creating quantum graph"): 

115 qgraph_filename = create_quantum_graph(config, out_prefix) 

116 

117 _LOG.info("Reading quantum graph from '%s'", qgraph_filename) 

118 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed reading quantum graph"): 

119 qgraph = QuantumGraph.loadUri(qgraph_filename) 

120 

121 if when_create.upper() == "QGRAPH_CMDLINE": 

122 if not os.path.exists(execution_butler_dir): 

123 raise OSError( 

124 f"Missing execution butler dir ({execution_butler_dir}) after " 

125 "creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE" 

126 ) 

127 elif when_create.upper() == "ACQUIRE": 

128 _create_execution_butler(config, qgraph_filename, execution_butler_dir, config["submitPath"]) 

129 

130 return qgraph_filename, qgraph, execution_butler_dir 

131 

132 

133def execute(command, filename): 

134 """Execute a command. 

135 

136 Parameters 

137 ---------- 

138 command : `str` 

139 String representing the command to execute. 

140 filename : `str` 

141 A file to which both stderr and stdout will be written to. 

142 

143 Returns 

144 ------- 

145 exit_code : `int` 

146 The exit code the command being executed finished with. 

147 """ 

148 buffer_size = 5000 

149 with open(filename, "w") as fh: 

150 print(command, file=fh) 

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

152 process = subprocess.Popen( 

153 shlex.split(command), shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT 

154 ) 

155 buffer = os.read(process.stdout.fileno(), buffer_size).decode() 

156 while process.poll is None or buffer: 

157 print(buffer, end="", file=fh) 

158 _LOG.info(buffer) 

159 buffer = os.read(process.stdout.fileno(), buffer_size).decode() 

160 process.stdout.close() 

161 process.wait() 

162 return process.returncode 

163 

164 

165def create_quantum_graph(config, out_prefix=""): 

166 """Create QuantumGraph from pipeline definition. 

167 

168 Parameters 

169 ---------- 

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

171 BPS configuration. 

172 out_prefix : `str`, optional 

173 Path in which to output QuantumGraph as well as the stdout/stderr 

174 from generating the QuantumGraph. Defaults to empty string so 

175 code will write the QuantumGraph and stdout/stderr to the current 

176 directory. 

177 

178 Returns 

179 ------- 

180 qgraph_filename : `str` 

181 Name of file containing generated QuantumGraph. 

182 """ 

183 # Create name of file to store QuantumGraph. 

184 qgraph_filename = os.path.join(out_prefix, config["qgraphFileTemplate"]) 

185 

186 # Get QuantumGraph generation command. 

187 search_opt = {"curvals": {"qgraphFile": qgraph_filename}} 

188 found, cmd = config.search("createQuantumGraph", opt=search_opt) 

189 if not found: 

190 _LOG.error("command for generating QuantumGraph not found") 

191 _LOG.info(cmd) 

192 

193 # Run QuantumGraph generation. 

194 out = os.path.join(out_prefix, "quantumGraphGeneration.out") 

195 status = execute(cmd, out) 

196 if status != 0: 

197 raise RuntimeError( 

198 f"QuantumGraph generation exited with non-zero exit code ({status})\n" 

199 f"Check {out} for more details." 

200 ) 

201 return qgraph_filename 

202 

203 

204def update_quantum_graph(config, qgraph_filename, out_prefix="", inplace=False): 

205 """Update output run in an existing quantum graph. 

206 

207 Parameters 

208 ---------- 

209 config : `BpsConfig` 

210 BPS configuration. 

211 qgraph_filename : `str` 

212 Name of file containing the quantum graph that needs to be updated. 

213 out_prefix : `str`, optional 

214 Path in which to output QuantumGraph as well as the stdout/stderr 

215 from generating the QuantumGraph. Defaults to empty string so 

216 code will write the QuantumGraph and stdout/stderr to the current 

217 directory. 

218 inplace : `bool`, optional 

219 If set to True, all updates of the graph will be done in place without 

220 creating a backup copy. Defaults to False. 

221 """ 

222 src_qgraph = Path(qgraph_filename) 

223 dest_qgraph = Path(qgraph_filename) 

224 

225 # If requested, create a backup copy of the quantum graph by adding 

226 # '_orig' suffix to its stem (the filename without the extension). 

227 if not inplace: 

228 _LOG.info("Backing up quantum graph from '%s'", qgraph_filename) 

229 src_qgraph = src_qgraph.parent / f"{src_qgraph.stem}_orig{src_qgraph.suffix}" 

230 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed backing up quantum graph"): 

231 shutil.copy2(qgraph_filename, src_qgraph) 

232 

233 # Get the command for updating the quantum graph. 

234 search_opt = {"curvals": {"inputQgraphFile": str(src_qgraph), "qgraphFile": str(dest_qgraph)}} 

235 found, cmd = config.search("updateQuantumGraph", opt=search_opt) 

236 if not found: 

237 _LOG.error("command for updating quantum graph not found") 

238 _LOG.info(cmd) 

239 

240 # Run the command to update the quantum graph. 

241 out = os.path.join(out_prefix, "quantumGraphUpdate.out") 

242 status = execute(cmd, out) 

243 if status != 0: 

244 raise RuntimeError( 

245 f"Updating quantum graph failed with non-zero exit code ({status})\n" 

246 f"Check {out} for more details." 

247 ) 

248 

249 

250def cluster_quanta(config, qgraph, name): 

251 """Call specified function to group quanta into clusters to be run 

252 together. 

253 

254 Parameters 

255 ---------- 

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

257 BPS configuration. 

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

259 Original full QuantumGraph for the run. 

260 name : `str` 

261 Name for the ClusteredQuantumGraph that will be generated. 

262 

263 Returns 

264 ------- 

265 graph : `lsst.ctrl.bps.ClusteredQuantumGraph` 

266 Generated ClusteredQuantumGraph. 

267 """ 

268 cluster_func = doImport(config["clusterAlgorithm"]) 

269 return cluster_func(config, qgraph, name)