Coverage for python/lsst/pipe/base/quantum_graph/aggregator/_scanner.py: 25%
153 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-04 01:47 -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 QuantumSuccessCaveats
43from ..._task_metadata import TaskMetadata
44from ...pipeline_graph import PipelineGraph, TaskImportMode
45from ...quantum_provenance_graph import ExceptionInfo
46from ...resource_usage import QuantumResourceUsage
47from .._multiblock import Compressor
48from .._predicted import (
49 PredictedDatasetModel,
50 PredictedQuantumDatasetsModel,
51 PredictedQuantumGraphReader,
52)
53from ._communicators import ScannerCommunicator
54from ._structs import IngestRequest, ScanReport, ScanResult, ScanStatus
57@dataclasses.dataclass
58class Scanner:
59 """A helper class for the provenance aggregator that reads metadata and log
60 files and scans for which outputs exist.
61 """
63 predicted_path: str
64 """Path to the predicted quantum graph."""
66 butler_path: str
67 """Path or alias to the central butler repository."""
69 comms: ScannerCommunicator
70 """Communicator object for this worker."""
72 reader: PredictedQuantumGraphReader = dataclasses.field(init=False)
73 """Reader for the predicted quantum graph."""
75 qbb: QuantumBackedButler = dataclasses.field(init=False)
76 """A quantum-backed butler used for log and metadata reads and existence
77 checks for other outputs (when necessary).
78 """
80 compressor: Compressor | None = None
81 """Object used to compress JSON blocks.
83 This is `None` until a compression dictionary is received from the writer
84 process.
85 """
87 init_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = dataclasses.field(init=False)
88 """Dictionary mapping init quantum IDs to their predicted models."""
90 def __post_init__(self) -> None:
91 if self.comms.config.mock_storage_classes:
92 import lsst.pipe.base.tests.mocks # noqa: F401
93 self.comms.log.verbose("Reading from predicted quantum graph.")
94 self.reader = self.comms.enter(
95 PredictedQuantumGraphReader.open(self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT)
96 )
97 self.reader.read_dimension_data()
98 self.reader.read_init_quanta()
99 self.comms.log.verbose("Initializing quantum-backed butler.")
100 self.qbb = self.make_qbb(self.butler_path, self.reader.pipeline_graph)
101 self.init_quanta = {q.quantum_id: q for q in self.reader.components.init_quanta.root}
103 @staticmethod
104 def make_qbb(butler_config: str, pipeline_graph: PipelineGraph) -> QuantumBackedButler:
105 """Make quantum-backed butler that can operate on the outputs of the
106 quantum graph.
108 Parameters
109 ----------
110 butler_config : `str`
111 Path or alias for the central butler repository that shares storage
112 with the quantum-backed butler.
113 pipeline_graph : `..pipeline_graph.PipelineGraph`
114 Graph of tasks and dataset types.
116 Returns
117 -------
118 qbb : `lsst.daf.butler.QuantumBackedButler`
119 Quantum-backed butler. This does not have the datastore records
120 needed to read overall-inputs.
121 """
122 return QuantumBackedButler.from_predicted(
123 butler_config,
124 predicted_inputs=[],
125 predicted_outputs=[],
126 dimensions=pipeline_graph.universe,
127 # We don't need the datastore records in the QG because we're
128 # only going to read metadata and logs, and those are never
129 # overall inputs.
130 datastore_records={},
131 dataset_types={node.name: node.dataset_type for node in pipeline_graph.dataset_types.values()},
132 )
134 @property
135 def pipeline_graph(self) -> PipelineGraph:
136 """Graph of tasks and dataset types."""
137 return self.reader.pipeline_graph
139 @staticmethod
140 def run(predicted_path: str, butler_path: str, comms: ScannerCommunicator) -> None:
141 """Run the scanner.
143 Parameters
144 ----------
145 predicted_path : `str`
146 Path to the predicted quantum graph.
147 butler_path : `str`
148 Path or alias to the central butler repository.
149 comms : `ScannerCommunicator`
150 Communicator for the scanner.
152 Notes
153 -----
154 This method is designed to run as the ``target`` in
155 `WorkerContext.make_worker`.
156 """
157 with comms:
158 scanner = Scanner(predicted_path, butler_path, comms)
159 scanner.loop()
161 def loop(self) -> None:
162 """Run the main loop for the scanner."""
163 self.comms.log.info("Scan request loop beginning.")
164 for quantum_id in self.comms.poll():
165 if self.compressor is None and (cdict_data := self.comms.get_compression_dict()) is not None:
166 self.compressor = zstandard.ZstdCompressor(
167 self.comms.config.zstd_level, zstandard.ZstdCompressionDict(cdict_data)
168 )
169 self.scan_quantum(quantum_id)
171 def scan_dataset(self, predicted: PredictedDatasetModel) -> bool:
172 """Scan for a dataset's existence.
174 Parameters
175 ----------
176 predicted : `.PredictedDatasetModel`
177 Information about the dataset from the predicted graph.
179 Returns
180 -------
181 exists : `bool``
182 Whether the dataset exists
183 """
184 ref = self.reader.components.make_dataset_ref(predicted)
185 return self.qbb.stored(ref)
187 def scan_quantum(self, quantum_id: uuid.UUID) -> ScanResult:
188 """Scan for a quantum's completion and error status, and its output
189 datasets' existence.
191 Parameters
192 ----------
193 quantum_id : `uuid.UUID`
194 Unique ID for the quantum.
196 Returns
197 -------
198 result : `ScanResult`
199 Scan result struct.
200 """
201 if (predicted_quantum := self.init_quanta.get(quantum_id)) is not None:
202 result = ScanResult(predicted_quantum.quantum_id, status=ScanStatus.INIT)
203 self.comms.log.debug("Created init scan for %s (%s)", quantum_id, predicted_quantum.task_label)
204 else:
205 self.reader.read_quantum_datasets([quantum_id])
206 predicted_quantum = self.reader.components.quantum_datasets[quantum_id]
207 self.comms.log.debug(
208 "Scanning %s (%s@%s)",
209 quantum_id,
210 predicted_quantum.task_label,
211 predicted_quantum.data_coordinate,
212 )
213 result = ScanResult(predicted_quantum.quantum_id, ScanStatus.INCOMPLETE)
214 del self.reader.components.quantum_datasets[quantum_id]
215 log_id = self._read_and_compress_log(predicted_quantum, result)
216 if not self.comms.config.assume_complete and not result.log:
217 self.comms.log.debug("Abandoning scan for %s; no log dataset.", quantum_id)
218 result.status = ScanStatus.ABANDONED
219 self.comms.report_scan(ScanReport(result.quantum_id, result.status))
220 return result
221 metadata_id = self._read_and_compress_metadata(predicted_quantum, result)
222 if result.metadata:
223 result.status = ScanStatus.SUCCESSFUL
224 result.existing_outputs.add(metadata_id)
225 elif self.comms.config.assume_complete:
226 result.status = ScanStatus.FAILED
227 else:
228 # We found the log dataset, but no metadata; this means the
229 # quantum failed, but a retry might still happen that could
230 # turn it into a success if we can't yet assume the run is
231 # complete.
232 self.comms.log.debug("Abandoning scan for %s.", quantum_id)
233 result.status = ScanStatus.ABANDONED
234 self.comms.report_scan(ScanReport(result.quantum_id, result.status))
235 return result
236 if result.log:
237 result.existing_outputs.add(log_id)
238 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
239 if predicted_output.dataset_id not in result.existing_outputs and self.scan_dataset(
240 predicted_output
241 ):
242 result.existing_outputs.add(predicted_output.dataset_id)
243 to_ingest = self._make_ingest_request(predicted_quantum, result)
244 self.comms.report_scan(ScanReport(result.quantum_id, result.status))
245 assert result.status is not ScanStatus.INCOMPLETE
246 assert result.status is not ScanStatus.ABANDONED
247 if self.comms.config.output_path is not None:
248 self.comms.request_write(result)
249 self.comms.request_ingest(to_ingest)
250 self.comms.log.debug("Finished scan for %s.", quantum_id)
251 return result
253 def _make_ingest_request(
254 self, predicted_quantum: PredictedQuantumDatasetsModel, result: ScanResult
255 ) -> IngestRequest:
256 """Make an ingest request from a quantum scan.
258 Parameters
259 ----------
260 predicted_quantum : `PredictedQuantumDatasetsModel`
261 Information about the predicted quantum.
262 result : `ScanResult`
263 Result of a quantum scan.
265 Returns
266 -------
267 ingest_request : `IngestRequest`
268 A request to be sent to the ingester.
269 """
270 predicted_outputs_by_id = {
271 d.dataset_id: d for d in itertools.chain.from_iterable(predicted_quantum.outputs.values())
272 }
273 to_ingest_predicted: list[PredictedDatasetModel] = []
274 to_ingest_refs: list[DatasetRef] = []
275 for dataset_id in result.existing_outputs:
276 predicted_output = predicted_outputs_by_id[dataset_id]
277 to_ingest_predicted.append(predicted_output)
278 to_ingest_refs.append(self.reader.components.make_dataset_ref(predicted_output))
279 to_ingest_records = self.qbb._datastore.export_predicted_records(to_ingest_refs)
280 return IngestRequest(result.quantum_id, to_ingest_predicted, to_ingest_records)
282 def _read_and_compress_metadata(
283 self, predicted_quantum: PredictedQuantumDatasetsModel, result: ScanResult
284 ) -> uuid.UUID:
285 """Attempt to read the metadata dataset for a quantum to extract
286 provenance information from it.
288 Parameters
289 ----------
290 predicted_quantum : `PredictedQuantumDatasetsModel`
291 Information about the predicted quantum.
292 result : `ScanResult`
293 Result object to be modified in-place.
295 Returns
296 -------
297 dataset_id : `uuid.UUID`
298 UUID of the metadata dataset.
299 """
300 assert not result.metadata, "We shouldn't be scanning again if we already read the metadata."
301 (predicted_dataset,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
302 ref = self.reader.components.make_dataset_ref(predicted_dataset)
303 try:
304 # This assumes QBB metadata writes are atomic, which should be the
305 # case. If it's not we'll probably get pydantic validation errors
306 # here.
307 content: TaskMetadata = self.qbb.get(ref, storageClass="TaskMetadata")
308 except FileNotFoundError:
309 if not self.comms.config.assume_complete:
310 return ref.id
311 else:
312 try:
313 # Int conversion guards against spurious conversion to
314 # float that can apparently sometimes happen in
315 # TaskMetadata.
316 result.caveats = QuantumSuccessCaveats(int(content["quantum"]["caveats"]))
317 except LookupError:
318 pass
319 try:
320 result.exception = ExceptionInfo._from_metadata(
321 content[predicted_quantum.task_label]["failure"]
322 )
323 except LookupError:
324 pass
325 try:
326 result.existing_outputs = {
327 uuid.UUID(id_str) for id_str in ensure_iterable(content["quantum"].getArray("outputs"))
328 }
329 except LookupError:
330 pass
331 result.resource_usage = QuantumResourceUsage.from_task_metadata(content)
332 result.metadata = content.model_dump_json().encode()
333 if self.compressor is not None:
334 result.metadata = self.compressor.compress(result.metadata)
335 result.is_compressed = True
336 return ref.id
338 def _read_and_compress_log(
339 self, predicted_quantum: PredictedQuantumDatasetsModel, result: ScanResult
340 ) -> uuid.UUID:
341 """Attempt to read the log dataset for a quantum to test for the
342 quantum's completion (the log is always written last) and aggregate
343 the log content in the provenance quantum graph.
345 Parameters
346 ----------
347 predicted_quantum : `PredictedQuantumDatasetsModel`
348 Information about the predicted quantum.
349 result : `ScanResult`
350 Result object to be modified in-place.
352 Returns
353 -------
354 dataset_id : `uuid.UUID`
355 UUID of the log dataset.
356 """
357 (predicted_dataset,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
358 ref = self.reader.components.make_dataset_ref(predicted_dataset)
359 try:
360 # This assumes QBB log writes are atomic, which should be the case.
361 # If it's not we'll probably get pydantic validation errors here.
362 content: ButlerLogRecords = self.qbb.get(ref)
363 except FileNotFoundError:
364 if not self.comms.config.assume_complete:
365 return ref.id
366 else:
367 result.log = content.model_dump_json().encode()
368 if self.compressor is not None:
369 result.log = self.compressor.compress(result.log)
370 result.is_compressed = True
371 return ref.id