Coverage for python / lsst / ctrl / bps / transform.py: 8%

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <https://www.gnu.org/licenses/>. 

27 

28"""Driver for the transformation of a QuantumGraph into a generic workflow.""" 

29 

30import copy 

31import dataclasses 

32import logging 

33import math 

34import os 

35import re 

36 

37from lsst.ctrl.bps import ClusteredQuantumGraph 

38from lsst.pipe.base import QuantumGraph 

39from lsst.utils.logging import VERBOSE 

40from lsst.utils.timer import timeMethod 

41 

42from . import ( 

43 DEFAULT_MEM_RETRIES, 

44 BpsConfig, 

45 GenericWorkflow, 

46 GenericWorkflowExec, 

47 GenericWorkflowFile, 

48 GenericWorkflowJob, 

49) 

50from .bps_utils import WhenToSaveQuantumGraphs, create_job_quantum_graph_filename, save_qg_subgraph 

51 

52# All available job attributes. 

53_ATTRS_ALL = frozenset([field.name for field in dataclasses.fields(GenericWorkflowJob)]) 

54 

55# Job attributes that need to be set to their maximal value in the cluster. 

56_ATTRS_MAX = frozenset( 

57 { 

58 "memory_multiplier", 

59 "number_of_retries", 

60 "request_cpus", 

61 "request_memory", 

62 "request_memory_max", 

63 } 

64) 

65 

66# Job attributes that need to be set to sum of their values in the cluster. 

67_ATTRS_SUM = frozenset( 

68 { 

69 "request_disk", 

70 "request_walltime", 

71 } 

72) 

73 

74# Job attributes do not fall into a specific category 

75_ATTRS_MISC = frozenset( 

76 { 

77 "label", # taskDef labels aren't same in job and may not match job label 

78 "cmdvals", 

79 "profile", 

80 "attrs", 

81 } 

82) 

83 

84# Attributes that need to be the same for each quanta in the cluster. 

85_ATTRS_UNIVERSAL = frozenset(_ATTRS_ALL - (_ATTRS_MAX | _ATTRS_MISC | _ATTRS_SUM)) 

86 

87_LOG = logging.getLogger(__name__) 

88 

89 

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

91def transform( 

92 config: BpsConfig, cqgraph: ClusteredQuantumGraph, prefix: str 

93) -> tuple[GenericWorkflow, BpsConfig]: 

94 """Transform a ClusteredQuantumGraph to a GenericWorkflow. 

95 

96 Parameters 

97 ---------- 

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

99 BPS configuration. 

100 cqgraph : `lsst.ctrl.bps.ClusteredQuantumGraph` 

101 A clustered quantum graph to transform into a generic workflow. 

102 prefix : `str` 

103 Root path for any output files. 

104 

105 Returns 

106 ------- 

107 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

108 The generic workflow transformed from the clustered quantum graph. 

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

110 Configuration to accompany GenericWorkflow. 

111 """ 

112 if cqgraph.name is not None: 

113 name = cqgraph.name 

114 else: 

115 _, name = config.search("uniqProcName", opt={"required": True}) 

116 

117 generic_workflow = create_generic_workflow(config, cqgraph, name, prefix) 

118 generic_workflow_config = create_generic_workflow_config(config, prefix) 

119 

120 return generic_workflow, generic_workflow_config 

121 

122 

123def add_workflow_init_nodes(config, qgraph, generic_workflow): 

124 """Add nodes to workflow graph that perform initialization steps. 

125 

126 Assumes that all of the initialization should be executed prior to any 

127 of the current workflow. 

128 

129 Parameters 

130 ---------- 

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

132 BPS configuration. 

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

134 The quantum graph the generic workflow represents. 

135 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

136 Generic workflow to which the initialization steps should be added. 

137 """ 

138 # Create a workflow graph that will have task and file nodes necessary for 

139 # initializing the pipeline execution 

140 init_workflow = create_init_workflow(config, qgraph, generic_workflow.get_file("runQgraphFile")) 

141 _LOG.debug("init_workflow nodes = %s", init_workflow.nodes()) 

142 generic_workflow.add_workflow_source(init_workflow) 

143 

144 

145def create_init_workflow( 

146 config: BpsConfig, qgraph: QuantumGraph, qgraph_gwfile: GenericWorkflowFile 

147) -> GenericWorkflow: 

