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

361 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-22 09:52 +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 copy 

31import datetime 

32import getpass 

33import logging 

34import shutil 

35from collections.abc import Iterable, Mapping, Sequence 

36from types import SimpleNamespace 

37from typing import TYPE_CHECKING 

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: 

69 from lsst.daf.butler import ( 

70 Config, 

71 DatasetType, 

72 DatastoreRecordData, 

73 DimensionUniverse, 

74 LimitedButler, 

75 Quantum, 

76 Registry, 

77 ) 

78 from lsst.pipe.base import TaskDef, TaskFactory 

79 

80 

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

82# Local non-exported definitions -- 

83# ---------------------------------- 

84 

85_LOG = logging.getLogger(__name__) 

86 

87 

88class _OutputChainedCollectionInfo: 

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

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

91 

92 Parameters 

93 ---------- 

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

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

96 name : `str` 

97 Name of the collection given on the command line. 

98 """ 

99 

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

101 self.name = name 

102 try: 

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

104 self.exists = True 

105 except MissingCollectionError: 

106 self.chain = () 

107 self.exists = False 

108 

109 def __str__(self) -> str: 

110 return self.name 

111 

112 name: str 

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

114 """ 

115 

116 exists: bool 

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

118 """ 

119 

120 chain: tuple[str, ...] 

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

122 

123 Empty if the collection does not already exist. 

124 """ 

125 

126 

127class _OutputRunCollectionInfo: 

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

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

130 

131 Parameters 

132 ---------- 

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

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

135 name : `str` 

136 Name of the collection given on the command line. 

137 """ 

138 

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

140 self.name = name 

141 try: 

142 actualType = registry.getCollectionType(name) 

143 if actualType is not CollectionType.RUN: 

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

145 self.exists = True 

146 except MissingCollectionError: 

147 self.exists = False 

148 

149 name: str 

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

151 """ 

152 

153 exists: bool 

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

155 """ 

156 

157 

158class _ButlerFactory: 

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

160 and output collections. 

161 

162 Parameters 

163 ---------- 

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

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

166 

167 args : `types.SimpleNamespace` 

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

169 either at construction or in later methods. 

170 

171 ``output`` 

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

173 input/output collection. 

174 

175 ``output_run`` 

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

177 collection. 

178 

179 ``extend_run`` 

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

181 and be extended. 

182 

183 ``replace_run`` 

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

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

186 replaced with a new one. 

187 

188 ``prune_replaced`` 

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

190 ``replace_run``). 

191 

192 ``inputs`` 

193 Input collections of any type; see 

194 :ref:`daf_butler_ordered_collection_searches` for details. 

195 

196 ``butler_config`` 

197 Path to a data repository root or configuration file. 

198 

199 writeable : `bool` 

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

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

202 is necessary. 

203 

204 Raises 

205 ------ 

206 ValueError 

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

208 """ 

209 

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

211 if args.output is not None: 

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

213 else: 

214 self.output = None 

215 if args.output_run is not None: 

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

217 elif self.output is not None: 

218 if args.extend_run: 

219 if not self.output.chain: 

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

221 runName = self.output.chain[0] 

222 else: 

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

224 self.outputRun = _OutputRunCollectionInfo(registry, runName) 

225 elif not writeable: 

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

227 self.outputRun = None 

228 else: 

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

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

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

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

233 # collection. 

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

235 

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

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

238 data repository. 

239 

240 Parameters 

241 ---------- 

242 args : `types.SimpleNamespace` 

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

244 construction parameter of the same name. 

