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

345 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-25 02:56 -0800

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 TYPE_CHECKING, Optional, Tuple 

38 

39from astropy.table import Table 

40from lsst.daf.butler import ( 

41 Butler, 

42 CollectionType, 

43 DatasetId, 

44 DatasetRef, 

45 DatastoreCacheManager, 

46 QuantumBackedButler, 

47) 

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

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

50from lsst.pipe.base import ( 

51 GraphBuilder, 

52 Instrument, 

53 Pipeline, 

54 PipelineDatasetTypes, 

55 QuantumGraph, 

56 buildExecutionButler, 

57) 

58from lsst.utils import doImportType 

59from lsst.utils.threads import disable_implicit_threading 

60 

61from . import util 

62from .dotTools import graph2dot, pipeline2dot 

63from .executionGraphFixup import ExecutionGraphFixup 

64from .mpGraphExecutor import MPGraphExecutor 

65from .preExecInit import PreExecInit, PreExecInitLimited 

66from .singleQuantumExecutor import SingleQuantumExecutor 

67 

68if TYPE_CHECKING: 68 ↛ 69line 68 didn't jump to line 69, because the condition on line 68 was never true

69 from lsst.daf.butler import DatastoreRecordData, Registry 

70 from lsst.pipe.base import TaskDef, TaskFactory 

71 

72 

73# ---------------------------------- 

74# Local non-exported definitions -- 

75# ---------------------------------- 

76 

77_LOG = logging.getLogger(__name__) 

78 

79 

80class _OutputChainedCollectionInfo: 

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

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

83 

84 Parameters 

85 ---------- 

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

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

88 name : `str` 

89 Name of the collection given on the command line. 

90 """ 

91 

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

93 self.name = name 

94 try: 

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

96 self.exists = True 

97 except MissingCollectionError: 

98 self.chain = () 

99 self.exists = False 

100 

101 def __str__(self) -> str: 

102 return self.name 

103 

104 name: str 

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

106 """ 

107 

108 exists: bool 

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

110 """ 

111 

112 chain: Tuple[str, ...] 

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

114 

115 Empty if the collection does not already exist. 

116 """ 

117 

118 

119class _OutputRunCollectionInfo: 

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

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

122 

123 Parameters 

124 ---------- 

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

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

127 name : `str` 

128 Name of the collection given on the command line. 

129 """ 

130 

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

132 self.name = name 

133 try: 

134 actualType = registry.getCollectionType(name) 

135 if actualType is not CollectionType.RUN: 

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

137 self.exists = True 

138 except MissingCollectionError: 

139 self.exists = False 

140 

141 name: str 

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

143 """ 

144 

145 exists: bool 

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

147 """ 

148 

149 

150class _ButlerFactory: 

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

152 and output collections. 

153 

154 Parameters 

155 ---------- 

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

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

158 

159 args : `types.SimpleNamespace` 

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

161 either at construction or in later methods. 

162 

163 ``output`` 

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

165 input/output collection. 

166 

167 ``output_run`` 

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

169 collection. 

170 

171 ``extend_run`` 

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

173 and be extended. 

174 

175 ``replace_run`` 

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

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

178 replaced with a new one. 

179 

180 ``prune_replaced`` 

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

182 ``replace_run``). 

183 

184 ``inputs`` 

185 Input collections of any type; see 

186 :ref:`daf_butler_ordered_collection_searches` for details. 

187 

188 ``butler_config`` 

189 Path to a data repository root or configuration file. 

190 

191 writeable : `bool` 

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

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

194 

195 Raises 

196 ------ 

197 ValueError 

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

199 """ 

200 

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

202 if args.output is not None: 

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

204 else: 

205 self.output = None 

206 if args.output_run is not None: 

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

208 elif self.output is not None: 

209 if args.extend_run: 

210 if not self.output.chain: 

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

212 runName = self.output.chain[0] 

213 else: 

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

215 self.outputRun = _OutputRunCollectionInfo(registry, runName) 

216 elif not writeable: 

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

218 self.outputRun = None 

219 else: 

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

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

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

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

224 # collection. 

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

226 

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

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

229 data repository. 

230 

231 Parameters 

232 ---------- 