148 """Create workflow for running initialization job(s). 

149 

150 Parameters 

151 ---------- 

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

153 BPS configuration. 

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

155 The quantum graph the generic workflow represents. 

156 qgraph_gwfile : `lsst.ctrl.bps.GenericWorkflowFile` 

157 File object for the full run QuantumGraph file. 

158 

159 Returns 

160 ------- 

161 init_workflow : `lsst.ctrl.bps.GenericWorkflow` 

162 GenericWorkflow consisting of job(s) to initialize workflow. 

163 """ 

164 _LOG.debug("creating init subgraph") 

165 _LOG.debug("creating init task input(s)") 

166 search_opt = { 

167 "curvals": {"curr_pipetask": "pipetaskInit"}, 

168 "replaceVars": False, 

169 "expandEnvVars": False, 

170 "replaceEnvVars": True, 

171 "required": False, 

172 } 

173 found, value = config.search("computeSite", opt=search_opt) 

174 if found: 

175 search_opt["curvals"]["curr_site"] = value 

176 found, value = config.search("computeCloud", opt=search_opt) 

177 if found: 

178 search_opt["curvals"]["curr_cloud"] = value 

179 

180 init_workflow = GenericWorkflow("init") 

181 init_workflow.add_file(qgraph_gwfile) 

182 

183 # create job for executing --init-only 

184 gwjob = GenericWorkflowJob("pipetaskInit", "pipetaskInit") 

185 

186 job_values = _get_job_values(config, search_opt, "runQuantumCommand") 

187 job_values["name"] = "pipetaskInit" 

188 job_values["label"] = "pipetaskInit" 

189 

190 # Adjust job attributes values if necessary. 

191 _handle_job_values(job_values, gwjob) 

192 

193 init_workflow.add_job(gwjob) 

194 init_workflow.add_job_inputs(gwjob.name, [qgraph_gwfile]) 

195 _enhance_command(config, init_workflow, gwjob, {}) 

196 

197 return init_workflow 

198 

199 

200def _enhance_command(config, generic_workflow, gwjob, cached_job_values): 

201 """Enhance command line with env and file placeholders 

202 and gather command line values. 

203 

204 Parameters 

205 ---------- 

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

207 BPS configuration. 

208 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

209 Generic workflow that contains the job. 

210 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

211 Generic workflow job to which the updated executable, arguments, 

212 and values should be saved. 

213 cached_job_values : `dict` [`str`, dict[`str`, `~typing.Any`]] 

214 Cached values common across jobs with same label. Updated if values 

215 aren't already saved for given gwjob's label. 

216 """ 

217 _LOG.debug("gwjob given to _enhance_command: %s", gwjob) 

218 

219 search_opt = { 

220 "curvals": {"curr_pipetask": gwjob.label}, 

221 "replaceVars": False, 

222 "expandEnvVars": False, 

223 "replaceEnvVars": True, 

224 "required": False, 

225 } 

226 

227 if gwjob.label not in cached_job_values: 

228 cached_job_values[gwjob.label] = {} 

229 # Allowing whenSaveJobQgraph and useLazyCommands per pipetask label. 

230 key = "whenSaveJobQgraph" 

231 _, when_save = config.search(key, opt=search_opt) 

232 cached_job_values[gwjob.label][key] = WhenToSaveQuantumGraphs[when_save.upper()] 

233 

234 key = "useLazyCommands" 

235 search_opt["default"] = True 

236 _, cached_job_values[gwjob.label][key] = config.search(key, opt=search_opt) 

237 del search_opt["default"] 

238 

239 # Change qgraph variable to match whether using run or per-job qgraph 

240 # Note: these are lookup keys, not actual physical filenames. 

241 if cached_job_values[gwjob.label]["whenSaveJobQgraph"] == WhenToSaveQuantumGraphs.NEVER: 

242 gwjob.arguments = gwjob.arguments.replace("{qgraphFile}", "{runQgraphFile}") 

243 elif gwjob.name == "pipetaskInit": 

244 gwjob.arguments = gwjob.arguments.replace("{qgraphFile}", "{runQgraphFile}") 

245 else: # Needed unique file keys for per-job QuantumGraphs 

246 gwjob.arguments = gwjob.arguments.replace("{qgraphFile}", f"{{qgraphFile_{gwjob.name}}}") 

247 

248 # Replace files with special placeholders 

249 for gwfile in generic_workflow.get_job_inputs(gwjob.name): 

250 gwjob.arguments = gwjob.arguments.replace(f"{{{gwfile.name}}}", f"<FILE:{gwfile.name}>") 

251 for gwfile in generic_workflow.get_job_outputs(gwjob.name): 

