Coverage for python/lsst/pipe/base/single_quantum_executor.py: 11%
200 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-06 01:36 -0800
1# This file is part of pipe_base.
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28__all__ = ["SingleQuantumExecutor"]
30import logging
31import time
32import uuid
33from collections import defaultdict
34from collections.abc import Callable, Mapping
35from itertools import chain
36from typing import Any
38from lsst.daf.butler import (
39 Butler,
40 ButlerMetrics,
41 DatasetRef,
42 DatasetType,
43 LimitedButler,
44 NamedKeyDict,
45 Quantum,
46)
47from lsst.utils.timer import logInfo
49from ._quantumContext import ExecutionResources, QuantumContext
50from ._status import AnnotatedPartialOutputsError, InvalidQuantumError, NoWorkFound, QuantumSuccessCaveats
51from .connections import AdjustQuantumHelper
52from .log_capture import LogCapture
53from .pipeline_graph import TaskNode
54from .pipelineTask import PipelineTask
55from .quantum_graph_executor import QuantumExecutor
56from .quantum_reports import QuantumReport
57from .task import _TASK_FULL_METADATA_TYPE, _TASK_METADATA_TYPE
58from .taskFactory import TaskFactory
60_LOG = logging.getLogger(__name__)
63class SingleQuantumExecutor(QuantumExecutor):
64 """Executor class which runs one Quantum at a time.
66 Parameters
67 ----------
68 butler : `~lsst.daf.butler.Butler` or `None`, optional
69 Data butler, `None` means that a limited butler should be used instead.
70 task_factory : `.TaskFactory`, optional
71 Instance of a task factory. Defaults to a new instance of
72 `lsst.pipe.base.TaskFactory`.
73 skip_existing_in : `str` or `~collections.abc.Iterable` [ `str` ]
74 Expressions representing the collections to search for existing output
75 datasets. See :ref:`daf_butler_ordered_collection_searches` for allowed
76 types. This class only checks for the presence of butler output run in
77 the list of collections. If the output run is present in the list then
78 the quanta whose complete outputs exist in the output run will be
79 skipped. `None` or empty string/sequence disables skipping.
80 clobber_outputs : `bool`, optional
81 If `True`, then outputs from a quantum that exist in output run
82 collection will be removed prior to executing a quantum. If
83 ``skip_existing_in`` contains output run, then only partial outputs
84 from a quantum will be removed. Only used when ``butler`` is not
85 `None`.
86 enable_lsst_debug : `bool`, optional
87 Enable debugging with ``lsstDebug`` facility for a task.
88 limited_butler_factory : `Callable`, optional
89 A method that creates a `~lsst.daf.butler.LimitedButler` instance for a
90 given Quantum. This parameter must be defined if ``butler`` is `None`.
91 If ``butler`` is not `None` then this parameter is ignored.
92 resources : `.ExecutionResources`, optional
93 The resources available to this quantum when executing.
94 skip_existing : `bool`, optional
95 If `True`, skip quanta whose metadata datasets are already stored.
96 Unlike ``skip_existing_in``, this works with limited butlers as well as
97 full butlers. Always set to `True` if ``skip_existing_in`` matches
98 ``butler.run``.
99 assume_no_existing_outputs : `bool`, optional
100 If `True`, assume preexisting outputs are impossible (e.g. because this
101 is known by higher-level code to be a new ``RUN`` collection), and do
102 not look for them. This causes the ``skip_existing`` and
103 ``clobber_outputs`` options to be ignored, but unlike just setting both
104 of those to `False`, it also avoids all dataset existence checks.
105 raise_on_partial_outputs : `bool`, optional
106 If `True` raise exceptions chained by `.AnnotatedPartialOutputsError`
107 immediately, instead of considering the partial result a success and
108 continuing to run downstream tasks.
109 job_metadata : `~collections.abc.Mapping`
110 Mapping with extra metadata to embed within the quantum metadata under
111 the "job" key. This is intended to correspond to information common
112 to all quanta being executed in a single process, such as the time
113 taken to load the quantum graph in a BPS job.
114 """
116 def __init__(
117 self,
118 *,
119 butler: Butler | None = None,
120 task_factory: TaskFactory | None = None,
121 skip_existing_in: Any = None,
122 clobber_outputs: bool = False,
123 enable_lsst_debug: bool = False,
124 limited_butler_factory: Callable[[Quantum], LimitedButler] | None = None,
125 resources: ExecutionResources | None = None,
126 skip_existing: bool = False,
127 assume_no_existing_outputs: bool = False,
128 raise_on_partial_outputs: bool = True,
129 job_metadata: Mapping[str, int | str | float] | None = None,
130 ):
131 self._butler = butler
132 self._task_factory = task_factory if task_factory is not None else TaskFactory()
133 self._clobber_outputs = clobber_outputs
134 self._enable_lsst_debug = enable_lsst_debug
135 self._limited_butler_factory = limited_butler_factory
136 self._resources = resources
137 self._assume_no_existing_outputs = assume_no_existing_outputs
138 self._raise_on_partial_outputs = raise_on_partial_outputs
139 self._job_metadata = job_metadata
141 if self._butler is None:
142 assert limited_butler_factory is not None, "limited_butler_factory is needed when butler is None"
144 # Find whether output run is in skip_existing_in.
145 self._skip_existing = skip_existing
146 if self._butler is not None and skip_existing_in and not self._skip_existing:
147 self._skip_existing = self._butler.run in self._butler.collections.query(
148 skip_existing_in, flatten_chains=True
149 )
151 def execute(
152 self, task_node: TaskNode, /, quantum: Quantum, quantum_id: uuid.UUID | None = None
153 ) -> tuple[Quantum, QuantumReport | None]:
154 # Docstring inherited from QuantumExecutor.execute
155 assert quantum.dataId is not None, "Quantum DataId cannot be None"
157 if self._butler is not None:
158 self._butler.registry.refresh()
160 result = self._execute(task_node, quantum, quantum_id=quantum_id)
161 report = QuantumReport(quantumId=quantum_id, dataId=quantum.dataId, taskLabel=task_node.label)
162 return result, report
164 def _execute(
165 self, task_node: TaskNode, /, quantum: Quantum, quantum_id: uuid.UUID | None = None
166 ) -> Quantum:
167 """Execute the quantum.
169 Internal implementation of `execute()`.
170 """
171 startTime = time.time()
173 # Make a limited butler instance if needed.
174 limited_butler: LimitedButler
175 if self._butler is not None:
176 limited_butler = self._butler
177 else:
178 # We check this in constructor, but mypy needs this check here.
179 assert self._limited_butler_factory is not None
180 limited_butler = self._limited_butler_factory(quantum)
182 if self._butler is not None:
183 log_capture = LogCapture.from_full(self._butler)
184 else:
185 log_capture = LogCapture.from_limited(limited_butler)
186 with log_capture.capture_logging(task_node, quantum) as captureLog:
187 # Save detailed resource usage before task start to metadata.
188 quantumMetadata = _TASK_METADATA_TYPE()
189 logInfo(None, "prep", metadata=quantumMetadata) # type: ignore[arg-type]
191 _LOG.info(
192 "Preparing execution of quantum for label=%s dataId=%s.", task_node.label, quantum.dataId
193 )
195 # check whether to skip or delete old outputs, if it returns True
196 # or raises an exception do not try to store logs, as they may be
197 # already in butler.
198 captureLog.store = False
199 if self._check_existing_outputs(quantum, task_node, limited_butler):
200 _LOG.info(
201 "Skipping already-successful quantum for label=%s dataId=%s.",
202 task_node.label,
203 quantum.dataId,
204 )
205 return quantum
206 captureLog.store = True
208 try:
209 quantum = self._updated_quantum_inputs(quantum, task_node, limited_butler)
210 except NoWorkFound as exc:
211 _LOG.info(
212 "Nothing to do for task '%s' on quantum %s; saving metadata and skipping: %s",
213 task_node.label,
214 quantum.dataId,
215 str(exc),
216 )
217 quantumMetadata["caveats"] = QuantumSuccessCaveats.from_adjust_quantum_no_work().value
218 # Make empty metadata that looks something like what a
219 # do-nothing task would write (but we don't bother with empty
220 # nested PropertySets for subtasks). This is slightly
221 # duplicative with logic in pipe_base that we can't easily call
222 # from here; we'll fix this on DM-29761.
223 logInfo(None, "end", metadata=quantumMetadata) # type: ignore[arg-type]
224 fullMetadata = _TASK_FULL_METADATA_TYPE()
225 fullMetadata[task_node.label] = _TASK_METADATA_TYPE()
226 fullMetadata["quantum"] = quantumMetadata
227 if self._job_metadata is not None:
228 fullMetadata["job"] = self._job_metadata
229 self._write_metadata(quantum, fullMetadata, task_node, limited_butler)
230 return quantum
232 # enable lsstDebug debugging
233 if self._enable_lsst_debug:
234 try:
235 _LOG.debug("Will try to import debug.py")
236 import debug # type: ignore # noqa:F401
237 except ImportError:
238 _LOG.warning("No 'debug' module found.")
240 # Ensure that we are executing a frozen config
241 task_node.config.freeze()
242 logInfo(None, "init", metadata=quantumMetadata) # type: ignore[arg-type]
243 init_input_refs = list(quantum.initInputs.values())
245 _LOG.info(
246 "Constructing task and executing quantum for label=%s dataId=%s.",
247 task_node.label,
248 quantum.dataId,
249 )
250 task = self._task_factory.makeTask(task_node, limited_butler, init_input_refs)
251 logInfo(None, "start", metadata=quantumMetadata) # type: ignore[arg-type]
252 try:
253 caveats, outputsPut, butler_metrics = self._run_quantum(
254 task, quantum, task_node, limited_butler, quantum_id=quantum_id
255 )
256 except Exception as e:
257 _LOG.error(
258 "Execution of task '%s' on quantum %s failed. Exception %s: %s",
259 task_node.label,
260 quantum.dataId,
261 e.__class__.__name__,
262 str(e),
263 )
264 raise
265 else:
266 quantumMetadata["butler_metrics"] = butler_metrics.model_dump()
267 quantumMetadata["caveats"] = caveats.value
268 # Stringify the UUID for easier compatibility with
269 # PropertyList.
270 quantumMetadata["outputs"] = [str(output) for output in outputsPut]
271 logInfo(None, "end", metadata=quantumMetadata) # type: ignore[arg-type]
272 fullMetadata = task.getFullMetadata()
273 fullMetadata["quantum"] = quantumMetadata
274 if self._job_metadata is not None:
275 fullMetadata["job"] = self._job_metadata
276 self._write_metadata(quantum, fullMetadata, task_node, limited_butler)
277 stopTime = time.time()
278 _LOG.info(
279 "Execution of task '%s' on quantum %s took %.3f seconds",
280 task_node.label,
281 quantum.dataId,
282 stopTime - startTime,
283 )
284 return quantum
286 def _check_existing_outputs(
287 self, quantum: Quantum, task_node: TaskNode, /, limited_butler: LimitedButler
288 ) -> bool:
289 """Decide whether this quantum needs to be executed.
291 If only partial outputs exist then they are removed if
292 ``clobberOutputs`` is True, otherwise an exception is raised.
294 The ``LimitedButler`` is used for everything, and should be set to
295 ``self.butler`` if no separate ``LimitedButler`` is available.
297 Parameters
298 ----------
299 quantum : `~lsst.daf.butler.Quantum`
300 Quantum to check for existing outputs.
301 task_node : `~.pipeline_graph.TaskNode`
302 Task definition structure.
303 limited_butler : `~lsst.daf.butler.LimitedButler`
304 Butler to use for querying and clobbering.
306 Returns
307 -------
308 exist : `bool`
309 `True` if ``self.skipExisting`` is defined, and a previous
310 execution of this quanta appears to have completed successfully
311 (either because metadata was written or all datasets were written).
312 `False` otherwise.
314 Raises
315 ------
316 RuntimeError
317 Raised if some outputs exist and some not.
318 """
319 if self._assume_no_existing_outputs:
320 return False
322 if self._skip_existing:
323 _LOG.debug(
324 "Checking existence of metadata from previous execution of label=%s dataId=%s.",
325 task_node.label,
326 quantum.dataId,
327 )
328 # Metadata output exists; this is sufficient to assume the previous
329 # run was successful and should be skipped.
330 [metadata_ref] = quantum.outputs[task_node.metadata_output.dataset_type_name]
331 if metadata_ref is not None:
332 if limited_butler.stored(metadata_ref):
333 return True
335 # Find and prune (partial) outputs if `self.clobberOutputs` is set.
336 _LOG.debug(
337 "Looking for existing outputs in the way for label=%s dataId=%s.", task_node.label, quantum.dataId
338 )
339 ref_dict = limited_butler.stored_many(chain.from_iterable(quantum.outputs.values()))
340 existingRefs = [ref for ref, exists in ref_dict.items() if exists]
341 missingRefs = [ref for ref, exists in ref_dict.items() if not exists]
342 if existingRefs:
343 if not missingRefs:
344 # Full outputs exist.
345 if self._skip_existing:
346 return True
347 elif self._clobber_outputs:
348 _LOG.info("Removing complete outputs for quantum %s: %s", quantum, existingRefs)
349 limited_butler.pruneDatasets(existingRefs, disassociate=True, unstore=True, purge=True)
350 else:
351 raise RuntimeError(
352 f"Complete outputs exists for a quantum {quantum} "
353 "and neither clobberOutputs nor skipExisting is set: "
354 f"existingRefs={existingRefs}"
355 )
356 else:
357 # Partial outputs from a failed quantum.
358 _LOG.debug(
359 "Partial outputs exist for quantum %s existingRefs=%s missingRefs=%s",
360 quantum,
361 existingRefs,
362 missingRefs,
363 )
364 if self._clobber_outputs:
365 # only prune
366 _LOG.info("Removing partial outputs for task %s: %s", task_node.label, existingRefs)
367 limited_butler.pruneDatasets(existingRefs, disassociate=True, unstore=True, purge=True)
368 return False
369 else:
370 raise RuntimeError(
371 "Registry inconsistency while checking for existing quantum outputs:"
372 f" quantum={quantum} existingRefs={existingRefs}"
373 f" missingRefs={missingRefs}"
374 )
376 # By default always execute.
377 return False
379 def _updated_quantum_inputs(
380 self, quantum: Quantum, task_node: TaskNode, /, limited_butler: LimitedButler
381 ) -> Quantum:
382 """Update quantum with extra information, returns a new updated
383 Quantum.
385 Some methods may require input DatasetRefs to have non-None
386 ``dataset_id``, but in case of intermediate dataset it may not be
387 filled during QuantumGraph construction. This method will retrieve
388 missing info from registry.
390 Parameters
391 ----------
392 quantum : `~lsst.daf.butler.Quantum`
393 Single Quantum instance.
394 task_node : `~.pipeline_graph.TaskNode`
395 Task definition structure.
396 limited_butler : `~lsst.daf.butler.LimitedButler`
397 Butler to use for querying.
399 Returns
400 -------
401 update : `~lsst.daf.butler.Quantum`
402 Updated Quantum instance.
403 """
404 anyChanges = False
405 updatedInputs: defaultdict[DatasetType, list] = defaultdict(list)
406 for key, refsForDatasetType in quantum.inputs.items():
407 _LOG.debug(
408 "Checking existence of input '%s' for label=%s dataId=%s.",
409 key.name,
410 task_node.label,
411 quantum.dataId,
412 )
413 toCheck = []
414 newRefsForDatasetType = updatedInputs[key]
415 for ref in refsForDatasetType:
416 if self._should_assume_exists(quantum, ref):
417 newRefsForDatasetType.append(ref)
418 else:
419 toCheck.append(ref)
420 if not toCheck:
421 _LOG.debug(
422 "Assuming overall input '%s' is present without checks for label=%s dataId=%s.",
423 key.name,
424 task_node.label,
425 quantum.dataId,
426 )
427 continue
428 stored = limited_butler.stored_many(toCheck)
429 for ref in toCheck:
430 if stored[ref]:
431 newRefsForDatasetType.append(ref)
432 else:
433 # This should only happen if a predicted intermediate was
434 # not actually produced upstream, but
435 # datastore misconfigurations can unfortunately also land
436 # us here.
437 _LOG.info("No dataset artifact found for %s", ref)
438 continue
439 if len(newRefsForDatasetType) != len(refsForDatasetType):
440 anyChanges = True
441 # If we removed any input datasets, let the task check if it has enough
442 # to proceed and/or prune related datasets that it also doesn't
443 # need/produce anymore. It will raise NoWorkFound if it can't run,
444 # which we'll let propagate up. This is exactly what we run during QG
445 # generation, because a task shouldn't care whether an input is missing
446 # because some previous task didn't produce it, or because it just
447 # wasn't there during QG generation.
448 namedUpdatedInputs = NamedKeyDict[DatasetType, list[DatasetRef]](updatedInputs.items())
449 helper = AdjustQuantumHelper(namedUpdatedInputs, quantum.outputs)
450 if anyChanges:
451 _LOG.debug("Running adjustQuantum for label=%s dataId=%s.", task_node.label, quantum.dataId)
452 assert quantum.dataId is not None, "Quantum DataId cannot be None"
453 helper.adjust_in_place(task_node.get_connections(), label=task_node.label, data_id=quantum.dataId)
454 return Quantum(
455 taskName=quantum.taskName,
456 taskClass=quantum.taskClass,
457 dataId=quantum.dataId,
458 initInputs=quantum.initInputs,
459 inputs=helper.inputs,
460 outputs=helper.outputs,
461 )
463 def _run_quantum(
464 self,
465 task: PipelineTask,
466 quantum: Quantum,
467 task_node: TaskNode,
468 /,
469 limited_butler: LimitedButler,
470 quantum_id: uuid.UUID | None = None,
471 ) -> tuple[QuantumSuccessCaveats, list[uuid.UUID], ButlerMetrics]:
472 """Execute task on a single quantum.
474 Parameters
475 ----------
476 task : `PipelineTask`
477 Task object.
478 quantum : `~lsst.daf.butler.Quantum`
479 Single Quantum instance.
480 task_node : `~.pipeline_graph.TaskNode`
481 Task definition structure.
482 limited_butler : `~lsst.daf.butler.LimitedButler`
483 Butler to use for dataset I/O.
484 quantum_id : `uuid.UUID` or `None`, optional
485 ID of the quantum being executed.
487 Returns
488 -------
489 flags : `QuantumSuccessCaveats`
490 Flags that describe qualified successes.
491 ids_put : list[ `uuid.UUID` ]
492 Record of all the dataset IDs that were written by this quantum
493 being executed.
494 metrics : `lsst.daf.butler.ButlerMetrics`
495 Butler metrics recorded for this quantum.
496 """
497 flags = QuantumSuccessCaveats.NO_CAVEATS
499 # Create a butler that operates in the context of a quantum
500 butlerQC = QuantumContext(limited_butler, quantum, resources=self._resources, quantum_id=quantum_id)
502 # Get the input and output references for the task
503 inputRefs, outputRefs = task_node.get_connections().buildDatasetRefs(quantum)
505 # Call task runQuantum() method.
506 try:
507 with limited_butler.record_metrics() as butler_metrics:
508 task.runQuantum(butlerQC, inputRefs, outputRefs)
509 except NoWorkFound as err:
510 # Not an error, just an early exit.
511 _LOG.info(
512 "Task '%s' on quantum %s exited early with no work found: %s.",
513 task_node.label,
514 quantum.dataId,
515 str(err),
516 )
517 flags |= err.FLAGS
518 except AnnotatedPartialOutputsError as caught:
519 error: BaseException
520 if caught.__cause__ is None:
521 _LOG.error(
522 "Incorrect use of AnnotatedPartialOutputsError: no chained exception found.",
523 task_node.label,
524 quantum.dataId,
525 )
526 error = caught
527 else:
528 error = caught.__cause__
529 if self._raise_on_partial_outputs:
530 # Note: this is a real edge case that required some
531 # experimentation: without 'from None' below, this raise would
532 # produce a "while one exception was being handled, another was
533 # raised" traceback involving AnnotatedPartialOutputsError.
534 # With the 'from None', we get just the error chained to it, as
535 # desired.
536 raise error from None
537 else:
538 _LOG.error(
539 "Task '%s' on quantum %s exited with partial outputs; "
540 "considering this a qualified success and proceeding.",
541 task_node.label,
542 quantum.dataId,
543 )
544 _LOG.error(error, exc_info=error)
545 flags |= caught.FLAGS
546 if not butlerQC.outputsPut:
547 flags |= QuantumSuccessCaveats.ALL_OUTPUTS_MISSING
548 if not butlerQC.outputsPut == butlerQC.allOutputs:
549 flags |= QuantumSuccessCaveats.ANY_OUTPUTS_MISSING
550 ids_put = [output[2] for output in butlerQC.outputsPut]
551 return flags, ids_put, butler_metrics
553 def _write_metadata(
554 self, quantum: Quantum, metadata: Any, task_node: TaskNode, /, limited_butler: LimitedButler
555 ) -> None:
556 # DatasetRef has to be in the Quantum outputs, can lookup by name
557 try:
558 [ref] = quantum.outputs[task_node.metadata_output.dataset_type_name]
559 except LookupError as exc:
560 raise InvalidQuantumError(
561 "Quantum outputs is missing metadata dataset type "
562 f"{task_node.metadata_output.dataset_type_name};"
563 " this could happen due to inconsistent options between QuantumGraph generation"
564 " and execution"
565 ) from exc
566 limited_butler.put(metadata, ref)
568 def _should_assume_exists(self, quantum: Quantum, ref: DatasetRef) -> bool | None:
569 """Report whether the given dataset can be assumed to exist because
570 some previous check reported that it did.
572 If this is `True` for a dataset does not in fact exist anymore, that's
573 an unexpected problem that we want to raise as an exception, and
574 definitely not a case where some predicted output just wasn't produced.
575 We can't always tell the difference, but in this case we can.
577 Parameters
578 ----------
579 quantum : `Quantum`
580 Quantum being processed.
581 ref : `lsst.daf.butler.DatasetRef`
582 Reference to the input dataset.
584 Returns
585 -------
586 exists : `bool` or `None`
587 `True` if this dataset is definitely an overall input, `False` if
588 some other quantum in the graph is expected to produce it, and
589 `None` if the answer could not be determined.
590 """
591 if quantum.datastore_records:
592 for datastore_record_data in quantum.datastore_records.values():
593 if ref.id in datastore_record_data.records:
594 return True
595 return False
596 return None