Coverage for python/lsst/ctrl/mpexec/cmdLineFwk.py: 14%
350 statements
« prev ^ index » next coverage.py v7.2.5, created at 2023-05-06 02:45 -0700
« prev ^ index » next coverage.py v7.2.5, created at 2023-05-06 02:45 -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/>.
22"""Module defining CmdLineFwk class and related methods.
23"""
25from __future__ import annotations
27__all__ = ["CmdLineFwk"]
29import atexit
30import copy
31import datetime
32import getpass
33import logging
34import shutil
35from collections.abc import Iterable, Sequence
36from types import SimpleNamespace
37from typing import TYPE_CHECKING, Optional, Tuple
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
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
68if TYPE_CHECKING: 68 ↛ 69line 68 didn't jump to line 69, because the condition on line 68 was never true
69 from lsst.daf.butler import DatastoreRecordData, LimitedButler, Quantum, Registry
70 from lsst.pipe.base import TaskDef, TaskFactory
73# ----------------------------------
74# Local non-exported definitions --
75# ----------------------------------
77_LOG = logging.getLogger(__name__)
80class _OutputChainedCollectionInfo:
81 """A helper class for handling command-line arguments related to an output
82 `~lsst.daf.butler.CollectionType.CHAINED` collection.
84 Parameters
85 ----------
86 registry : `lsst.daf.butler.Registry`
87 Butler registry that collections will be added to and/or queried from.
88 name : `str`
89 Name of the collection given on the command line.
90 """
92 def __init__(self, registry: Registry, name: str):
93 self.name = name
94 try:
95 self.chain = tuple(registry.getCollectionChain(name))
96 self.exists = True
97 except MissingCollectionError:
98 self.chain = ()
99 self.exists = False
101 def __str__(self) -> str:
102 return self.name
104 name: str
105 """Name of the collection provided on the command line (`str`).
106 """
108 exists: bool
109 """Whether this collection already exists in the registry (`bool`).
110 """
112 chain: Tuple[str, ...]
113 """The definition of the collection, if it already exists (`tuple`[`str`]).
115 Empty if the collection does not already exist.
116 """
119class _OutputRunCollectionInfo:
120 """A helper class for handling command-line arguments related to an output
121 `~lsst.daf.butler.CollectionType.RUN` collection.
123 Parameters
124 ----------
125 registry : `lsst.daf.butler.Registry`
126 Butler registry that collections will be added to and/or queried from.
127 name : `str`
128 Name of the collection given on the command line.
129 """
131 def __init__(self, registry: Registry, name: str):
132 self.name = name
133 try:
134 actualType = registry.getCollectionType(name)
135 if actualType is not CollectionType.RUN:
136 raise TypeError(f"Collection '{name}' exists but has type {actualType.name}, not RUN.")
137 self.exists = True
138 except MissingCollectionError:
139 self.exists = False
141 name: str
142 """Name of the collection provided on the command line (`str`).
143 """
145 exists: bool
146 """Whether this collection already exists in the registry (`bool`).
147 """
150class _ButlerFactory:
151 """A helper class for processing command-line arguments related to input
152 and output collections.
154 Parameters
155 ----------
156 registry : `lsst.daf.butler.Registry`
157 Butler registry that collections will be added to and/or queried from.
159 args : `types.SimpleNamespace`
160 Parsed command-line arguments. The following attributes are used,
161 either at construction or in later methods.
163 ``output``
164 The name of a `~lsst.daf.butler.CollectionType.CHAINED`
165 input/output collection.
167 ``output_run``
168 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output
169 collection.
171 ``extend_run``
172 A boolean indicating whether ``output_run`` should already exist
173 and be extended.
175 ``replace_run``
176 A boolean indicating that (if `True`) ``output_run`` should already
177 exist but will be removed from the output chained collection and
178 replaced with a new one.
180 ``prune_replaced``
181 A boolean indicating whether to prune the replaced run (requires
182 ``replace_run``).
184 ``inputs``
185 Input collections of any type; see
186 :ref:`daf_butler_ordered_collection_searches` for details.
188 ``butler_config``
189 Path to a data repository root or configuration file.
191 writeable : `bool`
192 If `True`, a `Butler` is being initialized in a context where actual
193 writes should happens, and hence no output run is necessary.
195 Raises
196 ------
197 ValueError
198 Raised if ``writeable is True`` but there are no output collections.
199 """
201 def __init__(self, registry: Registry, args: SimpleNamespace, writeable: bool):
202 if args.output is not None:
203 self.output = _OutputChainedCollectionInfo(registry, args.output)
204 else:
205 self.output = None
206 if args.output_run is not None:
207 self.outputRun = _OutputRunCollectionInfo(registry, args.output_run)
208 elif self.output is not None:
209 if args.extend_run:
210 if not self.output.chain:
211 raise ValueError("Cannot use --extend-run option with non-existing or empty output chain")
212 runName = self.output.chain[0]
213 else:
214 runName = "{}/{}".format(self.output, Instrument.makeCollectionTimestamp())
215 self.outputRun = _OutputRunCollectionInfo(registry, runName)
216 elif not writeable:
217 # If we're not writing yet, ok to have no output run.
218 self.outputRun = None
219 else:
220 raise ValueError("Cannot write without at least one of (--output, --output-run).")
221 # Recursively flatten any input CHAINED collections. We do this up
222 # front so we can tell if the user passes the same inputs on subsequent
223 # calls, even though we also flatten when we define the output CHAINED
224 # collection.
225 self.inputs = tuple(registry.queryCollections(args.input, flattenChains=True)) if args.input else ()
227 def check(self, args: SimpleNamespace) -> None:
228 """Check command-line options for consistency with each other and the
229 data repository.
231 Parameters
232 ----------
233 args : `types.SimpleNamespace`
234 Parsed command-line arguments. See class documentation for the
235 construction parameter of the same name.
236 """
237 assert not (args.extend_run and args.replace_run), "In mutually-exclusive group in ArgumentParser."
238 if self.inputs and self.output is not None and self.output.exists:
239 # Passing the same inputs that were used to initialize the output
240 # collection is allowed; this means they must _end_ with the same
241 # collections, because we push new runs to the front of the chain.
242 for c1, c2 in zip(self.inputs[::-1], self.output.chain[::-1]):
243 if c1 != c2:
244 raise ValueError(
245 f"Output CHAINED collection {self.output.name!r} exists, but it ends with "
246 "a different sequence of input collections than those given: "
247 f"{c1!r} != {c2!r} in inputs={self.inputs} vs "
248 f"{self.output.name}={self.output.chain}."
249 )
250 if len(self.inputs) > len(self.output.chain):
251 nNew = len(self.inputs) - len(self.output.chain)
252 raise ValueError(
253 f"Cannot add new input collections {self.inputs[:nNew]} after "
254 "output collection is first created."
255 )
256 if args.extend_run:
257 if self.outputRun is None:
258 raise ValueError("Cannot --extend-run when no output collection is given.")
259 elif not self.outputRun.exists:
260 raise ValueError(
261 f"Cannot --extend-run; output collection '{self.outputRun.name}' does not exist."
262 )
263 if not args.extend_run and self.outputRun is not None and self.outputRun.exists:
264 raise ValueError(
265 f"Output run '{self.outputRun.name}' already exists, but --extend-run was not given."
266 )
267 if args.prune_replaced and not args.replace_run:
268 raise ValueError("--prune-replaced requires --replace-run.")
269 if args.replace_run and (self.output is None or not self.output.exists):
270 raise ValueError("--output must point to an existing CHAINED collection for --replace-run.")
272 @classmethod
273 def _makeReadParts(cls, args: SimpleNamespace) -> tuple[Butler, Sequence[str], _ButlerFactory]:
274 """Common implementation for `makeReadButler` and
275 `makeButlerAndCollections`.
277 Parameters
278 ----------
279 args : `types.SimpleNamespace`
280 Parsed command-line arguments. See class documentation for the
281 construction parameter of the same name.
283 Returns
284 -------
285 butler : `lsst.daf.butler.Butler`
286 A read-only butler constructed from the repo at
287 ``args.butler_config``, but with no default collections.
288 inputs : `Sequence` [ `str` ]
289 A collection search path constructed according to ``args``.
290 self : `_ButlerFactory`
291 A new `_ButlerFactory` instance representing the processed version
292 of ``args``.
293 """
294 butler = Butler(args.butler_config, writeable=False)
295 self = cls(butler.registry, args, writeable=False)
296 self.check(args)
297 if self.output and self.output.exists:
298 if args.replace_run:
299 replaced = self.output.chain[0]
300 inputs = list(self.output.chain[1:])
301 _LOG.debug(
302 "Simulating collection search in '%s' after removing '%s'.", self.output.name, replaced
303 )
304 else:
305 inputs = [self.output.name]
306 else:
307 inputs = list(self.inputs)
308 if args.extend_run:
309 assert self.outputRun is not None, "Output collection has to be specified."
310 inputs.insert(0, self.outputRun.name)
311 collSearch = CollectionWildcard.from_expression(inputs).require_ordered()
312 return butler, collSearch, self
314 @classmethod
315 def makeReadButler(cls, args: SimpleNamespace) -> Butler:
316 """Construct a read-only butler according to the given command-line
317 arguments.
319 Parameters
320 ----------
321 args : `types.SimpleNamespace`
322 Parsed command-line arguments. See class documentation for the
323 construction parameter of the same name.
325 Returns
326 -------
327 butler : `lsst.daf.butler.Butler`
328 A read-only butler initialized with the collections specified by
329 ``args``.
330 """
331 cls.defineDatastoreCache() # Ensure that this butler can use a shared cache.
332 butler, inputs, _ = cls._makeReadParts(args)
333 _LOG.debug("Preparing butler to read from %s.", inputs)
334 return Butler(butler=butler, collections=inputs)
336 @classmethod
337 def makeButlerAndCollections(cls, args: SimpleNamespace) -> Tuple[Butler, Sequence[str], Optional[str]]:
338 """Return a read-only registry, a collection search path, and the name
339 of the run to be used for future writes.
341 Parameters
342 ----------
343 args : `types.SimpleNamespace`
344 Parsed command-line arguments. See class documentation for the
345 construction parameter of the same name.
347 Returns
348 -------
349 butler : `lsst.daf.butler.Butler`
350 A read-only butler that collections will be added to and/or queried
351 from.
352 inputs : `Sequence` [ `str` ]
353 Collections to search for datasets.
354 run : `str` or `None`
355 Name of the output `~lsst.daf.butler.CollectionType.RUN` collection
356 if it already exists, or `None` if it does not.
357 """
358 butler, inputs, self = cls._makeReadParts(args)
359 run: Optional[str] = None
360 if args.extend_run:
361 assert self.outputRun is not None, "Output collection has to be specified."
362 if self.outputRun is not None:
363 run = self.outputRun.name
364 _LOG.debug("Preparing registry to read from %s and expect future writes to '%s'.", inputs, run)
365 return butler, inputs, run
367 @staticmethod
368 def defineDatastoreCache() -> None:
369 """Define where datastore cache directories should be found.
371 Notes
372 -----
373 All the jobs should share a datastore cache if applicable. This
374 method asks for a shared fallback cache to be defined and then
375 configures an exit handler to clean it up.
376 """
377 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset()
378 if defined:
379 atexit.register(shutil.rmtree, cache_dir, ignore_errors=True)
380 _LOG.debug("Defining shared datastore cache directory to %s", cache_dir)
382 @classmethod
383 def makeWriteButler(cls, args: SimpleNamespace, taskDefs: Optional[Iterable[TaskDef]] = None) -> Butler:
384 """Return a read-write butler initialized to write to and read from
385 the collections specified by the given command-line arguments.
387 Parameters
388 ----------
389 args : `types.SimpleNamespace`
390 Parsed command-line arguments. See class documentation for the
391 construction parameter of the same name.
392 taskDefs : iterable of `TaskDef`, optional
393 Definitions for tasks in a pipeline. This argument is only needed
394 if ``args.replace_run`` is `True` and ``args.prune_replaced`` is
395 "unstore".
397 Returns
398 -------
399 butler : `lsst.daf.butler.Butler`
400 A read-write butler initialized according to the given arguments.
401 """
402 cls.defineDatastoreCache() # Ensure that this butler can use a shared cache.
403 butler = Butler(args.butler_config, writeable=True)
404 self = cls(butler.registry, args, writeable=True)
405 self.check(args)
406 assert self.outputRun is not None, "Output collection has to be specified." # for mypy
407 if self.output is not None:
408 chainDefinition = list(self.output.chain if self.output.exists else self.inputs)
409 if args.replace_run:
410 replaced = chainDefinition.pop(0)
411 if args.prune_replaced == "unstore":
412 # Remove datasets from datastore
413 with butler.transaction():
414 refs: Iterable[DatasetRef] = butler.registry.queryDatasets(..., collections=replaced)
415 # we want to remove regular outputs but keep
416 # initOutputs, configs, and versions.
417 if taskDefs is not None:
418 initDatasetNames = set(PipelineDatasetTypes.initOutputNames(taskDefs))
419 refs = [ref for ref in refs if ref.datasetType.name not in initDatasetNames]
420 butler.pruneDatasets(refs, unstore=True, disassociate=False)
421 elif args.prune_replaced == "purge":
422 # Erase entire collection and all datasets, need to remove
423 # collection from its chain collection first.
424 with butler.transaction():
425 butler.registry.setCollectionChain(self.output.name, chainDefinition, flatten=True)
426 butler.removeRuns([replaced], unstore=True)
427 elif args.prune_replaced is not None:
428 raise NotImplementedError(f"Unsupported --prune-replaced option '{args.prune_replaced}'.")
429 if not self.output.exists:
430 butler.registry.registerCollection(self.output.name, CollectionType.CHAINED)
431 if not args.extend_run:
432 butler.registry.registerCollection(self.outputRun.name, CollectionType.RUN)
433 chainDefinition.insert(0, self.outputRun.name)
434 butler.registry.setCollectionChain(self.output.name, chainDefinition, flatten=True)
435 _LOG.debug(
436 "Preparing butler to write to '%s' and read from '%s'=%s",
437 self.outputRun.name,
438 self.output.name,
439 chainDefinition,
440 )
441 butler.registry.defaults = RegistryDefaults(run=self.outputRun.name, collections=self.output.name)
442 else:
443 inputs = (self.outputRun.name,) + self.inputs
444 _LOG.debug("Preparing butler to write to '%s' and read from %s.", self.outputRun.name, inputs)
445 butler.registry.defaults = RegistryDefaults(run=self.outputRun.name, collections=inputs)
446 return butler
448 output: Optional[_OutputChainedCollectionInfo]
449 """Information about the output chained collection, if there is or will be
450 one (`_OutputChainedCollectionInfo` or `None`).
451 """
453 outputRun: Optional[_OutputRunCollectionInfo]
454 """Information about the output run collection, if there is or will be
455 one (`_OutputRunCollectionInfo` or `None`).
456 """
458 inputs: Tuple[str, ...]
459 """Input collections provided directly by the user (`tuple` [ `str` ]).
460 """
463# ------------------------
464# Exported definitions --
465# ------------------------
468class CmdLineFwk:
469 """PipelineTask framework which executes tasks from command line.
471 In addition to executing tasks this activator provides additional methods
472 for task management like dumping configuration or execution chain.
473 """
475 MP_TIMEOUT = 3600 * 24 * 30 # Default timeout (sec) for multiprocessing
477 def __init__(self) -> None:
478 pass
480 def makePipeline(self, args: SimpleNamespace) -> Pipeline:
481 """Build a pipeline from command line arguments.
483 Parameters
484 ----------
485 args : `types.SimpleNamespace`
486 Parsed command line
488 Returns
489 -------
490 pipeline : `~lsst.pipe.base.Pipeline`
491 """
492 if args.pipeline:
493 pipeline = Pipeline.from_uri(args.pipeline)
494 else:
495 pipeline = Pipeline("anonymous")
497 # loop over all pipeline actions and apply them in order
498 for action in args.pipeline_actions:
499 if action.action == "add_instrument":
500 pipeline.addInstrument(action.value)
502 elif action.action == "new_task":
503 pipeline.addTask(action.value, action.label)
505 elif action.action == "delete_task":
506 pipeline.removeTask(action.label)
508 elif action.action == "config":
509 # action value string is "field=value", split it at '='
510 field, _, value = action.value.partition("=")
511 pipeline.addConfigOverride(action.label, field, value)
513 elif action.action == "configfile":
514 pipeline.addConfigFile(action.label, action.value)
516 else:
517 raise ValueError(f"Unexpected pipeline action: {action.action}")
519 if args.save_pipeline:
520 pipeline.write_to_uri(args.save_pipeline)
522 if args.pipeline_dot:
523 pipeline2dot(pipeline, args.pipeline_dot)
525 return pipeline
527 def makeGraph(self, pipeline: Pipeline, args: SimpleNamespace) -> Optional[QuantumGraph]:
528 """Build a graph from command line arguments.
530 Parameters
531 ----------
532 pipeline : `~lsst.pipe.base.Pipeline`
533 Pipeline, can be empty or ``None`` if graph is read from a file.
534 args : `types.SimpleNamespace`
535 Parsed command line
537 Returns
538 -------
539 graph : `~lsst.pipe.base.QuantumGraph` or `None`
540 If resulting graph is empty then `None` is returned.
541 """
543 # make sure that --extend-run always enables --skip-existing
544 if args.extend_run:
545 args.skip_existing = True
547 butler, collections, run = _ButlerFactory.makeButlerAndCollections(args)
549 if args.skip_existing and run:
550 args.skip_existing_in += (run,)
552 if args.qgraph:
553 # click passes empty tuple as default value for qgraph_node_id
554 nodes = args.qgraph_node_id or None
555 qgraph = QuantumGraph.loadUri(
556 args.qgraph, butler.registry.dimensions, nodes=nodes, graphID=args.qgraph_id
557 )
559 # pipeline can not be provided in this case
560 if pipeline:
561 raise ValueError("Pipeline must not be given when quantum graph is read from file.")
562 if args.show_qgraph_header:
563 print(QuantumGraph.readHeader(args.qgraph))
564 else:
565 # make execution plan (a.k.a. DAG) for pipeline
566 graphBuilder = GraphBuilder(
567 butler.registry,
568 skipExistingIn=args.skip_existing_in,
569 clobberOutputs=args.clobber_outputs,
570 datastore=butler.datastore if args.qgraph_datastore_records else None,
571 )
572 # accumulate metadata
573 metadata = {
574 "input": args.input,
575 "output": args.output,
576 "butler_argument": args.butler_config,
577 "output_run": run,
578 "extend_run": args.extend_run,
579 "skip_existing_in": args.skip_existing_in,
580 "skip_existing": args.skip_existing,
581 "data_query": args.data_query,
582 "user": getpass.getuser(),
583 "time": f"{datetime.datetime.now()}",
584 }
585 assert run is not None, "Butler output run collection must be defined"
586 qgraph = graphBuilder.makeGraph(
587 pipeline,
588 collections,
589 run,
590 args.data_query,
591 metadata=metadata,
592 datasetQueryConstraint=args.dataset_query_constraint,
593 )
594 if args.show_qgraph_header:
595 qgraph.buildAndPrintHeader()
597 # Count quanta in graph; give a warning if it's empty and return None.
598 nQuanta = len(qgraph)
599 if nQuanta == 0:
600 return None
601 else:
602 if _LOG.isEnabledFor(logging.INFO):
603 qg_task_table = self._generateTaskTable(qgraph)
604 qg_task_table_formatted = "\n".join(qg_task_table.pformat_all())
605 _LOG.info(
606 "QuantumGraph contains %d quanta for %d tasks, graph ID: %r\n%s",
607 nQuanta,
608 len(qgraph.taskGraph),
609 qgraph.graphID,
610 qg_task_table_formatted,
611 )
613 if args.save_qgraph:
614 qgraph.saveUri(args.save_qgraph)
616 if args.save_single_quanta:
617 for quantumNode in qgraph:
618 sqgraph = qgraph.subset(quantumNode)
619 uri = args.save_single_quanta.format(quantumNode)
620 sqgraph.saveUri(uri)
622 if args.qgraph_dot:
623 graph2dot(qgraph, args.qgraph_dot)
625 if args.execution_butler_location:
626 butler = Butler(args.butler_config)
627 newArgs = copy.deepcopy(args)
629 def builderShim(butler: Butler) -> Butler:
630 newArgs.butler_config = butler._config
631 # Calling makeWriteButler is done for the side effects of
632 # calling that method, maining parsing all the args into
633 # collection names, creating collections, etc.
634 newButler = _ButlerFactory.makeWriteButler(newArgs)
635 return newButler
637 # Include output collection in collections for input
638 # files if it exists in the repo.
639 all_inputs = args.input
640 if args.output is not None:
641 try:
642 all_inputs += (next(iter(butler.registry.queryCollections(args.output))),)
643 except MissingCollectionError:
644 pass
646 _LOG.debug("Calling buildExecutionButler with collections=%s", all_inputs)
647 buildExecutionButler(
648 butler,
649 qgraph,
650 args.execution_butler_location,
651 run,
652 butlerModifier=builderShim,
653 collections=all_inputs,
654 clobber=args.clobber_execution_butler,
655 datastoreRoot=args.target_datastore_root,
656 transfer=args.transfer,
657 )
659 return qgraph
661 def runPipeline(
662 self,
663 graph: QuantumGraph,
664 taskFactory: TaskFactory,
665 args: SimpleNamespace,
666 butler: Optional[Butler] = None,
667 ) -> None:
668 """Execute complete QuantumGraph.
670 Parameters
671 ----------
672 graph : `QuantumGraph`
673 Execution graph.
674 taskFactory : `~lsst.pipe.base.TaskFactory`
675 Task factory
676 args : `types.SimpleNamespace`
677 Parsed command line
678 butler : `~lsst.daf.butler.Butler`, optional
679 Data Butler instance, if not defined then new instance is made
680 using command line options.
681 """
682 # make sure that --extend-run always enables --skip-existing
683 if args.extend_run:
684 args.skip_existing = True
686 if not args.enable_implicit_threading:
687 disable_implicit_threading()
689 # Make butler instance. QuantumGraph should have an output run defined,
690 # but we ignore it here and let command line decide actual output run.
691 if butler is None:
692 butler = _ButlerFactory.makeWriteButler(args, graph.iterTaskGraph())
694 if args.skip_existing:
695 args.skip_existing_in += (butler.run,)
697 # Enable lsstDebug debugging. Note that this is done once in the
698 # main process before PreExecInit and it is also repeated before
699 # running each task in SingleQuantumExecutor (which may not be
700 # needed if `multipocessing` always uses fork start method).
701 if args.enableLsstDebug:
702 try:
703 _LOG.debug("Will try to import debug.py")
704 import debug # type: ignore # noqa:F401
705 except ImportError:
706 _LOG.warn("No 'debug' module found.")
708 # Save all InitOutputs, configs, etc.
709 preExecInit = PreExecInit(butler, taskFactory, extendRun=args.extend_run, mock=args.mock)
710 preExecInit.initialize(
711 graph,
712 saveInitOutputs=not args.skip_init_writes,
713 registerDatasetTypes=args.register_dataset_types,
714 saveVersions=not args.no_versions,
715 )
717 if not args.init_only:
718 graphFixup = self._importGraphFixup(args)
719 quantumExecutor = SingleQuantumExecutor(
720 butler,
721 taskFactory,
722 skipExistingIn=args.skip_existing_in,
723 clobberOutputs=args.clobber_outputs,
724 enableLsstDebug=args.enableLsstDebug,
725 exitOnKnownError=args.fail_fast,
726 mock=args.mock,
727 mock_configs=args.mock_configs,
728 )
730 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
731 executor = MPGraphExecutor(
732 numProc=args.processes,
733 timeout=timeout,
734 startMethod=args.start_method,
735 quantumExecutor=quantumExecutor,
736 failFast=args.fail_fast,
737 pdb=args.pdb,
738 executionGraphFixup=graphFixup,
739 )
740 # Have to reset connection pool to avoid sharing connections with
741 # forked processes.
742 butler.registry.resetConnectionPool()
743 try:
744 with util.profile(args.profile, _LOG):
745 executor.execute(graph)
746 finally:
747 if args.summary:
748 report = executor.getReport()
749 if report:
750 with open(args.summary, "w") as out:
751 # Do not save fields that are not set.
752 out.write(report.json(exclude_none=True, indent=2))
754 def _generateTaskTable(self, qgraph: QuantumGraph) -> Table:
755 """Generate astropy table listing the number of quanta per task for a
756 given quantum graph.
758 Parameters
759 ----------
760 qgraph : `lsst.pipe.base.graph.graph.QuantumGraph`
761 A QuantumGraph object.
763 Returns
764 -------
765 qg_task_table : `astropy.table.table.Table`
766 An astropy table containing columns: Quanta and Tasks.
767 """
768 qg_quanta, qg_tasks = [], []
769 for task_def in qgraph.iterTaskGraph():
770 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def)
771 qg_quanta.append(num_qnodes)
772 qg_tasks.append(task_def.label)
773 qg_task_table = Table(dict(Quanta=qg_quanta, Tasks=qg_tasks))
774 return qg_task_table
776 def _importGraphFixup(self, args: SimpleNamespace) -> Optional[ExecutionGraphFixup]:
777 """Import/instantiate graph fixup object.
779 Parameters
780 ----------
781 args : `types.SimpleNamespace`
782 Parsed command line.
784 Returns
785 -------
786 fixup : `ExecutionGraphFixup` or `None`
788 Raises
789 ------
790 ValueError
791 Raised if import fails, method call raises exception, or returned
792 instance has unexpected type.
793 """
794 if args.graph_fixup:
795 try:
796 factory = doImportType(args.graph_fixup)
797 except Exception as exc:
798 raise ValueError("Failed to import graph fixup class/method") from exc
799 try:
800 fixup = factory()
801 except Exception as exc:
802 raise ValueError("Failed to make instance of graph fixup") from exc
803 if not isinstance(fixup, ExecutionGraphFixup):
804 raise ValueError("Graph fixup is not an instance of ExecutionGraphFixup class")
805 return fixup
806 return None
808 def preExecInitQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
809 # Load quantum graph. We do not really need individual Quanta here,
810 # but we need datastore records for initInputs, and those are only
811 # available from Quanta, so load the whole thing.
812 qgraph = QuantumGraph.loadUri(args.qgraph, graphID=args.qgraph_id)
813 universe = qgraph.universe
815 # Collect all init input/output dataset IDs.
816 predicted_inputs: set[DatasetId] = set()
817 predicted_outputs: set[DatasetId] = set()
818 for taskDef in qgraph.iterTaskGraph():
819 if (refs := qgraph.initInputRefs(taskDef)) is not None:
820 predicted_inputs.update(ref.getCheckedId() for ref in refs)
821 if (refs := qgraph.initOutputRefs(taskDef)) is not None:
822 predicted_outputs.update(ref.getCheckedId() for ref in refs)
823 predicted_outputs.update(ref.getCheckedId() for ref in qgraph.globalInitOutputRefs())
824 # remove intermediates from inputs
825 predicted_inputs -= predicted_outputs
827 # Very inefficient way to extract datastore records from quantum graph,
828 # we have to scan all quanta and look at their datastore records.
829 datastore_records: dict[str, DatastoreRecordData] = {}
830 for quantum_node in qgraph:
831 for store_name, records in quantum_node.quantum.datastore_records.items():
832 subset = records.subset(predicted_inputs)
833 if subset is not None:
834 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset)
836 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()}
838 # Make butler from everything.
839 butler = QuantumBackedButler.from_predicted(
840 config=args.butler_config,
841 predicted_inputs=predicted_inputs,
842 predicted_outputs=predicted_outputs,
843 dimensions=universe,
844 datastore_records=datastore_records,
845 search_paths=args.config_search_path,
846 dataset_types=dataset_types,
847 )
849 # Save all InitOutputs, configs, etc.
850 preExecInit = PreExecInitLimited(butler, task_factory)
851 preExecInit.initialize(qgraph)
853 def runGraphQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
854 # Load quantum graph.
855 nodes = args.qgraph_node_id or None
856 qgraph = QuantumGraph.loadUri(args.qgraph, nodes=nodes, graphID=args.qgraph_id)
858 if qgraph.metadata is None:
859 raise ValueError("QuantumGraph is missing metadata, cannot ")
861 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()}
863 def _butler_factory(quantum: Quantum) -> LimitedButler:
864 """Factory method to create QuantumBackedButler instances."""
865 return QuantumBackedButler.initialize(
866 config=args.butler_config,
867 quantum=quantum,
868 dimensions=qgraph.universe,
869 dataset_types=dataset_types,
870 )
872 # make special quantum executor
873 quantumExecutor = SingleQuantumExecutor(
874 butler=None,
875 taskFactory=task_factory,
876 enableLsstDebug=args.enableLsstDebug,
877 exitOnKnownError=args.fail_fast,
878 limited_butler_factory=_butler_factory,
879 )
881 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
882 executor = MPGraphExecutor(
883 numProc=args.processes,
884 timeout=timeout,
885 startMethod=args.start_method,
886 quantumExecutor=quantumExecutor,
887 failFast=args.fail_fast,
888 pdb=args.pdb,
889 )
890 try:
891 with util.profile(args.profile, _LOG):
892 executor.execute(qgraph)
893 finally:
894 if args.summary:
895 report = executor.getReport()
896 if report:
897 with open(args.summary, "w") as out:
898 # Do not save fields that are not set.
899 out.write(report.json(exclude_none=True, indent=2))