Coverage for python/lsst/ctrl/mpexec/cmdLineFwk.py: 15%

303 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-13 16:33 -0700

1# This file is part of ctrl_mpexec. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://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"""Module defining CmdLineFwk class and related methods. 

23""" 

24 

25from __future__ import annotations 

26 

27__all__ = ["CmdLineFwk"] 

28 

29import atexit 

30import copy 

31import datetime 

32import getpass 

33import logging 

34import shutil 

35from collections.abc import Iterable, Sequence 

36from types import SimpleNamespace 

37from typing import Optional, Tuple 

38 

39from astropy.table import Table 

40from lsst.daf.butler import Butler, CollectionType, DatasetRef, DatastoreCacheManager, Registry 

41from lsst.daf.butler.registry import MissingCollectionError, RegistryDefaults 

42from lsst.daf.butler.registry.wildcards import CollectionWildcard 

43from lsst.pipe.base import ( 

44 GraphBuilder, 

45 Instrument, 

46 Pipeline, 

47 PipelineDatasetTypes, 

48 QuantumGraph, 

49 TaskDef, 

50 TaskFactory, 

51 buildExecutionButler, 

52) 

53from lsst.utils import doImportType 

54from lsst.utils.threads import disable_implicit_threading 

55 

56from . import util 

57from .dotTools import graph2dot, pipeline2dot 

58from .executionGraphFixup import ExecutionGraphFixup 

59from .mpGraphExecutor import MPGraphExecutor 

60from .preExecInit import PreExecInit 

61from .singleQuantumExecutor import SingleQuantumExecutor 

62 

63# ---------------------------------- 

64# Local non-exported definitions -- 

65# ---------------------------------- 

66 

67_LOG = logging.getLogger(__name__) 

68 

69 

70class _OutputChainedCollectionInfo: 

71 """A helper class for handling command-line arguments related to an output 

72 `~lsst.daf.butler.CollectionType.CHAINED` collection. 

73 

74 Parameters 

75 ---------- 

76 registry : `lsst.daf.butler.Registry` 

77 Butler registry that collections will be added to and/or queried from. 

78 name : `str` 

79 Name of the collection given on the command line. 

80 """ 

81 

82 def __init__(self, registry: Registry, name: str): 

83 self.name = name 

84 try: 

85 self.chain = tuple(registry.getCollectionChain(name)) 

86 self.exists = True 

87 except MissingCollectionError: 

88 self.chain = () 

89 self.exists = False 

90 

91 def __str__(self) -> str: 

92 return self.name 

93 

94 name: str 

95 """Name of the collection provided on the command line (`str`). 

96 """ 

97 

98 exists: bool 

99 """Whether this collection already exists in the registry (`bool`). 

100 """ 

101 

102 chain: Tuple[str, ...] 

103 """The definition of the collection, if it already exists (`tuple`[`str`]). 

104 

105 Empty if the collection does not already exist. 

106 """ 

107 

108 

109class _OutputRunCollectionInfo: 

110 """A helper class for handling command-line arguments related to an output 

111 `~lsst.daf.butler.CollectionType.RUN` collection. 

112 

113 Parameters 

114 ---------- 

115 registry : `lsst.daf.butler.Registry` 

116 Butler registry that collections will be added to and/or queried from. 

117 name : `str` 

118 Name of the collection given on the command line. 

119 """ 

120 

121 def __init__(self, registry: Registry, name: str): 

122 self.name = name 

123 try: 

124 actualType = registry.getCollectionType(name) 

125 if actualType is not CollectionType.RUN: 

126 raise TypeError(f"Collection '{name}' exists but has type {actualType.name}, not RUN.") 

127 self.exists = True 

128 except MissingCollectionError: 

129 self.exists = False 

130 

131 name: str 

132 """Name of the collection provided on the command line (`str`). 

133 """ 

134 

135 exists: bool 

136 """Whether this collection already exists in the registry (`bool`). 

137 """ 

138 

139 

140class _ButlerFactory: 

141 """A helper class for processing command-line arguments related to input 

142 and output collections. 

143 

144 Parameters 

145 ---------- 

146 registry : `lsst.daf.butler.Registry` 

147 Butler registry that collections will be added to and/or queried from. 

148 

149 args : `types.SimpleNamespace` 

