Coverage for python/lsst/ctrl/mpexec/cmdLineFwk.py: 14%
357 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-06 02:54 -0700
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-06 02:54 -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, Mapping, 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:
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
81# ----------------------------------
82# Local non-exported definitions --
83# ----------------------------------
85_LOG = logging.getLogger(__name__)
88class _OutputChainedCollectionInfo:
89 """A helper class for handling command-line arguments related to an output
90 `~lsst.daf.butler.CollectionType.CHAINED` collection.
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 """
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
109 def __str__(self) -> str:
110 return self.name
112 name: str
113 """Name of the collection provided on the command line (`str`).
114 """
116 exists: bool
117 """Whether this collection already exists in the registry (`bool`).
118 """
120 chain: Tuple[str, ...]
121 """The definition of the collection, if it already exists (`tuple`[`str`]).
123 Empty if the collection does not already exist.
124 """
127class _OutputRunCollectionInfo:
128 """A helper class for handling command-line arguments related to an output
129 `~lsst.daf.butler.CollectionType.RUN` collection.
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 """
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
149 name: str
150 """Name of the collection provided on the command line (`str`).
151 """
153 exists: bool
154 """Whether this collection already exists in the registry (`bool`).
155 """
158class _ButlerFactory:
159 """A helper class for processing command-line arguments related to input
160 and output collections.
162 Parameters
163 ----------
164 registry : `lsst.daf.butler.Registry`
165 Butler registry that collections will be added to and/or queried from.
167 args : `types.SimpleNamespace`
168 Parsed command-line arguments. The following attributes are used,
169 either at construction or in later methods.
171 ``output``
172 The name of a `~lsst.daf.butler.CollectionType.CHAINED`
173 input/output collection.
175 ``output_run``
176 The name of a `~lsst.daf.butler.CollectionType.RUN` input/output
177 collection.
179 ``extend_run``
180 A boolean indicating whether ``output_run`` should already exist
181 and be extended.
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.
188 ``prune_replaced``
189 A boolean indicating whether to prune the replaced run (requires
190 ``replace_run``).
192 ``inputs``
193 Input collections of any type; see
194 :ref:`daf_butler_ordered_collection_searches` for details.
196 ``butler_config``
197 Path to a data repository root or configuration file.
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.
203 Raises
204 ------
205 ValueError
206 Raised if ``writeable is True`` but there are no output collections.
207 """
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 ()
235 def check(self, args: SimpleNamespace) -> None:
236 """Check command-line options for consistency with each other and the
237 data repository.
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.")
280 @classmethod
281 def _makeReadParts(cls, args: SimpleNamespace) -> tuple[Butler, Sequence[str], _ButlerFactory]:
282 """Common implementation for `makeReadButler` and
283 `makeButlerAndCollections`.
285 Parameters
286 ----------
287 args : `types.SimpleNamespace`
288 Parsed command-line arguments. See class documentation for the
289 construction parameter of the same name.
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
322 @classmethod
323 def makeReadButler(cls, args: SimpleNamespace) -> Butler:
324 """Construct a read-only butler according to the given command-line
325 arguments.
327 Parameters
328 ----------
329 args : `types.SimpleNamespace`
330 Parsed command-line arguments. See class documentation for the
331 construction parameter of the same name.
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)
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.
349 Parameters
350 ----------
351 args : `types.SimpleNamespace`
352 Parsed command-line arguments. See class documentation for the
353 construction parameter of the same name.
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
375 @staticmethod
376 def defineDatastoreCache() -> None:
377 """Define where datastore cache directories should be found.
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)
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.
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".
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
456 output: Optional[_OutputChainedCollectionInfo]
457 """Information about the output chained collection, if there is or will be
458 one (`_OutputChainedCollectionInfo` or `None`).
459 """
461 outputRun: Optional[_OutputRunCollectionInfo]
462 """Information about the output run collection, if there is or will be
463 one (`_OutputRunCollectionInfo` or `None`).
464 """
466 inputs: Tuple[str, ...]
467 """Input collections provided directly by the user (`tuple` [ `str` ]).
468 """
471class _QBBFactory:
472 """Class which is a callable for making QBB instances."""
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
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 )
491# ------------------------
492# Exported definitions --
493# ------------------------
496class CmdLineFwk:
497 """PipelineTask framework which executes tasks from command line.
499 In addition to executing tasks this activator provides additional methods
500 for task management like dumping configuration or execution chain.
501 """
503 MP_TIMEOUT = 3600 * 24 * 30 # Default timeout (sec) for multiprocessing
505 def __init__(self) -> None:
506 pass
508 def makePipeline(self, args: SimpleNamespace) -> Pipeline:
509 """Build a pipeline from command line arguments.
511 Parameters
512 ----------
513 args : `types.SimpleNamespace`
514 Parsed command line
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")
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)
530 elif action.action == "new_task":
531 pipeline.addTask(action.value, action.label)
533 elif action.action == "delete_task":
534 pipeline.removeTask(action.label)
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)
541 elif action.action == "configfile":
542 pipeline.addConfigFile(action.label, action.value)
544 else:
545 raise ValueError(f"Unexpected pipeline action: {action.action}")
547 if args.save_pipeline:
548 pipeline.write_to_uri(args.save_pipeline)
550 if args.pipeline_dot:
551 pipeline2dot(pipeline, args.pipeline_dot)
553 return pipeline
555 def makeGraph(self, pipeline: Pipeline, args: SimpleNamespace) -> Optional[QuantumGraph]:
556 """Build a graph from command line arguments.
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
565 Returns
566 -------
567 graph : `~lsst.pipe.base.QuantumGraph` or `None`
568 If resulting graph is empty then `None` is returned.
569 """
571 # make sure that --extend-run always enables --skip-existing
572 if args.extend_run:
573 args.skip_existing = True
575 butler, collections, run = _ButlerFactory.makeButlerAndCollections(args)
577 if args.skip_existing and run:
578 args.skip_existing_in += (run,)
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 )
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 # make execution plan (a.k.a. DAG) for pipeline
594 graphBuilder = GraphBuilder(
595 butler.registry,
596 skipExistingIn=args.skip_existing_in,
597 clobberOutputs=args.clobber_outputs,
598 datastore=butler.datastore if args.qgraph_datastore_records else None,
599 )
600 # accumulate metadata
601 metadata = {
602 "input": args.input,
603 "output": args.output,
604 "butler_argument": args.butler_config,
605 "output_run": run,
606 "extend_run": args.extend_run,
607 "skip_existing_in": args.skip_existing_in,
608 "skip_existing": args.skip_existing,
609 "data_query": args.data_query,
610 "user": getpass.getuser(),
611 "time": f"{datetime.datetime.now()}",
612 }
613 assert run is not None, "Butler output run collection must be defined"
614 qgraph = graphBuilder.makeGraph(
615 pipeline,
616 collections,
617 run,
618 args.data_query,
619 metadata=metadata,
620 datasetQueryConstraint=args.dataset_query_constraint,
621 )
622 if args.show_qgraph_header:
623 qgraph.buildAndPrintHeader()
625 # Count quanta in graph; give a warning if it's empty and return None.
626 nQuanta = len(qgraph)
627 if nQuanta == 0:
628 return None
629 else:
630 if _LOG.isEnabledFor(logging.INFO):
631 qg_task_table = self._generateTaskTable(qgraph)
632 qg_task_table_formatted = "\n".join(qg_task_table.pformat_all())
633 _LOG.info(
634 "QuantumGraph contains %d quanta for %d tasks, graph ID: %r\n%s",
635 nQuanta,
636 len(qgraph.taskGraph),
637 qgraph.graphID,
638 qg_task_table_formatted,
639 )
641 if args.save_qgraph:
642 qgraph.saveUri(args.save_qgraph)
644 if args.save_single_quanta:
645 for quantumNode in qgraph:
646 sqgraph = qgraph.subset(quantumNode)
647 uri = args.save_single_quanta.format(quantumNode)
648 sqgraph.saveUri(uri)
650 if args.qgraph_dot:
651 graph2dot(qgraph, args.qgraph_dot)
653 if args.execution_butler_location:
654 butler = Butler(args.butler_config)
655 newArgs = copy.deepcopy(args)
657 def builderShim(butler: Butler) -> Butler:
658 newArgs.butler_config = butler._config
659 # Calling makeWriteButler is done for the side effects of
660 # calling that method, maining parsing all the args into
661 # collection names, creating collections, etc.
662 newButler = _ButlerFactory.makeWriteButler(newArgs)
663 return newButler
665 # Include output collection in collections for input
666 # files if it exists in the repo.
667 all_inputs = args.input
668 if args.output is not None:
669 try:
670 all_inputs += (next(iter(butler.registry.queryCollections(args.output))),)
671 except MissingCollectionError:
672 pass
674 _LOG.debug("Calling buildExecutionButler with collections=%s", all_inputs)
675 buildExecutionButler(
676 butler,
677 qgraph,
678 args.execution_butler_location,
679 run,
680 butlerModifier=builderShim,
681 collections=all_inputs,
682 clobber=args.clobber_execution_butler,
683 datastoreRoot=args.target_datastore_root,
684 transfer=args.transfer,
685 )
687 return qgraph
689 def runPipeline(
690 self,
691 graph: QuantumGraph,
692 taskFactory: TaskFactory,
693 args: SimpleNamespace,
694 butler: Optional[Butler] = None,
695 ) -> None:
696 """Execute complete QuantumGraph.
698 Parameters
699 ----------
700 graph : `QuantumGraph`
701 Execution graph.
702 taskFactory : `~lsst.pipe.base.TaskFactory`
703 Task factory
704 args : `types.SimpleNamespace`
705 Parsed command line
706 butler : `~lsst.daf.butler.Butler`, optional
707 Data Butler instance, if not defined then new instance is made
708 using command line options.
709 """
710 # Check that output run defined on command line is consistent with
711 # quantum graph.
712 if args.output_run and graph.metadata:
713 graph_output_run = graph.metadata.get("output_run", args.output_run)
714 if graph_output_run != args.output_run:
715 raise ValueError(
716 f"Output run defined on command line ({args.output_run}) has to be "
717 f"identical to graph metadata ({graph_output_run}). "
718 "To update graph metadata run `pipetask update-graph-run` command."
719 )
721 # Make sure that --extend-run always enables --skip-existing,
722 # clobbering should be disabled if --extend-run is not specified.
723 if args.extend_run:
724 args.skip_existing = True
725 else:
726 args.clobber_outputs = False
728 if not args.enable_implicit_threading:
729 disable_implicit_threading()
731 # Make butler instance. QuantumGraph should have an output run defined,
732 # but we ignore it here and let command line decide actual output run.
733 if butler is None:
734 butler = _ButlerFactory.makeWriteButler(args, graph.iterTaskGraph())
736 if args.skip_existing:
737 args.skip_existing_in += (butler.run,)
739 # Enable lsstDebug debugging. Note that this is done once in the
740 # main process before PreExecInit and it is also repeated before
741 # running each task in SingleQuantumExecutor (which may not be
742 # needed if `multipocessing` always uses fork start method).
743 if args.enableLsstDebug:
744 try:
745 _LOG.debug("Will try to import debug.py")
746 import debug # type: ignore # noqa:F401
747 except ImportError:
748 _LOG.warn("No 'debug' module found.")
750 # Save all InitOutputs, configs, etc.
751 preExecInit = PreExecInit(butler, taskFactory, extendRun=args.extend_run, mock=args.mock)
752 preExecInit.initialize(
753 graph,
754 saveInitOutputs=not args.skip_init_writes,
755 registerDatasetTypes=args.register_dataset_types,
756 saveVersions=not args.no_versions,
757 )
759 if not args.init_only:
760 graphFixup = self._importGraphFixup(args)
761 quantumExecutor = SingleQuantumExecutor(
762 butler,
763 taskFactory,
764 skipExistingIn=args.skip_existing_in,
765 clobberOutputs=args.clobber_outputs,
766 enableLsstDebug=args.enableLsstDebug,
767 exitOnKnownError=args.fail_fast,
768 mock=args.mock,
769 mock_configs=args.mock_configs,
770 )
772 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
773 executor = MPGraphExecutor(
774 numProc=args.processes,
775 timeout=timeout,
776 startMethod=args.start_method,
777 quantumExecutor=quantumExecutor,
778 failFast=args.fail_fast,
779 pdb=args.pdb,
780 executionGraphFixup=graphFixup,
781 )
782 # Have to reset connection pool to avoid sharing connections with
783 # forked processes.
784 butler.registry.resetConnectionPool()
785 try:
786 with util.profile(args.profile, _LOG):
787 executor.execute(graph)
788 finally:
789 if args.summary:
790 report = executor.getReport()
791 if report:
792 with open(args.summary, "w") as out:
793 # Do not save fields that are not set.
794 out.write(report.json(exclude_none=True, indent=2))
796 def _generateTaskTable(self, qgraph: QuantumGraph) -> Table:
797 """Generate astropy table listing the number of quanta per task for a
798 given quantum graph.
800 Parameters
801 ----------
802 qgraph : `lsst.pipe.base.graph.graph.QuantumGraph`
803 A QuantumGraph object.
805 Returns
806 -------
807 qg_task_table : `astropy.table.table.Table`
808 An astropy table containing columns: Quanta and Tasks.
809 """
810 qg_quanta, qg_tasks = [], []
811 for task_def in qgraph.iterTaskGraph():
812 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def)
813 qg_quanta.append(num_qnodes)
814 qg_tasks.append(task_def.label)
815 qg_task_table = Table(dict(Quanta=qg_quanta, Tasks=qg_tasks))
816 return qg_task_table
818 def _importGraphFixup(self, args: SimpleNamespace) -> Optional[ExecutionGraphFixup]:
819 """Import/instantiate graph fixup object.
821 Parameters
822 ----------
823 args : `types.SimpleNamespace`
824 Parsed command line.
826 Returns
827 -------
828 fixup : `ExecutionGraphFixup` or `None`
830 Raises
831 ------
832 ValueError
833 Raised if import fails, method call raises exception, or returned
834 instance has unexpected type.
835 """
836 if args.graph_fixup:
837 try:
838 factory = doImportType(args.graph_fixup)
839 except Exception as exc:
840 raise ValueError("Failed to import graph fixup class/method") from exc
841 try:
842 fixup = factory()
843 except Exception as exc:
844 raise ValueError("Failed to make instance of graph fixup") from exc
845 if not isinstance(fixup, ExecutionGraphFixup):
846 raise ValueError("Graph fixup is not an instance of ExecutionGraphFixup class")
847 return fixup
848 return None
850 def preExecInitQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
851 # Load quantum graph. We do not really need individual Quanta here,
852 # but we need datastore records for initInputs, and those are only
853 # available from Quanta, so load the whole thing.
854 qgraph = QuantumGraph.loadUri(args.qgraph, graphID=args.qgraph_id)
855 universe = qgraph.universe
857 # Collect all init input/output dataset IDs.
858 predicted_inputs: set[DatasetId] = set()
859 predicted_outputs: set[DatasetId] = set()
860 for taskDef in qgraph.iterTaskGraph():
861 if (refs := qgraph.initInputRefs(taskDef)) is not None:
862 predicted_inputs.update(ref.id for ref in refs)
863 if (refs := qgraph.initOutputRefs(taskDef)) is not None:
864 predicted_outputs.update(ref.id for ref in refs)
865 predicted_outputs.update(ref.id for ref in qgraph.globalInitOutputRefs())
866 # remove intermediates from inputs
867 predicted_inputs -= predicted_outputs
869 # Very inefficient way to extract datastore records from quantum graph,
870 # we have to scan all quanta and look at their datastore records.
871 datastore_records: dict[str, DatastoreRecordData] = {}
872 for quantum_node in qgraph:
873 for store_name, records in quantum_node.quantum.datastore_records.items():
874 subset = records.subset(predicted_inputs)
875 if subset is not None:
876 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset)
878 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()}
880 # Make butler from everything.
881 butler = QuantumBackedButler.from_predicted(
882 config=args.butler_config,
883 predicted_inputs=predicted_inputs,
884 predicted_outputs=predicted_outputs,
885 dimensions=universe,
886 datastore_records=datastore_records,
887 search_paths=args.config_search_path,
888 dataset_types=dataset_types,
889 )
891 # Save all InitOutputs, configs, etc.
892 preExecInit = PreExecInitLimited(butler, task_factory)
893 preExecInit.initialize(qgraph)
895 def runGraphQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
896 # Load quantum graph.
897 nodes = args.qgraph_node_id or None
898 qgraph = QuantumGraph.loadUri(args.qgraph, nodes=nodes, graphID=args.qgraph_id)
900 if qgraph.metadata is None:
901 raise ValueError("QuantumGraph is missing metadata, cannot ")
903 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()}
905 _butler_factory = _QBBFactory(
906 butler_config=args.butler_config,
907 dimensions=qgraph.universe,
908 dataset_types=dataset_types,
909 )
911 # make special quantum executor
912 quantumExecutor = SingleQuantumExecutor(
913 butler=None,
914 taskFactory=task_factory,
915 enableLsstDebug=args.enableLsstDebug,
916 exitOnKnownError=args.fail_fast,
917 limited_butler_factory=_butler_factory,
918 )
920 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
921 executor = MPGraphExecutor(
922 numProc=args.processes,
923 timeout=timeout,
924 startMethod=args.start_method,
925 quantumExecutor=quantumExecutor,
926 failFast=args.fail_fast,
927 pdb=args.pdb,
928 )
929 try:
930 with util.profile(args.profile, _LOG):
931 executor.execute(qgraph)
932 finally:
933 if args.summary:
934 report = executor.getReport()
935 if report:
936 with open(args.summary, "w") as out:
937 # Do not save fields that are not set.
938 out.write(report.json(exclude_none=True, indent=2))