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

373 statements  

« prev     ^ index     » next       coverage.py v7.3.0, created at 2023-08-25 09:44 +0000

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 contextlib 

31import copy 

32import datetime 

33import getpass 

34import logging 

35import shutil 

36from collections.abc import Iterable, Mapping, Sequence 

37from types import SimpleNamespace 

38 

39import astropy.units as u 

40from astropy.table import Table 

41from lsst.daf.butler import ( 

42 Butler, 

43 CollectionType, 

44 Config, 

45 DatasetId, 

46 DatasetRef, 

47 DatasetType, 

48 DatastoreCacheManager, 

49 DatastoreRecordData, 

50 DimensionUniverse, 

51 LimitedButler, 

52 Quantum, 

53 QuantumBackedButler, 

54 Registry, 

55) 

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

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

58from lsst.pipe.base import ( 

59 ExecutionResources, 

60 GraphBuilder, 

61 Instrument, 

62 Pipeline, 

63 PipelineDatasetTypes, 

64 QuantumGraph, 

65 TaskDef, 

66 TaskFactory, 

67 buildExecutionButler, 

68) 

69from lsst.utils import doImportType 

70from lsst.utils.threads import disable_implicit_threading 

71 

72from . import util 

73from .dotTools import graph2dot, pipeline2dot 

74from .executionGraphFixup import ExecutionGraphFixup 

75from .mpGraphExecutor import MPGraphExecutor 

76from .preExecInit import PreExecInit, PreExecInitLimited 

77from .singleQuantumExecutor import SingleQuantumExecutor 

78 

79# ---------------------------------- 

80# Local non-exported definitions -- 

81# ---------------------------------- 

82 

83_LOG = logging.getLogger(__name__) 

84 

85 

86class _OutputChainedCollectionInfo: 

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

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

89 

90 Parameters 

91 ---------- 

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

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

94 name : `str` 

95 Name of the collection given on the command line. 

96 """ 

97 

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

99 self.name = name 

100 try: 

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

102 self.exists = True 

103 except MissingCollectionError: 

104 self.chain = () 

105 self.exists = False 

106 

107 def __str__(self) -> str: 

108 return self.name 

109 

110 name: str 

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

112 """ 

113 

114 exists: bool 

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

116 """ 

117 

118 chain: tuple[str, ...] 

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

120 

121 Empty if the collection does not already exist. 

122 """ 

123 

124 

125class _OutputRunCollectionInfo: 

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

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

128 

129 Parameters 

130 ---------- 

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

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

133 name : `str` 

134 Name of the collection given on the command line. 

135 """ 

136 

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

138 self.name = name 

139 try: 

140 actualType = registry.getCollectionType(name) 

141 if actualType is not CollectionType.RUN: 

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

143 self.exists = True 

144 except MissingCollectionError: 

145 self.exists = False 

146 

147 name: str 

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

149 """ 

150 

151 exists: bool 

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

153 """ 

154 

155 

156class _ButlerFactory: 

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

158 and output collections. 

159 

160 Parameters 

161 ---------- 

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

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

164 

165 args : `types.SimpleNamespace` 

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

167 either at construction or in later methods. 

168 

169 ``output`` 

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

171 input/output collection. 

172 

173 ``output_run`` 

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

175 collection. 

176 

177 ``extend_run`` 

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

179 and be extended. 

180 

181 ``replace_run`` 

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

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

184 replaced with a new one. 

185 

186 ``prune_replaced`` 

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

188 ``replace_run``). 

189 

190 ``inputs`` 

191 Input collections of any type; see 

192 :ref:`daf_butler_ordered_collection_searches` for details. 

193 

194 ``butler_config`` 

195 Path to a data repository root or configuration file. 

196 

197 writeable : `bool` 

198 If `True`, a `~lsst.daf.butler.Butler` is being initialized in a 

199 context where actual writes should happens, and hence no output run 

200 is necessary. 

201 

202 Raises 

203 ------ 

204 ValueError 

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

206 """ 

207 

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

209 if args.output is not None: 

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

211 else: 

212 self.output = None 

213 if args.output_run is not None: 

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

215 elif self.output is not None: 

