Coverage for python/lsst/ctrl/mpexec/cmdLineFwk.py: 14%
361 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-11 09:04 +0000
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-11 09:04 +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, Mapping, Sequence
36from types import SimpleNamespace
37from typing import TYPE_CHECKING
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 `~lsst.daf.butler.Butler` is being initialized in a
201 context where actual writes should happens, and hence no output run
202 is necessary.
204 Raises
205 ------
206 ValueError
207 Raised if ``writeable is True`` but there are no output collections.
208 """
210 def __init__(self, registry: Registry, args: SimpleNamespace, writeable: bool):
211 if args.output is not None:
212 self.output = _OutputChainedCollectionInfo(registry, args.output)
213 else:
214 self.output = None
215 if args.output_run is not None:
216 self.outputRun = _OutputRunCollectionInfo(registry, args.output_run)
217 elif self.output is not None:
218 if args.extend_run:
219 if not self.output.chain:
220 raise ValueError("Cannot use --extend-run option with non-existing or empty output chain")
221 runName = self.output.chain[0]
222 else:
223 runName = "{}/{}".format(self.output, Instrument.makeCollectionTimestamp())
224 self.outputRun = _OutputRunCollectionInfo(registry, runName)
225 elif not writeable:
226 # If we're not writing yet, ok to have no output run.
227 self.outputRun = None
228 else:
229 raise ValueError("Cannot write without at least one of (--output, --output-run).")
230 # Recursively flatten any input CHAINED collections. We do this up
231 # front so we can tell if the user passes the same inputs on subsequent
232 # calls, even though we also flatten when we define the output CHAINED
233 # collection.
234 self.inputs = tuple(registry.queryCollections(args.input, flattenChains=True)) if args.input else ()
236 def check(self, args: SimpleNamespace) -> None:
237 """Check command-line options for consistency with each other and the
238 data repository.
240 Parameters
241 ----------
242 args : `types.SimpleNamespace`
243 Parsed command-line arguments. See class documentation for the
244 construction parameter of the same name.
245 """
246 assert not (args.extend_run and args.replace_run), "In mutually-exclusive group in ArgumentParser."
247 if self.inputs and self.output is not None and self.output.exists:
248 # Passing the same inputs that were used to initialize the output
249 # collection is allowed; this means they must _end_ with the same
250 # collections, because we push new runs to the front of the chain.
251 for c1, c2 in zip(self.inputs[::-1], self.output.chain[::-1]):
252 if c1 != c2:
253 raise ValueError(
254 f"Output CHAINED collection {self.output.name!r} exists, but it ends with "
255 "a different sequence of input collections than those given: "
256 f"{c1!r} != {c2!r} in inputs={self.inputs} vs "
257 f"{self.output.name}={self.output.chain}."
258 )
259 if len(self.inputs) > len(self.output.chain):
260 nNew = len(self.inputs) - len(self.output.chain)
261 raise ValueError(
262 f"Cannot add new input collections {self.inputs[:nNew]} after "
263 "output collection is first created."
264 )
265 if args.extend_run:
266 if self.outputRun is None:
267 raise ValueError("Cannot --extend-run when no output collection is given.")
268 elif not self.outputRun.exists:
269 raise ValueError(
270 f"Cannot --extend-run; output collection '{self.outputRun.name}' does not exist."
271 )
272 if not args.extend_run and self.outputRun is not None and self.outputRun.exists:
273 raise ValueError(
274 f"Output run '{self.outputRun.name}' already exists, but --extend-run was not given."
275 )
276 if args.prune_replaced and not args.replace_run:
277 raise ValueError("--prune-replaced requires --replace-run.")
278 if args.replace_run and (self.output is None or not self.output.exists):
279 raise ValueError("--output must point to an existing CHAINED collection for --replace-run.")
281 @classmethod
282 def _makeReadParts(cls, args: SimpleNamespace) -> tuple[Butler, Sequence[str], _ButlerFactory]:
283 """Parse arguments to support implementations of `makeReadButler` and
284 `makeButlerAndCollections`.
286 Parameters
287 ----------
288 args : `types.SimpleNamespace`
289 Parsed command-line arguments. See class documentation for the
290 construction parameter of the same name.
292 Returns
293 -------
294 butler : `lsst.daf.butler.Butler`
295 A read-only butler constructed from the repo at
296 ``args.butler_config``, but with no default collections.
297 inputs : `Sequence` [ `str` ]
298 A collection search path constructed according to ``args``.
299 self : `_ButlerFactory`
300 A new `_ButlerFactory` instance representing the processed version
301 of ``args``.
302 """
303 butler = Butler(args.butler_config, writeable=False)
304 self = cls(butler.registry, args, writeable=False)
305 self.check(args)
306 if self.output and self.output.exists:
307 if args.replace_run:
308 replaced = self.output.chain[0]
309 inputs = list(self.output.chain[1:])
310 _LOG.debug(
311 "Simulating collection search in '%s' after removing '%s'.", self.output.name, replaced
312 )
313 else:
314 inputs = [self.output.name]
315 else:
316 inputs = list(self.inputs)
317 if args.extend_run:
318 assert self.outputRun is not None, "Output collection has to be specified."
319 inputs.insert(0, self.outputRun.name)
320 collSearch = CollectionWildcard.from_expression(inputs).require_ordered()
321 return butler, collSearch, self
323 @classmethod
324 def makeReadButler(cls, args: SimpleNamespace) -> Butler:
325 """Construct a read-only butler according to the given command-line
326 arguments.
328 Parameters
329 ----------
330 args : `types.SimpleNamespace`
331 Parsed command-line arguments. See class documentation for the
332 construction parameter of the same name.
334 Returns
335 -------
336 butler : `lsst.daf.butler.Butler`
337 A read-only butler initialized with the collections specified by
338 ``args``.
339 """
340 cls.defineDatastoreCache() # Ensure that this butler can use a shared cache.
341 butler, inputs, _ = cls._makeReadParts(args)
342 _LOG.debug("Preparing butler to read from %s.", inputs)
343 return Butler(butler=butler, collections=inputs)
345 @classmethod
346 def makeButlerAndCollections(cls, args: SimpleNamespace) -> tuple[Butler, Sequence[str], str | None]:
347 """Return a read-only registry, a collection search path, and the name
348 of the run to be used for future writes.
350 Parameters
351 ----------
352 args : `types.SimpleNamespace`
353 Parsed command-line arguments. See class documentation for the
354 construction parameter of the same name.
356 Returns
357 -------
358 butler : `lsst.daf.butler.Butler`
359 A read-only butler that collections will be added to and/or queried
360 from.
361 inputs : `Sequence` [ `str` ]
362 Collections to search for datasets.
363 run : `str` or `None`
364 Name of the output `~lsst.daf.butler.CollectionType.RUN` collection
365 if it already exists, or `None` if it does not.
366 """
367 butler, inputs, self = cls._makeReadParts(args)
368 run: str | None = None
369 if args.extend_run:
370 assert self.outputRun is not None, "Output collection has to be specified."
371 if self.outputRun is not None:
372 run = self.outputRun.name
373 _LOG.debug("Preparing registry to read from %s and expect future writes to '%s'.", inputs, run)
374 return butler, inputs, run
376 @staticmethod
377 def defineDatastoreCache() -> None:
378 """Define where datastore cache directories should be found.
380 Notes
381 -----
382 All the jobs should share a datastore cache if applicable. This
383 method asks for a shared fallback cache to be defined and then
384 configures an exit handler to clean it up.
385 """
386 defined, cache_dir = DatastoreCacheManager.set_fallback_cache_directory_if_unset()
387 if defined:
388 atexit.register(shutil.rmtree, cache_dir, ignore_errors=True)
389 _LOG.debug("Defining shared datastore cache directory to %s", cache_dir)
391 @classmethod
392 def makeWriteButler(cls, args: SimpleNamespace, taskDefs: Iterable[TaskDef] | None = None) -> Butler:
393 """Return a read-write butler initialized to write to and read from
394 the collections specified by the given command-line arguments.
396 Parameters
397 ----------
398 args : `types.SimpleNamespace`
399 Parsed command-line arguments. See class documentation for the
400 construction parameter of the same name.
401 taskDefs : iterable of `TaskDef`, optional
402 Definitions for tasks in a pipeline. This argument is only needed
403 if ``args.replace_run`` is `True` and ``args.prune_replaced`` is
404 "unstore".
406 Returns
407 -------
408 butler : `lsst.daf.butler.Butler`
409 A read-write butler initialized according to the given arguments.
410 """
411 cls.defineDatastoreCache() # Ensure that this butler can use a shared cache.
412 butler = Butler(args.butler_config, writeable=True)
413 self = cls(butler.registry, args, writeable=True)
414 self.check(args)
415 assert self.outputRun is not None, "Output collection has to be specified." # for mypy
416 if self.output is not None:
417 chainDefinition = list(self.output.chain if self.output.exists else self.inputs)
418 if args.replace_run:
419 replaced = chainDefinition.pop(0)
420 if args.prune_replaced == "unstore":
421 # Remove datasets from datastore
422 with butler.transaction():
423 refs: Iterable[DatasetRef] = butler.registry.queryDatasets(..., collections=replaced)
424 # we want to remove regular outputs but keep
425 # initOutputs, configs, and versions.
426 if taskDefs is not None:
427 initDatasetNames = set(PipelineDatasetTypes.initOutputNames(taskDefs))
428 refs = [ref for ref in refs if ref.datasetType.name not in initDatasetNames]
429 butler.pruneDatasets(refs, unstore=True, disassociate=False)
430 elif args.prune_replaced == "purge":
431 # Erase entire collection and all datasets, need to remove
432 # collection from its chain collection first.
433 with butler.transaction():
434 butler.registry.setCollectionChain(self.output.name, chainDefinition, flatten=True)
435 butler.removeRuns([replaced], unstore=True)
436 elif args.prune_replaced is not None:
437 raise NotImplementedError(f"Unsupported --prune-replaced option '{args.prune_replaced}'.")
438 if not self.output.exists:
439 butler.registry.registerCollection(self.output.name, CollectionType.CHAINED)
440 if not args.extend_run:
441 butler.registry.registerCollection(self.outputRun.name, CollectionType.RUN)
442 chainDefinition.insert(0, self.outputRun.name)
443 butler.registry.setCollectionChain(self.output.name, chainDefinition, flatten=True)
444 _LOG.debug(
445 "Preparing butler to write to '%s' and read from '%s'=%s",
446 self.outputRun.name,
447 self.output.name,
448 chainDefinition,
449 )
450 butler.registry.defaults = RegistryDefaults(run=self.outputRun.name, collections=self.output.name)
451 else:
452 inputs = (self.outputRun.name,) + self.inputs
453 _LOG.debug("Preparing butler to write to '%s' and read from %s.", self.outputRun.name, inputs)
454 butler.registry.defaults = RegistryDefaults(run=self.outputRun.name, collections=inputs)
455 return butler
457 output: _OutputChainedCollectionInfo | None
458 """Information about the output chained collection, if there is or will be
459 one (`_OutputChainedCollectionInfo` or `None`).
460 """
462 outputRun: _OutputRunCollectionInfo | None
463 """Information about the output run collection, if there is or will be
464 one (`_OutputRunCollectionInfo` or `None`).
465 """
467 inputs: tuple[str, ...]
468 """Input collections provided directly by the user (`tuple` [ `str` ]).
469 """
472class _QBBFactory:
473 """Class which is a callable for making QBB instances."""
475 def __init__(
476 self, butler_config: Config, dimensions: DimensionUniverse, dataset_types: Mapping[str, DatasetType]
477 ):
478 self.butler_config = butler_config
479 self.dimensions = dimensions
480 self.dataset_types = dataset_types
482 def __call__(self, quantum: Quantum) -> LimitedButler:
483 """Return freshly initialized `~lsst.daf.butler.QuantumBackedButler`.
485 Factory method to create QuantumBackedButler instances.
486 """
487 return QuantumBackedButler.initialize(
488 config=self.butler_config,
489 quantum=quantum,
490 dimensions=self.dimensions,
491 dataset_types=self.dataset_types,
492 )
495# ------------------------
496# Exported definitions --
497# ------------------------
500class CmdLineFwk:
501 """PipelineTask framework which executes tasks from command line.
503 In addition to executing tasks this activator provides additional methods
504 for task management like dumping configuration or execution chain.
505 """
507 MP_TIMEOUT = 3600 * 24 * 30 # Default timeout (sec) for multiprocessing
509 def __init__(self) -> None:
510 pass
512 def makePipeline(self, args: SimpleNamespace) -> Pipeline:
513 """Build a pipeline from command line arguments.
515 Parameters
516 ----------
517 args : `types.SimpleNamespace`
518 Parsed command line
520 Returns
521 -------
522 pipeline : `~lsst.pipe.base.Pipeline`
523 """
524 if args.pipeline:
525 pipeline = Pipeline.from_uri(args.pipeline)
526 else:
527 pipeline = Pipeline("anonymous")
529 # loop over all pipeline actions and apply them in order
530 for action in args.pipeline_actions:
531 if action.action == "add_instrument":
532 pipeline.addInstrument(action.value)
534 elif action.action == "new_task":
535 pipeline.addTask(action.value, action.label)
537 elif action.action == "delete_task":
538 pipeline.removeTask(action.label)
540 elif action.action == "config":
541 # action value string is "field=value", split it at '='
542 field, _, value = action.value.partition("=")
543 pipeline.addConfigOverride(action.label, field, value)
545 elif action.action == "configfile":
546 pipeline.addConfigFile(action.label, action.value)
548 else:
549 raise ValueError(f"Unexpected pipeline action: {action.action}")
551 if args.save_pipeline:
552 pipeline.write_to_uri(args.save_pipeline)
554 if args.pipeline_dot:
555 pipeline2dot(pipeline, args.pipeline_dot)
557 return pipeline
559 def makeGraph(self, pipeline: Pipeline, args: SimpleNamespace) -> QuantumGraph | None:
560 """Build a graph from command line arguments.
562 Parameters
563 ----------
564 pipeline : `~lsst.pipe.base.Pipeline`
565 Pipeline, can be empty or ``None`` if graph is read from a file.
566 args : `types.SimpleNamespace`
567 Parsed command line
569 Returns
570 -------
571 graph : `~lsst.pipe.base.QuantumGraph` or `None`
572 If resulting graph is empty then `None` is returned.
573 """
574 # make sure that --extend-run always enables --skip-existing
575 if args.extend_run:
576 args.skip_existing = True
578 butler, collections, run = _ButlerFactory.makeButlerAndCollections(args)
580 if args.skip_existing and run:
581 args.skip_existing_in += (run,)
583 if args.qgraph:
584 # click passes empty tuple as default value for qgraph_node_id
585 nodes = args.qgraph_node_id or None
586 qgraph = QuantumGraph.loadUri(args.qgraph, butler.dimensions, nodes=nodes, graphID=args.qgraph_id)
588 # pipeline can not be provided in this case
589 if pipeline:
590 raise ValueError("Pipeline must not be given when quantum graph is read from file.")
591 if args.show_qgraph_header:
592 print(QuantumGraph.readHeader(args.qgraph))
593 else:
594 task_defs = list(pipeline.toExpandedPipeline())
595 if args.mock:
596 from lsst.pipe.base.tests.mocks import mock_task_defs
598 task_defs = mock_task_defs(task_defs, unmocked_dataset_types=args.unmocked_dataset_types)
599 # make execution plan (a.k.a. DAG) for pipeline
600 graphBuilder = GraphBuilder(
601 butler.registry,
602 skipExistingIn=args.skip_existing_in,
603 clobberOutputs=args.clobber_outputs,
604 datastore=butler.datastore if args.qgraph_datastore_records else None,
605 )
606 # accumulate metadata
607 metadata = {
608 "input": args.input,
609 "output": args.output,
610 "butler_argument": args.butler_config,
611 "output_run": run,
612 "extend_run": args.extend_run,
613 "skip_existing_in": args.skip_existing_in,
614 "skip_existing": args.skip_existing,
615 "data_query": args.data_query,
616 "user": getpass.getuser(),
617 "time": f"{datetime.datetime.now()}",
618 }
619 assert run is not None, "Butler output run collection must be defined"
620 qgraph = graphBuilder.makeGraph(
621 task_defs,
622 collections,
623 run,
624 args.data_query,
625 metadata=metadata,
626 datasetQueryConstraint=args.dataset_query_constraint,
627 dataId=pipeline.get_data_id(butler.dimensions),
628 )
629 if args.show_qgraph_header:
630 qgraph.buildAndPrintHeader()
632 # Count quanta in graph; give a warning if it's empty and return None.
633 nQuanta = len(qgraph)
634 if nQuanta == 0:
635 return None
636 else:
637 if _LOG.isEnabledFor(logging.INFO):
638 qg_task_table = self._generateTaskTable(qgraph)
639 qg_task_table_formatted = "\n".join(qg_task_table.pformat_all())
640 _LOG.info(
641 "QuantumGraph contains %d quanta for %d tasks, graph ID: %r\n%s",
642 nQuanta,
643 len(qgraph.taskGraph),
644 qgraph.graphID,
645 qg_task_table_formatted,
646 )
648 if args.save_qgraph:
649 qgraph.saveUri(args.save_qgraph)
651 if args.save_single_quanta:
652 for quantumNode in qgraph:
653 sqgraph = qgraph.subset(quantumNode)
654 uri = args.save_single_quanta.format(quantumNode)
655 sqgraph.saveUri(uri)
657 if args.qgraph_dot:
658 graph2dot(qgraph, args.qgraph_dot)
660 if args.execution_butler_location:
661 butler = Butler(args.butler_config)
662 newArgs = copy.deepcopy(args)
664 def builderShim(butler: Butler) -> Butler:
665 newArgs.butler_config = butler._config
666 # Calling makeWriteButler is done for the side effects of
667 # calling that method, maining parsing all the args into
668 # collection names, creating collections, etc.
669 newButler = _ButlerFactory.makeWriteButler(newArgs)
670 return newButler
672 # Include output collection in collections for input
673 # files if it exists in the repo.
674 all_inputs = args.input
675 if args.output is not None:
676 try:
677 all_inputs += (next(iter(butler.registry.queryCollections(args.output))),)
678 except MissingCollectionError:
679 pass
681 _LOG.debug("Calling buildExecutionButler with collections=%s", all_inputs)
682 buildExecutionButler(
683 butler,
684 qgraph,
685 args.execution_butler_location,
686 run,
687 butlerModifier=builderShim,
688 collections=all_inputs,
689 clobber=args.clobber_execution_butler,
690 datastoreRoot=args.target_datastore_root,
691 transfer=args.transfer,
692 )
694 return qgraph
696 def runPipeline(
697 self,
698 graph: QuantumGraph,
699 taskFactory: TaskFactory,
700 args: SimpleNamespace,
701 butler: Butler | None = None,
702 ) -> None:
703 """Execute complete QuantumGraph.
705 Parameters
706 ----------
707 graph : `~lsst.pipe.base.QuantumGraph`
708 Execution graph.
709 taskFactory : `~lsst.pipe.base.TaskFactory`
710 Task factory
711 args : `types.SimpleNamespace`
712 Parsed command line
713 butler : `~lsst.daf.butler.Butler`, optional
714 Data Butler instance, if not defined then new instance is made
715 using command line options.
716 """
717 # Check that output run defined on command line is consistent with
718 # quantum graph.
719 if args.output_run and graph.metadata:
720 graph_output_run = graph.metadata.get("output_run", args.output_run)
721 if graph_output_run != args.output_run:
722 raise ValueError(
723 f"Output run defined on command line ({args.output_run}) has to be "
724 f"identical to graph metadata ({graph_output_run}). "
725 "To update graph metadata run `pipetask update-graph-run` command."
726 )
728 # Make sure that --extend-run always enables --skip-existing,
729 # clobbering should be disabled if --extend-run is not specified.
730 if args.extend_run:
731 args.skip_existing = True
732 else:
733 args.clobber_outputs = False
735 if not args.enable_implicit_threading:
736 disable_implicit_threading()
738 # Make butler instance. QuantumGraph should have an output run defined,
739 # but we ignore it here and let command line decide actual output run.
740 if butler is None:
741 butler = _ButlerFactory.makeWriteButler(args, graph.iterTaskGraph())
743 if args.skip_existing:
744 args.skip_existing_in += (butler.run,)
746 # Enable lsstDebug debugging. Note that this is done once in the
747 # main process before PreExecInit and it is also repeated before
748 # running each task in SingleQuantumExecutor (which may not be
749 # needed if `multipocessing` always uses fork start method).
750 if args.enableLsstDebug:
751 try:
752 _LOG.debug("Will try to import debug.py")
753 import debug # type: ignore # noqa:F401
754 except ImportError:
755 _LOG.warn("No 'debug' module found.")
757 # Save all InitOutputs, configs, etc.
758 preExecInit = PreExecInit(butler, taskFactory, extendRun=args.extend_run)
759 preExecInit.initialize(
760 graph,
761 saveInitOutputs=not args.skip_init_writes,
762 registerDatasetTypes=args.register_dataset_types,
763 saveVersions=not args.no_versions,
764 )
766 if not args.init_only:
767 graphFixup = self._importGraphFixup(args)
768 quantumExecutor = SingleQuantumExecutor(
769 butler,
770 taskFactory,
771 skipExistingIn=args.skip_existing_in,
772 clobberOutputs=args.clobber_outputs,
773 enableLsstDebug=args.enableLsstDebug,
774 exitOnKnownError=args.fail_fast,
775 )
777 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
778 executor = MPGraphExecutor(
779 numProc=args.processes,
780 timeout=timeout,
781 startMethod=args.start_method,
782 quantumExecutor=quantumExecutor,
783 failFast=args.fail_fast,
784 pdb=args.pdb,
785 executionGraphFixup=graphFixup,
786 )
787 # Have to reset connection pool to avoid sharing connections with
788 # forked processes.
789 butler.registry.resetConnectionPool()
790 try:
791 with util.profile(args.profile, _LOG):
792 executor.execute(graph)
793 finally:
794 if args.summary:
795 report = executor.getReport()
796 if report:
797 with open(args.summary, "w") as out:
798 # Do not save fields that are not set.
799 out.write(report.json(exclude_none=True, indent=2))
801 def _generateTaskTable(self, qgraph: QuantumGraph) -> Table:
802 """Generate astropy table listing the number of quanta per task for a
803 given quantum graph.
805 Parameters
806 ----------
807 qgraph : `lsst.pipe.base.graph.graph.QuantumGraph`
808 A QuantumGraph object.
810 Returns
811 -------
812 qg_task_table : `astropy.table.table.Table`
813 An astropy table containing columns: Quanta and Tasks.
814 """
815 qg_quanta, qg_tasks = [], []
816 for task_def in qgraph.iterTaskGraph():
817 num_qnodes = qgraph.getNumberOfQuantaForTask(task_def)
818 qg_quanta.append(num_qnodes)
819 qg_tasks.append(task_def.label)
820 qg_task_table = Table(dict(Quanta=qg_quanta, Tasks=qg_tasks))
821 return qg_task_table
823 def _importGraphFixup(self, args: SimpleNamespace) -> ExecutionGraphFixup | None:
824 """Import/instantiate graph fixup object.
826 Parameters
827 ----------
828 args : `types.SimpleNamespace`
829 Parsed command line.
831 Returns
832 -------
833 fixup : `ExecutionGraphFixup` or `None`
835 Raises
836 ------
837 ValueError
838 Raised if import fails, method call raises exception, or returned
839 instance has unexpected type.
840 """
841 if args.graph_fixup:
842 try:
843 factory = doImportType(args.graph_fixup)
844 except Exception as exc:
845 raise ValueError("Failed to import graph fixup class/method") from exc
846 try:
847 fixup = factory()
848 except Exception as exc:
849 raise ValueError("Failed to make instance of graph fixup") from exc
850 if not isinstance(fixup, ExecutionGraphFixup):
851 raise ValueError("Graph fixup is not an instance of ExecutionGraphFixup class")
852 return fixup
853 return None
855 def preExecInitQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
856 # Load quantum graph. We do not really need individual Quanta here,
857 # but we need datastore records for initInputs, and those are only
858 # available from Quanta, so load the whole thing.
859 qgraph = QuantumGraph.loadUri(args.qgraph, graphID=args.qgraph_id)
860 universe = qgraph.universe
862 # Collect all init input/output dataset IDs.
863 predicted_inputs: set[DatasetId] = set()
864 predicted_outputs: set[DatasetId] = set()
865 for taskDef in qgraph.iterTaskGraph():
866 if (refs := qgraph.initInputRefs(taskDef)) is not None:
867 predicted_inputs.update(ref.id for ref in refs)
868 if (refs := qgraph.initOutputRefs(taskDef)) is not None:
869 predicted_outputs.update(ref.id for ref in refs)
870 predicted_outputs.update(ref.id for ref in qgraph.globalInitOutputRefs())
871 # remove intermediates from inputs
872 predicted_inputs -= predicted_outputs
874 # Very inefficient way to extract datastore records from quantum graph,
875 # we have to scan all quanta and look at their datastore records.
876 datastore_records: dict[str, DatastoreRecordData] = {}
877 for quantum_node in qgraph:
878 for store_name, records in quantum_node.quantum.datastore_records.items():
879 subset = records.subset(predicted_inputs)
880 if subset is not None:
881 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset)
883 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()}
885 # Make butler from everything.
886 butler = QuantumBackedButler.from_predicted(
887 config=args.butler_config,
888 predicted_inputs=predicted_inputs,
889 predicted_outputs=predicted_outputs,
890 dimensions=universe,
891 datastore_records=datastore_records,
892 search_paths=args.config_search_path,
893 dataset_types=dataset_types,
894 )
896 # Save all InitOutputs, configs, etc.
897 preExecInit = PreExecInitLimited(butler, task_factory)
898 preExecInit.initialize(qgraph)
900 def runGraphQBB(self, task_factory: TaskFactory, args: SimpleNamespace) -> None:
901 # Load quantum graph.
902 nodes = args.qgraph_node_id or None
903 qgraph = QuantumGraph.loadUri(args.qgraph, nodes=nodes, graphID=args.qgraph_id)
905 if qgraph.metadata is None:
906 raise ValueError("QuantumGraph is missing metadata, cannot ")
908 dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()}
910 _butler_factory = _QBBFactory(
911 butler_config=args.butler_config,
912 dimensions=qgraph.universe,
913 dataset_types=dataset_types,
914 )
916 # make special quantum executor
917 quantumExecutor = SingleQuantumExecutor(
918 butler=None,
919 taskFactory=task_factory,
920 enableLsstDebug=args.enableLsstDebug,
921 exitOnKnownError=args.fail_fast,
922 limited_butler_factory=_butler_factory,
923 )
925 timeout = self.MP_TIMEOUT if args.timeout is None else args.timeout
926 executor = MPGraphExecutor(
927 numProc=args.processes,
928 timeout=timeout,
929 startMethod=args.start_method,
930 quantumExecutor=quantumExecutor,
931 failFast=args.fail_fast,
932 pdb=args.pdb,
933 )
934 try:
935 with util.profile(args.profile, _LOG):
936 executor.execute(qgraph)
937 finally:
938 if args.summary:
939 report = executor.getReport()
940 if report:
941 with open(args.summary, "w") as out:
942 # Do not save fields that are not set.
943 out.write(report.json(exclude_none=True, indent=2))