Coverage for python/lsst/pipe/base/quantum_graph/aggregator/

1# This file is part of pipe_base.

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.

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 ProvenanceInitQuantumModel, ProvenanceQuantumAttemptModel, ProvenanceQuantumModel

54from ._communicators import ScannerCommunicator

55from ._structs import IngestRequest, InProgressScan, ScanReport, ScanStatus, WriteRequest

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}

103

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.

108

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.

116

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 )

134

135 @property

136 def pipeline_graph(self) -> PipelineGraph:

137 """Graph of tasks and dataset types."""

138 return self.reader.pipeline_graph

139

140 @staticmethod

141 def run(predicted_path: str, butler_path: str, comms: ScannerCommunicator) -> None:

142 """Run the scanner.

143

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.

152

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()

161

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)

171

172 def scan_dataset(self, predicted: PredictedDatasetModel) -> bool:

173 """Scan for a dataset's existence.

174

175 Parameters

176 ----------

177 predicted : `.PredictedDatasetModel`

178 Information about the dataset from the predicted graph.

179

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)

187

188 def scan_quantum(self, quantum_id: uuid.UUID) -> InProgressScan:

189 """Scan for a quantum's completion and error status, and its output

190 datasets' existence.

191

192 Parameters

193 ----------

194 quantum_id : `uuid.UUID`

195 Unique ID for the quantum.

196

197 Returns

198 -------

199 result : `InProgressScan`

200 Scan result struct.

201 """

202 if (predicted_quantum := self.init_quanta.get(quantum_id)) is not None:

203 result = InProgressScan(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 = InProgressScan(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

234 if len(result.logs.attempts) < len(result.attempts):

235 # Logs were not found for this attempt; must have been a hard error

236 # that kept the `finally` block from running or otherwise

237 # interrupted the writing of the logs.

238 result.logs.attempts.append(None)

239 if result.status is ScanStatus.SUCCESSFUL:

240 # But we found the metadata! Either that hard error happened

241 # at a very unlucky time (in between those two writes), or

242 # something even weirder happened.

243 result.attempts[-1].status = QuantumAttemptStatus.LOGS_MISSING

244 else:

245 result.attempts[-1].status = QuantumAttemptStatus.FAILED

246 if len(result.metadata.attempts) < len(result.attempts):

247 # Metadata missing usually just means a failure. In any case, the

248 # status will already be correct, either because it was set to a

249 # failure when we read the logs, or left at UNKNOWN if there were

250 # no logs. Note that scanners never process BLOCKED quanta at all.

251 result.metadata.attempts.append(None)

252 assert len(result.logs.attempts) == len(result.attempts) or len(result.metadata.attempts) == len(

253 result.attempts

254 ), (

255 "The only way we can add more than one quantum attempt is by "

256 "extracting info stored with the logs, and that always appends "

257 "a log attempt and a metadata attempt, so this must be a bug in "

258 "the scanner."

259 )

260 # Scan for output dataset existence, skipping any the metadata reported

261 # on as well as and the metadata and logs themselves (since we just

262 # checked those).

263 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):

264 if predicted_output.dataset_id not in result.outputs:

265 result.outputs[predicted_output.dataset_id] = self.scan_dataset(predicted_output)

266 to_ingest = self._make_ingest_request(predicted_quantum, result)

267 if self.comms.config.output_path is not None:

268 to_write = self._make_write_request(predicted_quantum, result)

269 self.comms.request_write(to_write)

270 self.comms.request_ingest(to_ingest)

271 self.comms.report_scan(ScanReport(result.quantum_id, result.status))

272 self.comms.log.debug("Finished scan for %s.", quantum_id)

273 return result

274

275 def _make_ingest_request(

276 self, predicted_quantum: PredictedQuantumDatasetsModel, result: InProgressScan

277 ) -> IngestRequest:

278 """Make an ingest request from a quantum scan.

279

280 Parameters

281 ----------

282 predicted_quantum : `PredictedQuantumDatasetsModel`

283 Information about the predicted quantum.

284 result : `InProgressScan`

285 Result of a quantum scan.

286

287 Returns

288 -------

289 ingest_request : `IngestRequest`

290 A request to be sent to the ingester.

291 """

292 predicted_outputs_by_id = {

293 d.dataset_id: d for d in itertools.chain.from_iterable(predicted_quantum.outputs.values())

294 }

295 to_ingest_predicted: list[PredictedDatasetModel] = []

296 to_ingest_refs: list[DatasetRef] = []

297 for dataset_id, was_produced in result.outputs.items():

298 if was_produced:

299 predicted_output = predicted_outputs_by_id[dataset_id]

300 to_ingest_predicted.append(predicted_output)

301 to_ingest_refs.append(self.reader.components.make_dataset_ref(predicted_output))

302 to_ingest_records = self.qbb._datastore.export_predicted_records(to_ingest_refs)

303 return IngestRequest(result.quantum_id, to_ingest_predicted, to_ingest_records)

304

305 def _make_write_request(

306 self, predicted_quantum: PredictedQuantumDatasetsModel, result: InProgressScan

307 ) -> WriteRequest:

308 """Make a write request from a quantum scan.

309

310 Parameters

311 ----------

312 predicted_quantum : `PredictedQuantumDatasetsModel`

313 Information about the predicted quantum.

314 result : `InProgressScan`

315 Result of a quantum scan.

316

317 Returns

318 -------

319 write_request : `WriteRequest`

320 A request to be sent to the writer.

321 """

322 quantum: ProvenanceInitQuantumModel | ProvenanceQuantumModel

323 if result.status is ScanStatus.INIT:

324 quantum = ProvenanceInitQuantumModel.from_predicted(predicted_quantum)

325 else:

326 quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum)