216 if args.extend_run: 

217 if not self.output.chain: 

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

219 runName = self.output.chain[0] 

220 else: 

221 runName = f"{self.output}/{Instrument.makeCollectionTimestamp()}" 

222 self.outputRun = _OutputRunCollectionInfo(registry, runName) 

223 elif not writeable: 

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

225 self.outputRun = None 

226 else: 

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

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

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

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

231 # collection. 

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

233 

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

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

236 data repository. 

237 

238 Parameters 

239 ---------- 

240 args : `types.SimpleNamespace` 

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

242 construction parameter of the same name. 

243 """ 

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

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

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

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

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

249 for c1, c2 in zip(self.inputs[::-1], self.output.chain[::-1], strict=False): 

250 if c1 != c2: 

251 raise ValueError( 

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

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

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

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

256 ) 

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

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

259 raise ValueError( 

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

261 "output collection is first created." 

262 ) 

263 if args.extend_run: 

264 if self.outputRun is None: 

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

266 elif not self.outputRun.exists: 

267 raise ValueError( 

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

269 ) 

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

271 raise ValueError( 

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

273 ) 

274 if args.prune_replaced and not args.replace_run: 

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

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

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

278 

279 @classmethod 

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

281 """Parse arguments to support implementations of `makeReadButler` and 

282 `makeButlerAndCollections`. 

283 

284 Parameters 

285 ---------- 

286 args : `types.SimpleNamespace` 

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

288 construction parameter of the same name. 

289 

290 Returns 

291 ------- 

292 butler : `lsst.daf.butler.Butler` 

293 A read-only butler constructed from the repo at 

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

295 inputs : `~collections.abc.Sequence` [ `str` ] 

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

297 self : `_ButlerFactory` 

298 A new `_ButlerFactory` instance representing the processed version 

299 of ``args``. 

300 """ 

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

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

303 self.check(args) 

304 if self.output and self.output.exists: 

305 if args.replace_run: 

306 replaced = self.output.chain[0] 

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

308 _LOG.debug( 

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

310 ) 

311 else: 

312 inputs = [self.output.name] 

313 else: 

314 inputs = list(self.inputs) 

315 if args.extend_run: 

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

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

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

319 return butler, collSearch, self 

320 

321 @classmethod 

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

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

324 arguments. 

325 

326 Parameters 

327 ---------- 

328 args : `types.SimpleNamespace` 

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

330 construction parameter of the same name. 

331 

332 Returns 

333 ------- 

334 butler : `lsst.daf.butler.Butler` 

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

336 ``args``. 

337 """ 

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

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

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

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

342 

343 @classmethod 

344 def makeButlerAndCollections(cls, args: SimpleNamespace) -> tuple[Butler, Sequence[str], str | None]: 

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

346 of the run to be used for future writes. 

347 

348 Parameters 

349 ---------- 

350 args : `types.SimpleNamespace` 

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

352 construction parameter of the same name. 

353 

354 Returns 

355 ------- 

356 butler : `lsst.daf.butler.Butler` 

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

358 from. 

359 inputs : `Sequence` [ `str` ] 

360 Collections to search for datasets. 

361 run : `str` or `None` 

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

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

364 """ 

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

366 run: str | None = None 

367 if args.extend_run: 

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

369 if self.outputRun is not None: 

370 run = self.outputRun.name 

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

372 return butler, inputs, run 

373 

374 @staticmethod 

375 def defineDatastoreCache() -> None: 

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

377 

378 Notes 

379 ----- 

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

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

382 configures an exit handler to clean it up. 

383 """ 

384 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

385 if defined: 

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

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

388 

389 @classmethod 

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

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

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

393 

394 Parameters 

395 ---------- 

396 args : `types.SimpleNamespace` 

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

398 construction parameter of the same name. 

399 taskDefs : iterable of `TaskDef`, optional 

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

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

402 "unstore". 

403 

404 Returns 

405 ------- 

406 butler : `lsst.daf.butler.Butler` 

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

