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

35from contextlib import AbstractContextManager

36from typing import Any, Literal, Self

38import zstandard

40from lsst.daf.butler import ButlerLogRecords, DatasetRef, QuantumBackedButler

41from lsst.utils.iteration import ensure_iterable

43from ... import automatic_connection_constants as acc

44from ..._status import ExceptionInfo, QuantumAttemptStatus, QuantumSuccessCaveats

45from ..._task_metadata import TaskMetadata

46from ...log_capture import _ExecutionLogRecordsExtra

47from ...pipeline_graph import PipelineGraph, TaskImportMode

48from ...resource_usage import QuantumResourceUsage

49from .._multiblock import Compressor

50from .._predicted import (

51 PredictedDatasetModel,

52 PredictedQuantumDatasetsModel,

53 PredictedQuantumGraphReader,

54)

55from .._provenance import ProvenanceInitQuantumModel, ProvenanceQuantumAttemptModel, ProvenanceQuantumModel

56from ._communicators import ScannerCommunicator

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

60@dataclasses.dataclass

61class Scanner(AbstractContextManager):

62 """A helper class for the provenance aggregator that reads metadata and log

63 files and scans for which outputs exist.

64 """

66 predicted_path: str

67 """Path to the predicted quantum graph."""

69 butler_path: str

70 """Path or alias to the central butler repository."""

72 comms: ScannerCommunicator

73 """Communicator object for this worker."""

75 reader: PredictedQuantumGraphReader = dataclasses.field(init=False)

76 """Reader for the predicted quantum graph."""

78 qbb: QuantumBackedButler = dataclasses.field(init=False)

79 """A quantum-backed butler used for log and metadata reads and existence

80 checks for other outputs (when necessary).

81 """

83 compressor: Compressor | None = None

84 """Object used to compress JSON blocks.

86 This is `None` until a compression dictionary is received from the writer

87 process.

88 """

90 init_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = dataclasses.field(init=False)

91 """Dictionary mapping init quantum IDs to their predicted models."""

93 def __post_init__(self) -> None:

94 if self.comms.config.mock_storage_classes:

95 import lsst.pipe.base.tests.mocks # noqa: F401

96 self.comms.log.verbose("Reading from predicted quantum graph.")

97 self.reader = self.comms.enter(

98 PredictedQuantumGraphReader.open(self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT)

99 )

100 self.reader.read_dimension_data()

101 self.reader.read_init_quanta()

102 self.comms.log.verbose("Initializing quantum-backed butler.")

103 self.qbb = self.make_qbb(self.butler_path, self.reader.pipeline_graph)

104 self.init_quanta = {q.quantum_id: q for q in self.reader.components.init_quanta.root}

105

106 def __enter__(self) -> Self:

107 return self

108

109 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Literal[False]:

110 try:

111 self.qbb.close()

112 except Exception:

113 self.comms.log.exception("An exception occurred during Ingester exit")

114 return False

115

116 @staticmethod

117 def make_qbb(butler_config: str, pipeline_graph: PipelineGraph) -> QuantumBackedButler:

118 """Make quantum-backed butler that can operate on the outputs of the

119 quantum graph.

120

121 Parameters

122 ----------

123 butler_config : `str`

124 Path or alias for the central butler repository that shares storage

125 with the quantum-backed butler.

126 pipeline_graph : `..pipeline_graph.PipelineGraph`

127 Graph of tasks and dataset types.

128

129 Returns

130 -------

131 qbb : `lsst.daf.butler.QuantumBackedButler`

132 Quantum-backed butler. This does not have the datastore records

133 needed to read overall-inputs.

134 """

135 return QuantumBackedButler.from_predicted(

136 butler_config,

137 predicted_inputs=[],

138 predicted_outputs=[],

139 dimensions=pipeline_graph.universe,

140 # We don't need the datastore records in the QG because we're

141 # only going to read metadata and logs, and those are never

142 # overall inputs.

143 datastore_records={},

144 dataset_types={node.name: node.dataset_type for node in pipeline_graph.dataset_types.values()},

145 )

146

147 @property

148 def pipeline_graph(self) -> PipelineGraph:

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

150 return self.reader.pipeline_graph

151

152 @staticmethod

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

154 """Run the scanner.

155

156 Parameters

157 ----------

158 predicted_path : `str`

159 Path to the predicted quantum graph.

160 butler_path : `str`

161 Path or alias to the central butler repository.

162 comms : `ScannerCommunicator`

163 Communicator for the scanner.

164

165 Notes

166 -----

167 This method is designed to run as the ``target`` in

168 `WorkerContext.make_worker`.

169 """