252 gwjob.arguments = gwjob.arguments.replace(f"{{{gwfile.name}}}", f"<FILE:{gwfile.name}>") 

253 

254 # Save dict of other values needed to complete command line. 

255 # (Be careful to not replace env variables as they may 

256 # be different in compute job.) 

257 search_opt["replaceVars"] = True 

258 

259 for key in re.findall(r"{([^}]+)}", gwjob.arguments): 

260 if key not in gwjob.cmdvals: 

261 if key not in cached_job_values[gwjob.label]: 

262 _, cached_job_values[gwjob.label][key] = config.search(key, opt=search_opt) 

263 gwjob.cmdvals[key] = cached_job_values[gwjob.label][key] 

264 

265 # backwards compatibility 

266 if not cached_job_values[gwjob.label]["useLazyCommands"]: 

267 if "bpsUseShared" not in cached_job_values[gwjob.label]: 

268 key = "bpsUseShared" 

269 search_opt["default"] = True 

270 _, cached_job_values[gwjob.label][key] = config.search(key, opt=search_opt) 

271 del search_opt["default"] 

272 

273 gwjob.arguments = _fill_arguments( 

274 cached_job_values[gwjob.label]["bpsUseShared"], generic_workflow, gwjob.arguments, gwjob.cmdvals 

275 ) 

276 

277 

278def _fill_arguments(use_shared, generic_workflow, arguments, cmdvals): 

279 """Replace placeholders in command line string in job. 

280 

281 Parameters 

282 ---------- 

283 use_shared : `bool` 

284 Whether using shared filesystem. 

285 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

286 Generic workflow containing the job. 

287 arguments : `str` 

288 String containing placeholders. 

289 cmdvals : `dict` [`str`, `~typing.Any`] 

290 Any command line values that can be used to replace placeholders. 

291 

292 Returns 

293 ------- 

294 arguments : `str` 

295 Command line with FILE and ENV placeholders replaced. 

296 """ 

297 # Replace file placeholders 

298 for file_key in re.findall(r"<FILE:([^>]+)>", arguments): 

299 gwfile = generic_workflow.get_file(file_key) 

300 if not gwfile.wms_transfer: 

301 # Must assume full URI if in command line and told WMS is not 

302 # responsible for transferring file. 

303 uri = gwfile.src_uri 

304 elif use_shared: 

305 if gwfile.job_shared: 

306 # Have shared filesystems and jobs can share file. 

307 uri = gwfile.src_uri 

308 else: 

309 uri = os.path.basename(gwfile.src_uri) 

310 else: # Using push transfer 

311 uri = os.path.basename(gwfile.src_uri) 

312 

313 arguments = arguments.replace(f"<FILE:{file_key}>", uri) 

314 

315 # Replace env placeholder with submit-side values 

316 arguments = re.sub(r"<ENV:([^>]+)>", r"$\1", arguments) 

317 arguments = os.path.expandvars(arguments) 

318 

319 # Replace remaining vars 

320 arguments = arguments.format(**cmdvals) 

321 

322 return arguments 

323 

324 

325def _get_qgraph_gwfile(config, save_qgraph_per_job, gwjob, run_qgraph_file, prefix): 

326 """Get qgraph location to be used by job. 

327 

328 Parameters 

329 ---------- 

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

331 Bps configuration. 

332 save_qgraph_per_job : `lsst.ctrl.bps.bps_utils.WhenToSaveQuantumGraphs` 

333 What submission stage to save per-job qgraph files (or NEVER) 

334 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

335 Job for which determining QuantumGraph file. 

336 run_qgraph_file : `lsst.ctrl.bps.GenericWorkflowFile` 

337 File representation of the full run QuantumGraph. 

338 prefix : `str` 

339 Path prefix for any files written. 

340 

341 Returns 

342 ------- 

343 gwfile : `lsst.ctrl.bps.GenericWorkflowFile` 

344 Representation of butler location (may not include filename). 

345 """ 

346 qgraph_gwfile = None 

347 if save_qgraph_per_job != WhenToSaveQuantumGraphs.NEVER: 

348 qgraph_gwfile = GenericWorkflowFile( 

349 f"qgraphFile_{gwjob.name}", 

350 src_uri=create_job_quantum_graph_filename(config, gwjob, prefix), 

351 wms_transfer=True, 

352 job_access_remote=True, 

353 job_shared=True, 

354 ) 

355 else: 

356 qgraph_gwfile = run_qgraph_file 

357 

358 return qgraph_gwfile 