233 args : `types.SimpleNamespace` 

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

235 construction parameter of the same name. 

236 """ 

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

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

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

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

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

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

243 if c1 != c2: 

244 raise ValueError( 

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

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

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

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

249 ) 

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

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

252 raise ValueError( 

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

254 "output collection is first created." 

255 ) 

256 if args.extend_run: 

257 if self.outputRun is None: 

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

259 elif not self.outputRun.exists: 

260 raise ValueError( 

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

262 ) 

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

264 raise ValueError( 

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

266 ) 

267 if args.prune_replaced and not args.replace_run: 

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

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

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

271 

272 @classmethod 

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

274 """Common implementation for `makeReadButler` and 

275 `makeButlerAndCollections`. 

276 

277 Parameters 

278 ---------- 

279 args : `types.SimpleNamespace` 

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

281 construction parameter of the same name. 

282 

283 Returns 

284 ------- 

285 butler : `lsst.daf.butler.Butler` 

286 A read-only butler constructed from the repo at 

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

288 inputs : `Sequence` [ `str` ] 

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

290 self : `_ButlerFactory` 

291 A new `_ButlerFactory` instance representing the processed version 

292 of ``args``. 

293 """ 

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

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

296 self.check(args) 

297 if self.output and self.output.exists: 

298 if args.replace_run: 

299 replaced = self.output.chain[0] 

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

301 _LOG.debug( 

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

303 ) 

304 else: 

305 inputs = [self.output.name] 

306 else: 

307 inputs = list(self.inputs) 

308 if args.extend_run: 

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

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

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

312 return butler, collSearch, self 

313 

314 @classmethod 

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

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

317 arguments. 

318 

319 Parameters 

320 ---------- 

321 args : `types.SimpleNamespace` 

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

323 construction parameter of the same name. 

324 

325 Returns 

326 ------- 

327 butler : `lsst.daf.butler.Butler` 

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

329 ``args``. 

330 """ 

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

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

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

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

335 

336 @classmethod 

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

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

339 of the run to be used for future writes. 

340 

341 Parameters 

342 ---------- 

343 args : `types.SimpleNamespace` 

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

345 construction parameter of the same name. 

346 

347 Returns 

348 ------- 

349 butler : `lsst.daf.butler.Butler` 

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

351 from. 

352 inputs : `Sequence` [ `str` ] 

353 Collections to search for datasets. 

354 run : `str` or `None` 

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

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

357 """ 

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

359 run: Optional[str] = None 

360 if args.extend_run: 

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

362 if self.outputRun is not None: 

363 run = self.outputRun.name 

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

365 return butler, inputs, run 

366 

367 @staticmethod 

368 def defineDatastoreCache() -> None: 

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

370 

371 Notes 

372 ----- 

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

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

375 configures an exit handler to clean it up. 

376 """ 

377 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

378 if defined: 

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

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

381 

382 @classmethod 

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

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

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

386 

387 Parameters 

388 ---------- 

389 args : `types.SimpleNamespace` 

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

391 construction parameter of the same name. 

392 taskDefs : iterable of `TaskDef`, optional 

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

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

395 "unstore". 

396 

397 Returns 

398 ------- 

399 butler : `lsst.daf.butler.Butler` 

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

401 """ 

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

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

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

405 self.check(args) 

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

407 if self.output is not None: 

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

409 if args.replace_run: 

410 replaced = chainDefinition.pop(0) 

411 if args.prune_replaced == "unstore": 

412 # Remove datasets from datastore 

413 with butler.transaction(): 

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

415 # we want to remove regular outputs but keep 

416 # initOutputs, configs, and versions. 

417 if taskDefs is not None: 

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

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

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

421 elif args.prune_replaced == "purge": 

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

423 # collection from its chain collection first. 

424 with butler.transaction(): 

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

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

427 elif args.prune_replaced is not None: 

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

429 if not self.output.exists: 

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

431 if not args.extend_run: 

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

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

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

435 _LOG.debug( 

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

437 self.outputRun.name, 

438 self.output.name, 

439 chainDefinition, 

440 ) 

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

442 else: 

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

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

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

446 return butler 

447 

448 output: Optional[_OutputChainedCollectionInfo] 

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

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

451 """ 