408 """ 

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

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

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

412 self.check(args) 

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

414 if self.output is not None: 

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

416 if args.replace_run: 

417 replaced = chainDefinition.pop(0) 

418 if args.prune_replaced == "unstore": 

419 # Remove datasets from datastore 

420 with butler.transaction(): 

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

422 # we want to remove regular outputs but keep 

423 # initOutputs, configs, and versions. 

424 if taskDefs is not None: 

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

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

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

428 elif args.prune_replaced == "purge": 

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

430 # collection from its chain collection first. 

431 with butler.transaction(): 

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

433 butler.removeRuns([replaced], unstore=True) 

434 elif args.prune_replaced is not None: 

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

436 if not self.output.exists: 

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

438 if not args.extend_run: 

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

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

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

442 _LOG.debug( 

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

444 self.outputRun.name, 

445 self.output.name, 

446 chainDefinition, 

447 ) 

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

449 else: 

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

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

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

453 return butler 

454 

455 output: _OutputChainedCollectionInfo | None 

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

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

458 """ 

459 

460 outputRun: _OutputRunCollectionInfo | None 

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

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

463 """ 

464 

465 inputs: tuple[str, ...] 

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

467 """ 

468 

469 

470class _QBBFactory: 

471 """Class which is a callable for making QBB instances.""" 

472 

473 def __init__( 

474 self, butler_config: Config, dimensions: DimensionUniverse, dataset_types: Mapping[str, DatasetType] 

475 ): 

476 self.butler_config = butler_config 

477 self.dimensions = dimensions 

478 self.dataset_types = dataset_types 

479 

480 def __call__(self, quantum: Quantum) -> LimitedButler: 

481 """Return freshly initialized `~lsst.daf.butler.QuantumBackedButler`. 

482 

483 Factory method to create QuantumBackedButler instances. 

484 """ 

485 return QuantumBackedButler.initialize( 

486 config=self.butler_config, 

487 quantum=quantum, 

488 dimensions=self.dimensions, 

489 dataset_types=self.dataset_types, 

490 ) 

491 

492 

493# ------------------------ 

494# Exported definitions -- 

495# ------------------------ 

496 

497 

498class CmdLineFwk: 

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

500 

501 In addition to executing tasks this activator provides additional methods 

502 for task management like dumping configuration or execution chain. 

503 """ 

504 

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

506 

507 def __init__(self) -> None: 

508 pass 

509 

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

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

512 

513 Parameters 

514 ---------- 

515 args : `types.SimpleNamespace` 

516 Parsed command line 

517 

518 Returns 

519 ------- 

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

521 """ 

522 if args.pipeline: 

523 pipeline = Pipeline.from_uri(args.pipeline) 

524 else: 

525 pipeline = Pipeline("anonymous") 

526 

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

528 for action in args.pipeline_actions: 

529 if action.action == "add_instrument": 

530 pipeline.addInstrument(action.value) 

531 

532 elif action.action == "new_task": 

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

534 

535 elif action.action == "delete_task": 

536 pipeline.removeTask(action.label) 

537 

538 elif action.action == "config": 

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

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

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

542 

543 elif action.action == "configfile": 

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

545 

546 else: 

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

548 

549 if args.save_pipeline: 

550 pipeline.write_to_uri(args.save_pipeline) 

551 

552 if args.pipeline_dot: 

553 pipeline2dot(pipeline, args.pipeline_dot) 

554 

555 return pipeline 

556 

557 def makeGraph(self, pipeline: Pipeline, args: SimpleNamespace) -> QuantumGraph | None: 

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

559 

560 Parameters 

561 ---------- 

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

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

564 args : `types.SimpleNamespace` 

565 Parsed command line 

566 

567 Returns 

568 ------- 

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

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

