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

Shortcuts on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

82 statements  

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.daf.butler import DimensionUniverse 

34from lsst.pipe.base.graph import QuantumGraph 

35from lsst.utils import doImport 

36from lsst.ctrl.bps.bps_utils import _create_execution_butler 

37from lsst.utils.logging import VERBOSE 

38from lsst.utils.timer import time_this, timeMethod 

39 

40 

41_LOG = logging.getLogger(__name__) 

42 

43 

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

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

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

47 

48 Parameters 

49 ---------- 

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

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

52 out_prefix : `str`, optional 

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

54 the QuantumGraph. Default value is empty string. 

55 

56 Returns 

57 ------- 

58 qgraph_filename : `str` 

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

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

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

62 result of running code that generates it. 

63 execution_butler_dir : `str` or None 

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

65 created during this submission step. 

66 """ 

67 # consistently name execution butler directory 

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

69 if not execution_butler_dir.startswith('/'): 

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

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

72 

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

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

75 if found and input_qgraph_filename: 

76 if out_prefix is not None: 

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

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

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

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

81 shutil.copy2(input_qgraph_filename, qgraph_filename) 

82 else: 

83 # Use QuantumGraph file in original given location. 

84 qgraph_filename = input_qgraph_filename 

85 

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

87 # butler if not providing QuantumGraph) 

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

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

90 if not found: 

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

92 

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

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

95 with time_this(log=_LOG, level=logging.INFO, prefix=None, 

96 msg="Completed copying execution butler"): 

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) 

110 

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

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

113 raise OSError(f"Missing execution butler dir ({execution_butler_dir}) after " 

114 f"creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE") 

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

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

117 

118 return qgraph_filename, qgraph, execution_butler_dir 

119 

120 

121def execute(command, filename): 

122 """Execute a command. 

123 

124 Parameters 

125 ---------- 

126 command : `str` 

127 String representing the command to execute. 

128 filename : `str` 

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

130 

131 Returns 

132 ------- 

133 exit_code : `int` 

134 The exit code the command being executed finished with. 

135 """ 

136 buffer_size = 5000 

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

138 print(command, file=fh) 

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

140 process = subprocess.Popen( 

141 shlex.split(command), shell=False, stdout=subprocess.PIPE, 

142 stderr=subprocess.STDOUT 

143 ) 

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

145 while process.poll is None or buffer: 

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

147 _LOG.info(buffer) 

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

149 process.stdout.close() 

150 process.wait() 

151 return process.returncode 

152 

153 

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

155 """Create QuantumGraph from pipeline definition. 

156 

157 Parameters 

158 ---------- 

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

160 BPS configuration. 

161 out_prefix : `str`, optional 

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

163 from generating the QuantumGraph. Defaults to empty string so 

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

165 directory. 

166 

167 Returns 

168 ------- 

169 qgraph_filename : `str` 

170 Name of file containing generated QuantumGraph. 

171 """ 

172 # Create name of file to store QuantumGraph. 

173 qgraph_filename = os.path.join(out_prefix, config['qgraphFileTemplate']) 

174 

175 # Get QuantumGraph generation command. 

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

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

178 if not found: 

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

180 _LOG.info(cmd) 

181 

182 # Run QuantumGraph generation. 

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

184 status = execute(cmd, out) 

185 if status != 0: 

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

187 f"Check {out} for more details.") 

188 return qgraph_filename 

189 

190 

191def read_quantum_graph(qgraph_filename): 

192 """Read the QuantumGraph from disk. 

193 

194 Parameters 

195 ---------- 

196 qgraph_filename : `str` 

197 Name of file containing QuantumGraph to be used for workflow 

198 generation. 

199 

200 Returns 

201 ------- 

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

203 The QuantumGraph read from a file. 

204 

205 Raises 

206 ------ 

207 RuntimeError 

208 If the QuantumGraph contains 0 Quanta. 

209 """ 

210 qgraph = QuantumGraph.loadUri(qgraph_filename, DimensionUniverse()) 

211 if len(qgraph) == 0: 

212 raise RuntimeError("QuantumGraph is empty") 

213 return qgraph 

214 

215 

216def cluster_quanta(config, qgraph, name): 

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

218 together. 

219 

220 Parameters 

221 ---------- 

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

223 BPS configuration. 

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

225 Original full QuantumGraph for the run. 

226 name : `str` 

227 Name for the ClusteredQuantumGraph that will be generated. 

228 

229 Returns 

230 ------- 

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

232 Generated ClusteredQuantumGraph. 

233 """ 

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

235 return cluster_func(config, qgraph, name)