452 

453 outputRun: Optional[_OutputRunCollectionInfo] 

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

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

456 """ 

457 

458 inputs: Tuple[str, ...] 

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

460 """ 

461 

462 

463# ------------------------ 

464# Exported definitions -- 

465# ------------------------ 

466 

467 

468class CmdLineFwk: 

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

470 

471 In addition to executing tasks this activator provides additional methods 

472 for task management like dumping configuration or execution chain. 

473 """ 

474 

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

476 

477 def __init__(self) -> None: 

478 pass 

479 

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

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

482 

483 Parameters 

484 ---------- 

485 args : `types.SimpleNamespace` 

486 Parsed command line 

487 

488 Returns 

489 ------- 

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

491 """ 

492 if args.pipeline: 

493 pipeline = Pipeline.from_uri(args.pipeline) 

494 else: 

495 pipeline = Pipeline("anonymous") 

496 

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

498 for action in args.pipeline_actions: 

499 if action.action == "add_instrument": 

500 

501 pipeline.addInstrument(action.value) 

502 

503 elif action.action == "new_task": 

504 

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

506 

507 elif action.action == "delete_task": 

508 

509 pipeline.removeTask(action.label) 

510 

511 elif action.action == "config": 

512 

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

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

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

516 

517 elif action.action == "configfile": 

518 

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

520 

521 else: 

522 

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

524 

525 if args.save_pipeline: 

526 pipeline.write_to_uri(args.save_pipeline) 

527 

528 if args.pipeline_dot: 

529 pipeline2dot(pipeline, args.pipeline_dot) 

530 

531 return pipeline 

532 

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

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

535 

536 Parameters 

537 ---------- 

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

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

540 args : `types.SimpleNamespace` 

541 Parsed command line 

542 

543 Returns 

544 ------- 

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

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

547 """ 

548 

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

550 if args.extend_run: 

551 args.skip_existing = True 

552 

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

554 

555 if args.skip_existing and run: 

556 args.skip_existing_in += (run,) 

557 

558 if args.qgraph: 

559 # click passes empty tuple as default value for qgraph_node_id 

560 nodes = args.qgraph_node_id or None 

561 qgraph = QuantumGraph.loadUri( 

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

563 ) 

564 

565 # pipeline can not be provided in this case 

566 if pipeline: 

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

568 if args.show_qgraph_header: 

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

570 else: 

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

572 graphBuilder = GraphBuilder( 

573 butler.registry, 

574 skipExistingIn=args.skip_existing_in, 

575 clobberOutputs=args.clobber_outputs, 

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

577 ) 

578 # accumulate metadata 

579 metadata = { 

580 "input": args.input, 

581 "output": args.output, 

582 "butler_argument": args.butler_config, 

583 "output_run": run, 

584 "extend_run": args.extend_run, 

585 "skip_existing_in": args.skip_existing_in, 

586 "skip_existing": args.skip_existing, 

587 "data_query": args.data_query, 

588 "user": getpass.getuser(), 

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

590 } 

591 qgraph = graphBuilder.makeGraph( 

592 pipeline, 

593 collections, 

594 run, 

595 args.data_query, 

596 metadata=metadata, 

597 datasetQueryConstraint=args.dataset_query_constraint, 

598 resolveRefs=True, 

599 ) 

600 if args.show_qgraph_header: 

601 qgraph.buildAndPrintHeader() 

602 

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

604 nQuanta = len(qgraph) 

605 if nQuanta == 0: 

606 return None 

607 else: 

608 if _LOG.isEnabledFor(logging.INFO): 

609 qg_task_table = self._generateTaskTable(qgraph) 

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

611 _LOG.info( 

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

613 nQuanta, 

614 len(qgraph.taskGraph), 

615 qgraph.graphID, 

616 qg_task_table_formatted, 

617 ) 

618 

619 if args.save_qgraph: 

620 qgraph.saveUri(args.save_qgraph) 

621 

622 if args.save_single_quanta: 

623 for quantumNode in qgraph: 

624 sqgraph = qgraph.subset(quantumNode) 

625 uri = args.save_single_quanta.format(quantumNode) 

626 sqgraph.saveUri(uri) 

627 

628 if args.qgraph_dot: 

629 graph2dot(qgraph, args.qgraph_dot) 

630 

631 if args.execution_butler_location: 

632 butler = Butler(args.butler_config) 

633 newArgs = copy.deepcopy(args) 

634 

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

636 newArgs.butler_config = butler._config 

637 # Calling makeWriteButler is done for the side effects of 

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

639 # collection names, creating collections, etc. 

640 newButler = _ButlerFactory.makeWriteButler(newArgs) 

641 return newButler 

642 

643 # Include output collection in collections for input 

644 # files if it exists in the repo. 

645 all_inputs = args.input 

646 if args.output is not None: 

647 try: 

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

649 except MissingCollectionError: 

650 pass 

651 

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

653 buildExecutionButler( 

654 butler, 

655 qgraph, 

656 args.execution_butler_location, 

657 run, 

658 butlerModifier=builderShim, 

659 collections=all_inputs, 

660 clobber=args.clobber_execution_butler, 

661 datastoreRoot=args.target_datastore_root, 

662 transfer=args.transfer, 

663 ) 

664 

665 return qgraph 

666 

667 def runPipeline( 

668 self, 

669 graph: QuantumGraph, 

670 taskFactory: TaskFactory, 

671 args: SimpleNamespace, 

672 butler: Optional[Butler] = None, 

673 ) -> None: 

674 """Execute complete QuantumGraph. 