327 quantum.attempts = result.attempts

328 request = WriteRequest(

329 result.quantum_id,

330 result.status,

331 existing_outputs={

332 dataset_id for dataset_id, was_produced in result.outputs.items() if was_produced

333 },

334 quantum=quantum.model_dump_json().encode(),

335 logs=result.logs.model_dump_json().encode() if result.logs.attempts else b"",

336 metadata=result.metadata.model_dump_json().encode() if result.metadata.attempts else b"",

337 )

338 if self.compressor is not None:

339 request.quantum = self.compressor.compress(request.quantum)

340 request.logs = self.compressor.compress(request.logs) if request.logs else b""

341 request.metadata = self.compressor.compress(request.metadata) if request.metadata else b""

342 request.is_compressed = True

343 return request

344

345 def _read_metadata(

346 self,

347 predicted_quantum: PredictedQuantumDatasetsModel,

348 result: InProgressScan,

349 last_attempt: ProvenanceQuantumAttemptModel,

350 ) -> bool:

351 """Attempt to read the metadata dataset for a quantum to extract

352 provenance information from it.

353

354 Parameters

355 ----------

356 predicted_quantum : `PredictedQuantumDatasetsModel`

357 Information about the predicted quantum.

358 result : `InProgressScan`

359 Result object to be modified in-place.

360 last_attempt : `ScanningProvenanceQuantumAttemptModel`

361 Structure to fill in with information about the last attempt to

362 run this quantum.

363

364 Returns

365 -------

366 complete : `bool`

367 Whether the quantum is complete.

368 """

369 (predicted_dataset,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]

370 ref = self.reader.components.make_dataset_ref(predicted_dataset)

371 try:

372 # This assumes QBB metadata writes are atomic, which should be the

373 # case. If it's not we'll probably get pydantic validation errors

374 # here.

375 metadata: TaskMetadata = self.qbb.get(ref, storageClass="TaskMetadata")

376 except FileNotFoundError:

377 result.outputs[ref.id] = False

378 if self.comms.config.assume_complete:

379 result.status = ScanStatus.FAILED

380 else:

381 result.status = ScanStatus.ABANDONED

382 return False

383 else:

384 result.status = ScanStatus.SUCCESSFUL

385 result.outputs[ref.id] = True

386 last_attempt.status = QuantumAttemptStatus.SUCCESSFUL

387 try:

388 # Int conversion guards against spurious conversion to

389 # float that can apparently sometimes happen in

390 # TaskMetadata.

391 last_attempt.caveats = QuantumSuccessCaveats(int(metadata["quantum"]["caveats"]))

392 except LookupError:

393 pass

394 try:

395 last_attempt.exception = ExceptionInfo._from_metadata(

396 metadata[predicted_quantum.task_label]["failure"]

397 )

398 except LookupError:

399 pass

400 try:

401 for id_str in ensure_iterable(metadata["quantum"].getArray("outputs")):

402 result.outputs[uuid.UUID(id_str)]

403 except LookupError:

404 pass

405 else:

406 # If the metadata told us what it wrote, anything not in that

407 # list was not written.

408 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):

409 result.outputs.setdefault(predicted_output.dataset_id, False)

410 last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)

411 result.metadata.attempts.append(metadata)

412 return True

413

414 def _read_log(

415 self,

416 predicted_quantum: PredictedQuantumDatasetsModel,

417 result: InProgressScan,

418 last_attempt: ProvenanceQuantumAttemptModel,

419 ) -> bool:

420 """Attempt to read the log dataset for a quantum to test for the

421 quantum's completion (the log is always written last) and aggregate

422 the log content in the provenance quantum graph.

423

424 Parameters

425 ----------

426 predicted_quantum : `PredictedQuantumDatasetsModel`

427 Information about the predicted quantum.

428 result : `InProgressScan`

429 Result object to be modified in-place.

430 last_attempt : `ScanningProvenanceQuantumAttemptModel`

431 Structure to fill in with information about the last attempt to

432 run this quantum.

433

434 Returns

435 -------

436 complete : `bool`

437 Whether the quantum is complete.

438 """

439 (predicted_dataset,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]

440 ref = self.reader.components.make_dataset_ref(predicted_dataset)

441 try:

442 # This assumes QBB log writes are atomic, which should be the case.

443 # If it's not we'll probably get pydantic validation errors here.

444 log_records: ButlerLogRecords = self.qbb.get(ref)

445 except FileNotFoundError:

446 result.outputs[ref.id] = False

447 if self.comms.config.assume_complete:

448 result.status = ScanStatus.FAILED

449 else:

450 result.status = ScanStatus.ABANDONED

451 return False

452 else:

453 # Set the attempt's run status to FAILED, since the default is

454 # UNKNOWN (i.e. logs *and* metadata are missing) and we now know

455 # the logs exist. This will usually get replaced by SUCCESSFUL

456 # when we look for metadata next.

457 last_attempt.status = QuantumAttemptStatus.FAILED

458 result.outputs[ref.id] = True

459 if log_records.extra:

460 log_extra = _ExecutionLogRecordsExtra.model_validate(log_records.extra)

461 self._extract_from_log_extra(log_extra, result, last_attempt=last_attempt)

462 result.logs.attempts.append(list(log_records))

463 return True

464

465 def _extract_from_log_extra(

466 self,

467 log_extra: _ExecutionLogRecordsExtra,

468 result: InProgressScan,

469 last_attempt: ProvenanceQuantumAttemptModel | None,

470 ) -> None:

471 for previous_attempt_log_extra in log_extra.previous_attempts:

472 self._extract_from_log_extra(previous_attempt_log_extra, result, last_attempt=None)

473 quantum_attempt: ProvenanceQuantumAttemptModel

474 if last_attempt is None:

475 # This is not the last attempt, so it must be a failure.

476 quantum_attempt = ProvenanceQuantumAttemptModel(

477 attempt=len(result.attempts), status=QuantumAttemptStatus.FAILED

478 )

479 # We also need to get the logs from this extra provenance, since

480 # they won't be the main section of the log records.

481 result.logs.attempts.append(log_extra.logs)

482 # The special last attempt is only appended after we attempt to

483 # read metadata later, but we have to append this one now.

484 result.attempts.append(quantum_attempt)

485 else:

486 assert not log_extra.logs, "Logs for the last attempt should not be stored in the extra JSON."

487 quantum_attempt = last_attempt

488 if log_extra.exception is not None or log_extra.metadata is not None or last_attempt is None:

489 # We won't be getting a separate metadata dataset, so anything we

490 # might get from the metadata has to come from this extra

491 # provenance in the logs.

492 quantum_attempt.exception = log_extra.exception

493 if log_extra.metadata is not None:

494 quantum_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(log_extra.metadata)

495 result.metadata.attempts.append(log_extra.metadata)

496 else:

497 result.metadata.attempts.append(None)

498 # Regardless of whether this is the last attempt or not, we can only

499 # get the previous_process_quanta from the log extra.

500 quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)

Coverage for python/lsst/pipe/base/quantum_graph/aggregator/_scanner.py: 20%

195 statements