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

96 statements  

« prev     ^ index     » next       coverage.py v7.3.0, created at 2023-08-23 10:45 +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"""Driver to execute steps outside of BPS that need to be done first 

23including running QuantumGraph generation and reading the QuantumGraph 

24into memory. 

25""" 

26 

27import logging 

28import os 

29import shlex 

30import shutil 

31import subprocess 

32from pathlib import Path 

33 

34from lsst.ctrl.bps.bps_utils import _create_execution_butler 

35from lsst.pipe.base.graph import QuantumGraph 

36from lsst.utils import doImport 

37from lsst.utils.logging import VERBOSE 

38from lsst.utils.timer import time_this, timeMethod 

39 

40_LOG = logging.getLogger(__name__) 

41 

42 

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

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

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

46 

47 Parameters 

48 ---------- 

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

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

51 out_prefix : `str`, optional 

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

53 the QuantumGraph. Default value is empty string. 

54 

55 Returns 

56 ------- 

57 qgraph_filename : `str` 

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

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

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

61 result of running code that generates it. 

62 execution_butler_dir : `str` or None 

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

64 created during this submission step. 

65 """ 

66 # consistently name execution butler directory 

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

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

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

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

71 

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

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

74 if found and input_qgraph_filename: 

75 if out_prefix is not None: 

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

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

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

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

80 shutil.copy2(input_qgraph_filename, qgraph_filename) 

81 else: 

82 # Use QuantumGraph file in original given location. 

83 qgraph_filename = input_qgraph_filename 

84 

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

86 if "finalJob" in config: 

87 update_quantum_graph(config, qgraph_filename, out_prefix) 

88 

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

90 # butler if not providing QuantumGraph) 

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

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

93 if not found: 

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

95 

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

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

98 with time_this( 

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

100 ): 

101 shutil.copytree(user_exec_butler_dir, execution_butler_dir) 

102 else: 

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

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

105 

106 # Run command to create the QuantumGraph. 

107 _LOG.info("Creating quantum graph") 

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

109 qgraph_filename = create_quantum_graph(config, out_prefix) 

110 

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

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

113 qgraph = QuantumGraph.loadUri(qgraph_filename) 

114 

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

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

117 raise OSError( 

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

119 "creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE" 

120 ) 

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

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

123 

124 return qgraph_filename, qgraph, execution_butler_dir 

125 

126 

127def execute(command, filename): 

128 """Execute a command. 

129 

130 Parameters 

131 ---------- 

132 command : `str` 

133 String representing the command to execute. 

134 filename : `str` 

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

136 

137 Returns 

138 ------- 

139 exit_code : `int` 

140 The exit code the command being executed finished with. 

141 """ 

142 buffer_size = 5000 

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

144 print(command, file=fh) 

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

146 process = subprocess.Popen( 

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

148 ) 

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

150 while process.poll is None or buffer: 

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

152 _LOG.info(buffer) 

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

154 process.stdout.close() 

155 process.wait() 

156 return process.returncode 

157 

158 

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

160 """Create QuantumGraph from pipeline definition. 

161 

162 Parameters 

163 ---------- 

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

165 BPS configuration. 

166 out_prefix : `str`, optional 

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

168 from generating the QuantumGraph. Defaults to empty string so 

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

170 directory. 

171 

172 Returns 

173 ------- 

174 qgraph_filename : `str` 

175 Name of file containing generated QuantumGraph. 

176 """ 

177 # Create name of file to store QuantumGraph. 

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

179 

180 # Get QuantumGraph generation command. 

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

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

183 if not found: 

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

185 _LOG.info(cmd) 

186 

187 # Run QuantumGraph generation. 

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

189 status = execute(cmd, out) 

190 if status != 0: 

191 raise RuntimeError( 

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

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

194 ) 

195 return qgraph_filename 

196 

197 

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

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

200 

201 Parameters 

202 ---------- 

203 config : `BpsConfig` 

204 BPS configuration. 

205 qgraph_filename : `str` 

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

207 out_prefix : `str`, optional 

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

209 from generating the QuantumGraph. Defaults to empty string so 

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

211 directory. 

212 inplace : `bool`, optional 

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

214 creating a backup copy. Defaults to False. 

215 """ 

216 src_qgraph = Path(qgraph_filename) 

217 dest_qgraph = Path(qgraph_filename) 

218 

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

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

221 if not inplace: 

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

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

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

225 shutil.copy2(qgraph_filename, src_qgraph) 

226 

227 # Get the command for updating the quantum graph. 

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

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

230 if not found: 

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

232 _LOG.info(cmd) 

233 

234 # Run the command to update the quantum graph. 

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

236 status = execute(cmd, out) 

237 if status != 0: 

238 raise RuntimeError( 

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

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

241 ) 

242 

243 

244def cluster_quanta(config, qgraph, name): 

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

246 together. 

247 

248 Parameters 

249 ---------- 

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

251 BPS configuration. 

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

253 Original full QuantumGraph for the run. 

254 name : `str` 

255 Name for the ClusteredQuantumGraph that will be generated. 

256 

257 Returns 

258 ------- 

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

260 Generated ClusteredQuantumGraph. 

261 """ 

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

263 return cluster_func(config, qgraph, name)