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

361 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-09 02:48 -0700

1# This file is part of ctrl_mpexec. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22"""Module defining CmdLineFwk class and related methods. 

23""" 

24 

25from __future__ import annotations 

26 

27__all__ = ["CmdLineFwk"] 

28 

29import atexit 

30import copy 

31import datetime 

32import getpass 

33import logging 

34import shutil 

35from collections.abc import Iterable, Mapping, 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: 

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 `Butler` is being initialized in a context where actual 

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

202 

203 Raises 

204 ------ 

205 ValueError 

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

207 """ 

208 

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

210 if args.output is not None: 

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

212 else: 

213 self.output = None 

214 if args.output_run is not None: 

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

216 elif self.output is not None: 

217 if args.extend_run: 

218 if not self.output.chain: 

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

220 runName = self.output.chain[0] 

221 else: 

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

223 self.outputRun = _OutputRunCollectionInfo(registry, runName) 

224 elif not writeable: 

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

226 self.outputRun = None 

227 else: 

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

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

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

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

232 # collection. 

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

234 

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

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

237 data repository. 

238 

239 Parameters 

240 ---------- 

241 args : `types.SimpleNamespace` 

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

243 construction parameter of the same name. 

244 """ 

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

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

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

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

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

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

251 if c1 != c2: 

252 raise ValueError( 

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

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

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

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

257 ) 

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

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

260 raise ValueError( 

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

262 "output collection is first created." 

263 ) 

264 if args.extend_run: 

265 if self.outputRun is None: 

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

267 elif not self.outputRun.exists: 

268 raise ValueError( 

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

270 ) 

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

272 raise ValueError( 

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

274 ) 

275 if args.prune_replaced and not args.replace_run: 

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

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

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

279 

280 @classmethod 

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

282 """Common implementation for `makeReadButler` and 

283 `makeButlerAndCollections`. 

284 

285 Parameters 

286 ---------- 

287 args : `types.SimpleNamespace` 

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

289 construction parameter of the same name. 

290 

291 Returns 

292 ------- 

293 butler : `lsst.daf.butler.Butler` 

294 A read-only butler constructed from the repo at 

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

296 inputs : `Sequence` [ `str` ] 

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

298 self : `_ButlerFactory` 

299 A new `_ButlerFactory` instance representing the processed version 

300 of ``args``. 

301 """ 

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

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

304 self.check(args) 

305 if self.output and self.output.exists: 

306 if args.replace_run: 

307 replaced = self.output.chain[0] 

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

309 _LOG.debug( 

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

311 ) 

312 else: 

313 inputs = [self.output.name] 

314 else: 

315 inputs = list(self.inputs) 

316 if args.extend_run: 

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

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

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

320 return butler, collSearch, self 

321 

322 @classmethod 

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

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

325 arguments. 

326 

327 Parameters 

328 ---------- 

329 args : `types.SimpleNamespace` 

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

331 construction parameter of the same name. 

332 

333 Returns 

334 ------- 

335 butler : `lsst.daf.butler.Butler` 

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

337 ``args``. 

338 """ 

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

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

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

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

343 

344 @classmethod 

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

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

347 of the run to be used for future writes. 

348 

349 Parameters 

350 ---------- 

351 args : `types.SimpleNamespace` 

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

353 construction parameter of the same name. 

354 

355 Returns 

356 ------- 

357 butler : `lsst.daf.butler.Butler` 

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

359 from. 

360 inputs : `Sequence` [ `str` ] 

361 Collections to search for datasets. 

362 run : `str` or `None` 

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

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

365 """ 

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

367 run: Optional[str] = None 

368 if args.extend_run: 

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

370 if self.outputRun is not None: 

371 run = self.outputRun.name 

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

373 return butler, inputs, run 

374 

375 @staticmethod 

376 def defineDatastoreCache() -> None: 

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

378 

379 Notes 

380 ----- 

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

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

383 configures an exit handler to clean it up. 

384 """ 

385 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset() 

386 if defined: 

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

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

389 

390 @classmethod 

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

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

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

394 

395 Parameters 

396 ---------- 

397 args : `types.SimpleNamespace` 

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

399 construction parameter of the same name. 

400 taskDefs : iterable of `TaskDef`, optional 

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

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

403 "unstore". 

404 

405 Returns 

406 ------- 

407 butler : `lsst.daf.butler.Butler` 

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