245 """ 

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

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

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

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

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

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

252 if c1 != c2: 

253 raise ValueError( 

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

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

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

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

258 ) 

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

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

261 raise ValueError( 

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

263 "output collection is first created." 

264 ) 

265 if args.extend_run: 

266 if self.outputRun is None: 

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

268 elif not self.outputRun.exists: 

269 raise ValueError( 

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

271 ) 

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

273 raise ValueError( 

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

275 ) 

276 if args.prune_replaced and not args.replace_run: 

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

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

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

280 

281 @classmethod 

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

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

284 `makeButlerAndCollections`. 

285 

286 Parameters 

287 ---------- 

288 args : `types.SimpleNamespace` 

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

290 construction parameter of the same name. 

291 

292 Returns 

293 ------- 

294 butler : `lsst.daf.butler.Butler` 

295 A read-only butler constructed from the repo at 

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

297 inputs : `Sequence` [ `str` ] 

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

299 self : `_ButlerFactory` 

300 A new `_ButlerFactory` instance representing the processed version 

301 of ``args``. 

302 """ 

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

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

305 self.check(args) 

306 if self.output and self.output.exists: 

307 if args.replace_run: 

308 replaced = self.output.chain[0] 

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

310 _LOG.debug( 

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

312 ) 

313 else: 

314 inputs = [self.output.name] 

315 else: 

316 inputs = list(self.inputs) 

317 if args.extend_run: 

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

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

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

321 return butler, collSearch, self 

322 

323 @classmethod 

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

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

326 arguments. 

327 

328 Parameters 

329 ---------- 

330 args : `types.SimpleNamespace` 

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

332 construction parameter of the same name. 

333 

334 Returns 

335 ------- 

336 butler : `lsst.daf.butler.Butler` 

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

338 ``args``. 

339 """ 

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

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

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

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

344 

345 @classmethod 

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

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

348 of the run to be used for future writes. 

349 

350 Parameters 

351 ---------- 

352 args : `types.SimpleNamespace` 

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

354 construction parameter of the same name. 

355 

356 Returns 

357 ------- 

358 butler : `lsst.daf.butler.Butler` 

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

360 from. 

361 inputs : `Sequence` [ `str` ] 

362 Collections to search for datasets. 

363 run : `str` or `None` 

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

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

366 """ 

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

368 run: str | None = None 

369 if args.extend_run: 

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

371 if self.outputRun is not None: 

372 run = self.outputRun.name 

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

374 return butler, inputs, run 

375 

376 @staticmethod 

377 def defineDatastoreCache() -> None: 

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

379 

380 Notes 

381 ----- 

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

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

384 configures an exit handler to clean it up. 

385 """ 

386 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

387 if defined: 

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

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

390 

391 @classmethod 

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

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

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

395 

396 Parameters 

397 ---------- 

398 args : `types.SimpleNamespace` 

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

400 construction parameter of the same name. 

401 taskDefs : iterable of `TaskDef`, optional 

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

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

404 "unstore". 

405 

406 Returns 

407 ------- 

408 butler : `lsst.daf.butler.Butler` 

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

410 """ 

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

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

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

414 self.check(args) 

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

416 if self.output is not None: 

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

418 if args.replace_run: 

419 replaced = chainDefinition.pop(0) 

420 if args.prune_replaced == "unstore": 

421 # Remove datasets from datastore 

422 with butler.transaction(): 

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

424 # we want to remove regular outputs but keep 

425 # initOutputs, configs, and versions. 

426 if taskDefs is not None: 

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

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

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

430 elif args.prune_replaced == "purge": 

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

432 # collection from its chain collection first. 

433 with butler.transaction(): 

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

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

436 elif args.prune_replaced is not None: 

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

438 if not self.output.exists: 

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

440 if not args.extend_run: 

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

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

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

444 _LOG.debug( 

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

446 self.outputRun.name, 

447 self.output.name, 

448 chainDefinition, 

449 ) 

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

451 else: 

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

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

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

455 return butler 

456 

457 output: _OutputChainedCollectionInfo | None 

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

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

460 """ 

461 

462 outputRun: _OutputRunCollectionInfo | None 

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

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

465 """ 

466 

467 inputs: tuple[str, ...] 

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

