Hide keyboard shortcuts

Hot-keys 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

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 <http://www.gnu.org/licenses/>. 

21 

22"""Driver functions for each subcommand. 

23 

24Driver functions ensure that ensure all setup work is done before running 

25the subcommand method. 

26""" 

27 

28 

29__all__ = [ 

30 "acquire_qgraph_driver", 

31 "cluster_qgraph_driver", 

32 "transform_driver", 

33 "prepare_driver", 

34 "submit_driver", 

35 "report_driver", 

36 "cancel_driver", 

37] 

38 

39 

40import getpass 

41import logging 

42import os 

43import re 

44import pickle 

45import time 

46import shutil 

47 

48 

49from lsst.obs.base import Instrument 

50 

51from . import BPS_SEARCH_ORDER, BpsConfig 

52from .bps_draw import draw_networkx_dot 

53from .pre_transform import acquire_quantum_graph, cluster_quanta 

54from .transform import transform 

55from .prepare import prepare 

56from .submit import submit 

57from .cancel import cancel 

58from .report import report 

59 

60 

61_LOG = logging.getLogger(__name__) 

62 

63 

64def _init_submission_driver(config_file, **kwargs): 

65 """Initialize runtime environment. 

66 

67 Parameters 

68 ---------- 

69 config_file : `str` 

70 Name of the configuration file. 

71 

72 Returns 

73 ------- 

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

75 Batch Processing Service configuration. 

76 """ 

77 config = BpsConfig(config_file, BPS_SEARCH_ORDER) 

78 

79 # Override config with command-line values 

80 # Handle diffs between pipetask argument names vs bps yaml 

81 translation = {"input": "inCollection", 

82 "output_run": "outCollection", 

83 "qgraph": "qgraphFile", 

84 "pipeline": "pipelineYaml"} 

85 for key, value in kwargs.items(): 

86 # Don't want to override config with None or empty string values. 

87 if value: 

88 # pipetask argument parser converts some values to list, 

89 # but bps will want string. 

90 if not isinstance(value, str): 

91 value = ",".join(value) 

92 new_key = translation.get(key, re.sub(r"_(\S)", lambda match: match.group(1).upper(), key)) 

93 config[f".bps_cmdline.{new_key}"] = value 

94 

95 # Set some initial values 

96 config[".bps_defined.timestamp"] = Instrument.makeCollectionTimestamp() 

97 if "operator" not in config: 

98 config[".bps_defined.operator"] = getpass.getuser() 

99 

100 if "uniqProcName" not in config: 

101 config[".bps_defined.uniqProcName"] = config["outCollection"].replace("/", "_") 

102 

103 # make submit directory to contain all outputs 

104 submit_path = config["submitPath"] 

105 os.makedirs(submit_path, exist_ok=True) 

106 config[".bps_defined.submitPath"] = submit_path 

107 

108 # save copy of configs (orig and expanded config) 

109 shutil.copy2(config_file, submit_path) 

110 with open(f"{submit_path}/{config['uniqProcName']}_config.yaml", "w") as fh: 

111 config.dump(fh) 

112 

113 return config 

114 

115 

116def acquire_qgraph_driver(config_file, **kwargs): 

117 """Read a quantum graph from a file or create one from pipeline definition. 

118 

119 Parameters 

120 ---------- 

121 config_file : `str` 

122 Name of the configuration file. 

123 

124 Returns 

125 ------- 

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

127 Updated configuration. 

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

129 A graph representing quanta. 

130 """ 

131 stime = time.time() 

132 config = _init_submission_driver(config_file, **kwargs) 

133 submit_path = config[".bps_defined.submitPath"] 

134 _LOG.info("Acquiring QuantumGraph (it will be created from pipeline definition if needed)") 

135 qgraph_file, qgraph, execution_butler_dir = acquire_quantum_graph(config, out_prefix=submit_path) 

136 config[".bps_defined.executionButlerDir"] = execution_butler_dir 

137 _LOG.info("Run QuantumGraph file %s", qgraph_file) 

138 config[".bps_defined.runQgraphFile"] = qgraph_file 

139 _LOG.info("Acquiring QuantumGraph took %.2f seconds", time.time() - stime) 

140 

141 return config, qgraph 

142 

143 

144def cluster_qgraph_driver(config_file, **kwargs): 

145 """Group quanta into clusters. 

146 

147 Parameters 

148 ---------- 

149 config_file : `str` 

150 Name of the configuration file. 

151 

152 Returns 

153 ------- 

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

155 Updated configuration. 

156 clustered_qgraph : `lsst.ctrl.bps.ClusteredQuantumGraph` 

157 A graph representing clustered quanta. 

158 """ 

159 stime = time.time() 

160 config, qgraph = acquire_qgraph_driver(config_file, **kwargs) 

161 _LOG.info("Clustering quanta") 

162 clustered_qgraph = cluster_quanta(config, qgraph, config["uniqProcName"]) 

163 _LOG.info("Clustering quanta took %.2f seconds", time.time() - stime) 

164 

165 submit_path = config[".bps_defined.submitPath"] 

