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

83 statements  

« prev     ^ index     » next       coverage.py v7.2.2, created at 2023-03-23 16:51 -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 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 

32 

33from lsst.ctrl.bps.bps_utils import _create_execution_butler 

34from lsst.daf.butler import 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 # Copy Execution Butler if user provided (shouldn't provide execution 

86 # butler if not providing QuantumGraph) 

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

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

89 if not found: 

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

91 

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

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

94 with time_this( 

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

96 ): 

97 shutil.copytree(user_exec_butler_dir, execution_butler_dir) 

98 else: 

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

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

101 

102 # Run command to create the QuantumGraph. 

103 _LOG.info("Creating quantum graph") 

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

105 qgraph_filename = create_quantum_graph(config, out_prefix) 

106 

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

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

109 qgraph = read_quantum_graph(qgraph_filename, config["butlerConfig"]) 

110 

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

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

113 raise OSError( 

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

115 f"creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE" 

116 ) 

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

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

119 

120 return qgraph_filename, qgraph, execution_butler_dir 

121 

122 

123def execute(command, filename): 

124 """Execute a command. 

125 

126 Parameters 

127 ---------- 

128 command : `str` 

129 String representing the command to execute. 

130 filename : `str` 

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

132 

133 Returns 

134 ------- 

135 exit_code : `int` 

136 The exit code the command being executed finished with. 

137 """ 

138 buffer_size = 5000 

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

140 print(command, file=fh) 

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

142 process = subprocess.Popen( 

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

144 ) 

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

146 while process.poll is None or buffer: 

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

148 _LOG.info(buffer) 

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

150 process.stdout.close() 

151 process.wait() 

152 return process.returncode 

153 

154 

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

156 """Create QuantumGraph from pipeline definition. 

157 

158 Parameters 

159 ---------- 

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

161 BPS configuration. 

162 out_prefix : `str`, optional 

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

164 from generating the QuantumGraph. Defaults to empty string so 

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

166 directory. 

167 

168 Returns 

169 ------- 

170 qgraph_filename : `str` 

171 Name of file containing generated QuantumGraph. 

172 """ 

173 # Create name of file to store QuantumGraph. 

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

175 

176 # Get QuantumGraph generation command. 

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

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

179 if not found: 

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

181 _LOG.info(cmd) 

182 

183 # Run QuantumGraph generation. 

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

185 status = execute(cmd, out) 

186 if status != 0: 

187 raise RuntimeError( 

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

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

190 ) 

191 return qgraph_filename 

192 

193 

194def read_quantum_graph(qgraph_filename, butler_uri): 

195 """Read the QuantumGraph from disk. 

196 

197 Parameters 

198 ---------- 

199 qgraph_filename : `str` 

200 Name of file containing QuantumGraph to be used for workflow 

201 generation. 

202 butler_uri : `str` 

203 Location of butler repository that can be used to create a 

204 butler object. 

205 

206 Returns 

207 ------- 

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

209 The QuantumGraph read from a file. 

210 

211 Raises 

212 ------ 

213 RuntimeError 

214 If the QuantumGraph contains 0 Quanta. 

215 """ 

216 # Get the DimensionUniverse from the butler repository 

217 butler = Butler(butler_uri, writeable=False) 

218 qgraph = QuantumGraph.loadUri(qgraph_filename, butler.registry.dimensions) 

219 if len(qgraph) == 0: 

220 raise RuntimeError("QuantumGraph is empty") 

221 return qgraph 

222 

223 

224def cluster_quanta(config, qgraph, name): 

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

226 together. 

227 

228 Parameters 

229 ---------- 

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

231 BPS configuration. 

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

233 Original full QuantumGraph for the run. 

234 name : `str` 

235 Name for the ClusteredQuantumGraph that will be generated. 

236 

237 Returns 

238 ------- 

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

240 Generated ClusteredQuantumGraph. 

241 """ 

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

243 return cluster_func(config, qgraph, name)