409 """ 

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

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

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

413 self.check(args) 

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

415 if self.output is not None: 

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

417 if args.replace_run: 

418 replaced = chainDefinition.pop(0) 

419 if args.prune_replaced == "unstore": 

420 # Remove datasets from datastore 

421 with butler.transaction(): 

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

423 # we want to remove regular outputs but keep 

424 # initOutputs, configs, and versions. 

425 if taskDefs is not None: 

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

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

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

429 elif args.prune_replaced == "purge": 

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

431 # collection from its chain collection first. 

432 with butler.transaction(): 

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

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

435 elif args.prune_replaced is not None: 

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

437 if not self.output.exists: 

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

439 if not args.extend_run: 

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

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

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

443 _LOG.debug( 

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

445 self.outputRun.name, 

446 self.output.name, 

447 chainDefinition, 

448 ) 

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

450 else: 

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

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

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

454 return butler 

455 

456 output: Optional[_OutputChainedCollectionInfo] 

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

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

459 """ 

460 

461 outputRun: Optional[_OutputRunCollectionInfo] 

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

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

464 """ 

465 

466 inputs: Tuple[str, ...] 

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

468 """ 

469 

470 

471class _QBBFactory: 

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

473 

474 def __init__( 

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

476 ): 

477 self.butler_config = butler_config 

478 self.dimensions = dimensions 

479 self.dataset_types = dataset_types 

480 

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

482 """Factory method to create QuantumBackedButler instances.""" 

483 return QuantumBackedButler.initialize( 

484 config=self.butler_config, 

485 quantum=quantum, 

486 dimensions=self.dimensions, 

487 dataset_types=self.dataset_types, 

488 ) 

489 

490 

491# ------------------------ 

492# Exported definitions -- 

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

494 

495 

496class CmdLineFwk: 

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

498 

499 In addition to executing tasks this activator provides additional methods 

500 for task management like dumping configuration or execution chain. 

501 """ 

502 

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

504 

505 def __init__(self) -> None: 

506 pass 

507 

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

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

510 

511 Parameters 

512 ---------- 

513 args : `types.SimpleNamespace` 

514 Parsed command line 

515 

516 Returns 

517 ------- 

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

519 """ 

520 if args.pipeline: 

521 pipeline = Pipeline.from_uri(args.pipeline) 

522 else: 

523 pipeline = Pipeline("anonymous") 

524 

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

526 for action in args.pipeline_actions: 

527 if action.action == "add_instrument": 

528 pipeline.addInstrument(action.value) 

529 

530 elif action.action == "new_task": 

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

532 

533 elif action.action == "delete_task": 

534 pipeline.removeTask(action.label) 

535 

536 elif action.action == "config": 

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

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

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

540 

541 elif action.action == "configfile": 

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

543 

544 else: 

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

546 

547 if args.save_pipeline: 

548 pipeline.write_to_uri(args.save_pipeline) 

549 

550 if args.pipeline_dot: 

551 pipeline2dot(pipeline, args.pipeline_dot) 

552 

553 return pipeline 

554 

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

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

557 

558 Parameters 

559 ---------- 

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

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

562 args : `types.SimpleNamespace` 

563 Parsed command line 

564 

565 Returns 

566 ------- 

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

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

