Coverage for python/lsst/ctrl/mpexec/cmdLineFwk.py: 14%
345 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-02-01 10:07 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2023-02-01 10:07 +0000
1# This file is part of ctrl_mpexec.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
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, 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.pruneCollection(replaced, purge=True, 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":
501 pipeline.addInstrument(action.value)
503 elif action.action == "new_task":
505 pipeline.addTask(action.value, action.label)
507 elif action.action == "delete_task":
509 pipeline.removeTask(action.label)
511 elif action.action == "config":
513 # action value string is "field=value", split it at '='
514 field, _, value = action.value.partition("=")
515 pipeline.addConfigOverride(action.label, field, value)
517 elif action.action == "configfile":
519 pipeline.addConfigFile(action.label, action.value)
521 else:
523 raise ValueError(f"Unexpected pipeline action: {action.action}")
525 if args.save_pipeline:
526 pipeline.write_to_uri(args.save_pipeline)
528 if args.pipeline_dot:
529 pipeline2dot(pipeline, args.pipeline_dot)
531 return pipeline
533 def makeGraph(self, pipeline: Pipeline, args: SimpleNamespace) -> Optional[QuantumGraph]:
534 """Build a graph from command line arguments.
536 Parameters
537 ----------
538 pipeline : `~lsst.pipe.base.Pipeline`
539 Pipeline, can be empty or ``None`` if graph is read from a file.
540 args : `types.SimpleNamespace`
541 Parsed command line
543 Returns
544 -------
545 graph : `~lsst.pipe.base.QuantumGraph` or `None`
546 If resulting graph is empty then `None` is returned.
547 """
549 # make sure that --extend-run always enables --skip-existing
550 if args.extend_run:
551 args.skip_existing = True
553 butler, collections, run = _ButlerFactory.makeButlerAndCollections(args)
555 if args.skip_existing and run:
556 args.skip_existing_in += (run,)
558 if args.qgraph:
559 # click passes empty tuple as default value for qgraph_node_id
560 nodes = args.qgraph_node_id or None
561 qgraph = QuantumGraph.loadUri(
562 args.qgraph, butler.registry.dimensions, nodes=nodes, graphID=args.qgraph_id
563 )
565 # pipeline can not be provided in this case
566 if pipeline:
567 raise ValueError("Pipeline must not be given when quantum graph is read from file.")
568 if args.show_qgraph_header:
569 print(QuantumGraph.readHeader(args.qgraph))
570 else:
571 # make execution plan (a.k.a. DAG) for pipeline
572 graphBuilder = GraphBuilder(
573 butler.registry,
574 skipExistingIn=args.skip_existing_in,
575 clobberOutputs=args.clobber_outputs,
576 datastore=butler.datastore if args.qgraph_datastore_records else None,
577 )
578 # accumulate metadata
579 metadata = {
580 "input": args.input,
581 "output": args.output,
582 "butler_argument": args.butler_config,
583 "output_run": run,
584 "extend_run": args.extend_run,
585 "skip_existing_in": args.skip_existing_in,
586 "skip_existing": args.skip_existing,
587 "data_query": args.data_query,
588 "user": getpass.getuser(),
589 "time": f"{datetime.datetime.now()}",
590 }
591 qgraph = graphBuilder.makeGraph(
592 pipeline,
593 collections,
594 run,
595 args.data_query,
596 metadata=metadata,
597 datasetQueryConstraint=args.dataset_query_constraint,
598 resolveRefs=True,
599 )
600 if args.show_qgraph_header:
601 qgraph.buildAndPrintHeader()
603 # Count quanta in graph; give a warning if it's empty and return None.
604 nQuanta = len(qgraph)
605 if nQuanta == 0:
606 return None
607 else:
608 if _LOG.isEnabledFor(logging.INFO):
609 qg_task_table = self._generateTaskTable(qgraph)
610 qg_task_table_formatted = "\n".join(qg_task_table.pformat_all())
611 _LOG.info(
612 "QuantumGraph contains %d quanta for %d tasks, graph ID: %r\n%s",
613 nQuanta,
614 len(qgraph.taskGraph),
615 qgraph.graphID,
616 qg_task_table_formatted,
617 )
619 if args.save_qgraph:
620 qgraph.saveUri(args.save_qgraph)
622 if args.save_single_quanta:
623 for quantumNode in qgraph:
624 sqgraph = qgraph.subset(quantumNode)
625 uri = args.save_single_quanta.format(quantumNode)
626 sqgraph.saveUri(uri)
628 if args.qgraph_dot:
629 graph2dot(qgraph, args.qgraph_dot)
631 if args.execution_butler_location:
632 butler = Butler(args.butler_config)
633 newArgs = copy.deepcopy(args)
635 def builderShim(butler: Butler) -> Butler:
636 newArgs.butler_config = butler._config
637 # Calling makeWriteButler is done for the side effects of
638 # calling that method, maining parsing all the args into
639 # collection names, creating collections, etc.
640 newButler = _ButlerFactory.makeWriteButler(newArgs)
641 return newButler
643 # Include output collection in collections for input
644 # files if it exists in the repo.
645 all_inputs = args.input
646 if args.output is not None:
647 try:
648 all_inputs += (next(iter(butler.registry.queryCollections(args.output))),)
649 except MissingCollectionError:
650 pass
652 _LOG.debug("Calling buildExecutionButler with collections=%s", all_inputs)
653 buildExecutionButler(
654 butler,
655 qgraph,
656 args.execution_butler_location,
657 run,
658 butlerModifier=builderShim,
659 collections=all_inputs,
660 clobber=args.clobber_execution_butler,
661 datastoreRoot=args.target_datastore_root,
662 transfer=args.transfer,
663 )
665 return qgraph
667 def runPipeline(
668 self,
669 graph: QuantumGraph,
670 taskFactory: TaskFactory,
671 args: SimpleNamespace,
672 butler: Optional[Butler] = None,
673 ) -> None:
674 """Execute complete QuantumGraph.
676 Parameters
677 ----------
678 graph : `QuantumGraph`
679 Execution graph.
680 taskFactory : `~lsst.pipe.base.TaskFactory`
681 Task factory
682 args : `types.SimpleNamespace`
683 Parsed command line
684 butler : `~lsst.daf.butler.Butler`, optional
685 Data Butler instance, if not defined then new instance is made
686 using command line options.
687 """
688 # make sure that --extend-run always enables --skip-existing
689 if args.extend_run:
690 args.skip_existing = True
692 if not args.enable_implicit_threading:
693 disable_implicit_threading()
695 # Make butler instance. QuantumGraph should have an output run defined,
696 # but we ignore it here and let command line decide actual output run.
697 if butler is None:
698 butler = _ButlerFactory.makeWriteButler(args, graph.iterTaskGraph())
700 if args.skip_existing:
701 args.skip_existing_in += (butler.run,)
703 # Enable lsstDebug debugging. Note that this is done once in the
704 # main process before PreExecInit and it is also repeated before
705 # running each task in SingleQuantumExecutor (which may not be
706 # needed if `multipocessing` always uses fork start method).
707 if args.enableLsstDebug:
708 try:
709 _LOG.debug("Will try to import debug.py")
710 import debug # type: ignore # noqa:F401
711 except ImportError:
712 _LOG.warn("No 'debug' module found.")
714 # Save all InitOutputs, configs, etc.
715 preExecInit = PreExecInit(butler, taskFactory, extendRun=args.extend_run, mock=args.mock)
716 preExecInit.initialize(
717 graph,
718 saveInitOutputs=not args.skip_init_writes,
719 registerDatasetTypes=args.register_dataset_types,
720 saveVersions=not args.no_versions,
721 )
723 if not args.init_only:
724 graphFixup = self._importGraphFixup(args)
725 quantumExecutor = SingleQuantumExecutor(
726 butler,
727 taskFactory,
728 skipExistingIn=args.skip_existing_in,
729 clobberOutputs=args.clobber_outputs,
730 enableLsstDebug=args.enableLsstDebug,
731 exitOnKnownError=args.fail_fast,
732 mock=args.mock,
733 mock_configs=args.mock_configs,
734 )
736 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
737 executor = MPGraphExecutor(
738 numProc=args.processes,
739 timeout=timeout,
740 startMethod=args.start_method,
741 quantumExecutor=quantumExecutor,
742 failFast=args.fail_fast,
743 pdb=args.pdb,
744 executionGraphFixup=graphFixup,
745 )
746 # Have to reset connection pool to avoid sharing connections with
747 # forked processes.
748 butler.registry.resetConnectionPool()
749 try:
750 with util.profile(args.profile, _LOG):
751 executor.execute(graph)
752 finally:
753 if args.summary:
754 report = executor.getReport()
755 if report:
756 with open(args.summary, "w") as out:
757 # Do not save fields that are not set.
758 out.write(report.json(exclude_none=True, indent=2))
760 def _generateTaskTable(self, qgraph: QuantumGraph) -> Table:
761 """Generate astropy table listing the number of quanta per task for a
762 given quantum graph.
764 Parameters
765 ----------
766 qgraph : `lsst.pipe.base.graph.graph.QuantumGraph`
767 A QuantumGraph object.
769 Returns
770 -------
771 qg_task_table : `astropy.table.table.Table`
772 An astropy table containing columns: Quanta and Tasks.
773 """
774 qg_quanta, qg_tasks = [], []
775 for task_def in qgraph.iterTaskGraph():
776 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def)
777 qg_quanta.append(num_qnodes)
778 qg_tasks.append(task_def.label)
779 qg_task_table = Table(dict(Quanta=qg_quanta, Tasks=qg_tasks))
780 return qg_task_table
782 def _importGraphFixup(self, args: SimpleNamespace) -> Optional[ExecutionGraphFixup]:
783 """Import/instantiate graph fixup object.
785 Parameters
786 ----------
787 args : `types.SimpleNamespace`
788 Parsed command line.
790 Returns
791 -------
792 fixup : `ExecutionGraphFixup` or `None`
794 Raises
795 ------
796 ValueError
797 Raised if import fails, method call raises exception, or returned
798 instance has unexpected type.
799 """
800 if args.graph_fixup:
801 try:
802 factory = doImportType(args.graph_fixup)
803 except Exception as exc:
804 raise ValueError("Failed to import graph fixup class/method") from exc
805 try:
806 fixup = factory()
807 except Exception as exc:
808 raise ValueError("Failed to make instance of graph fixup") from exc
809 if not isinstance(fixup, ExecutionGraphFixup):
810 raise ValueError("Graph fixup is not an instance of ExecutionGraphFixup class")
811 return fixup
812 return None
814 def preExecInitQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
816 # Load quantum graph. We do not really need individual Quanta here,
817 # but we need datastore records for initInputs, and those are only
818 # available from Quanta, so load the whole thing.
819 qgraph = QuantumGraph.loadUri(args.qgraph, graphID=args.qgraph_id)
820 universe = qgraph.universe
822 # Collect all init input/output dataset IDs.
823 predicted_inputs: set[DatasetId] = set()
824 predicted_outputs: set[DatasetId] = set()
825 for taskDef in qgraph.iterTaskGraph():
826 if (refs := qgraph.initInputRefs(taskDef)) is not None:
827 predicted_inputs.update(ref.getCheckedId() for ref in refs)
828 if (refs := qgraph.initOutputRefs(taskDef)) is not None:
829 predicted_outputs.update(ref.getCheckedId() for ref in refs)
830 predicted_outputs.update(ref.getCheckedId() for ref in qgraph.globalInitOutputRefs())
831 # remove intermediates from inputs
832 predicted_inputs -= predicted_outputs
834 # Very inefficient way to extract datastore records from quantum graph,
835 # we have to scan all quanta and look at their datastore records.
836 datastore_records: dict[str, DatastoreRecordData] = {}
837 for quantum_node in qgraph:
838 for store_name, records in quantum_node.quantum.datastore_records.items():
839 subset = records.subset(predicted_inputs)
840 if subset is not None:
841 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset)
843 # Make butler from everything.
844 butler = QuantumBackedButler.from_predicted(
845 config=args.butler_config,
846 predicted_inputs=predicted_inputs,
847 predicted_outputs=predicted_outputs,
848 dimensions=universe,
849 datastore_records=datastore_records,
850 search_paths=args.config_search_path,
851 )
853 # Save all InitOutputs, configs, etc.
854 preExecInit = PreExecInitLimited(butler, task_factory)
855 preExecInit.initialize(qgraph)
857 def runGraphQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
859 # Load quantum graph.
860 nodes = args.qgraph_node_id or None
861 qgraph = QuantumGraph.loadUri(args.qgraph, nodes=nodes, graphID=args.qgraph_id)
863 if qgraph.metadata is None:
864 raise ValueError("QuantumGraph is missing metadata, cannot ")
866 # make special quantum executor
867 quantumExecutor = SingleQuantumExecutor(
868 butler=None,
869 taskFactory=task_factory,
870 enableLsstDebug=args.enableLsstDebug,
871 exitOnKnownError=args.fail_fast,
872 butler_config=args.butler_config,
873 universe=qgraph.universe,
874 )
876 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
877 executor = MPGraphExecutor(
878 numProc=args.processes,
879 timeout=timeout,
880 startMethod=args.start_method,
881 quantumExecutor=quantumExecutor,
882 failFast=args.fail_fast,
883 pdb=args.pdb,
884 )
885 try:
886 with util.profile(args.profile, _LOG):
887 executor.execute(qgraph)
888 finally:
889 if args.summary:
890 report = executor.getReport()
891 if report:
892 with open(args.summary, "w") as out:
893 # Do not save fields that are not set.
894 out.write(report.json(exclude_none=True, indent=2))