469 """ 

470 

471 

472class _QBBFactory: 

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

474 

475 def __init__( 

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

477 ): 

478 self.butler_config = butler_config 

479 self.dimensions = dimensions 

480 self.dataset_types = dataset_types 

481 

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

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

484 

485 Factory method to create QuantumBackedButler instances. 

486 """ 

487 return QuantumBackedButler.initialize( 

488 config=self.butler_config, 

489 quantum=quantum, 

490 dimensions=self.dimensions, 

491 dataset_types=self.dataset_types, 

492 ) 

493 

494 

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

496# Exported definitions -- 

497# ------------------------ 

498 

499 

500class CmdLineFwk: 

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

502 

503 In addition to executing tasks this activator provides additional methods 

504 for task management like dumping configuration or execution chain. 

505 """ 

506 

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

508 

509 def __init__(self) -> None: 

510 pass 

511 

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

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

514 

515 Parameters 

516 ---------- 

517 args : `types.SimpleNamespace` 

518 Parsed command line 

519 

520 Returns 

521 ------- 

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

523 """ 

524 if args.pipeline: 

525 pipeline = Pipeline.from_uri(args.pipeline) 

526 else: 

527 pipeline = Pipeline("anonymous") 

528 

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

530 for action in args.pipeline_actions: 

531 if action.action == "add_instrument": 

532 pipeline.addInstrument(action.value) 

533 

534 elif action.action == "new_task": 

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

536 

537 elif action.action == "delete_task": 

538 pipeline.removeTask(action.label) 

539 

540 elif action.action == "config": 

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

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

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

544 

545 elif action.action == "configfile": 

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

547 

548 else: 

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

550 

551 if args.save_pipeline: 

552 pipeline.write_to_uri(args.save_pipeline) 

553 

554 if args.pipeline_dot: 

555 pipeline2dot(pipeline, args.pipeline_dot) 

556 

557 return pipeline 

558 

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

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

561 

562 Parameters 

563 ---------- 

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

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

566 args : `types.SimpleNamespace` 

567 Parsed command line 

568 

569 Returns 

570 ------- 

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

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