150 Parsed command-line arguments. The following attributes are used, 

151 either at construction or in later methods. 

152 

153 ``output`` 

154 The name of a `~lsst.daf.butler.CollectionType.CHAINED` 

155 input/output collection. 

156 

157 ``output_run`` 

158 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output 

159 collection. 

160 

161 ``extend_run`` 

162 A boolean indicating whether ``output_run`` should already exist 

163 and be extended. 

164 

165 ``replace_run`` 

166 A boolean indicating that (if `True`) ``output_run`` should already 

167 exist but will be removed from the output chained collection and 

168 replaced with a new one. 

169 

170 ``prune_replaced`` 

171 A boolean indicating whether to prune the replaced run (requires 

172 ``replace_run``). 

173 

174 ``inputs`` 

175 Input collections of any type; see 

176 :ref:`daf_butler_ordered_collection_searches` for details. 

177 

178 ``butler_config`` 

179 Path to a data repository root or configuration file. 

180 

181 writeable : `bool` 

182 If `True`, a `Butler` is being initialized in a context where actual 

183 writes should happens, and hence no output run is necessary. 

184 

185 Raises 

186 ------ 

187 ValueError 

188 Raised if ``writeable is True`` but there are no output collections. 

189 """ 

190 

191 def __init__(self, registry: Registry, args: SimpleNamespace, writeable: bool): 

192 if args.output is not None: 

193 self.output = _OutputChainedCollectionInfo(registry, args.output) 

194 else: 

195 self.output = None 

196 if args.output_run is not None: 

197 self.outputRun = _OutputRunCollectionInfo(registry, args.output_run) 

198 elif self.output is not None: 

199 if args.extend_run: 

200 if not self.output.chain: 

201 raise ValueError("Cannot use --extend-run option with non-existing or empty output chain") 

202 runName = self.output.chain[0] 

203 else: 

204 runName = "{}/{}".format(self.output, Instrument.makeCollectionTimestamp()) 

205 self.outputRun = _OutputRunCollectionInfo(registry, runName) 

206 elif not writeable: 

207 # If we're not writing yet, ok to have no output run. 

208 self.outputRun = None 

209 else: 

210 raise ValueError("Cannot write without at least one of (--output, --output-run).") 

211 # Recursively flatten any input CHAINED collections. We do this up 

212 # front so we can tell if the user passes the same inputs on subsequent 

213 # calls, even though we also flatten when we define the output CHAINED 

214 # collection. 

215 self.inputs = tuple(registry.queryCollections(args.input, flattenChains=True)) if args.input else () 

216 

217 def check(self, args: SimpleNamespace) -> None: 

218 """Check command-line options for consistency with each other and the 

219 data repository. 

220 

221 Parameters 

222 ---------- 

223 args : `types.SimpleNamespace` 

224 Parsed command-line arguments. See class documentation for the 

225 construction parameter of the same name. 

226 """ 

227 assert not (args.extend_run and args.replace_run), "In mutually-exclusive group in ArgumentParser." 

228 if self.inputs and self.output is not None and self.output.exists: 

229 # Passing the same inputs that were used to initialize the output 

230 # collection is allowed; this means they must _end_ with the same 

231 # collections, because we push new runs to the front of the chain. 

232 for c1, c2 in zip(self.inputs[::-1], self.output.chain[::-1]): 

233 if c1 != c2: 

234 raise ValueError( 

235 f"Output CHAINED collection {self.output.name!r} exists, but it ends with " 

236 "a different sequence of input collections than those given: " 

237 f"{c1!r} != {c2!r} in inputs={self.inputs} vs " 

238 f"{self.output.name}={self.output.chain}." 

239 ) 

240 if len(self.inputs) > len(self.output.chain): 

241 nNew = len(self.inputs) - len(self.output.chain) 

242 raise ValueError( 

243 f"Cannot add new input collections {self.inputs[:nNew]} after " 

244 "output collection is first created." 

245 ) 

246 if args.extend_run: 

247 if self.outputRun is None: 

248 raise ValueError("Cannot --extend-run when no output collection is given.") 

249 elif not self.outputRun.exists: 

250 raise ValueError( 

251 f"Cannot --extend-run; output collection '{self.outputRun.name}' does not exist." 

252 ) 

253 if not args.extend_run and self.outputRun is not None and self.outputRun.exists: 

254 raise ValueError( 

255 f"Output run '{self.outputRun.name}' already exists, but --extend-run was not given." 

256 ) 

257 if args.prune_replaced and not args.replace_run: 

258 raise ValueError("--prune-replaced requires --replace-run.") 

259 if args.replace_run and (self.output is None or not self.output.exists): 

260 raise ValueError("--output must point to an existing CHAINED collection for --replace-run.") 

261 

262 @classmethod 

263 def _makeReadParts(cls, args: SimpleNamespace) -> tuple[Butler, Sequence[str], _ButlerFactory]: 

264 """Common implementation for `makeReadButler` and 

