Coverage for python/lsst/pipe/base/quantum_graph/aggregator/_scanner.py: 21%
192 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/>.
28from __future__ import annotations
30__all__ = ("Scanner",)
32import dataclasses
33import itertools
34import uuid
36import zstandard
38from lsst.daf.butler import ButlerLogRecords, DatasetRef, QuantumBackedButler
39from lsst.utils.iteration import ensure_iterable
41from ... import automatic_connection_constants as acc
42from ..._status import ExceptionInfo, QuantumAttemptStatus, QuantumSuccessCaveats
43from ..._task_metadata import TaskMetadata
44from ...log_capture import _ExecutionLogRecordsExtra
45from ...pipeline_graph import PipelineGraph, TaskImportMode
46from ...resource_usage import QuantumResourceUsage
47from .._multiblock import Compressor
48from .._predicted import (
49 PredictedDatasetModel,
50 PredictedQuantumDatasetsModel,
51 PredictedQuantumGraphReader,
52)
53from .._provenance import ProvenanceQuantumAttemptModel
54from ._communicators import ScannerCommunicator
55from ._structs import IngestRequest, ScanReport, ScanResult, ScanStatus
58@dataclasses.dataclass
59class Scanner:
60 """A helper class for the provenance aggregator that reads metadata and log
61 files and scans for which outputs exist.
62 """
64 predicted_path: str
65 """Path to the predicted quantum graph."""
67 butler_path: str
68 """Path or alias to the central butler repository."""
70 comms: ScannerCommunicator
71 """Communicator object for this worker."""
73 reader: PredictedQuantumGraphReader = dataclasses.field(init=False)
74 """Reader for the predicted quantum graph."""
76 qbb: QuantumBackedButler = dataclasses.field(init=False)
77 """A quantum-backed butler used for log and metadata reads and existence
78 checks for other outputs (when necessary).
79 """
81 compressor: Compressor | None = None
82 """Object used to compress JSON blocks.
84 This is `None` until a compression dictionary is received from the writer
85 process.
86 """
88 init_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = dataclasses.field(init=False)
89 """Dictionary mapping init quantum IDs to their predicted models."""
91 def __post_init__(self) -> None:
92 if self.comms.config.mock_storage_classes:
93 import lsst.pipe.base.tests.mocks # noqa: F401
94 self.comms.log.verbose("Reading from predicted quantum graph.")
95 self.reader = self.comms.enter(
96 PredictedQuantumGraphReader.open(self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT)
97 )
98 self.reader.read_dimension_data()
99 self.reader.read_init_quanta()
100 self.comms.log.verbose("Initializing quantum-backed butler.")
101 self.qbb = self.make_qbb(self.butler_path, self.reader.pipeline_graph)
102 self.init_quanta = {q.quantum_id: q for q in self.reader.components.init_quanta.root}
104 @staticmethod
105 def make_qbb(butler_config: str, pipeline_graph: PipelineGraph) -> QuantumBackedButler:
106 """Make quantum-backed butler that can operate on the outputs of the
107 quantum graph.
109 Parameters
110 ----------
111 butler_config : `str`
112 Path or alias for the central butler repository that shares storage
113 with the quantum-backed butler.
114 pipeline_graph : `..pipeline_graph.PipelineGraph`
115 Graph of tasks and dataset types.
117 Returns
118 -------
119 qbb : `lsst.daf.butler.QuantumBackedButler`
120 Quantum-backed butler. This does not have the datastore records
121 needed to read overall-inputs.
122 """
123 return QuantumBackedButler.from_predicted(
124 butler_config,
125 predicted_inputs=[],
126 predicted_outputs=[],
127 dimensions=pipeline_graph.universe,
128 # We don't need the datastore records in the QG because we're
129 # only going to read metadata and logs, and those are never
130 # overall inputs.
131 datastore_records={},
132 dataset_types={node.name: node.dataset_type for node in pipeline_graph.dataset_types.values()},
133 )
135 @property
136 def pipeline_graph(self) -> PipelineGraph:
137 """Graph of tasks and dataset types."""
138 return self.reader.pipeline_graph
140 @staticmethod
141 def run(predicted_path: str, butler_path: str, comms: ScannerCommunicator) -> None:
142 """Run the scanner.
144 Parameters
145 ----------
146 predicted_path : `str`
147 Path to the predicted quantum graph.
148 butler_path : `str`
149 Path or alias to the central butler repository.
150 comms : `ScannerCommunicator`
151 Communicator for the scanner.
153 Notes
154 -----
155 This method is designed to run as the ``target`` in
156 `WorkerContext.make_worker`.
157 """
158 with comms:
159 scanner = Scanner(predicted_path, butler_path, comms)
160 scanner.loop()
162 def loop(self) -> None:
163 """Run the main loop for the scanner."""
164 self.comms.log.info("Scan request loop beginning.")
165 for quantum_id in self.comms.poll():
166 if self.compressor is None and (cdict_data := self.comms.get_compression_dict()) is not None:
167 self.compressor = zstandard.ZstdCompressor(
168 self.comms.config.zstd_level, zstandard.ZstdCompressionDict(cdict_data)
169 )
170 self.scan_quantum(quantum_id)
172 def scan_dataset(self, predicted: PredictedDatasetModel) -> bool:
173 """Scan for a dataset's existence.
175 Parameters
176 ----------
177 predicted : `.PredictedDatasetModel`
178 Information about the dataset from the predicted graph.
180 Returns
181 -------
182 exists : `bool``
183 Whether the dataset exists.
184 """
185 ref = self.reader.components.make_dataset_ref(predicted)
186 return self.qbb.stored(ref)
188 def scan_quantum(self, quantum_id: uuid.UUID) -> ScanResult:
189 """Scan for a quantum's completion and error status, and its output
190 datasets' existence.
192 Parameters
193 ----------
194 quantum_id : `uuid.UUID`
195 Unique ID for the quantum.
197 Returns
198 -------
199 result : `ScanResult`
200 Scan result struct.
201 """
202 if (predicted_quantum := self.init_quanta.get(quantum_id)) is not None:
203 result = ScanResult(predicted_quantum.quantum_id, status=ScanStatus.INIT)
204 self.comms.log.debug("Created init scan for %s (%s)", quantum_id, predicted_quantum.task_label)
205 else:
206 self.reader.read_quantum_datasets([quantum_id])
207 predicted_quantum = self.reader.components.quantum_datasets[quantum_id]
208 self.comms.log.debug(
209 "Scanning %s (%s@%s)",
210 quantum_id,
211 predicted_quantum.task_label,
212 predicted_quantum.data_coordinate,
213 )
214 result = ScanResult(predicted_quantum.quantum_id, ScanStatus.INCOMPLETE)
215 del self.reader.components.quantum_datasets[quantum_id]
216 last_attempt = ProvenanceQuantumAttemptModel()
217 if not self._read_log(predicted_quantum, result, last_attempt):
218 self.comms.log.debug("Abandoning scan for %s; no log dataset.", quantum_id)
219 self.comms.report_scan(ScanReport(result.quantum_id, result.status))
220 return result
221 if not self._read_metadata(predicted_quantum, result, last_attempt):
222 # We found the log dataset, but no metadata; this means the
223 # quantum failed, but a retry might still happen that could
224 # turn it into a success if we can't yet assume the run is
225 # complete.
226 self.comms.log.debug("Abandoning scan for %s.", quantum_id)
227 self.comms.report_scan(ScanReport(result.quantum_id, result.status))
228 return result
229 last_attempt.attempt = len(result.attempts)
230 result.attempts.append(last_attempt)
231 assert result.status is not ScanStatus.INCOMPLETE
232 assert result.status is not ScanStatus.ABANDONED
233 assert result.log_model is not None, "Only set to None after converting to JSON."
234 assert result.metadata_model is not None, "Only set to None after converting to JSON."
236 if len(result.log_model.attempts) < len(result.attempts):
237 # Logs were not found for this attempt; must have been a hard error
238 # that kept the `finally` block from running or otherwise
239 # interrupted the writing of the logs.
240 result.log_model.attempts.append(None)
241 if result.status is ScanStatus.SUCCESSFUL:
242 # But we found the metadata! Either that hard error happened
243 # at a very unlucky time (in between those two writes), or
244 # something even weirder happened.
245 result.attempts[-1].status = QuantumAttemptStatus.LOGS_MISSING
246 else:
247 result.attempts[-1].status = QuantumAttemptStatus.FAILED
248 if len(result.metadata_model.attempts) < len(result.attempts):
249 # Metadata missing usually just means a failure. In any case, the
250 # status will already be correct, either because it was set to a
251 # failure when we read the logs, or left at UNKNOWN if there were
252 # no logs. Note that scanners never process BLOCKED quanta at all.
253 result.metadata_model.attempts.append(None)
254 assert len(result.log_model.attempts) == len(result.attempts) or len(
255 result.metadata_model.attempts
256 ) == len(result.attempts), (
257 "The only way we can add more than one quantum attempt is by "
258 "extracting info stored with the logs, and that always appends "
259 "a log attempt and a metadata attempt, so this must be a bug in "
260 "the scanner."
261 )
262 # Now that we're done gathering the log and metadata information into
263 # models, dump them to JSON and delete the originals.
264 result.log_content = result.log_model.model_dump_json().encode()
265 result.log_model = None
266 result.metadata_content = result.metadata_model.model_dump_json().encode()
267 result.metadata_model = None
268 if self.compressor is not None:
269 if result.log_content is not None:
270 result.log_content = self.compressor.compress(result.log_content)
271 if result.metadata_content is not None:
272 result.metadata_content = self.compressor.compress(result.metadata_content)
273 result.is_compressed = True
274 # Scan for output dataset existence, skipping any the metadata reported
275 # as having been definitively written, as well as and the metadata and
276 # logs themselves (since we just checked those).
277 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
278 if predicted_output.dataset_id not in result.existing_outputs and self.scan_dataset(
279 predicted_output
280 ):
281 result.existing_outputs.add(predicted_output.dataset_id)
282 to_ingest = self._make_ingest_request(predicted_quantum, result)
283 self.comms.report_scan(ScanReport(result.quantum_id, result.status))
284 if self.comms.config.output_path is not None:
285 self.comms.request_write(result)
286 self.comms.request_ingest(to_ingest)
287 self.comms.log.debug("Finished scan for %s.", quantum_id)
288 return result
290 def _make_ingest_request(
291 self, predicted_quantum: PredictedQuantumDatasetsModel, result: ScanResult
292 ) -> IngestRequest:
293 """Make an ingest request from a quantum scan.
295 Parameters
296 ----------
297 predicted_quantum : `PredictedQuantumDatasetsModel`
298 Information about the predicted quantum.
299 result : `ScanResult`
300 Result of a quantum scan.
302 Returns
303 -------
304 ingest_request : `IngestRequest`
305 A request to be sent to the ingester.
306 """
307 predicted_outputs_by_id = {
308 d.dataset_id: d for d in itertools.chain.from_iterable(predicted_quantum.outputs.values())
309 }
310 to_ingest_predicted: list[PredictedDatasetModel] = []
311 to_ingest_refs: list[DatasetRef] = []
312 for dataset_id in result.existing_outputs:
313 predicted_output = predicted_outputs_by_id[dataset_id]
314 to_ingest_predicted.append(predicted_output)
315 to_ingest_refs.append(self.reader.components.make_dataset_ref(predicted_output))
316 to_ingest_records = self.qbb._datastore.export_predicted_records(to_ingest_refs)
317 return IngestRequest(result.quantum_id, to_ingest_predicted, to_ingest_records)
319 def _read_metadata(
320 self,
321 predicted_quantum: PredictedQuantumDatasetsModel,
322 result: ScanResult,
323 last_attempt: ProvenanceQuantumAttemptModel,
324 ) -> bool:
325 """Attempt to read the metadata dataset for a quantum to extract
326 provenance information from it.
328 Parameters
329 ----------
330 predicted_quantum : `PredictedQuantumDatasetsModel`
331 Information about the predicted quantum.
332 result : `ScanResult`
333 Result object to be modified in-place.
334 last_attempt : `ScanningProvenanceQuantumAttemptModel`
335 Structure to fill in with information about the last attempt to
336 run this quantum.
338 Returns
339 -------
340 complete : `bool`
341 Whether the quantum is complete.
342 """
343 (predicted_dataset,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
344 ref = self.reader.components.make_dataset_ref(predicted_dataset)
345 try:
346 # This assumes QBB metadata writes are atomic, which should be the
347 # case. If it's not we'll probably get pydantic validation errors
348 # here.
349 metadata: TaskMetadata = self.qbb.get(ref, storageClass="TaskMetadata")
350 except FileNotFoundError:
351 if self.comms.config.assume_complete:
352 result.status = ScanStatus.FAILED
353 else:
354 result.status = ScanStatus.ABANDONED
355 return False
356 else:
357 result.status = ScanStatus.SUCCESSFUL
358 result.existing_outputs.add(ref.id)
359 last_attempt.status = QuantumAttemptStatus.SUCCESSFUL
360 try:
361 # Int conversion guards against spurious conversion to
362 # float that can apparently sometimes happen in
363 # TaskMetadata.
364 last_attempt.caveats = QuantumSuccessCaveats(int(metadata["quantum"]["caveats"]))
365 except LookupError:
366 pass
367 try:
368 last_attempt.exception = ExceptionInfo._from_metadata(
369 metadata[predicted_quantum.task_label]["failure"]
370 )
371 except LookupError:
372 pass
373 try:
374 result.existing_outputs.update(
375 uuid.UUID(id_str) for id_str in ensure_iterable(metadata["quantum"].getArray("outputs"))
376 )
377 except LookupError:
378 pass
379 last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)
380 assert result.metadata_model is not None, "Only set to None after converting to JSON."
381 result.metadata_model.attempts.append(metadata)
382 return True
384 def _read_log(
385 self,
386 predicted_quantum: PredictedQuantumDatasetsModel,
387 result: ScanResult,
388 last_attempt: ProvenanceQuantumAttemptModel,
389 ) -> bool:
390 """Attempt to read the log dataset for a quantum to test for the
391 quantum's completion (the log is always written last) and aggregate
392 the log content in the provenance quantum graph.
394 Parameters
395 ----------
396 predicted_quantum : `PredictedQuantumDatasetsModel`
397 Information about the predicted quantum.
398 result : `ScanResult`
399 Result object to be modified in-place.
400 last_attempt : `ScanningProvenanceQuantumAttemptModel`
401 Structure to fill in with information about the last attempt to
402 run this quantum.
404 Returns
405 -------
406 complete : `bool`
407 Whether the quantum is complete.
408 """
409 (predicted_dataset,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
410 ref = self.reader.components.make_dataset_ref(predicted_dataset)
411 try:
412 # This assumes QBB log writes are atomic, which should be the case.
413 # If it's not we'll probably get pydantic validation errors here.
414 log_records: ButlerLogRecords = self.qbb.get(ref)
415 except FileNotFoundError:
416 if self.comms.config.assume_complete:
417 result.status = ScanStatus.FAILED
418 else:
419 result.status = ScanStatus.ABANDONED
420 return False
421 else:
422 # Set the attempt's run status to FAILED, since the default is
423 # UNKNOWN (i.e. logs *and* metadata are missing) and we now know
424 # the logs exist. This will usually get replaced by SUCCESSFUL
425 # when we look for metadata next.
426 last_attempt.status = QuantumAttemptStatus.FAILED
427 result.existing_outputs.add(ref.id)
428 if log_records.extra:
429 log_extra = _ExecutionLogRecordsExtra.model_validate(log_records.extra)
430 self._extract_from_log_extra(log_extra, result, last_attempt=last_attempt)
431 assert result.log_model is not None, "Only set to None after converting to JSON."
432 result.log_model.attempts.append(list(log_records))
433 return True
435 def _extract_from_log_extra(
436 self,
437 log_extra: _ExecutionLogRecordsExtra,
438 result: ScanResult,
439 last_attempt: ProvenanceQuantumAttemptModel | None,
440 ) -> None:
441 for previous_attempt_log_extra in log_extra.previous_attempts:
442 self._extract_from_log_extra(previous_attempt_log_extra, result, last_attempt=None)
443 quantum_attempt: ProvenanceQuantumAttemptModel
444 if last_attempt is None:
445 # This is not the last attempt, so it must be a failure.
446 quantum_attempt = ProvenanceQuantumAttemptModel(
447 attempt=len(result.attempts), status=QuantumAttemptStatus.FAILED
448 )
449 # We also need to get the logs from this extra provenance, since
450 # they won't be the main section of the log records.
451 assert result.log_model is not None, "Only set to None after converting to JSON."
452 result.log_model.attempts.append(log_extra.logs)
453 # The special last attempt is only appended after we attempt to
454 # read metadata later, but we have to append this one now.
455 result.attempts.append(quantum_attempt)
456 else:
457 assert not log_extra.logs, "Logs for the last attempt should not be stored in the extra JSON."
458 quantum_attempt = last_attempt
459 if log_extra.exception is not None or log_extra.metadata is not None or last_attempt is None:
460 # We won't be getting a separate metadata dataset, so anything we
461 # might get from the metadata has to come from this extra
462 # provenance in the logs.
463 quantum_attempt.exception = log_extra.exception
464 assert result.metadata_model is not None, "Only set to None after converting to JSON."
465 if log_extra.metadata is not None:
466 quantum_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(log_extra.metadata)
467 result.metadata_model.attempts.append(log_extra.metadata)
468 else:
469 result.metadata_model.attempts.append(None)
470 # Regardless of whether this is the last attempt or not, we can only
471 # get the previous_process_quanta from the log extra.
472 quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)