166 _, save_clustered_qgraph = config.search("saveClusteredQgraph", opt={"default": False}) 

167 if save_clustered_qgraph: 

168 with open(os.path.join(submit_path, "bps_clustered_qgraph.pickle"), "wb") as outfh: 

169 pickle.dump(clustered_qgraph, outfh) 

170 _, save_dot = config.search("saveDot", opt={"default": False}) 

171 if save_dot: 

172 draw_networkx_dot(clustered_qgraph, os.path.join(submit_path, "bps_clustered_qgraph.dot")) 

173 return config, clustered_qgraph 

174 

175 

176def transform_driver(config_file, **kwargs): 

177 """Create a workflow for a specific workflow management system. 

178 

179 Parameters 

180 ---------- 

181 config_file : `str` 

182 Name of the configuration file. 

183 

184 Returns 

185 ------- 

186 generic_workflow_config : `lsst.ctrl.bps.BpsConfig` 

187 Configuration to use when creating the workflow. 

188 generic_workflow : `lsst.ctrl.bps.BaseWmsWorkflow` 

189 Representation of the abstract/scientific workflow specific to a given 

190 workflow management system. 

191 """ 

192 stime = time.time() 

193 config, clustered_qgraph = cluster_qgraph_driver(config_file, **kwargs) 

194 submit_path = config[".bps_defined.submitPath"] 

195 _LOG.info("Creating Generic Workflow") 

196 generic_workflow, generic_workflow_config = transform(config, clustered_qgraph, submit_path) 

197 _LOG.info("Creating Generic Workflow took %.2f seconds", time.time() - stime) 

198 _LOG.info("Generic Workflow name %s", generic_workflow.name) 

199 

200 _, save_workflow = config.search("saveGenericWorkflow", opt={"default": False}) 

201 if save_workflow: 

202 with open(os.path.join(submit_path, "bps_generic_workflow.pickle"), "wb") as outfh: 

203 generic_workflow.save(outfh, "pickle") 

204 _, save_dot = config.search("saveDot", opt={"default": False}) 

205 if save_dot: 

206 with open(os.path.join(submit_path, "bps_generic_workflow.dot"), "w") as outfh: 

207 generic_workflow.draw(outfh, "dot") 

208 return generic_workflow_config, generic_workflow 

209 

210 

211def prepare_driver(config_file, **kwargs): 

212 """Create a representation of the generic workflow. 

213 

214 Parameters 

215 ---------- 

216 config_file : `str` 

217 Name of the configuration file. 

218 

219 Returns 

220 ------- 

221 wms_config : `lsst.ctrl.bps.BpsConfig` 

222 Configuration to use when creating the workflow. 

223 workflow : `lsst.ctrl.bps.BaseWmsWorkflow` 

224 Representation of the abstract/scientific workflow specific to a given 

225 workflow management system. 

226 """ 

227 stime = time.time() 

228 generic_workflow_config, generic_workflow = transform_driver(config_file, **kwargs) 

229 submit_path = generic_workflow_config[".bps_defined.submitPath"] 

230 _LOG.info("Creating specific implementation of workflow") 

231 wms_workflow = prepare(generic_workflow_config, generic_workflow, submit_path) 

232 wms_workflow_config = generic_workflow_config 

233 _LOG.info("Creating specific implementation of workflow took %.2f seconds", time.time() - stime) 

234 print(f"Submit dir: {wms_workflow.submit_path}") 

235 return wms_workflow_config, wms_workflow 

236 

237 

238def submit_driver(config_file, **kwargs): 

239 """Submit workflow for execution. 

240 

241 Parameters 

242 ---------- 

243 config_file : `str` 

244 Name of the configuration file. 

245 """ 

246 wms_workflow_config, wms_workflow = prepare_driver(config_file, **kwargs) 

247 submit(wms_workflow_config, wms_workflow) 

248 print(f"Run Id: {wms_workflow.run_id}") 

249 

250 

251def report_driver(wms_service, run_id, user, hist_days, pass_thru): 

252 """Print out summary of jobs submitted for execution. 

253 

254 Parameters 

255 ---------- 

256 wms_service : `str` 

257 Name of the class. 

258 run_id : `str` 

259 A run id the report will be restricted to. 

260 user : `str` 

261 A user name the report will be restricted to. 

262 hist_days : int 

263 Number of days 

264 pass_thru : `str` 

265 A string to pass directly to the WMS service class. 

266 """ 

267 report(wms_service, run_id, user, hist_days, pass_thru) 

268 

269 

270def cancel_driver(wms_service, run_id, user, require_bps, pass_thru): 

271 """Cancel submitted workflows. 

272 

273 Parameters 

274 ---------- 

275 wms_service : `str` 

276 Name of the Workload Management System service class. 

277 run_id : `str` 

278 ID or path of job that should be canceled. 

279 user : `str` 

280 User whose submitted jobs should be canceled. 

281 require_bps : `bool` 

282 Whether to require given run_id/user to be a bps submitted job. 

283 pass_thru : `str` 

284 Information to pass through to WMS. 

285 """ 

286 cancel(wms_service, run_id, user, require_bps, pass_thru)