265 `makeButlerAndCollections`. 

266 

267 Parameters 

268 ---------- 

269 args : `types.SimpleNamespace` 

270 Parsed command-line arguments. See class documentation for the 

271 construction parameter of the same name. 

272 

273 Returns 

274 ------- 

275 butler : `lsst.daf.butler.Butler` 

276 A read-only butler constructed from the repo at 

277 ``args.butler_config``, but with no default collections. 

278 inputs : `Sequence` [ `str` ] 

279 A collection search path constructed according to ``args``. 

280 self : `_ButlerFactory` 

281 A new `_ButlerFactory` instance representing the processed version 

282 of ``args``. 

283 """ 

284 butler = Butler(args.butler_config, writeable=False) 

285 self = cls(butler.registry, args, writeable=False) 

286 self.check(args) 

287 if self.output and self.output.exists: 

288 if args.replace_run: 

289 replaced = self.output.chain[0] 

290 inputs = list(self.output.chain[1:]) 

291 _LOG.debug( 

292 "Simulating collection search in '%s' after removing '%s'.", self.output.name, replaced 

293 ) 

294 else: 

295 inputs = [self.output.name] 

296 else: 

297 inputs = list(self.inputs) 

298 if args.extend_run: 

299 assert self.outputRun is not None, "Output collection has to be specified." 

300 inputs.insert(0, self.outputRun.name) 

301 collSearch = CollectionWildcard.from_expression(inputs).require_ordered() 

302 return butler, collSearch, self 

303 

304 @classmethod 

305 def makeReadButler(cls, args: SimpleNamespace) -> Butler: 

306 """Construct a read-only butler according to the given command-line 

307 arguments. 

308 

309 Parameters 

310 ---------- 

311 args : `types.SimpleNamespace` 

312 Parsed command-line arguments. See class documentation for the 

313 construction parameter of the same name. 

314 

315 Returns 

316 ------- 

317 butler : `lsst.daf.butler.Butler` 

318 A read-only butler initialized with the collections specified by 

319 ``args``. 

320 """ 

321 cls.defineDatastoreCache() # Ensure that this butler can use a shared cache. 

322 butler, inputs, _ = cls._makeReadParts(args) 

323 _LOG.debug("Preparing butler to read from %s.", inputs) 

324 return Butler(butler=butler, collections=inputs) 

325 

326 @classmethod 

327 def makeButlerAndCollections(cls, args: SimpleNamespace) -> Tuple[Butler, Sequence[str], Optional[str]]: 

328 """Return a read-only registry, a collection search path, and the name 

329 of the run to be used for future writes. 

330 

331 Parameters 

332 ---------- 

333 args : `types.SimpleNamespace` 

334 Parsed command-line arguments. See class documentation for the 

335 construction parameter of the same name. 

336 

337 Returns 

338 ------- 

339 butler : `lsst.daf.butler.Butler` 

340 A read-only butler that collections will be added to and/or queried 

341 from. 

342 inputs : `Sequence` [ `str` ] 

343 Collections to search for datasets. 

344 run : `str` or `None` 

345 Name of the output `~lsst.daf.butler.CollectionType.RUN` collection 

346 if it already exists, or `None` if it does not. 

347 """ 

348 butler, inputs, self = cls._makeReadParts(args) 

349 run: Optional[str] = None 

350 if args.extend_run: 

351 assert self.outputRun is not None, "Output collection has to be specified." 

352 run = self.outputRun.name 

353 _LOG.debug("Preparing registry to read from %s and expect future writes to '%s'.", inputs, run) 

354 return butler, inputs, run 

355 

356 @staticmethod 

357 def defineDatastoreCache() -> None: 

358 """Define where datastore cache directories should be found. 