571 """ 

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

573 if args.extend_run: 

574 args.skip_existing = True 

575 

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

577 

578 if args.skip_existing and run: 

579 args.skip_existing_in += (run,) 

580 

581 if args.qgraph: 

582 # click passes empty tuple as default value for qgraph_node_id 

583 nodes = args.qgraph_node_id or None 

584 qgraph = QuantumGraph.loadUri(args.qgraph, butler.dimensions, nodes=nodes, graphID=args.qgraph_id) 

585 

586 # pipeline can not be provided in this case 

587 if pipeline: 

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

589 if args.show_qgraph_header: 

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

591 else: 

592 task_defs = list(pipeline.toExpandedPipeline()) 

593 if args.mock: 

594 from lsst.pipe.base.tests.mocks import mock_task_defs 

595 

596 task_defs = mock_task_defs( 

597 task_defs, 

598 unmocked_dataset_types=args.unmocked_dataset_types, 

599 force_failures=args.mock_failure, 

600 ) 

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

602 graphBuilder = GraphBuilder( 

603 butler.registry, 

604 skipExistingIn=args.skip_existing_in, 

605 clobberOutputs=args.clobber_outputs, 

606 datastore=butler._datastore if args.qgraph_datastore_records else None, 

607 ) 

608 # accumulate metadata 

609 metadata = { 

610 "input": args.input, 

611 "output": args.output, 

612 "butler_argument": args.butler_config, 

613 "output_run": run, 

614 "extend_run": args.extend_run, 

615 "skip_existing_in": args.skip_existing_in, 

616 "skip_existing": args.skip_existing, 

617 "data_query": args.data_query, 

618 "user": getpass.getuser(), 

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

620 } 

621 assert run is not None, "Butler output run collection must be defined" 

622 qgraph = graphBuilder.makeGraph( 

623 task_defs, 

624 collections, 

625 run, 

626 args.data_query, 

627 metadata=metadata, 

628 datasetQueryConstraint=args.dataset_query_constraint, 

629 dataId=pipeline.get_data_id(butler.dimensions), 

630 ) 

631 if args.show_qgraph_header: 

632 qgraph.buildAndPrintHeader() 

633 

634 if len(qgraph) == 0: 

635 # Nothing to do. 

636 return None 

637 self._summarize_qgraph(qgraph) 

638 

639 if args.save_qgraph: 

640 qgraph.saveUri(args.save_qgraph) 

641 

642 if args.save_single_quanta: 

643 for quantumNode in qgraph: 

644 sqgraph = qgraph.subset(quantumNode) 

645 uri = args.save_single_quanta.format(quantumNode) 

646 sqgraph.saveUri(uri) 

647 

648 if args.qgraph_dot: 

649 graph2dot(qgraph, args.qgraph_dot) 

650 

651 if args.execution_butler_location: 

652 butler = Butler(args.butler_config) 

653 newArgs = copy.deepcopy(args) 

654 

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

656 newArgs.butler_config = butler._config 

657 # Calling makeWriteButler is done for the side effects of 

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

659 # collection names, creating collections, etc. 

660 newButler = _ButlerFactory.makeWriteButler(newArgs) 

661 return newButler 

662 

663 # Include output collection in collections for input 

664 # files if it exists in the repo. 

665 all_inputs = args.input 

666 if args.output is not None: 

667 with contextlib.suppress(MissingCollectionError): 

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

669 

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

671 buildExecutionButler( 

672 butler, 

673 qgraph, 

674 args.execution_butler_location, 

675 run, 

676 butlerModifier=builderShim, 

677 collections=all_inputs, 

678 clobber=args.clobber_execution_butler, 

679 datastoreRoot=args.target_datastore_root, 

680 transfer=args.transfer, 

681 ) 

682 

683 return qgraph 

684 

685 def _make_execution_resources(self, args: SimpleNamespace) -> ExecutionResources: 

686 """Construct the execution resource class from arguments. 

687 

688 Parameters 

689 ---------- 

690 args : `types.SimpleNamespace` 

691 Parsed command line. 

692 

693 Returns 

694 ------- 

695 resources : `~lsst.pipe.base.ExecutionResources` 

696 The resources available to each quantum. 

697 """ 

698 return ExecutionResources( 

699 num_cores=args.cores_per_quantum, max_mem=args.memory_per_quantum, default_mem_units=u.MB 

700 ) 

701 

702 def runPipeline( 

703 self, 

704 graph: QuantumGraph, 

705 taskFactory: TaskFactory, 

706 args: SimpleNamespace, 

707 butler: Butler | None = None, 

708 ) -> None: 

709 """Execute complete QuantumGraph. 

710 

711 Parameters 

712 ---------- 

