Coverage for python/lsst/ctrl/bps/drivers.py: 19%

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

104 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 <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 shutil 

45 

46 

47from lsst.daf.butler.core.utils import time_this 

48from lsst.obs.base import Instrument 

49 

50from . import BPS_SEARCH_ORDER, BpsConfig 

51from .pre_transform import acquire_quantum_graph, cluster_quanta 

52from .transform import transform 

53from .prepare import prepare 

54from .submit import submit 

55from .cancel import cancel 

56from .report import report 

57 

58 

59_LOG = logging.getLogger(__name__) 

60 

61 

62def _init_submission_driver(config_file, **kwargs): 

63 """Initialize runtime environment. 

64 

65 Parameters 

66 ---------- 

67 config_file : `str` 

68 Name of the configuration file. 

69 

70 Returns 

71 ------- 

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

73 Batch Processing Service configuration. 

74 """ 

75 config = BpsConfig(config_file, BPS_SEARCH_ORDER) 

76 

77 # Override config with command-line values 

78 # Handle diffs between pipetask argument names vs bps yaml 

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

80 "output_run": "outputRun", 

81 "qgraph": "qgraphFile", 

82 "pipeline": "pipelineYaml"} 

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

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

85 if value: 

86 # pipetask argument parser converts some values to list, 

87 # but bps will want string. 

88 if not isinstance(value, str): 

89 value = ",".join(value) 

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

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

92 

93 # Set some initial values 

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

95 if "operator" not in config: 

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

97 

98 if "outCollection" in config: 

99 raise KeyError("outCollection is deprecated. Replace all outCollection references with outputRun.") 

100 

101 if "outputRun" not in config: 

102 raise KeyError("Must specify the output run collection using outputRun") 

103 

104 if "uniqProcName" not in config: 

105 config[".bps_defined.uniqProcName"] = config["outputRun"].replace("/", "_") 

106 

107 if "submitPath" not in config: 

108 raise KeyError("Must specify the submit-side run directory using submitPath") 

109 

110 # make submit directory to contain all outputs 

111 submit_path = config["submitPath"] 

112 os.makedirs(submit_path, exist_ok=True) 

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

114 

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

116 shutil.copy2(config_file, submit_path) 

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

118 config.dump(fh) 

119 

120 return config 

121 

122 

123def acquire_qgraph_driver(config_file, **kwargs): 

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

125 

126 Parameters 

127 ---------- 

128 config_file : `str` 

129 Name of the configuration file. 

130 

131 Returns 

132 ------- 

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

134 Updated configuration. 

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

136 A graph representing quanta. 

137 """ 

138 config = _init_submission_driver(config_file, **kwargs) 

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

140 

141 _LOG.info("Starting acquire stage (generating and/or reading quantum graph)") 

142 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Acquire stage completed"): 

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

144 

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

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

147 return config, qgraph 

148 

149 

150def cluster_qgraph_driver(config_file, **kwargs): 

151 """Group quanta into clusters. 

152 

153 Parameters 

154 ---------- 

155 config_file : `str` 

156 Name of the configuration file. 

157 

158 Returns 

159 ------- 

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

161 Updated configuration. 

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

163 A graph representing clustered quanta. 

164 """ 

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

166 

167 _LOG.info("Starting cluster stage (grouping quanta into jobs)") 

168 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Cluster stage completed"): 

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

170 

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

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

173 if save_clustered_qgraph: 

174 clustered_qgraph.save(os.path.join(submit_path, "bps_clustered_qgraph.pickle")) 

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

176 if save_dot: 

177 clustered_qgraph.draw(os.path.join(submit_path, "bps_clustered_qgraph.dot")) 

178 return config, clustered_qgraph 

179 

180 

181def transform_driver(config_file, **kwargs): 

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

183 

184 Parameters 

185 ---------- 

186 config_file : `str` 

187 Name of the configuration file. 

188 

189 Returns 

190 ------- 

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

192 Configuration to use when creating the workflow. 

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

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

195 workflow management system. 

196 """ 

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

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

199 

200 _LOG.info("Starting transform stage (creating generic workflow)") 

201 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Transform stage completed"): 

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

203 _LOG.info("Generic workflow name '%s'", generic_workflow.name) 

204 

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

206 if save_workflow: 

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

208 generic_workflow.save(outfh, "pickle") 

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

210 if save_dot: 

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

212 generic_workflow.draw(outfh, "dot") 

213 return generic_workflow_config, generic_workflow 

214 

215 

216def prepare_driver(config_file, **kwargs): 

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

218 

219 Parameters 

220 ---------- 

221 config_file : `str` 

222 Name of the configuration file. 

223 

224 Returns 

225 ------- 

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

227 Configuration to use when creating the workflow. 

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

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

230 workflow management system. 

231 """ 

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

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

234 

235 _LOG.info("Starting prepare stage (creating specific implementation of workflow)") 

236 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Prepare stage completed"): 

237 wms_workflow = prepare(generic_workflow_config, generic_workflow, submit_path) 

238 

239 wms_workflow_config = generic_workflow_config 

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

241 return wms_workflow_config, wms_workflow 

242 

243 

244def submit_driver(config_file, **kwargs): 

245 """Submit workflow for execution. 

246 

247 Parameters 

248 ---------- 

249 config_file : `str` 

250 Name of the configuration file. 

251 """ 

252 _LOG.info("Starting submission process") 

253 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed entire submission process"): 

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

255 

256 _LOG.info("Starting submit stage") 

257 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed submit stage"): 

258 submit(wms_workflow_config, wms_workflow) 

259 _LOG.info("Run '%s' submitted for execution with id '%s'", wms_workflow.name, wms_workflow.run_id) 

260 

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

262 

263 

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

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

266 

267 Parameters 

268 ---------- 

269 wms_service : `str` 

270 Name of the class. 

271 run_id : `str` 

272 A run id the report will be restricted to. 

273 user : `str` 

274 A user name the report will be restricted to. 

275 hist_days : int 

276 Number of days 

277 pass_thru : `str` 

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

279 """ 

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

281 

282 

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

284 """Cancel submitted workflows. 

285 

286 Parameters 

287 ---------- 

288 wms_service : `str` 

289 Name of the Workload Management System service class. 

290 run_id : `str` 

291 ID or path of job that should be canceled. 

292 user : `str` 

293 User whose submitted jobs should be canceled. 

294 require_bps : `bool` 

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

296 pass_thru : `str` 

297 Information to pass through to WMS. 

298 """ 

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