359 

360 Notes 

361 ----- 

362 All the jobs should share a datastore cache if applicable. This 

363 method asks for a shared fallback cache to be defined and then 

364 configures an exit handler to clean it up. 

365 """ 

366 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

367 if defined: 

368 atexit.register(shutil.rmtree, cache_dir, ignore_errors=True) 

369 _LOG.debug("Defining shared datastore cache directory to %s", cache_dir) 

370 

371 @classmethod 

372 def makeWriteButler(cls, args: SimpleNamespace, taskDefs: Optional[Iterable[TaskDef]] = None) -> Butler: 

373 """Return a read-write butler initialized to write to and read from 

374 the collections specified by the given command-line arguments. 

375 

376 Parameters 

377 ---------- 

378 args : `types.SimpleNamespace` 

379 Parsed command-line arguments. See class documentation for the 

380 construction parameter of the same name. 

381 taskDefs : iterable of `TaskDef`, optional 

382 Definitions for tasks in a pipeline. This argument is only needed 

383 if ``args.replace_run`` is `True` and ``args.prune_replaced`` is 

384 "unstore". 

385 

386 Returns 

387 ------- 

388 butler : `lsst.daf.butler.Butler` 

389 A read-write butler initialized according to the given arguments. 

390 """ 

391 cls.defineDatastoreCache() # Ensure that this butler can use a shared cache. 

392 butler = Butler(args.butler_config, writeable=True) 

393 self = cls(butler.registry, args, writeable=True) 

394 self.check(args) 

395 assert self.outputRun is not None, "Output collection has to be specified." # for mypy 

396 if self.output is not None: 

397 chainDefinition = list(self.output.chain if self.output.exists else self.inputs) 

398 if args.replace_run: 

399 replaced = chainDefinition.pop(0) 

400 if args.prune_replaced == "unstore": 

401 # Remove datasets from datastore 

402 with butler.transaction(): 

403 refs: Iterable[DatasetRef] = butler.registry.queryDatasets(..., collections=replaced) 

404 # we want to remove regular outputs but keep 

405 # initOutputs, configs, and versions. 

406 if taskDefs is not None: 

407 initDatasetNames = set(PipelineDatasetTypes.initOutputNames(taskDefs)) 

408 refs = [ref for ref in refs if ref.datasetType.name not in initDatasetNames] 

409 butler.pruneDatasets(refs, unstore=True, disassociate=False) 

410 elif args.prune_replaced == "purge": 

411 # Erase entire collection and all datasets, need to remove 

412 # collection from its chain collection first. 

413 with butler.transaction(): 

414 butler.registry.setCollectionChain(self.output.name, chainDefinition, flatten=True) 

415 butler.pruneCollection(replaced, purge=True, unstore=True) 

416 elif args.prune_replaced is not None: 

417 raise NotImplementedError(f"Unsupported --prune-replaced option '{args.prune_replaced}'.") 

418 if not self.output.exists: 

419 butler.registry.registerCollection(self.output.name, CollectionType.CHAINED) 

420 if not args.extend_run: 

421 butler.registry.registerCollection(self.outputRun.name, CollectionType.RUN) 

422 chainDefinition.insert(0, self.outputRun.name) 

423 butler.registry.setCollectionChain(self.output.name, chainDefinition, flatten=True) 

424 _LOG.debug( 

425 "Preparing butler to write to '%s' and read from '%s'=%s", 

426 self.outputRun.name, 

427 self.output.name, 

428 chainDefinition, 

429 ) 

430 butler.registry.defaults = RegistryDefaults(run=self.outputRun.name, collections=self.output.name) 

431 else: 

432 inputs = (self.outputRun.name,) + self.inputs 

433 _LOG.debug("Preparing butler to write to '%s' and read from %s.", self.outputRun.name, inputs) 

434 butler.registry.defaults = RegistryDefaults(run=self.outputRun.name, collections=inputs) 

435 return butler 

436 

437 output: Optional[_OutputChainedCollectionInfo] 

438 """Information about the output chained collection, if there is or will be 

439 one (`_OutputChainedCollectionInfo` or `None`). 

