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

1# This file is part of ctrl_bps.

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.

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

28"""Driver functions for each subcommand.

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

31the subcommand method.

32"""

34__all__ = [

35 "acquire_qgraph_driver",

36 "cancel_driver",

37 "cluster_qgraph_driver",

38 "ping_driver",

39 "prepare_driver",

40 "report_driver",

41 "restart_driver",

42 "status_driver",

43 "submit_driver",

44 "submitcmd_driver",

45 "transform_driver",

46]

49import logging

50import os

51from pathlib import Path

53from lsst.pipe.base.quantum_graph import PredictedQuantumGraph

54from lsst.utils.timer import time_this

55from lsst.utils.usage import get_peak_mem_usage

57from . import BPS_DEFAULTS, BPS_SEARCH_ORDER, DEFAULT_MEM_FMT, DEFAULT_MEM_UNIT, BpsConfig

58from .bps_utils import _dump_env_info, _dump_pkg_info, _make_id_link

59from .cancel import cancel

60from .construct import construct

61from .initialize import (

62 custom_job_validator,

63 init_submission,

64 out_collection_validator,

65 output_run_validator,

66 submit_path_validator,

67)

68from .ping import ping

69from .pre_transform import acquire_quantum_graph, cluster_quanta

70from .prepare import prepare

71from .report import BPS_POSTPROCESSORS, display_report, retrieve_report

72from .restart import restart

73from .status import status

74from .submit import submit

75from .transform import transform

77_LOG = logging.getLogger(__name__)

80def _init_submission_driver(config_file: str, **kwargs) -> BpsConfig:

81 """Initialize runtime environment.

83 Parameters

84 ----------

85 config_file : `str`

86 Name of the configuration file.

87 **kwargs : `~typing.Any`

88 Additional modifiers to the configuration.

90 Returns

91 -------

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

93 Batch Processing Service configuration.

94 """

95 validators = [submit_path_validator, output_run_validator, out_collection_validator]

96 _LOG.info("Initializing BPS configuration and creating submit directory")

97 with time_this(

98 log=_LOG,

99 level=logging.INFO,

100 prefix=None,

101 msg="BPS configuration initialized and submit directory created",

102 mem_usage=True,

103 mem_unit=DEFAULT_MEM_UNIT,

104 mem_fmt=DEFAULT_MEM_FMT,

105 ):

106 config = init_submission(config_file, validators=validators, **kwargs)

107 _log_mem_usage()

108

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

110 print(f"Submit dir: {submit_path}")

111 return config

112

113

114def acquire_qgraph_driver(config_file: str, **kwargs) -> tuple[BpsConfig, PredictedQuantumGraph]:

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

116

117 Parameters

118 ----------

119 config_file : `str`

120 Name of the configuration file.

121 **kwargs : `~typing.Any`

122 Additional modifiers to the configuration.

123

124 Returns

125 -------

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

127 Updated configuration.