359 

360 

361def _get_job_values(config, search_opt, cmd_line_key): 

362 """Gather generic workflow job values from the bps config. 

363 

364 Parameters 

365 ---------- 

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

367 Bps configuration. 

368 search_opt : `dict` [`str`, `~typing.Any`] 

369 Search options to be used when searching config. 

370 cmd_line_key : `str` or None 

371 Which command line key to search for (e.g., "runQuantumCommand"). 

372 

373 Returns 

374 ------- 

375 job_values : `dict` [ `str`, `~typing.Any` ]` 

376 A mapping between job attributes and their values. 

377 """ 

378 _LOG.debug("cmd_line_key=%s, search_opt=%s", cmd_line_key, search_opt) 

379 

380 # Create a dummy job to easily access the default values. 

381 default_gwjob = GenericWorkflowJob("default_job", "default_label") 

382 

383 job_values = {} 

384 for attr in _ATTRS_ALL: 

385 # Variable names in yaml are camel case instead of snake case. 

386 yaml_name = re.sub(r"_(\S)", lambda match: match.group(1).upper(), attr) 

387 found, value = config.search(yaml_name, opt=search_opt) 

388 if found: 

389 job_values[attr] = value 

390 else: 

391 job_values[attr] = getattr(default_gwjob, attr) 

392 

393 # Need to replace all config variables in environment values 

394 # While replacing variables, convert to plain dict 

395 if job_values["environment"]: 

396 old_searchobj = search_opt.get("searchobj", None) 

397 old_replace_vars = search_opt.get("replaceVars", None) 

398 job_env = job_values["environment"] 

399 search_opt["searchobj"] = job_env 

400 search_opt["replaceVars"] = True 

401 job_values["environment"] = {} 

402 for name in job_env: 

403 job_values["environment"][name] = str(config.search(name, search_opt)[1]) 

404 if old_searchobj is None: 

405 del search_opt["searchobj"] 

406 else: 

407 search_opt["searchobj"] = old_searchobj 

408 if old_replace_vars is None: 

409 del search_opt["replaceVars"] 

410 else: 

411 search_opt["replaceVars"] = old_replace_vars 

412 

413 # If the automatic memory scaling is enabled (i.e. the memory multiplier 

414 # is set and it is a positive number greater than 1.0), adjust number 

415 # of retries when necessary. If the memory multiplier is invalid, disable 

416 # automatic memory scaling. 

417 if job_values["memory_multiplier"] is not None: 

418 if math.ceil(float(job_values["memory_multiplier"])) > 1: 

419 if job_values["number_of_retries"] is None: 

420 job_values["number_of_retries"] = DEFAULT_MEM_RETRIES 

421 else: 

422 job_values["memory_multiplier"] = None 

423 

424 if cmd_line_key: 

425 found, cmdline = config.search(cmd_line_key, opt=search_opt) 

426 # Make sure cmdline isn't None as that could be sent in as a 

427 # default value in search_opt. 

428 if found and cmdline: 

429 cmd, args = cmdline.split(" ", 1) 

430 job_values["executable"] = GenericWorkflowExec(os.path.basename(cmd), cmd, False) 

431 if args: 

432 job_values["arguments"] = args 

433 

434 return job_values 

435 

436 

437def _handle_job_values(quantum_job_values, gwjob, attributes=_ATTRS_ALL): 

438 """Set the job attributes in the cluster to their correct values. 

439 

440 Parameters 

441 ---------- 

442 quantum_job_values : `dict` [`str`, Any] 

443 Job values for running single Quantum. 

444 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

445 Generic workflow job in which to store the universal values. 

446 attributes : `~collections.abc.Iterable` [`str`], optional 

447 Job attributes to be set in the job following different rules. 

448 The default value is _ATTRS_ALL. 

449 """ 

450 _LOG.debug("Call to _handle_job_values") 

451 _handle_job_values_universal(quantum_job_values, gwjob, attributes) 

452 _handle_job_values_max(quantum_job_values, gwjob, attributes) 

453 _handle_job_values_sum(quantum_job_values, gwjob, attributes) 

454 

455 

456def _handle_job_values_universal(quantum_job_values, gwjob, attributes=_ATTRS_UNIVERSAL): 

457 """Handle job attributes that must have the same value for every quantum 

458 in the cluster. 

459 

460 Parameters 

461 ---------- 

462 quantum_job_values : `dict` [`str`, Any] 

463 Job values for running single Quantum. 

464 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

465 Generic workflow job in which to store the universal values. 

466 attributes : `~collections.abc.Iterable` [`str`], optional 

467 Job attributes to be set in the job following different rules. 

468 The default value is _ATTRS_UNIVERSAL. 

469 """ 