440 """ 

441 

442 outputRun: Optional[_OutputRunCollectionInfo] 

443 """Information about the output run collection, if there is or will be 

444 one (`_OutputRunCollectionInfo` or `None`). 

445 """ 

446 

447 inputs: Tuple[str, ...] 

448 """Input collections provided directly by the user (`tuple` [ `str` ]). 

449 """ 

450 

451 

452# ------------------------ 

453# Exported definitions -- 

454# ------------------------ 

455 

456 

457class CmdLineFwk: 

458 """PipelineTask framework which executes tasks from command line. 

459 

460 In addition to executing tasks this activator provides additional methods 

461 for task management like dumping configuration or execution chain. 

462 """ 

463 

464 MP_TIMEOUT = 3600 * 24 * 30 # Default timeout (sec) for multiprocessing 

465 

466 def __init__(self) -> None: 

467 pass 

468 

469 def makePipeline(self, args: SimpleNamespace) -> Pipeline: 

470 """Build a pipeline from command line arguments. 

471 

472 Parameters 

473 ---------- 

474 args : `types.SimpleNamespace` 

475 Parsed command line 

476 

477 Returns 

478 ------- 

479 pipeline : `~lsst.pipe.base.Pipeline` 

480 """ 

481 if args.pipeline: 

482 pipeline = Pipeline.from_uri(args.pipeline) 

483 else: 

484 pipeline = Pipeline("anonymous") 

485 

486 # loop over all pipeline actions and apply them in order 

487 for action in args.pipeline_actions: 

488 if action.action == "add_instrument": 

489 pipeline.addInstrument(action.value) 

490 

491 elif action.action == "new_task": 

492 pipeline.addTask(action.value, action.label) 

493 

494 elif action.action == "delete_task": 

495 pipeline.removeTask(action.label) 

496 

497 elif action.action == "config": 

498 # action value string is "field=value", split it at '=' 

499 field, _, value = action.value.partition("=") 

500 pipeline.addConfigOverride(action.label, field, value) 

501 

502 elif action.action == "configfile": 

503 pipeline.addConfigFile(action.label, action.value) 

504 

505 else: 

506 raise ValueError(f"Unexpected pipeline action: {action.action}") 

507 

508 if args.save_pipeline: 

509 pipeline.write_to_uri(args.save_pipeline) 

510 

511 if args.pipeline_dot: 

512 pipeline2dot(pipeline, args.pipeline_dot) 

513 

514 return pipeline 

515 

516 def makeGraph(self, pipeline: Pipeline, args: SimpleNamespace) -> Optional[QuantumGraph]: 

517 """Build a graph from command line arguments. 

518 

519 Parameters 

520 ---------- 

521 pipeline : `~lsst.pipe.base.Pipeline` 

522 Pipeline, can be empty or ``None`` if graph is read from a file. 

523 args : `types.SimpleNamespace` 

524 Parsed command line 

525 

526 Returns 

527 ------- 

528 graph : `~lsst.pipe.base.QuantumGraph` or `None` 

529 If resulting graph is empty then `None` is returned. 

