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

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

80 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.daf.butler.core.utils import time_this 

35from lsst.pipe.base.graph import QuantumGraph 

36from lsst.utils import doImport 

37from lsst.ctrl.bps.bps_utils import _create_execution_butler 

38 

39 

40_LOG = logging.getLogger(__name__) 

41 

42 

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

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

45 

46 Parameters 

47 ---------- 

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

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

50 out_prefix : `str`, optional 

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

52 the QuantumGraph. Default value is empty string. 

53 

54 Returns 

55 ------- 

56 qgraph_filename : `str` 

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

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

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

60 result of running code that generates it. 

61 execution_butler_dir : `str` or None 

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

63 created during this submission step. 

64 """ 

65 # consistently name execution butler directory 

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

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

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

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

70 

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

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

73 if found and input_qgraph_filename: 

74 if out_prefix is not None: 

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

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

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

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

79 shutil.copy2(input_qgraph_filename, qgraph_filename) 

80 else: 

81 # Use QuantumGraph file in original given location. 

82 qgraph_filename = input_qgraph_filename 

83 

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

85 # butler if not providing QuantumGraph) 

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

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

88 if not found: 

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

90 

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

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

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

94 msg="Completed copying execution butler"): 

95 shutil.copytree(user_exec_butler_dir, execution_butler_dir) 

96 else: 

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

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

99 

100 # Run command to create the QuantumGraph. 

101 _LOG.info("Creating quantum graph") 

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

103 qgraph_filename = create_quantum_graph(config, out_prefix) 

104 

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

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

107 qgraph = read_quantum_graph(qgraph_filename) 

108 

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

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

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

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

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

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

115 

116 return qgraph_filename, qgraph, execution_butler_dir 

117 

118 

119def execute(command, filename): 

120 """Execute a command. 

121 

122 Parameters 

123 ---------- 

124 command : `str` 

125 String representing the command to execute. 

126 filename : `str` 

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

128 

129 Returns 

130 ------- 

131 exit_code : `int` 

132 The exit code the command being executed finished with. 

133 """ 

134 buffer_size = 5000 

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

136 print(command, file=fh) 

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

138 process = subprocess.Popen( 

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

140 stderr=subprocess.STDOUT 

141 ) 

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

143 while process.poll is None or buffer: 

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

145 _LOG.info(buffer) 

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

147 process.stdout.close() 

148 process.wait() 

149 return process.returncode 

150 

151 

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

153 """Create QuantumGraph from pipeline definition. 

154 

155 Parameters 

156 ---------- 

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

158 BPS configuration. 

159 out_prefix : `str`, optional 

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

161 from generating the QuantumGraph. Defaults to empty string so 

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

163 directory. 

164 

165 Returns 

166 ------- 

167 qgraph_filename : `str` 

168 Name of file containing generated QuantumGraph. 

169 """ 

170 # Create name of file to store QuantumGraph. 

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

172 

173 # Get QuantumGraph generation command. 

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

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

176 if not found: 

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

178 _LOG.info(cmd) 

179 

180 # Run QuantumGraph generation. 

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

182 status = execute(cmd, out) 

183 if status != 0: 

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

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

186 return qgraph_filename 

187 

188 

189def read_quantum_graph(qgraph_filename): 

190 """Read the QuantumGraph from disk. 

191 

192 Parameters 

193 ---------- 

194 qgraph_filename : `str` 

195 Name of file containing QuantumGraph to be used for workflow 

196 generation. 

197 

198 Returns 

199 ------- 

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

201 The QuantumGraph read from a file. 

202 

203 Raises 

204 ------ 

205 RuntimeError 

206 If the QuantumGraph contains 0 Quanta. 

207 """ 

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

209 if len(qgraph) == 0: 

210 raise RuntimeError("QuantumGraph is empty") 

211 return qgraph 

212 

213 

214def cluster_quanta(config, qgraph, name): 

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

216 together. 

217 

218 Parameters 

219 ---------- 

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

221 BPS configuration. 

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

223 Original full QuantumGraph for the run. 

224 name : `str` 

225 Name for the ClusteredQuantumGraph that will be generated. 

226 

227 Returns 

228 ------- 

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

230 Generated ClusteredQuantumGraph. 

231 """ 

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

233 return cluster_func(config, qgraph, name)