470 for attr in _ATTRS_UNIVERSAL & set(attributes): 

471 _LOG.debug( 

472 "Handling job %s (job=%s, quantum=%s)", 

473 attr, 

474 getattr(gwjob, attr), 

475 quantum_job_values.get(attr, "MISSING"), 

476 ) 

477 current_value = getattr(gwjob, attr) 

478 try: 

479 quantum_value = quantum_job_values[attr] 

480 except KeyError: 

481 continue 

482 else: 

483 if not current_value: 

484 setattr(gwjob, attr, quantum_value) 

485 elif current_value != quantum_value: 

486 _LOG.error( 

487 "Inconsistent value for %s in Cluster %s Quantum Number %s\n" 

488 "Current cluster value: %s\n" 

489 "Quantum value: %s", 

490 attr, 

491 gwjob.name, 

492 quantum_job_values.get("qgraphNodeId", "MISSING"), 

493 current_value, 

494 quantum_value, 

495 ) 

496 raise RuntimeError(f"Inconsistent value for {attr} in cluster {gwjob.name}.") 

497 

498 

499def _handle_job_values_max(quantum_job_values, gwjob, attributes=_ATTRS_MAX): 

500 """Handle job attributes that should be set to their maximum value in 

501 the in cluster. 

502 

503 Parameters 

504 ---------- 

505 quantum_job_values : `dict` [`str`, `~typing.Any`] 

506 Job values for running single Quantum. 

507 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

508 Generic workflow job in which to store the aggregate values. 

509 attributes : `~collections.abc.Iterable` [`str`], optional 

510 Job attributes to be set in the job following different rules. 

511 The default value is _ATTR_MAX. 

512 """ 

513 for attr in _ATTRS_MAX & set(attributes): 

514 current_value = getattr(gwjob, attr) 

515 try: 

516 quantum_value = quantum_job_values[attr] 

517 except KeyError: 

518 continue 

519 else: 

520 needs_update = False 

521 if current_value is None: 

522 if quantum_value is not None: 

523 needs_update = True 

524 else: 

525 if quantum_value is not None and current_value < quantum_value: 

526 needs_update = True 

527 if needs_update: 

528 setattr(gwjob, attr, quantum_value) 

529 

530 # When updating memory requirements for a job, check if memory 

531 # autoscaling is enabled. If it is, always use the memory 

532 # multiplier and the number of retries which comes with the 

533 # quantum. 

534 # 

535 # Note that as a result, the quantum with the biggest memory 

536 # requirements will determine whether the memory autoscaling 

537 # will be enabled (or disabled) depending on the value of its 

538 # memory multiplier. 

539 if attr == "request_memory": 

540 gwjob.memory_multiplier = quantum_job_values["memory_multiplier"] 

541 if gwjob.memory_multiplier is not None: 

542 gwjob.number_of_retries = quantum_job_values["number_of_retries"] 

543 

544 

545def _handle_job_values_sum(quantum_job_values, gwjob, attributes=_ATTRS_SUM): 

546 """Handle job attributes that are the sum of their values in the cluster. 

547 

548 Parameters 

549 ---------- 

550 quantum_job_values : `dict` [`str`, `~typing.Any`] 

551 Job values for running single Quantum. 

552 gwjob : `lsst.ctrl.bps.GenericWorkflowJob` 

553 Generic workflow job in which to store the aggregate values. 

554 attributes : `~collections.abc.Iterable` [`str`], optional 

555 Job attributes to be set in the job following different rules. 

556 The default value is _ATTRS_SUM. 

557 """ 

558 for attr in _ATTRS_SUM & set(attributes): 

559 current_value = getattr(gwjob, attr) 

560 if not current_value: 

561 setattr(gwjob, attr, quantum_job_values[attr]) 

562 else: 

563 setattr(gwjob, attr, current_value + quantum_job_values[attr]) 

564 

565 

566def create_generic_workflow( 

567 config: BpsConfig, cqgraph: ClusteredQuantumGraph, name: str, prefix: str 

568) -> GenericWorkflow: 