530 """ 

531 

532 # make sure that --extend-run always enables --skip-existing 

533 if args.extend_run: 

534 args.skip_existing = True 

535 

536 butler, collections, run = _ButlerFactory.makeButlerAndCollections(args) 

537 

538 if args.skip_existing and run: 

539 args.skip_existing_in += (run,) 

540 

541 if args.qgraph: 

542 # click passes empty tuple as default value for qgraph_node_id 

543 nodes = args.qgraph_node_id or None 

544 qgraph = QuantumGraph.loadUri( 

545 args.qgraph, butler.registry.dimensions, nodes=nodes, graphID=args.qgraph_id 

546 ) 

547 

548 # pipeline can not be provided in this case 

549 if pipeline: 

550 raise ValueError("Pipeline must not be given when quantum graph is read from file.") 

551 if args.show_qgraph_header: 

552 print(QuantumGraph.readHeader(args.qgraph)) 

553 else: 

554 # make execution plan (a.k.a. DAG) for pipeline 

555 graphBuilder = GraphBuilder( 

556 butler.registry, 

557 skipExistingIn=args.skip_existing_in, 

558 clobberOutputs=args.clobber_outputs, 

559 datastore=butler.datastore if args.qgraph_datastore_records else None, 

560 ) 

561 # accumulate metadata 

562 metadata = { 

563 "input": args.input, 

564 "output": args.output, 

565 "butler_argument": args.butler_config, 

566 "output_run": args.output_run, 

567 "extend_run": args.extend_run, 

568 "skip_existing_in": args.skip_existing_in, 

569 "skip_existing": args.skip_existing, 

570 "data_query": args.data_query, 

571 "user": getpass.getuser(), 

572 "time": f"{datetime.datetime.now()}", 

573 } 

574 qgraph = graphBuilder.makeGraph( 

575 pipeline, 

576 collections, 

577 run, 

578 args.data_query, 

579 metadata=metadata, 

580 datasetQueryConstraint=args.dataset_query_constraint, 

581 ) 

582 if args.show_qgraph_header: 

583 qgraph.buildAndPrintHeader() 

584 

585 # Count quanta in graph; give a warning if it's empty and return None. 

586 nQuanta = len(qgraph) 

587 if nQuanta == 0: 

588 return None 

589 else: 

590 if _LOG.isEnabledFor(logging.INFO): 

591 qg_task_table = self._generateTaskTable(qgraph) 

592 qg_task_table_formatted = "\n".join(qg_task_table.pformat_all()) 

593 _LOG.info( 

594 "QuantumGraph contains %d quanta for %d tasks, graph ID: %r\n%s", 

595 nQuanta, 

596 len(qgraph.taskGraph), 

597 qgraph.graphID, 

598 qg_task_table_formatted, 

599 ) 

600 

601 if args.save_qgraph: 

602 qgraph.saveUri(args.save_qgraph) 

603 

604 if args.save_single_quanta: 

605 for quantumNode in qgraph: 

606 sqgraph = qgraph.subset(quantumNode) 

607 uri = args.save_single_quanta.format(quantumNode) 

608 sqgraph.saveUri(uri) 

609 

610 if args.qgraph_dot: 

611 graph2dot(qgraph, args.qgraph_dot) 

612 

613 if args.execution_butler_location: 

614 butler = Butler(args.butler_config) 

615 newArgs = copy.deepcopy(args) 

616 

617 def builderShim(butler: Butler) -> Butler: 

618 newArgs.butler_config = butler._config 

619 # Calling makeWriteButler is done for the side effects of 

620 # calling that method, maining parsing all the args into 

621 # collection names, creating collections, etc. 

622 newButler = _ButlerFactory.makeWriteButler(newArgs) 

623 return newButler 

624 

625 # Include output collection in collections for input 

626 # files if it exists in the repo. 

627 all_inputs = args.input 

628 if args.output is not None: 

629 try: 

630 all_inputs += (next(iter(butler.registry.queryCollections(args.output))),) 

631 except MissingCollectionError: 

632 pass 

633 

634 _LOG.debug("Calling buildExecutionButler with collections=%s", all_inputs) 

635 buildExecutionButler( 

636 butler, 

637 qgraph, 

638 args.execution_butler_location, 

639 run, 

640 butlerModifier=builderShim, 

641 collections=all_inputs, 

642 clobber=args.clobber_execution_butler, 

643 datastoreRoot=args.target_datastore_root, 

644 transfer=args.transfer, 

645 ) 

646 

647 return qgraph 

648 

649 def runPipeline( 

650 self, 

651 graph: QuantumGraph, 

652 taskFactory: TaskFactory, 

653 args: SimpleNamespace, 

654 butler: Optional[Butler] = None, 

655 ) -> None: 

656 """Execute complete QuantumGraph. 

657 

658 Parameters 

659 ---------- 

660 graph : `QuantumGraph` 

661 Execution graph. 

662 taskFactory : `~lsst.pipe.base.TaskFactory` 

663 Task factory 

664 args : `types.SimpleNamespace` 

665 Parsed command line 

666 butler : `~lsst.daf.butler.Butler`, optional 

667 Data Butler instance, if not defined then new instance is made 

668 using command line options. 