569 """ 

570 

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

572 if args.extend_run: 

573 args.skip_existing = True 

574 

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

576 

577 if args.skip_existing and run: 

578 args.skip_existing_in += (run,) 

579 

580 if args.qgraph: 

581 # click passes empty tuple as default value for qgraph_node_id 

582 nodes = args.qgraph_node_id or None 

583 qgraph = QuantumGraph.loadUri( 

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

585 ) 

586 

587 # pipeline can not be provided in this case 

588 if pipeline: 

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

590 if args.show_qgraph_header: 

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

592 else: 

593 task_defs = list(pipeline.toExpandedPipeline()) 

594 if args.mock: 

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

596 

597 task_defs = mock_task_defs(task_defs, unmocked_dataset_types=args.unmocked_dataset_types) 

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

599 graphBuilder = GraphBuilder( 

600 butler.registry, 

601 skipExistingIn=args.skip_existing_in, 

602 clobberOutputs=args.clobber_outputs, 

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

604 ) 

605 # accumulate metadata 

606 metadata = { 

607 "input": args.input, 

608 "output": args.output, 

609 "butler_argument": args.butler_config, 

610 "output_run": run, 

611 "extend_run": args.extend_run, 

612 "skip_existing_in": args.skip_existing_in, 

613 "skip_existing": args.skip_existing, 

614 "data_query": args.data_query, 

615 "user": getpass.getuser(), 

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

617 } 

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

619 qgraph = graphBuilder.makeGraph( 

620 task_defs, 

621 collections, 

622 run, 

623 args.data_query, 

624 metadata=metadata, 

625 datasetQueryConstraint=args.dataset_query_constraint, 

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

627 ) 

628 if args.show_qgraph_header: 

629 qgraph.buildAndPrintHeader() 

630 

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

632 nQuanta = len(qgraph) 

633 if nQuanta == 0: 

634 return None 

635 else: 

636 if _LOG.isEnabledFor(logging.INFO): 

637 qg_task_table = self._generateTaskTable(qgraph) 

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

639 _LOG.info( 

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

641 nQuanta, 

642 len(qgraph.taskGraph), 

643 qgraph.graphID, 

644 qg_task_table_formatted, 

645 ) 

646 

647 if args.save_qgraph: 

648 qgraph.saveUri(args.save_qgraph) 

649 

650 if args.save_single_quanta: 

651 for quantumNode in qgraph: 

652 sqgraph = qgraph.subset(quantumNode) 

653 uri = args.save_single_quanta.format(quantumNode) 

654 sqgraph.saveUri(uri) 

655 

656 if args.qgraph_dot: 

657 graph2dot(qgraph, args.qgraph_dot) 

658 

659 if args.execution_butler_location: 

660 butler = Butler(args.butler_config) 

661 newArgs = copy.deepcopy(args) 

662 

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

664 newArgs.butler_config = butler._config 

665 # Calling makeWriteButler is done for the side effects of 

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

667 # collection names, creating collections, etc. 

668 newButler = _ButlerFactory.makeWriteButler(newArgs) 

669 return newButler 

670 

671 # Include output collection in collections for input 

672 # files if it exists in the repo. 

673 all_inputs = args.input 

674 if args.output is not None: 

675 try: 

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

677 except MissingCollectionError: 

678 pass 

679 

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

681 buildExecutionButler( 

682 butler, 

683 qgraph, 

684 args.execution_butler_location, 

685 run, 

686 butlerModifier=builderShim, 

687 collections=all_inputs, 

688 clobber=args.clobber_execution_butler, 

689 datastoreRoot=args.target_datastore_root, 

690 transfer=args.transfer, 

691 ) 

692 

693 return qgraph 

694 

695 def runPipeline( 

696 self, 

697 graph: QuantumGraph, 

698 taskFactory: TaskFactory, 

699 args: SimpleNamespace, 

700 butler: Optional[Butler] = None, 

701 ) -> None: 

702 """Execute complete QuantumGraph. 

703 

704 Parameters 

705 ---------- 

706 graph : `QuantumGraph` 

707 Execution graph. 

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

709 Task factory 

710 args : `types.SimpleNamespace` 

711 Parsed command line 

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

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

714 using command line options. 

715 """ 

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

717 # quantum graph. 

718 if args.output_run and graph.metadata: 

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

720 if graph_output_run != args.output_run: 

721 raise ValueError( 

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

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

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

725 ) 

726 

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

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

729 if args.extend_run: 

730 args.skip_existing = True 

731 else: 

732 args.clobber_outputs = False 

733 

734 if not args.enable_implicit_threading: 

735 disable_implicit_threading() 

736 

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

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

739 if butler is None: 

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

741 

742 if args.skip_existing: 

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

744 

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

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

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

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

749 if args.enableLsstDebug: 

750 try: 

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

752 import debug # type: ignore # noqa:F401 

753 except ImportError: 

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

755 

756 # Save all InitOutputs, configs, etc. 

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

758 preExecInit.initialize( 

759 graph, 

760 saveInitOutputs=not args.skip_init_writes, 

761 registerDatasetTypes=args.register_dataset_types, 

762 saveVersions=not args.no_versions, 

763 ) 

764 

765 if not args.init_only: 

766 graphFixup = self._importGraphFixup(args) 

767 quantumExecutor = SingleQuantumExecutor( 

768 butler, 

769 taskFactory, 

770 skipExistingIn=args.skip_existing_in, 

771 clobberOutputs=args.clobber_outputs, 

772 enableLsstDebug=args.enableLsstDebug, 

773 exitOnKnownError=args.fail_fast, 

774 ) 

775 

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

777 executor = MPGraphExecutor( 

778 numProc=args.processes, 

779 timeout=timeout, 

780 startMethod=args.start_method, 

781 quantumExecutor=quantumExecutor, 

782 failFast=args.fail_fast, 

783 pdb=args.pdb, 

784 executionGraphFixup=graphFixup, 

785 ) 

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

787 # forked processes. 

788 butler.registry.resetConnectionPool() 

789 try: 

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

791 executor.execute(graph) 

792 finally: 

793 if args.summary: 

794 report = executor.getReport() 

795 if report: 

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

797 # Do not save fields that are not set. 

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

799 

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

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

802 given quantum graph. 

803 

804 Parameters 

805 ---------- 

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

807 A QuantumGraph object. 

808 

809 Returns 

810 ------- 

811 qg_task_table : `astropy.table.table.Table` 

812 An astropy table containing columns: Quanta and Tasks. 

813 """ 