675 

676 Parameters 

677 ---------- 

678 graph : `QuantumGraph` 

679 Execution graph. 

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

681 Task factory 

682 args : `types.SimpleNamespace` 

683 Parsed command line 

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

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

686 using command line options. 

687 """ 

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

689 if args.extend_run: 

690 args.skip_existing = True 

691 

692 if not args.enable_implicit_threading: 

693 disable_implicit_threading() 

694 

695 # Make butler instance. QuantumGraph should have an output run defined, 

696 # but we ignore it here and let command line decide actual output run. 

697 if butler is None: 

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

699 

700 if args.skip_existing: 

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

702 

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

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

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

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

707 if args.enableLsstDebug: 

708 try: 

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

710 import debug # type: ignore # noqa:F401 

711 except ImportError: 

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

713 

714 # Save all InitOutputs, configs, etc. 

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

716 preExecInit.initialize( 

717 graph, 

718 saveInitOutputs=not args.skip_init_writes, 

719 registerDatasetTypes=args.register_dataset_types, 

720 saveVersions=not args.no_versions, 

721 ) 

722 

723 if not args.init_only: 

724 graphFixup = self._importGraphFixup(args) 

725 quantumExecutor = SingleQuantumExecutor( 

726 butler, 

727 taskFactory, 

728 skipExistingIn=args.skip_existing_in, 

729 clobberOutputs=args.clobber_outputs, 

730 enableLsstDebug=args.enableLsstDebug, 

731 exitOnKnownError=args.fail_fast, 

732 mock=args.mock, 

733 mock_configs=args.mock_configs, 

734 ) 

735 

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

737 executor = MPGraphExecutor( 

738 numProc=args.processes, 

739 timeout=timeout, 

740 startMethod=args.start_method, 

741 quantumExecutor=quantumExecutor, 

742 failFast=args.fail_fast, 

743 pdb=args.pdb, 

744 executionGraphFixup=graphFixup, 

745 ) 

746 # Have to reset connection pool to avoid sharing connections with 

747 # forked processes. 

748 butler.registry.resetConnectionPool() 

749 try: 

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

751 executor.execute(graph) 

752 finally: 

753 if args.summary: 

754 report = executor.getReport() 

755 if report: 

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

757 # Do not save fields that are not set. 

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

759 

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

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

762 given quantum graph. 

763 

764 Parameters 

765 ---------- 

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

767 A QuantumGraph object. 

768 

769 Returns 

770 ------- 

771 qg_task_table : `astropy.table.table.Table` 

772 An astropy table containing columns: Quanta and Tasks. 

773 """ 

774 qg_quanta, qg_tasks = [], [] 

775 for task_def in qgraph.iterTaskGraph(): 

776 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def) 