569 """Create a generic workflow from a ClusteredQuantumGraph such that it 

570 has information needed for WMS (e.g., command lines). 

571 

572 Parameters 

573 ---------- 

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

575 BPS configuration. 

576 cqgraph : `lsst.ctrl.bps.ClusteredQuantumGraph` 

577 ClusteredQuantumGraph for running a specific pipeline on a specific 

578 payload. 

579 name : `str` 

580 Name for the workflow (typically unique). 

581 prefix : `str` 

582 Root path for any output files. 

583 

584 Returns 

585 ------- 

586 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

587 Generic workflow for the given ClusteredQuantumGraph + config. 

588 """ 

589 # Determine whether saving per-job QuantumGraph files in the loop. 

590 _, when_save = config.search("whenSaveJobQgraph", {"default": WhenToSaveQuantumGraphs.TRANSFORM.name}) 

591 save_qgraph_per_job = WhenToSaveQuantumGraphs[when_save.upper()] 

592 

593 search_opt = {"replaceVars": False, "expandEnvVars": False, "replaceEnvVars": True, "required": False} 

594 

595 generic_workflow = GenericWorkflow(name) 

596 

597 # Save full run QuantumGraph for use by jobs 

598 generic_workflow.add_file( 

599 GenericWorkflowFile( 

600 "runQgraphFile", 

601 src_uri=config["runQgraphFile"], 

602 wms_transfer=True, 

603 job_access_remote=True, 

604 job_shared=True, 

605 ) 

606 ) 

607 

608 # Cache pipetask specific or more generic job values to minimize number 

609 # on config searches. 

610 cached_job_values = {} 

611 cached_pipetask_values = {} 

612 

613 for cluster in cqgraph.clusters(): 

614 _LOG.debug("Loop over clusters: %s, %s", cluster, type(cluster)) 

615 _LOG.debug( 

616 "cqgraph: name=%s, len=%s, label=%s, ids=%s", 

617 cluster.name, 

618 len(cluster.qgraph_node_ids), 

619 cluster.label, 

620 cluster.qgraph_node_ids, 

621 ) 

622 

623 gwjob = GenericWorkflowJob(cluster.name, cluster.label) 

624 

625 # First get job values from cluster or cluster config 

626 search_opt["curvals"] = {"curr_cluster": cluster.label} 

627 found, value = config.search("computeSite", opt=search_opt) 

628 if found: 

629 search_opt["curvals"]["curr_site"] = value 

630 found, value = config.search("computeCloud", opt=search_opt) 

631 if found: 

632 search_opt["curvals"]["curr_cloud"] = value 

633 

634 # If some config values are set for this cluster 

635 if cluster.label not in cached_job_values: 

636 _LOG.debug("config['cluster'][%s] = %s", cluster.label, config["cluster"][cluster.label]) 

637 cached_job_values[cluster.label] = {} 

638 

639 # Allowing whenSaveJobQgraph and useLazyCommands per cluster label. 

640 key = "whenSaveJobQgraph" 

641 _, when_save = config.search(key, opt=search_opt) 

642 cached_job_values[cluster.label][key] = WhenToSaveQuantumGraphs[when_save.upper()] 

643 

644 key = "useLazyCommands" 

645 search_opt["default"] = True 

646 _, cached_job_values[cluster.label][key] = config.search(key, opt=search_opt) 

647 del search_opt["default"] 

648 

649 if cluster.label in config["cluster"]: 

650 # Don't want to get global defaults here so only look in 

651 # cluster section. 

652 cached_job_values[cluster.label].update( 

653 _get_job_values(config["cluster"][cluster.label], search_opt, "runQuantumCommand") 

654 ) 

655 cluster_job_values = copy.copy(cached_job_values[cluster.label]) 

656 

657 cluster_job_values["name"] = cluster.name 

658 cluster_job_values["label"] = cluster.label 

659 cluster_job_values["quanta_counts"] = cluster.quanta_counts 

660 cluster_job_values["tags"] = cluster.tags 

661 _LOG.debug("cluster_job_values = %s", cluster_job_values) 

662 _handle_job_values(cluster_job_values, gwjob, cluster_job_values.keys()) 

663 

664 # For purposes of whether to continue searching for a value is whether 

665 # the value evaluates to False. 

666 unset_attributes = {attr for attr in _ATTRS_ALL if not getattr(gwjob, attr)} 

667 

668 _LOG.debug("unset_attributes=%s", unset_attributes) 

669 _LOG.debug("set=%s", _ATTRS_ALL - unset_attributes) 

670 

671 # For job info not defined at cluster level, attempt to get job info 

672 # either common or aggregate for all Quanta in cluster. 

673 for node_id in iter(cluster.qgraph_node_ids): 