573 """ 

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

575 if args.extend_run: 

576 args.skip_existing = True 

577 

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

579 

580 if args.skip_existing and run: 

581 args.skip_existing_in += (run,) 

582 

583 if args.qgraph: 

584 # click passes empty tuple as default value for qgraph_node_id 

585 nodes = args.qgraph_node_id or None 

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

587 

588 # pipeline can not be provided in this case 

589 if pipeline: 

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

591 if args.show_qgraph_header: 

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

593 else: 

594 task_defs = list(pipeline.toExpandedPipeline()) 

595 if args.mock: 

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

597 

598 task_defs = mock_task_defs( 

599 task_defs, 

600 unmocked_dataset_types=args.unmocked_dataset_types, 

601 force_failures=args.mock_failure, 

602 ) 

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

604 graphBuilder = GraphBuilder( 

605 butler.registry, 

606 skipExistingIn=args.skip_existing_in, 

607 clobberOutputs=args.clobber_outputs, 

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

609 ) 

610 # accumulate metadata 

611 metadata = { 

612 "input": args.input, 

613 "output": args.output, 

614 "butler_argument": args.butler_config, 

615 "output_run": run, 

616 "extend_run": args.extend_run, 

617 "skip_existing_in": args.skip_existing_in, 

618 "skip_existing": args.skip_existing, 

619 "data_query": args.data_query, 

620 "user": getpass.getuser(), 

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

622 } 

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

624 qgraph = graphBuilder.makeGraph( 

625 task_defs, 

626 collections, 

627 run, 

628 args.data_query, 

629 metadata=metadata, 

630 datasetQueryConstraint=args.dataset_query_constraint, 

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

632 ) 

633 if args.show_qgraph_header: 

634 qgraph.buildAndPrintHeader() 

635 

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

637 nQuanta = len(qgraph) 

638 if nQuanta == 0: 

639 return None 

640 else: 

641 if _LOG.isEnabledFor(logging.INFO): 

642 qg_task_table = self._generateTaskTable(qgraph) 

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

644 _LOG.info( 

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

646 nQuanta, 

647 len(qgraph.taskGraph), 

648 qgraph.graphID, 

649 qg_task_table_formatted, 

650 ) 

651 

652 if args.save_qgraph: 

653 qgraph.saveUri(args.save_qgraph) 

654 

655 if args.save_single_quanta: 

656 for quantumNode in qgraph: 

657 sqgraph = qgraph.subset(quantumNode) 

658 uri = args.save_single_quanta.format(quantumNode) 

659 sqgraph.saveUri(uri) 

660 

661 if args.qgraph_dot: 

662 graph2dot(qgraph, args.qgraph_dot) 

663 

664 if args.execution_butler_location: 

665 butler = Butler(args.butler_config) 

666 newArgs = copy.deepcopy(args) 

667 

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

669 newArgs.butler_config = butler._config 

670 # Calling makeWriteButler is done for the side effects of 

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

672 # collection names, creating collections, etc. 

673 newButler = _ButlerFactory.makeWriteButler(newArgs) 

674 return newButler 

675 

676 # Include output collection in collections for input 

677 # files if it exists in the repo. 

678 all_inputs = args.input 

679 if args.output is not None: 

680 try: 

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

682 except MissingCollectionError: 

683 pass 

684 

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

686 buildExecutionButler( 

687 butler, 

688 qgraph, 

689 args.execution_butler_location, 

690 run, 

691 butlerModifier=builderShim, 

692 collections=all_inputs, 

693 clobber=args.clobber_execution_butler, 

694 datastoreRoot=args.target_datastore_root, 

695 transfer=args.transfer, 

696 ) 

697 

698 return qgraph 

699 

700 def runPipeline( 

701 self, 

702 graph: QuantumGraph, 

703 taskFactory: TaskFactory, 

704 args: SimpleNamespace, 

705 butler: Butler | None = None, 

706 ) -> None: 

707 """Execute complete QuantumGraph. 

708 

709 Parameters 

710 ---------- 

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

712 Execution graph. 

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

714 Task factory 

715 args : `types.SimpleNamespace` 

716 Parsed command line 

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

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

719 using command line options. 

720 """ 

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

722 # quantum graph. 

723 if args.output_run and graph.metadata: 

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

725 if graph_output_run != args.output_run: 

726 raise ValueError( 

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

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

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

730 ) 

731 

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

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

734 if args.extend_run: 

735 args.skip_existing = True 

736 else: 

737 args.clobber_outputs = False 

738 

739 if not args.enable_implicit_threading: 

740 disable_implicit_threading() 

741 

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

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

744 if butler is None: 

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

746 

747 if args.skip_existing: 

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

749 

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

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

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

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

754 if args.enableLsstDebug: 

755 try: 

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

757 import debug # type: ignore # noqa:F401 

758 except ImportError: 

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

760 

761 # Save all InitOutputs, configs, etc. 

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

763 preExecInit.initialize( 

764 graph, 

765 saveInitOutputs=not args.skip_init_writes, 

766 registerDatasetTypes=args.register_dataset_types, 

767 saveVersions=not args.no_versions, 

768 ) 

769 

770 if not args.init_only: 

771 graphFixup = self._importGraphFixup(args) 

772 quantumExecutor = SingleQuantumExecutor( 

773 butler, 

774 taskFactory, 

775 skipExistingIn=args.skip_existing_in, 

776 clobberOutputs=args.clobber_outputs, 

777 enableLsstDebug=args.enableLsstDebug, 

778 exitOnKnownError=args.fail_fast, 

779 ) 

780 

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

782 executor = MPGraphExecutor( 

783 numProc=args.processes, 

784 timeout=timeout, 

785 startMethod=args.start_method, 

786 quantumExecutor=quantumExecutor, 

787 failFast=args.fail_fast, 

788 pdb=args.pdb, 

789 executionGraphFixup=graphFixup, 

790 ) 

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

792 # forked processes. 

793 butler.registry.resetConnectionPool() 

794 try: 

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

796 executor.execute(graph) 

797 finally: 

798 if args.summary: 

799 report = executor.getReport() 

800 if report: 

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

802 # Do not save fields that are not set. 

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

804 

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

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

807 given quantum graph. 

808 

809 Parameters 

810 ---------- 

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

812 A QuantumGraph object. 

813 

814 Returns 

815 ------- 

816 qg_task_table : `astropy.table.table.Table` 

817 An astropy table containing columns: Quanta and Tasks. 

818 """ 

