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