674 _LOG.debug("node_id=%s", node_id) 

675 quantum_info = cqgraph.get_quantum_info(node_id) 

676 

677 task_label = quantum_info["task_label"] 

678 if task_label not in cached_pipetask_values: 

679 search_opt["curvals"]["curr_pipetask"] = task_label 

680 cached_pipetask_values[task_label] = _get_job_values(config, search_opt, "runQuantumCommand") 

681 _handle_job_values(cached_pipetask_values[task_label], gwjob, unset_attributes) 

682 

683 # Update job with workflow attribute and profile values. 

684 qgraph_gwfile = _get_qgraph_gwfile( 

685 config, save_qgraph_per_job, gwjob, generic_workflow.get_file("runQgraphFile"), prefix 

686 ) 

687 

688 generic_workflow.add_job(gwjob) 

689 generic_workflow.add_job_inputs(gwjob.name, [qgraph_gwfile]) 

690 

691 gwjob.cmdvals["qgraphNodeId"] = ",".join( 

692 sorted([f"{node_id}" for node_id in cluster.qgraph_node_ids]) 

693 ) 

694 _enhance_command(config, generic_workflow, gwjob, cached_job_values) 

695 

696 # If writing per-job QuantumGraph files during TRANSFORM stage, 

697 # write it now while in memory. 

698 if save_qgraph_per_job == WhenToSaveQuantumGraphs.TRANSFORM: 

699 save_qg_subgraph(cqgraph.qgraph, qgraph_gwfile.src_uri, cluster.qgraph_node_ids) 

700 

701 # Create job dependencies. 

702 for parent in cqgraph.clusters(): 

703 for child in cqgraph.successors(parent): 

704 generic_workflow.add_job_relationships(parent.name, child.name) 

705 

706 # Add initial workflow. 

707 if config.get("runInit", "{default: False}"): 

708 add_workflow_init_nodes(config, cqgraph.qgraph, generic_workflow) 

709 

710 generic_workflow.run_attrs.update( 

711 { 

712 "bps_isjob": "True", 

713 "bps_project": config["project"], 

714 "bps_campaign": config["campaign"], 

715 "bps_run": generic_workflow.name, 

716 "bps_operator": config["operator"], 

717 "bps_payload": config["payloadName"], 

718 "bps_runsite": config["computeSite"], 

719 } 

720 ) 

721 

722 # Add final job 

723 add_final_job(config, generic_workflow, prefix) 

724 

725 if "ordering" in config: 

726 generic_workflow.add_special_job_ordering(config["ordering"]) 

727 

728 return generic_workflow 

729 

730 

731def create_generic_workflow_config(config, prefix): 

732 """Create generic workflow configuration. 

733 

734 Parameters 

735 ---------- 

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

737 Bps configuration. 

738 prefix : `str` 

739 Root path for any output files. 

740 

741 Returns 

742 ------- 

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

744 Configuration accompanying the GenericWorkflow. 

745 """ 

746 generic_workflow_config = BpsConfig(config) 

747 generic_workflow_config["workflowName"] = config["uniqProcName"] 

748 generic_workflow_config["workflowPath"] = prefix 

749 return generic_workflow_config 

750 

751 

752def add_final_job(config: BpsConfig, generic_workflow: GenericWorkflow, prefix: str) -> None: 

753 """Add final workflow job depending upon configuration. 

754 

755 Depending on configuration, the final job will be added as a special job 

756 which will always run regardless of the exit status of the workflow or 

757 a regular sink node which will only run if the workflow execution finished 

758 with no errors. 

759 

760 Parameters 

761 ---------- 

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

763 Bps configuration. 

764 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

765 Generic workflow to which attributes should be added. 

766 prefix : `str` 

767 Directory in which to output final script. 

768 """ 

769 _, when_run = config.search(".finalJob.whenRun") 

770 if when_run.upper() != "NEVER": 

771 gwjob = create_final_job(config, generic_workflow, prefix) 

772 if when_run.upper() == "ALWAYS": 

773 generic_workflow.add_final(gwjob) 

774 elif when_run.upper() == "SUCCESS": 

775 add_final_job_as_sink(generic_workflow, gwjob) 

776 else: 

777 raise ValueError(f"Invalid value for finalJob.whenRun: {when_run}") 

778 

779 

780def create_final_job(config: BpsConfig, generic_workflow: GenericWorkflow, prefix: str) -> GenericWorkflowJob: 