814 qg_quanta, qg_tasks = [], [] 

815 for task_def in qgraph.iterTaskGraph(): 

816 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def) 

817 qg_quanta.append(num_qnodes) 

818 qg_tasks.append(task_def.label) 

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

820 return qg_task_table 

821 

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

823 """Import/instantiate graph fixup object. 

824 

825 Parameters 

826 ---------- 

827 args : `types.SimpleNamespace` 

828 Parsed command line. 

829 

830 Returns 

831 ------- 

832 fixup : `ExecutionGraphFixup` or `None` 

833 

834 Raises 

835 ------ 

836 ValueError 

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

838 instance has unexpected type. 

839 """ 

840 if args.graph_fixup: 

841 try: 

842 factory = doImportType(args.graph_fixup) 

843 except Exception as exc: 

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

845 try: 

846 fixup = factory() 

847 except Exception as exc: 

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

849 if not isinstance(fixup, ExecutionGraphFixup): 

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

851 return fixup 

852 return None 

853 

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

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

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

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

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

859 universe = qgraph.universe 

860 

861 # Collect all init input/output dataset IDs. 

862 predicted_inputs: set[DatasetId] = set() 

863 predicted_outputs: set[DatasetId] = set() 

864 for taskDef in qgraph.iterTaskGraph(): 

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

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

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

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

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

870 # remove intermediates from inputs 

871 predicted_inputs -= predicted_outputs 

872 

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

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

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

876 for quantum_node in qgraph: 

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

878 subset = records.subset(predicted_inputs) 

879 if subset is not None: 

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

881 

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

883 

884 # Make butler from everything. 

885 butler = QuantumBackedButler.from_predicted( 

886 config=args.butler_config, 

887 predicted_inputs=predicted_inputs, 

888 predicted_outputs=predicted_outputs, 

889 dimensions=universe, 

890 datastore_records=datastore_records, 

891 search_paths=args.config_search_path, 

892 dataset_types=dataset_types, 

893 ) 

894 

895 # Save all InitOutputs, configs, etc. 

896 preExecInit = PreExecInitLimited(butler, task_factory) 

897 preExecInit.initialize(qgraph) 

898 

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

900 # Load quantum graph. 

901 nodes = args.qgraph_node_id or None 

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

903 

904 if qgraph.metadata is None: 

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

906 

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

908 

909 _butler_factory = _QBBFactory( 

910 butler_config=args.butler_config, 

911 dimensions=qgraph.universe, 

912 dataset_types=dataset_types, 

913 ) 

914 

915 # make special quantum executor 

916 quantumExecutor = SingleQuantumExecutor( 

917 butler=None, 

918 taskFactory=task_factory, 

919 enableLsstDebug=args.enableLsstDebug, 

920 exitOnKnownError=args.fail_fast, 

921 limited_butler_factory=_butler_factory, 

922 ) 

923 

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

925 executor = MPGraphExecutor( 

926 numProc=args.processes, 

927 timeout=timeout, 

928 startMethod=args.start_method, 

929 quantumExecutor=quantumExecutor, 

930 failFast=args.fail_fast, 

931 pdb=args.pdb, 

932 ) 

933 try: 

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

935 executor.execute(qgraph) 

936 finally: 

937 if args.summary: 

938 report = executor.getReport() 

939 if report: 

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

941 # Do not save fields that are not set. 

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