669 """ 

670 # make sure that --extend-run always enables --skip-existing 

671 if args.extend_run: 

672 args.skip_existing = True 

673 

674 if not args.enable_implicit_threading: 

675 disable_implicit_threading() 

676 

677 # make butler instance 

678 if butler is None: 

679 butler = _ButlerFactory.makeWriteButler(args, graph.iterTaskGraph()) 

680 

681 if args.skip_existing: 

682 args.skip_existing_in += (butler.run,) 

683 

684 # Enable lsstDebug debugging. Note that this is done once in the 

685 # main process before PreExecInit and it is also repeated before 

686 # running each task in SingleQuantumExecutor (which may not be 

687 # needed if `multipocessing` always uses fork start method). 

688 if args.enableLsstDebug: 

689 try: 

690 _LOG.debug("Will try to import debug.py") 

691 import debug # type: ignore # noqa:F401 

692 except ImportError: 

693 _LOG.warn("No 'debug' module found.") 

694 

695 # Save all InitOutputs, configs, etc. 

696 preExecInit = PreExecInit(butler, taskFactory, extendRun=args.extend_run, mock=args.mock) 

697 preExecInit.initialize( 

698 graph, 

699 saveInitOutputs=not args.skip_init_writes, 

700 registerDatasetTypes=args.register_dataset_types, 

701 saveVersions=not args.no_versions, 

702 ) 

703 

704 if not args.init_only: 

705 graphFixup = self._importGraphFixup(args) 

706 quantumExecutor = SingleQuantumExecutor( 

707 taskFactory, 

708 skipExistingIn=args.skip_existing_in, 

709 clobberOutputs=args.clobber_outputs, 

710 enableLsstDebug=args.enableLsstDebug, 

711 exitOnKnownError=args.fail_fast, 

712 mock=args.mock, 

713 mock_configs=args.mock_configs, 

714 ) 

715 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout 

716 executor = MPGraphExecutor( 

717 numProc=args.processes, 

718 timeout=timeout, 

719 startMethod=args.start_method, 

720 quantumExecutor=quantumExecutor, 

721 failFast=args.fail_fast, 

722 pdb=args.pdb, 

723 executionGraphFixup=graphFixup, 

724 ) 

725 try: 

726 with util.profile(args.profile, _LOG): 

727 executor.execute(graph, butler) 

728 finally: 

729 if args.summary: 

730 report = executor.getReport() 

731 if report: 

732 with open(args.summary, "w") as out: 

733 # Do not save fields that are not set. 

734 out.write(report.json(exclude_none=True, indent=2)) 

735 

736 def _generateTaskTable(self, qgraph: QuantumGraph) -> Table: 

737 """Generate astropy table listing the number of quanta per task for a 

738 given quantum graph. 

739 

740 Parameters 

741 ---------- 

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

743 A QuantumGraph object. 

744 

745 Returns 

746 ------- 

747 qg_task_table : `astropy.table.table.Table` 

748 An astropy table containing columns: Quanta and Tasks. 

749 """ 

750 qg_quanta, qg_tasks = [], [] 

751 for task_def in qgraph.iterTaskGraph(): 

752 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def) 

753 qg_quanta.append(num_qnodes) 

754 qg_tasks.append(task_def.label) 

755 qg_task_table = Table(dict(Quanta=qg_quanta, Tasks=qg_tasks)) 

756 return qg_task_table 

757 

758 def _importGraphFixup(self, args: SimpleNamespace) -> Optional[ExecutionGraphFixup]: 

759 """Import/instantiate graph fixup object. 

760 

761 Parameters 

762 ---------- 

763 args : `types.SimpleNamespace` 

764 Parsed command line. 

765 

766 Returns 

767 ------- 

768 fixup : `ExecutionGraphFixup` or `None` 

769 

770 Raises 

771 ------ 

772 ValueError 

773 Raised if import fails, method call raises exception, or returned 

774 instance has unexpected type. 

775 """ 

776 if args.graph_fixup: 

777 try: 

778 factory = doImportType(args.graph_fixup) 

779 except Exception as exc: 

780 raise ValueError("Failed to import graph fixup class/method") from exc 

781 try: 

782 fixup = factory() 

783 except Exception as exc: 

784 raise ValueError("Failed to make instance of graph fixup") from exc 

785 if not isinstance(fixup, ExecutionGraphFixup): 

786 raise ValueError("Graph fixup is not an instance of ExecutionGraphFixup class") 

787 return fixup 

788 return None