781 """Create the final workflow job. 

782 

783 Parameters 

784 ---------- 

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

786 Bps configuration. 

787 generic_workflow : `lsst.ctrl.bps.GenericWorkflow` 

788 Generic workflow to which attributes should be added. 

789 prefix : `str` 

790 Directory in which to output final script. 

791 

792 Returns 

793 ------- 

794 final_job : `lsst.ctrl.bps.GenericWorkflowJob` 

795 Final workflow job. 

796 """ 

797 job_name = "finalJob" 

798 gwjob = GenericWorkflowJob(job_name, job_name) 

799 

800 search_opt = {"searchobj": config[job_name], "curvals": {}, "default": None} 

801 found, value = config.search("computeSite", opt=search_opt) 

802 if found: 

803 search_opt["curvals"]["curr_site"] = value 

804 found, value = config.search("computeCloud", opt=search_opt) 

805 if found: 

806 search_opt["curvals"]["curr_cloud"] = value 

807 

808 # Set job attributes based on the values find in the config excluding 

809 # the ones in the _ATTRS_MISC group. The attributes in this group are 

810 # somewhat "special": 

811 # * HTCondor plugin, which uses 'attrs' and 'profile', has its own 

812 # mechanism for setting them, 

813 # * 'cmdvals' is being set internally, not via config. 

814 job_values = _get_job_values(config, search_opt, None) 

815 for attr in _ATTRS_ALL - _ATTRS_MISC: 

816 if not getattr(gwjob, attr) and job_values.get(attr, None): 

817 setattr(gwjob, attr, job_values[attr]) 

818 

819 # Create script and add command line to job. 

820 gwjob.executable, gwjob.arguments = create_final_command(config, prefix) 

821 

822 # Determine inputs from command line. 

823 for file_key in re.findall(r"<FILE:([^>]+)>", gwjob.arguments): 

824 gwfile = generic_workflow.get_file(file_key) 

825 generic_workflow.add_job_inputs(gwjob.name, gwfile) 

826 

827 _enhance_command(config, generic_workflow, gwjob, {}) 

828 return gwjob 

829 

830 

831def create_final_command(config: BpsConfig, prefix: str) -> tuple[GenericWorkflowExec, str]: 

832 """Create the command and shell script for the final job. 

833 

834 Parameters 

835 ---------- 

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

837 Bps configuration. 

838 prefix : `str` 

839 Directory in which to output final script. 

840 

841 Returns 

842 ------- 

843 executable : `lsst.ctrl.bps.GenericWorkflowExec` 

844 Executable object for the final script. 

845 arguments : `str` 

846 Command line needed to call the final script. 

847 

848 Raises 

849 ------ 

850 RuntimeError if no commands found. 

851 """ 

852 search_opt = { 

853 "replaceVars": True, 

854 "skipNames": ["butlerConfig", "qgraphFile"], 

855 "replaceEnvVars": False, 

856 "expandEnvVars": False, 

857 "searchobj": config["finalJob"], 

858 } 

859 

860 script_file = os.path.join(prefix, "final_job.bash") 

861 with open(script_file, "w", encoding="utf8") as fh: 

862 print("#!/bin/bash\n", file=fh) 

863 print("set -e", file=fh) 

864 print("set -x", file=fh) 

865 

866 print("qgraphFile=$1", file=fh) 

867 print("butlerConfig=$2", file=fh) 

868 

869 command_len = 0 # Make sure at least write one actual command 

870 i = 1 

871 found, command = config.search(f"command{i}", opt=search_opt) 

872 while found: 

873 # The files will be args to script, so change to shell vars 

874 command = command.replace("{qgraphFile}", "${qgraphFile}") 

875 command = command.replace("{butlerConfig}", "${butlerConfig}") 

876 

877 print(command, file=fh) 

878 command_len += len(command.strip()) 

879 

880 # Search for next command 

881 i += 1 

882 found, command = config.search(f"command{i}", opt=search_opt) 

883 if command_len == 0: 

884 raise RuntimeError( 

885 "No finalJob commands were found. Use NEVER for finalJob.whenRun to turn off finalJob" 

886 ) 

887 os.chmod(script_file, 0o755) 

888 executable = GenericWorkflowExec(os.path.basename(script_file), script_file, True) 

889 

890 _, orig_butler = config.search("butlerConfig") 

891 return executable, f"<FILE:runQgraphFile> {orig_butler}" 

892 

893 

894def add_final_job_as_sink(generic_workflow, final_job): 

895 """Add final job as the single sink for the workflow. 

896 

897 Parameters 

898 ----------