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

346 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-13 02:55 -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 # Execution butler builder relies on non-resolved output refs. 

592 resolve_refs = not args.execution_butler_location 

593 qgraph = graphBuilder.makeGraph( 

594 pipeline, 

595 collections, 

596 run, 

597 args.data_query, 

598 metadata=metadata, 

599 datasetQueryConstraint=args.dataset_query_constraint, 

600 resolveRefs=resolve_refs, 

601 ) 

602 if args.show_qgraph_header: 

603 qgraph.buildAndPrintHeader() 

604 

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

606 nQuanta = len(qgraph) 

607 if nQuanta == 0: 

608 return None 

609 else: 

610 if _LOG.isEnabledFor(logging.INFO): 

611 qg_task_table = self._generateTaskTable(qgraph) 

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

613 _LOG.info( 

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

615 nQuanta, 

616 len(qgraph.taskGraph), 

617 qgraph.graphID, 

618 qg_task_table_formatted, 

619 ) 

620 

621 if args.save_qgraph: 

622 qgraph.saveUri(args.save_qgraph) 

623 

624 if args.save_single_quanta: 

625 for quantumNode in qgraph: 

626 sqgraph = qgraph.subset(quantumNode) 

627 uri = args.save_single_quanta.format(quantumNode) 

628 sqgraph.saveUri(uri) 

629 

630 if args.qgraph_dot: 

631 graph2dot(qgraph, args.qgraph_dot) 

632 

633 if args.execution_butler_location: 

634 butler = Butler(args.butler_config) 

635 newArgs = copy.deepcopy(args) 

636 

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

638 newArgs.butler_config = butler._config 

639 # Calling makeWriteButler is done for the side effects of 

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

641 # collection names, creating collections, etc. 

642 newButler = _ButlerFactory.makeWriteButler(newArgs) 

643 return newButler 

644 

645 # Include output collection in collections for input 

646 # files if it exists in the repo. 

647 all_inputs = args.input 

648 if args.output is not None: 

649 try: 

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

651 except MissingCollectionError: 

652 pass 

653 

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

655 buildExecutionButler( 

656 butler, 

657 qgraph, 

658 args.execution_butler_location, 

659 run, 

660 butlerModifier=builderShim, 

661 collections=all_inputs, 

662 clobber=args.clobber_execution_butler, 

663 datastoreRoot=args.target_datastore_root, 

664 transfer=args.transfer, 

665 ) 

666 

667 return qgraph 

668 

669 def runPipeline( 

670 self, 

671 graph: QuantumGraph, 

672 taskFactory: TaskFactory, 

673 args: SimpleNamespace, 

674 butler: Optional[Butler] = None, 

675 ) -> None: 

676 """Execute complete QuantumGraph. 

677 

678 Parameters 

679 ---------- 

680 graph : `QuantumGraph` 

681 Execution graph. 

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

683 Task factory 

684 args : `types.SimpleNamespace` 

685 Parsed command line 

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

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

688 using command line options. 

689 """ 

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

691 if args.extend_run: 

692 args.skip_existing = True 

693 

694 if not args.enable_implicit_threading: 

695 disable_implicit_threading() 

696 

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

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

699 if butler is None: 

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

701 

702 if args.skip_existing: 

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

704 

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

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

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

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

709 if args.enableLsstDebug: 

710 try: 

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

712 import debug # type: ignore # noqa:F401 

713 except ImportError: 

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

715 

716 # Save all InitOutputs, configs, etc. 

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

718 preExecInit.initialize( 

719 graph, 

720 saveInitOutputs=not args.skip_init_writes, 

721 registerDatasetTypes=args.register_dataset_types, 

722 saveVersions=not args.no_versions, 

723 ) 

724 

725 if not args.init_only: 

726 graphFixup = self._importGraphFixup(args) 

727 quantumExecutor = SingleQuantumExecutor( 

728 butler, 

729 taskFactory, 

730 skipExistingIn=args.skip_existing_in, 

731 clobberOutputs=args.clobber_outputs, 

732 enableLsstDebug=args.enableLsstDebug, 

733 exitOnKnownError=args.fail_fast, 

734 mock=args.mock, 

735 mock_configs=args.mock_configs, 

736 ) 