819 qg_quanta, qg_tasks = [], [] 

820 for task_def in qgraph.iterTaskGraph(): 

821 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def) 

822 qg_quanta.append(num_qnodes) 

823 qg_tasks.append(task_def.label) 

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

825 return qg_task_table 

826 

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

828 """Import/instantiate graph fixup object. 

829 

830 Parameters 

831 ---------- 

832 args : `types.SimpleNamespace` 

833 Parsed command line. 

834 

835 Returns 

836 ------- 

837 fixup : `ExecutionGraphFixup` or `None` 

838 

839 Raises 

840 ------ 

841 ValueError 

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

843 instance has unexpected type. 

844 """ 

845 if args.graph_fixup: 

846 try: 

847 factory = doImportType(args.graph_fixup) 

848 except Exception as exc: 

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

850 try: 

851 fixup = factory() 

852 except Exception as exc: 

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

854 if not isinstance(fixup, ExecutionGraphFixup): 

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

856 return fixup 

857 return None 

858 

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

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

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

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

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

864 universe = qgraph.universe 

865 

866 # Collect all init input/output dataset IDs. 

867 predicted_inputs: set[DatasetId] = set() 

868 predicted_outputs: set[DatasetId] = set() 

869 for taskDef in qgraph.iterTaskGraph(): 

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

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

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

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

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

875 # remove intermediates from inputs 

876 predicted_inputs -= predicted_outputs 

877 

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

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

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

881 for quantum_node in qgraph: 

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

883 subset = records.subset(predicted_inputs) 

884 if subset is not None: 

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

886 

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

888 

889 # Make butler from everything. 

890 butler = QuantumBackedButler.from_predicted( 

891 config=args.butler_config, 

892 predicted_inputs=predicted_inputs, 

893 predicted_outputs=predicted_outputs, 

894 dimensions=universe, 

895 datastore_records=datastore_records, 

896 search_paths=args.config_search_path, 

897 dataset_types=dataset_types, 

898 ) 

899 

900 # Save all InitOutputs, configs, etc. 

901 preExecInit = PreExecInitLimited(butler, task_factory) 

902 preExecInit.initialize(qgraph) 

903 

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

905 # Load quantum graph. 

906 nodes = args.qgraph_node_id or None 

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

908 

909 if qgraph.metadata is None: 

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

911 

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

913 

914 _butler_factory = _QBBFactory( 

915 butler_config=args.butler_config, 

916 dimensions=qgraph.universe, 

917 dataset_types=dataset_types, 

918 ) 

919 

920 # make special quantum executor 

921 quantumExecutor = SingleQuantumExecutor( 

922 butler=None, 

923 taskFactory=task_factory, 

924 enableLsstDebug=args.enableLsstDebug, 

925 exitOnKnownError=args.fail_fast, 

926 limited_butler_factory=_butler_factory, 

927 ) 

928 

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

930 executor = MPGraphExecutor( 

931 numProc=args.processes, 

932 timeout=timeout, 

933 startMethod=args.start_method, 

934 quantumExecutor=quantumExecutor, 

935 failFast=args.fail_fast, 

936 pdb=args.pdb, 

937 ) 

938 try: 

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

940 executor.execute(qgraph) 

941 finally: 

942 if args.summary: 

943 report = executor.getReport() 

944 if report: 

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

946 # Do not save fields that are not set. 

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