713 graph : `~lsst.pipe.base.QuantumGraph` 

714 Execution graph. 

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

716 Task factory 

717 args : `types.SimpleNamespace` 

718 Parsed command line 

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

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

721 using command line options. 

722 """ 

723 # Check that output run defined on command line is consistent with 

724 # quantum graph. 

725 if args.output_run and graph.metadata: 

726 graph_output_run = graph.metadata.get("output_run", args.output_run) 

727 if graph_output_run != args.output_run: 

728 raise ValueError( 

729 f"Output run defined on command line ({args.output_run}) has to be " 

730 f"identical to graph metadata ({graph_output_run}). " 

731 "To update graph metadata run `pipetask update-graph-run` command." 

732 ) 

733 

734 # Make sure that --extend-run always enables --skip-existing, 

735 # clobbering should be disabled if --extend-run is not specified. 

736 if args.extend_run: 

737 args.skip_existing = True 

738 else: 

739 args.clobber_outputs = False 

740 

741 if not args.enable_implicit_threading: 

742 disable_implicit_threading() 

743 

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

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

746 if butler is None: 

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

748 

749 if args.skip_existing: 

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

751 

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

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

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

755 # needed if `multiprocessing` always uses fork start method). 

756 if args.enableLsstDebug: 

757 try: 

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

759 import debug # type: ignore # noqa:F401 

760 except ImportError: 

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

762 

763 # Save all InitOutputs, configs, etc. 

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

765 preExecInit.initialize( 

766 graph, 

767 saveInitOutputs=not args.skip_init_writes, 

768 registerDatasetTypes=args.register_dataset_types, 

769 saveVersions=not args.no_versions, 

770 ) 

771 

772 if not args.init_only: 

773 graphFixup = self._importGraphFixup(args) 

774 resources = self._make_execution_resources(args) 

775 quantumExecutor = SingleQuantumExecutor( 

776 butler, 

777 taskFactory, 

778 skipExistingIn=args.skip_existing_in, 

779 clobberOutputs=args.clobber_outputs, 

780 enableLsstDebug=args.enableLsstDebug, 

781 exitOnKnownError=args.fail_fast, 

782 resources=resources, 

783 ) 

784 

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

786 executor = MPGraphExecutor( 

787 numProc=args.processes, 

788 timeout=timeout, 

789 startMethod=args.start_method, 

790 quantumExecutor=quantumExecutor, 

791 failFast=args.fail_fast, 

792 pdb=args.pdb, 

793 executionGraphFixup=graphFixup, 

794 ) 

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

796 # forked processes. 

797 butler.registry.resetConnectionPool() 

798 try: 

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

800 executor.execute(graph) 

801 finally: 

802 if args.summary: 

803 report = executor.getReport() 

804 if report: 

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

806 # Do not save fields that are not set. 

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

808 

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

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

811 given quantum graph. 

812 

813 Parameters 

814 ---------- 

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

816 A QuantumGraph object. 

817 

818 Returns 

819 ------- 

820 qg_task_table : `astropy.table.table.Table` 

821 An astropy table containing columns: Quanta and Tasks. 

822 """ 

823 qg_quanta, qg_tasks = [], [] 

824 for task_def in qgraph.iterTaskGraph(): 

825 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def) 

826 qg_quanta.append(num_qnodes) 

827 qg_tasks.append(task_def.label) 

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

829 return qg_task_table 

830 

831 def _summarize_qgraph(self, qgraph: QuantumGraph) -> int: 

832 """Report a summary of the quanta in the graph. 

833 

834 Parameters 

835 ---------- 

836 qgraph : `lsst.pipe.base.QuantumGraph` 

837 The graph to be summarized. 

838 

839 Returns 

840 ------- 

841 n_quanta : `int` 

842 The number of quanta in the graph. 

843 """ 

844 n_quanta = len(qgraph) 

845 if n_quanta == 0: 

846 _LOG.info("QuantumGraph contains no quanta.") 

847 else: 

848 if _LOG.isEnabledFor(logging.INFO): 

849 qg_task_table = self._generateTaskTable(qgraph) 

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