737 

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

739 executor = MPGraphExecutor( 

740 numProc=args.processes, 

741 timeout=timeout, 

742 startMethod=args.start_method, 

743 quantumExecutor=quantumExecutor, 

744 failFast=args.fail_fast, 

745 pdb=args.pdb, 

746 executionGraphFixup=graphFixup, 

747 ) 

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

749 # forked processes. 

750 butler.registry.resetConnectionPool() 

751 try: 

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

753 executor.execute(graph) 

754 finally: 

755 if args.summary: 

756 report = executor.getReport() 

757 if report: 

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

759 # Do not save fields that are not set. 

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

761 

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

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

764 given quantum graph. 

765 

766 Parameters 

767 ---------- 

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

769 A QuantumGraph object. 

770 

771 Returns 

772 ------- 

773 qg_task_table : `astropy.table.table.Table` 

774 An astropy table containing columns: Quanta and Tasks. 

775 """ 

776 qg_quanta, qg_tasks = [], [] 

777 for task_def in qgraph.iterTaskGraph(): 

778 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def) 

779 qg_quanta.append(num_qnodes) 

780 qg_tasks.append(task_def.label) 

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

782 return qg_task_table 

783 

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

785 """Import/instantiate graph fixup object. 

786 

787 Parameters 

788 ---------- 

789 args : `types.SimpleNamespace` 

790 Parsed command line. 

791 

792 Returns 

793 ------- 

794 fixup : `ExecutionGraphFixup` or `None` 

795 

796 Raises 

797 ------ 

798 ValueError 

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

800 instance has unexpected type. 

801 """ 

802 if args.graph_fixup: 

803 try: 

804 factory = doImportType(args.graph_fixup) 

805 except Exception as exc: 

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

807 try: 

808 fixup = factory() 

809 except Exception as exc: 

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

811 if not isinstance(fixup, ExecutionGraphFixup): 

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

813 return fixup 

814 return None 

815 

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

817 

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

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

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

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

822 universe = qgraph.universe 

823 

824 # Collect all init input/output dataset IDs. 

825 predicted_inputs: set[DatasetId] = set() 

826 predicted_outputs: set[DatasetId] = set() 

827 for taskDef in qgraph.iterTaskGraph(): 

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

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

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

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

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

833 # remove intermediates from inputs 

834 predicted_inputs -= predicted_outputs 

835 

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

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

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

839 for quantum_node in qgraph: 

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

841 subset = records.subset(predicted_inputs) 

842 if subset is not None: 

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

844 

845 # Make butler from everything. 

846 butler = QuantumBackedButler.from_predicted( 

847 config=args.butler_config, 

848 predicted_inputs=predicted_inputs, 

849 predicted_outputs=predicted_outputs, 

850 dimensions=universe, 

851 datastore_records=datastore_records, 

852 search_paths=args.config_search_path, 

853 ) 

854 

855 # Save all InitOutputs, configs, etc. 

856 preExecInit = PreExecInitLimited(butler, task_factory) 

857 preExecInit.initialize(qgraph) 

858 

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

860 

861 # Load quantum graph. 

862 nodes = args.qgraph_node_id or None 

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

864 

865 if qgraph.metadata is None: 

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

867 

868 # make special quantum executor 

869 quantumExecutor = SingleQuantumExecutor( 

870 butler=None, 

871 taskFactory=task_factory, 

872 enableLsstDebug=args.enableLsstDebug, 

873 exitOnKnownError=args.fail_fast, 

874 butler_config=args.butler_config, 

875 universe=qgraph.universe, 

876 ) 

877 

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

879 executor = MPGraphExecutor( 

880 numProc=args.processes, 

881 timeout=timeout, 

882 startMethod=args.start_method, 

883 quantumExecutor=quantumExecutor, 

884 failFast=args.fail_fast, 

885 pdb=args.pdb, 

886 ) 

887 try: 

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

889 executor.execute(qgraph) 

890 finally: 

891 if args.summary: 

892 report = executor.getReport() 

893 if report: 

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

895 # Do not save fields that are not set. 

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