777 qg_quanta.append(num_qnodes) 

778 qg_tasks.append(task_def.label) 

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

780 return qg_task_table 

781 

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

783 """Import/instantiate graph fixup object. 

784 

785 Parameters 

786 ---------- 

787 args : `types.SimpleNamespace` 

788 Parsed command line. 

789 

790 Returns 

791 ------- 

792 fixup : `ExecutionGraphFixup` or `None` 

793 

794 Raises 

795 ------ 

796 ValueError 

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

798 instance has unexpected type. 

799 """ 

800 if args.graph_fixup: 

801 try: 

802 factory = doImportType(args.graph_fixup) 

803 except Exception as exc: 

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

805 try: 

806 fixup = factory() 

807 except Exception as exc: 

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

809 if not isinstance(fixup, ExecutionGraphFixup): 

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

811 return fixup 

812 return None 

813 

814 def preExecInitQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None: 

815 

816 # Load quantum graph. We do not really need individual Quanta here, 

817 # but we need datastore records for initInputs, and those are only 

818 # available from Quanta, so load the whole thing. 

819 qgraph = QuantumGraph.loadUri(args.qgraph, graphID=args.qgraph_id) 

820 universe = qgraph.universe 

821 

822 # Collect all init input/output dataset IDs. 

823 predicted_inputs: set[DatasetId] = set() 

824 predicted_outputs: set[DatasetId] = set() 

825 for taskDef in qgraph.iterTaskGraph(): 

826 if (refs := qgraph.initInputRefs(taskDef)) is not None: 

827 predicted_inputs.update(ref.getCheckedId() for ref in refs) 

828 if (refs := qgraph.initOutputRefs(taskDef)) is not None: 

829 predicted_outputs.update(ref.getCheckedId() for ref in refs) 

830 predicted_outputs.update(ref.getCheckedId() for ref in qgraph.globalInitOutputRefs()) 

831 # remove intermediates from inputs 

832 predicted_inputs -= predicted_outputs 

833 

834 # Very inefficient way to extract datastore records from quantum graph, 

835 # we have to scan all quanta and look at their datastore records. 

836 datastore_records: dict[str, DatastoreRecordData] = {} 

837 for quantum_node in qgraph: 

838 for store_name, records in quantum_node.quantum.datastore_records.items(): 

839 subset = records.subset(predicted_inputs) 

840 if subset is not None: 

841 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset) 

842 

843 # Make butler from everything. 

844 butler = QuantumBackedButler.from_predicted( 

845 config=args.butler_config, 

846 predicted_inputs=predicted_inputs, 

847 predicted_outputs=predicted_outputs, 

848 dimensions=universe, 

849 datastore_records=datastore_records, 

850 search_paths=args.config_search_path, 

851 ) 

852 

853 # Save all InitOutputs, configs, etc. 

854 preExecInit = PreExecInitLimited(butler, task_factory) 

855 preExecInit.initialize(qgraph) 

856 

857 def runGraphQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None: 

858 

859 # Load quantum graph. 

860 nodes = args.qgraph_node_id or None 

861 qgraph = QuantumGraph.loadUri(args.qgraph, nodes=nodes, graphID=args.qgraph_id) 

862 

863 if qgraph.metadata is None: 

864 raise ValueError("QuantumGraph is missing metadata, cannot ") 

865 

866 # make special quantum executor 

867 quantumExecutor = SingleQuantumExecutor( 

868 butler=None, 

869 taskFactory=task_factory, 

870 enableLsstDebug=args.enableLsstDebug, 

871 exitOnKnownError=args.fail_fast, 

872 butler_config=args.butler_config, 

873 universe=qgraph.universe, 

874 ) 

875 

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

877 executor = MPGraphExecutor( 

878 numProc=args.processes, 

879 timeout=timeout, 

880 startMethod=args.start_method, 

881 quantumExecutor=quantumExecutor, 

882 failFast=args.fail_fast, 

883 pdb=args.pdb, 

884 ) 

885 try: 

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

887 executor.execute(qgraph) 

888 finally: 

889 if args.summary: 

890 report = executor.getReport() 

891 if report: 

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

893 # Do not save fields that are not set. 

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