851 quanta_str = "quantum" if n_quanta == 1 else "quanta" 

852 n_tasks = len(qgraph.taskGraph) 

853 n_tasks_plural = "" if n_tasks == 1 else "s" 

854 _LOG.info( 

855 "QuantumGraph contains %d %s for %d task%s, graph ID: %r\n%s", 

856 n_quanta, 

857 quanta_str, 

858 n_tasks, 

859 n_tasks_plural, 

860 qgraph.graphID, 

861 qg_task_table_formatted, 

862 ) 

863 return n_quanta 

864 

865 def _importGraphFixup(self, args: SimpleNamespace) -> ExecutionGraphFixup | None: 

866 """Import/instantiate graph fixup object. 

867 

868 Parameters 

869 ---------- 

870 args : `types.SimpleNamespace` 

871 Parsed command line. 

872 

873 Returns 

874 ------- 

875 fixup : `ExecutionGraphFixup` or `None` 

876 

877 Raises 

878 ------ 

879 ValueError 

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

881 instance has unexpected type. 

882 """ 

883 if args.graph_fixup: 

884 try: 

885 factory = doImportType(args.graph_fixup) 

886 except Exception as exc: 

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

888 try: 

889 fixup = factory() 

890 except Exception as exc: 

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

892 if not isinstance(fixup, ExecutionGraphFixup): 

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

894 return fixup 

895 return None 

896 

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

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

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

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

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

902 universe = qgraph.universe 

903 

904 # Collect all init input/output dataset IDs. 

905 predicted_inputs: set[DatasetId] = set() 

906 predicted_outputs: set[DatasetId] = set() 

907 for taskDef in qgraph.iterTaskGraph(): 

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

909 predicted_inputs.update(ref.id for ref in refs) 

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

911 predicted_outputs.update(ref.id for ref in refs) 

912 predicted_outputs.update(ref.id for ref in qgraph.globalInitOutputRefs()) 

913 # remove intermediates from inputs 

914 predicted_inputs -= predicted_outputs 

915 

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

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

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

919 for quantum_node in qgraph: 

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

921 subset = records.subset(predicted_inputs) 

922 if subset is not None: 

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

924 

925 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()} 

926 

927 # Make butler from everything. 

928 butler = QuantumBackedButler.from_predicted( 

929 config=args.butler_config, 

930 predicted_inputs=predicted_inputs, 

931 predicted_outputs=predicted_outputs, 

932 dimensions=universe, 

933 datastore_records=datastore_records, 

934 search_paths=args.config_search_path, 

935 dataset_types=dataset_types, 

936 ) 

937 

938 # Save all InitOutputs, configs, etc. 

939 preExecInit = PreExecInitLimited(butler, task_factory) 

940 preExecInit.initialize(qgraph) 

941 

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

943 # Load quantum graph. 

944 nodes = args.qgraph_node_id or None 

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

946 

947 if qgraph.metadata is None: 

948 raise ValueError("QuantumGraph is missing metadata, cannot continue.") 

949 

950 self._summarize_qgraph(qgraph) 

951 

952 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()} 

953 

954 _butler_factory = _QBBFactory( 

955 butler_config=args.butler_config, 

956 dimensions=qgraph.universe, 

957 dataset_types=dataset_types, 

958 ) 

959 

960 # make special quantum executor 

961 resources = self._make_execution_resources(args) 

962 quantumExecutor = SingleQuantumExecutor( 

963 butler=None, 

964 taskFactory=task_factory, 

965 enableLsstDebug=args.enableLsstDebug, 

966 exitOnKnownError=args.fail_fast, 

967 limited_butler_factory=_butler_factory, 

968 resources=resources, 

969 ) 

970 

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

972 executor = MPGraphExecutor( 

973 numProc=args.processes, 

974 timeout=timeout, 

975 startMethod=args.start_method, 

976 quantumExecutor=quantumExecutor, 

977 failFast=args.fail_fast, 

978 pdb=args.pdb, 

979 ) 

980 try: 

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

982 executor.execute(qgraph) 

983 finally: 

984 if args.summary: 

985 report = executor.getReport() 

986 if report: 

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

988 # Do not save fields that are not set. 

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