128 qgraph : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`

129 A graph representing quanta.

130 """

131 config = _init_submission_driver(config_file, **kwargs)

132

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

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

135 with time_this(

136 log=_LOG,

137 level=logging.INFO,

138 prefix=None,

139 msg="Acquire stage completed",

140 mem_usage=True,

141 mem_unit=DEFAULT_MEM_UNIT,

142 mem_fmt=DEFAULT_MEM_FMT,

143 ):

144 qgraph_file, qgraph = acquire_quantum_graph(config, out_prefix=submit_path)

145 _log_mem_usage()

146

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

148 return config, qgraph

149

150

151def cluster_qgraph_driver(config_file, **kwargs):

152 """Group quanta into clusters.

153

154 Parameters

155 ----------

156 config_file : `str`

157 Name of the configuration file.

158 **kwargs : `~typing.Any`

159 Additional modifiers to the configuration.

160

161 Returns

162 -------

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

164 Updated configuration.

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

166 A graph representing clustered quanta.

167 """

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

169

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

171 with time_this(

172 log=_LOG,

173 level=logging.INFO,

174 prefix=None,

175 msg="Cluster stage completed",

176 mem_usage=True,

177 mem_unit=DEFAULT_MEM_UNIT,

178 mem_fmt=DEFAULT_MEM_FMT,

179 ):

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

181 _log_mem_usage()

182

183 _LOG.info("ClusteredQuantumGraph contains %d cluster(s)", len(clustered_qgraph))

184

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

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

187 if save_clustered_qgraph:

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

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

190 if save_dot:

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

192 return config, clustered_qgraph

193

194

195def transform_driver(config_file, **kwargs):

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

197

198 Parameters

199 ----------

200 config_file : `str`

201 Name of the configuration file.

202 **kwargs : `~typing.Any`

203 Additional modifiers to the configuration.

204

205 Returns

206 -------

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

208 Configuration to use when creating the workflow.

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

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

211 workflow management system.

212 """

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

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

215

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

217 with time_this(

218 log=_LOG,

219 level=logging.INFO,

220 prefix=None,

221 msg="Transform stage completed",

222 mem_usage=True,

223 mem_unit=DEFAULT_MEM_UNIT,

224 mem_fmt=DEFAULT_MEM_FMT,

225 ):

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

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

228 _log_mem_usage()

229

230 num_jobs = sum(generic_workflow.job_counts.values())

231 _LOG.info("GenericWorkflow contains %d job(s) (including final)", num_jobs)

232

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

234 if save_workflow:

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

236 generic_workflow.save(outfh, "pickle")

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

238 if save_dot:

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

240 generic_workflow.draw(outfh, "dot")

241 return generic_workflow_config, generic_workflow

242

243

244def prepare_driver(config_file, **kwargs):

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

246

247 Parameters

248 ----------

249 config_file : `str`

250 Name of the configuration file.

251 **kwargs : `~typing.Any`

252 Additional modifiers to the configuration.

253

254 Returns

255 -------

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

257 Configuration to use when creating the workflow.

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

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

260 workflow management system.

261 """

262 kwargs.setdefault("runWmsSubmissionChecks", True)

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

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

265

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

267 with time_this(

268 log=_LOG,

269 level=logging.INFO,

270 prefix=None,

271 msg="Prepare stage completed",

272 mem_usage=True,

273 mem_unit=DEFAULT_MEM_UNIT,

274 mem_fmt=DEFAULT_MEM_FMT,

275 ):

276 wms_workflow = prepare(generic_workflow_config, generic_workflow, submit_path)

277 _log_mem_usage()

278

279 wms_workflow_config = generic_workflow_config

280 return wms_workflow_config, wms_workflow

281

282

283def submit_driver(config_file, **kwargs):

284 """Submit workflow for execution.

285

286 Parameters

287 ----------

288 config_file : `str`

289 Name of the configuration file.

290 **kwargs : `~typing.Any`

291 Additional modifiers to the configuration.

292 """

293 kwargs.setdefault("runWmsSubmissionChecks", True)

294

295 _LOG.info(

296 "DISCLAIMER: All values regarding memory consumption reported below are approximate and may "

297 "not accurately reflect actual memory usage by the bps process."

298 )

299

300 remote_build = {}

301 config = BpsConfig(

302 config_file,

303 search_order=BPS_SEARCH_ORDER,

304 defaults=BPS_DEFAULTS,

305 wms_service_class_fqn=kwargs.get("wms_service"),

306 )

307 _, remote_build = config.search("remoteBuild", opt={"default": {}})

308 if remote_build:

309 if config["wmsServiceClass"] == "lsst.ctrl.bps.panda.PanDAService":

310 if not remote_build.search("enabled", opt={"default": False})[1]:

311 remote_build = {}

312 _LOG.info("The workflow is submitted to the local Data Facility.")

313 else:

314 _LOG.info(

315 "Remote submission is enabled. The workflow is submitted to a remote Data Facility."

316 )

317 _LOG.info("Initializing execution environment")

318 with time_this(

319 log=_LOG,

320 level=logging.INFO,

321 prefix=None,

322 msg="Initializing execution environment completed",

323 mem_usage=True,

324 mem_unit=DEFAULT_MEM_UNIT,

325 mem_fmt=DEFAULT_MEM_FMT,

326 ):

327 config = _init_submission_driver(config_file, **kwargs)

328 kwargs["remote_build"] = remote_build

329 kwargs["config_file"] = config_file

330 wms_workflow = None

331 else:

332 _LOG.info("The workflow is submitted to the local Data Facility.")

333

334 _LOG.info("Starting submission process")

335 with time_this(

336 log=_LOG,

337 level=logging.INFO,

338 prefix=None,

339 msg="Submission process completed",

340 mem_usage=True,

341 mem_unit=DEFAULT_MEM_UNIT,

342 mem_fmt=DEFAULT_MEM_FMT,

343 ):

344 if not remote_build:

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

346 else:

347 wms_workflow_config = config

348

349 _LOG.info("Starting submit stage")

350 with time_this(

351 log=_LOG,

352 level=logging.INFO,

353 prefix=None,

354 msg="Submit stage completed",

355 mem_usage=True,

356 mem_unit=DEFAULT_MEM_UNIT,

357 mem_fmt=DEFAULT_MEM_FMT,

358 ):

359 workflow = submit(wms_workflow_config, wms_workflow, **kwargs)

360 if not wms_workflow:

361 wms_workflow = workflow

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

363 _log_mem_usage()

364

365 _make_id_link(wms_workflow_config, wms_workflow.run_id)

366

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

368 print(f"Run Name: {wms_workflow.name}")

369

370

371def restart_driver(wms_service, run_id):

372 """Restart a failed workflow.

373

374 Parameters

375 ----------

376 wms_service : `str`

377 Name of the class.

378 run_id : `str`

379 Id or path of workflow that need to be restarted.

380 """

381 if wms_service is None:

382 default_config = BpsConfig({}, defaults=BPS_DEFAULTS)

383 wms_service = default_config["wmsServiceClass"]

384

385 new_run_id, run_name, message = restart(wms_service, run_id)

386 if new_run_id is not None:

387 path = Path(run_id)

388 if path.exists():

389 _dump_env_info(f"{run_id}/{run_name}.env.info.yaml")

390 _dump_pkg_info(f"{run_id}/{run_name}.pkg.info.yaml")

391 config = BpsConfig(f"{run_id}/{run_name}_config.yaml")

392 _make_id_link(config, new_run_id)

393

394 print(f"Run Id: {new_run_id}")

395 print(f"Run Name: {run_name}")

396 else:

397 if message:

398 print(f"Restart failed: {message}")

399 else:

400 print("Restart failed: Unknown error")

401

402

403def report_driver(wms_service, run_id, user, hist_days, pass_thru, is_global=False, return_exit_codes=False):

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

405

406 Parameters

407 ----------

408 wms_service : `str`

409 Name of the class.

410 run_id : `str`

411 A run id the report will be restricted to.

412 user : `str`

413 A user the report will be restricted to.

414 hist_days : `float`

415 Number of days.

416 pass_thru : `str`

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

418 is_global : `bool`, optional

419 If set, all available job queues will be queried for job information.

420 Defaults to False which means that only a local job queue will be

421 queried for information.

422

423 Only applicable in the context of a WMS using distributed job queues

424 (e.g., HTCondor).

425 return_exit_codes : `bool`, optional

426 If set, return exit codes related to jobs with a

427 non-success status. Defaults to False, which means that only

428 the summary state is returned.

429

430 Only applicable in the context of a WMS with associated

431 handlers to return exit codes from jobs.

432 """

433 if wms_service is None:

434 default_config = BpsConfig(BPS_DEFAULTS)

435 wms_service = os.environ.get("BPS_WMS_SERVICE_CLASS", default_config["wmsServiceClass"])

436

437 # When reporting on single run:

438 # * increase history until better mechanism for handling completed jobs is

439 # available.

440 # * massage the retrieved reports using BPS report postprocessors.

441 if run_id:

442 hist_days = max(hist_days, 2)

443 postprocessors = BPS_POSTPROCESSORS

444 else:

445 postprocessors = None

446

447 runs, messages = retrieve_report(

448 wms_service,

449 run_id=run_id,

450 user=user,

451 hist=hist_days,

452 pass_thru=pass_thru,

453 is_global=is_global,

454 return_exit_codes=return_exit_codes,

455 postprocessors=postprocessors,

456 )

457

458 if runs or messages:

459 display_report(

460 runs,

461 messages,

462 is_detailed=bool(run_id),

463 is_global=is_global,

464 return_exit_codes=return_exit_codes,

465 )

466 else:

467 if run_id:

468 print(

469 f"No records found for job id '{run_id}'. "

470 f"Hints: Double check id, retry with a larger --hist value (currently: {hist_days}), "

471 "and/or use --global to search all job queues."

472 )

473

474

475def status_driver(wms_service: str, run_id: str, hist_days: float, is_global: bool = False) -> int:

476 """Print out status of workflow submitted for execution.

477

478 Parameters

479 ----------

480 wms_service : `str`

481 Name of the class.

482 run_id : `str`

483 A run id the report will be restricted to.

484 hist_days : `float`

485 Number of days.

486 is_global : `bool`, optional

487 If set, all available job queues will be queried for job information.

488 Defaults to False which means that only a local job queue will be

489 queried for information.

490

491 Only applicable in the context of a WMS using distributed job queues

492 (e.g., HTCondor).

493

494 Returns

495 -------

496 state : `int`

497 Status of submitted workflow.

498 """

499 if wms_service is None:

500 default_config = BpsConfig(BPS_DEFAULTS)

501 wms_service = os.environ.get("BPS_WMS_SERVICE_CLASS", default_config["wmsServiceClass"])

502

503 state, message = status(

504 wms_service,

505 run_id=run_id,

506 hist=hist_days,

507 is_global=is_global,

508 )

509

510 _LOG.info("status: %s", state.name)

511 if message:

512 _LOG.warning(message)

513

514 return state.value

515

516

517def cancel_driver(wms_service, run_id, user, require_bps, pass_thru, is_global=False):

518 """Cancel submitted workflows.

519

520 Parameters

521 ----------

522 wms_service : `str`

523 Name of the Workload Management System service class.

524 run_id : `str`

525 ID or path of job that should be canceled.

526 user : `str`

527 User whose submitted jobs should be canceled.

528 require_bps : `bool`

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

530 pass_thru : `str`

531 Information to pass through to WMS.

532 is_global : `bool`, optional

533 If set, all available job queues will be checked for jobs to cancel.

534 Defaults to False which means that only a local job queue will be

535 checked.

536

537 Only applicable in the context of a WMS using distributed job queues

538 (e.g., HTCondor).

539 """

540 if wms_service is None:

541 default_config = BpsConfig({}, defaults=BPS_DEFAULTS)

542 wms_service = default_config["wmsServiceClass"]

543 cancel(wms_service, run_id, user, require_bps, pass_thru, is_global=is_global)

544

545

546def ping_driver(wms_service=None, pass_thru=None):

547 """Check whether WMS services are up, reachable, and any authentication,

548 if needed, succeeds.

549

550 The services to be checked are those needed for submit, report, cancel,

551 restart, but ping cannot guarantee whether jobs would actually run

552 successfully.

553

554 Parameters

555 ----------

556 wms_service : `str`, optional

557 Name of the Workload Management System service class.

558 pass_thru : `str`, optional

559 Information to pass through to WMS.

560

561 Returns

562 -------

563 success : `int`

564 Whether services are up and usable (0) or not (non-zero).

565 """

566 if wms_service is None:

567 default_config = BpsConfig({}, defaults=BPS_DEFAULTS)

568 wms_service = default_config["wmsServiceClass"]

569 status, message = ping(wms_service, pass_thru)

570

571 if message:

572 if not status:

573 _LOG.info(message)

574 else:

575 _LOG.error(message)

576

577 # Log overall status message

578 if not status:

579 _LOG.info("Ping successful.")

580 else:

581 _LOG.error("Ping failed (%d).", status)

582

583 return status

584

585

586def submitcmd_driver(config_file: str, **kwargs) -> None:

587 """Submit a command for execution.

588

589 Parameters

590 ----------

591 config_file : `str`

592 Name of the configuration file.

593 **kwargs : `~typing.Any`

594 Additional modifiers to the configuration.

595 """

596 validators = [submit_path_validator, custom_job_validator]

597 _LOG.info("Initializing BPS configuration and creating submit directory")

598 with time_this(

599 log=_LOG,

600 level=logging.INFO,

601 prefix=None,

602 msg="BPS configuration initialized and submit directory created",

603 mem_usage=True,

604 mem_unit=DEFAULT_MEM_UNIT,

605 mem_fmt=DEFAULT_MEM_FMT,

606 ):

607 config = init_submission(config_file, validators=validators, **kwargs)

608 _log_mem_usage()

609

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

611

612 _LOG.info("Starting construction stage (creating generic workflow)")

613 with time_this(

614 log=_LOG,

615 level=logging.INFO,

616 prefix=None,

617 msg="Construction stage completed",

618 mem_usage=True,

619 mem_unit=DEFAULT_MEM_UNIT,

620 mem_fmt=DEFAULT_MEM_FMT,

621 ):

622 generic_workflow, generic_workflow_config = construct(config)

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

624 _log_mem_usage()

625

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

627 if save_workflow:

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

629 generic_workflow.save(outfh, "pickle")

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

631 if save_dot:

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

633 generic_workflow.draw(outfh, "dot")

634

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

636 with time_this(

637 log=_LOG,

638 level=logging.INFO,

639 prefix=None,

640 msg="Prepare stage completed",

641 mem_usage=True,

642 mem_unit=DEFAULT_MEM_UNIT,

643 mem_fmt=DEFAULT_MEM_FMT,

644 ):

645 wms_workflow = prepare(generic_workflow_config, generic_workflow, submit_path)

646 _log_mem_usage()

647

648 wms_workflow_config = generic_workflow_config

649

650 if kwargs.get("dry_run", False):

651 return

652

653 _LOG.info("Starting submit stage")

654 with time_this(

655 log=_LOG,

656 level=logging.INFO,

657 prefix=None,

658 msg="Submit stage completed",

659 mem_usage=True,

660 mem_unit=DEFAULT_MEM_UNIT,

661 mem_fmt=DEFAULT_MEM_FMT,

662 ):

663 submit(wms_workflow_config, wms_workflow, **kwargs)

664 _log_mem_usage()

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

666 print(f"Run Name: {wms_workflow.name}")

667

668

669def _log_mem_usage() -> None:

670 """Log memory usage."""

671 if _LOG.isEnabledFor(logging.INFO):

672 _LOG.info(

673 "Peak memory usage for bps process %s (main), %s (largest child process)",

674 *tuple(f"{val.to(DEFAULT_MEM_UNIT):{DEFAULT_MEM_FMT}}" for val in get_peak_mem_usage()),

675 )