170 with comms, Scanner(predicted_path, butler_path, comms) as scanner:

171 scanner.loop()

172

173 def loop(self) -> None:

174 """Run the main loop for the scanner."""

175 self.comms.log.info("Scan request loop beginning.")

176 for quantum_id in self.comms.poll():

177 if self.compressor is None and (cdict_data := self.comms.get_compression_dict()) is not None:

178 self.compressor = zstandard.ZstdCompressor(

179 self.comms.config.zstd_level, zstandard.ZstdCompressionDict(cdict_data)

180 )

181 self.scan_quantum(quantum_id)

182

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

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

185

186 Parameters

187 ----------

188 predicted : `.PredictedDatasetModel`

189 Information about the dataset from the predicted graph.

190

191 Returns

192 -------

193 exists : `bool``

194 Whether the dataset exists.

195 """

196 ref = self.reader.components.make_dataset_ref(predicted)

197 return self.qbb.stored(ref)

198

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

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

201 datasets' existence.

202

203 Parameters

204 ----------

205 quantum_id : `uuid.UUID`

206 Unique ID for the quantum.

207

208 Returns

209 -------

210 result : `InProgressScan`

211 Scan result struct.

212 """

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

214 result = InProgressScan(predicted_quantum.quantum_id, status=ScanStatus.INIT)

215 self.comms.log.debug("Created init scan for %s (%s)", quantum_id, predicted_quantum.task_label)

216 else:

217 self.reader.read_quantum_datasets([quantum_id])

218 predicted_quantum = self.reader.components.quantum_datasets[quantum_id]

219 self.comms.log.debug(

220 "Scanning %s (%s@%s)",

221 quantum_id,

222 predicted_quantum.task_label,

223 predicted_quantum.data_coordinate,

224 )

225 result = InProgressScan(predicted_quantum.quantum_id, ScanStatus.INCOMPLETE)

226 del self.reader.components.quantum_datasets[quantum_id]

227 last_attempt = ProvenanceQuantumAttemptModel()

228 if not self._read_log(predicted_quantum, result, last_attempt):

229 self.comms.log.debug("Abandoning scan for %s; no log dataset.", quantum_id)

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

231 return result

232 if not self._read_metadata(predicted_quantum, result, last_attempt):

233 # We found the log dataset, but no metadata; this means the

234 # quantum failed, but a retry might still happen that could

235 # turn it into a success if we can't yet assume the run is

236 # complete.

237 self.comms.log.debug("Abandoning scan for %s.", quantum_id)

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

239 return result

240 last_attempt.attempt = len(result.attempts)

241 result.attempts.append(last_attempt)

242 assert result.status is not ScanStatus.INCOMPLETE

243 assert result.status is not ScanStatus.ABANDONED

244

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

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

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

248 # interrupted the writing of the logs.

249 result.logs.attempts.append(None)

250 if result.status is ScanStatus.SUCCESSFUL:

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

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

253 # something even weirder happened.

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

255 else:

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

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

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

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

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

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

262 result.metadata.attempts.append(None)

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

264 result.attempts

265 ), (

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

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

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

269 "the scanner."

270 )

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

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

273 # checked those).

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

275 if predicted_output.dataset_id not in result.outputs:

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

277 to_ingest = self._make_ingest_request(predicted_quantum, result)

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

279 to_write = self._make_write_request(predicted_quantum, result)

280 self.comms.request_write(to_write)

281 self.comms.request_ingest(to_ingest)

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

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

284 return result

285

286 def _make_ingest_request(

287 self, predicted_quantum: PredictedQuantumDatasetsModel, result: InProgressScan

288 ) -> IngestRequest:

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

290

291 Parameters

292 ----------

293 predicted_quantum : `PredictedQuantumDatasetsModel`

294 Information about the predicted quantum.

295 result : `InProgressScan`

296 Result of a quantum scan.

297

298 Returns

299 -------

300 ingest_request : `IngestRequest`

301 A request to be sent to the ingester.

302 """

303 predicted_outputs_by_id = {

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

305 }

306 to_ingest_predicted: list[PredictedDatasetModel] = []

307 to_ingest_refs: list[DatasetRef] = []

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

309 if was_produced:

310 predicted_output = predicted_outputs_by_id[dataset_id]

311 to_ingest_predicted.append(predicted_output)

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

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

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

315

316 def _make_write_request(

317 self, predicted_quantum: PredictedQuantumDatasetsModel, result: InProgressScan

318 ) -> WriteRequest:

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

320

321 Parameters

322 ----------

323 predicted_quantum : `PredictedQuantumDatasetsModel`

324 Information about the predicted quantum.

325 result : `InProgressScan`

326 Result of a quantum scan.

327

328 Returns

329 -------

330 write_request : `WriteRequest`

331 A request to be sent to the writer.

332 """

333 quantum: ProvenanceInitQuantumModel | ProvenanceQuantumModel

334 if result.status is ScanStatus.INIT:

335 quantum = ProvenanceInitQuantumModel.from_predicted(predicted_quantum)

336 else:

337 quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum)

338 quantum.attempts = result.attempts

339 request = WriteRequest(

340 result.quantum_id,

341 result.status,

342 existing_outputs={

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

344 },

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

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

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

348 )

349 if self.compressor is not None:

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

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

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

353 request.is_compressed = True

354 return request

355

356 def _read_metadata(

357 self,

358 predicted_quantum: PredictedQuantumDatasetsModel,

359 result: InProgressScan,

360 last_attempt: ProvenanceQuantumAttemptModel,

361 ) -> bool:

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

363 provenance information from it.

364

365 Parameters

366 ----------

367 predicted_quantum : `PredictedQuantumDatasetsModel`

368 Information about the predicted quantum.

369 result : `InProgressScan`

370 Result object to be modified in-place.

371 last_attempt : `ScanningProvenanceQuantumAttemptModel`

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

373 run this quantum.

374

375 Returns

376 -------

377 complete : `bool`

378 Whether the quantum is complete.

379 """

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

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

382 try:

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

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

385 # here.

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

387 except FileNotFoundError:

388 result.outputs[ref.id] = False

389 if self.comms.config.assume_complete:

390 result.status = ScanStatus.FAILED

391 else:

392 result.status = ScanStatus.ABANDONED

393 return False

394 else:

395 result.status = ScanStatus.SUCCESSFUL

396 result.outputs[ref.id] = True

397 last_attempt.status = QuantumAttemptStatus.SUCCESSFUL

398 try:

399 # Int conversion guards against spurious conversion to

400 # float that can apparently sometimes happen in

401 # TaskMetadata.

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

403 except LookupError:

404 pass

405 try:

406 last_attempt.exception = ExceptionInfo._from_metadata(

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

408 )

409 except LookupError:

410 pass

411 try:

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

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

414 except LookupError:

415 pass

416 else:

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

418 # list was not written.

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

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

421 last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)

422 result.metadata.attempts.append(metadata)

423 return True

424

425 def _read_log(

426 self,

427 predicted_quantum: PredictedQuantumDatasetsModel,

428 result: InProgressScan,

429 last_attempt: ProvenanceQuantumAttemptModel,

430 ) -> bool:

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

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

433 the log content in the provenance quantum graph.

434

435 Parameters

436 ----------

437 predicted_quantum : `PredictedQuantumDatasetsModel`

438 Information about the predicted quantum.

439 result : `InProgressScan`

440 Result object to be modified in-place.

441 last_attempt : `ScanningProvenanceQuantumAttemptModel`

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

443 run this quantum.

444

445 Returns

446 -------

447 complete : `bool`

448 Whether the quantum is complete.

449 """

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

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

452 try:

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

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

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

456 except FileNotFoundError:

457 result.outputs[ref.id] = False

458 if self.comms.config.assume_complete:

459 result.status = ScanStatus.FAILED

460 else:

461 result.status = ScanStatus.ABANDONED

462 return False

463 else:

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

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

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

467 # when we look for metadata next.

468 last_attempt.status = QuantumAttemptStatus.FAILED

469 result.outputs[ref.id] = True

470 if log_records.extra:

471 log_extra = _ExecutionLogRecordsExtra.model_validate(log_records.extra)

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

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

474 return True

475

476 def _extract_from_log_extra(

477 self,

478 log_extra: _ExecutionLogRecordsExtra,

479 result: InProgressScan,

480 last_attempt: ProvenanceQuantumAttemptModel | None,

481 ) -> None:

482 for previous_attempt_log_extra in log_extra.previous_attempts:

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

484 quantum_attempt: ProvenanceQuantumAttemptModel

485 if last_attempt is None:

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

487 quantum_attempt = ProvenanceQuantumAttemptModel(

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

489 )

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

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

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

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

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

495 result.attempts.append(quantum_attempt)

496 else:

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

498 quantum_attempt = last_attempt

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

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

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

502 # provenance in the logs.

503 quantum_attempt.exception = log_extra.exception

504 if log_extra.metadata is not None:

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

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

507 else:

508 result.metadata.attempts.append(None)

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

510 # get the previous_process_quanta from the log extra.

511 quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)

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

204 statements