Coverage for tests/test_aggregator.py: 100%

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# (https://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 <https://www.gnu.org/licenses/>. 

27 

28from __future__ import annotations 

29 

30import dataclasses 

31import itertools 

32import json 

33import os 

34import re 

35import tempfile 

36import time 

37import unittest.mock 

38import uuid 

39from collections.abc import Iterator 

40from contextlib import contextmanager 

41from typing import Any, cast 

42 

43import astropy.table 

44import click.testing 

45import numpy as np 

46import pydantic 

47from click.testing import CliRunner, Result 

48 

49import lsst.utils.tests 

50from lsst.daf.butler import Butler, ButlerLogRecords, QuantumBackedButler 

51from lsst.pex.config import Config 

52from lsst.pipe.base import ( 

53 AlgorithmError, 

54 QuantumAttemptStatus, 

55 QuantumSuccessCaveats, 

56 TaskMetadata, 

57) 

58from lsst.pipe.base import automatic_connection_constants as acc 

59from lsst.pipe.base.cli.cmd.commands import ( 

60 aggregate_graph as aggregate_graph_cli, 

61) 

62from lsst.pipe.base.cli.cmd.commands import ( 

63 provenance_report as provenance_report_cli, 

64) 

65from lsst.pipe.base.graph_walker import GraphWalker 

66from lsst.pipe.base.pipeline_graph import Edge 

67from lsst.pipe.base.quantum_graph import ( 

68 FORMAT_VERSION, 

69 PredictedDatasetInfo, 

70 PredictedQuantumGraph, 

71 PredictedQuantumInfo, 

72 ProvenanceDatasetInfo, 

73 ProvenanceQuantumGraph, 

74 ProvenanceQuantumInfo, 

75 ProvenanceQuantumReport, 

76 ProvenanceReport, 

77 ProvenanceTaskMetadataModel, 

78) 

79from lsst.pipe.base.quantum_graph.aggregator import AggregatorConfig, FatalWorkerError, aggregate_graph 

80from lsst.pipe.base.quantum_graph.ingest_graph import ingest_graph 

81from lsst.pipe.base.resource_usage import QuantumResourceUsage 

82from lsst.pipe.base.single_quantum_executor import SingleQuantumExecutor 

83from lsst.pipe.base.tests.mocks import ( 

84 DirectButlerRepo, 

85 DynamicConnectionConfig, 

86 DynamicTestPipelineTaskConfig, 

87) 

88from lsst.pipe.base.tests.util import patch_deterministic_uuid4 

89from lsst.utils.packages import Packages 

90 

91 

92@dataclasses.dataclass 

93class PrepInfo: 

94 """Struct of objects used in an aggregator test.""" 

95 

96 butler: Butler 

97 butler_path: str 

98 predicted: PredictedQuantumGraph 

99 predicted_path: str 

100 config: AggregatorConfig 

101 

102 

103class AggregatorTestCase(unittest.TestCase): 

104 """Unit tests for `lsst.pipe.base.quantum_graph.aggregator`.""" 

105 

106 @staticmethod 

107 @contextmanager 

108 def make_test_repo() -> Iterator[PrepInfo]: 

109 """Make a test data repository and predicted quantum graph. 

110 

111 Returns 

112 ------- 

113 prep_info : `PrepInfo` 

114 Objects used in aggregator tests. 

115 

116 Notes 

117 ----- 

118 The pipeline graph used by this task looks like this: 

119 

120 ■ calibrate: {detector, visit} 

121 ╭─┤ 

122 ■ │ consolidate: {visit} 

123 │ 

124 ■ resample: {patch, visit} 

125 │ 

126 ■ coadd: {band, patch} 

127 

128 The data can be visualized via:: 

129 

130 python -m lsst.daf.butler.tests.registry_data.spatial 

131 

132 One of the 'calibrate' quanta (visit=2, detector=2) is configured to 

133 fail with `lsst.pipe.base.AnnotatedPartialOutputsError`. This lets us 

134 test both success-with-caveats and failures, depending on how we 

135 configure the executor. This ``{visit: 2, detector: 2}`` data ID is 

136 the only one that overlaps ``{tract: 1, patch: 1}`` and 

137 ``{tract: 0, patch: 5}``, so it should chain to the 'resample' and 

138 'coadd' tasks, too. 

139 """ 

140 with patch_deterministic_uuid4(100): 

141 with DirectButlerRepo.make_temporary("base.yaml", "spatial.yaml") as (helper, root): 

142 calibrate_config = DynamicTestPipelineTaskConfig() 

143 calibrate_config.fail_exception = "lsst.pipe.base.AnnotatedPartialOutputsError" 

144 calibrate_config.fail_condition = "visit=2 AND detector=2" 

145 helper.add_task( 

146 "calibrate", 

147 config=calibrate_config, 

148 dimensions=["visit", "detector"], 

149 inputs={ 

150 "input_image": DynamicConnectionConfig( 

151 dataset_type_name="raw", 

152 dimensions=["visit", "detector"], 

153 ) 

154 }, 

155 prerequisite_inputs={ 

156 "refcat": DynamicConnectionConfig( 

157 dataset_type_name="references", 

158 dimensions=["htm7"], 

159 multiple=True, 

160 ) 

161 }, 

162 init_outputs={ 

163 "output_schema": DynamicConnectionConfig( 

164 dataset_type_name="source_schema", 

165 ) 

166 }, 

167 outputs={ 

168 "output_image": DynamicConnectionConfig( 

169 dataset_type_name="image", 

170 dimensions=["visit", "detector"], 

171 ), 

172 "output_table": DynamicConnectionConfig( 

173 dataset_type_name="source_detector", 

174 dimensions=["visit", "detector"], 

175 ), 

176 }, 

177 ) 

178 helper.add_task( 

179 "consolidate", 

180 dimensions=["visit"], 

181 init_inputs={ 

182 "input_schema": DynamicConnectionConfig( 

183 dataset_type_name="source_schema", 

184 ) 

185 }, 

186 inputs={ 

187 "input_table": DynamicConnectionConfig( 

188 dataset_type_name="source_detector", 

189 dimensions=["visit", "detector"], 

190 multiple=True, 

191 ) 

192 }, 

193 outputs={ 

194 "output_table": DynamicConnectionConfig( 

195 dataset_type_name="source", 

196 dimensions=["visit"], 

197 ) 

198 }, 

199 ) 

200 helper.add_task( 

201 "resample", 

202 dimensions=["patch", "visit"], 

203 inputs={ 

204 "input_image": DynamicConnectionConfig( 

205 dataset_type_name="image", 

206 dimensions=["visit", "detector"], 

207 multiple=True, 

208 ) 

209 }, 

210 outputs={ 

211 "output_image": DynamicConnectionConfig( 

212 dataset_type_name="warp", 

213 dimensions=["patch", "visit"], 

214 ) 

215 }, 

216 ) 

217 helper.add_task( 

218 "coadd", 

219 dimensions=["patch", "band"], 

220 inputs={ 

221 "input_image": DynamicConnectionConfig( 

222 dataset_type_name="warp", 

223 dimensions=["patch", "visit"], 

224 multiple=True, 

225 ) 

226 }, 

227 outputs={ 

228 "output_image": DynamicConnectionConfig( 

229 dataset_type_name="coadd", 

230 dimensions=["patch", "band"], 

231 ), 

232 }, 

233 ) 

234 pqgc = helper.make_quantum_graph_builder().finish(output="out_chain") 

235 # We use the butler root for various QG files just because it's 

236 # a convenient temporary directory. 

237 predicted_path = os.path.join(root, "predicted.qg") 

238 pqgc.write(predicted_path) 

239 config = AggregatorConfig( 

240 output_path=os.path.join(root, "provenance.qg"), 

241 # Set these small to see logic paths that otherwise only 

242 # affect large graphs. 

243 ingest_batch_size=10, 

244 zstd_dict_size=256, 

245 zstd_dict_n_inputs=16, 

246 ) 

247 yield PrepInfo( 

248 butler=helper.butler.clone(collections="out_chain"), 

249 butler_path=root, 

250 predicted=pqgc.assemble(), 

251 predicted_path=predicted_path, 

252 config=config, 

253 ) 

254 

255 def iter_graph_execution( 

256 self, 

257 repo: str, 

258 qg: PredictedQuantumGraph, 

259 raise_on_partial_outputs: bool, 

260 is_retry: bool = False, 

261 ) -> Iterator[uuid.UUID]: 

262 """Return an iterator that executes and yields quanta one by one. 

263 

264 Parameters 

265 ---------- 

266 repo : `str` 

267 Butler repository path or alias. 

268 qg : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` 

269 Predicted quantum graph. Must have datastore records attached, 

270 since execution uses a quantum-backed butler. 

271 raise_on_partial_outputs : `bool` 

272 Whether to raise on `lsst.pipe.base.AnnotatedPartialOutputsError` 

273 or treat it as a success with caveats. 

274 is_retry : `bool`, optional 

275 If `True`, this is a retry attempt and hence some outputs may 

276 already be present; skip successes and reprocess failures. 

277 

278 Returns 

279 ------- 

280 quanta : `~collections.abc.Iterator` [`uuid.UUID`] 

281 An iterator over all executed quantum IDs (not blocked ones). 

282 """ 

283 qbb = qg.make_init_qbb(repo) 

284 self.enterContext(qbb) 

285 qg.init_output_run(qbb) 

286 sqe = SingleQuantumExecutor( 

287 limited_butler_factory=lambda quantum: QuantumBackedButler.initialize( 

288 repo, 

289 quantum, 

290 qg.pipeline_graph.universe, 

291 ), 

292 assume_no_existing_outputs=not is_retry, 

293 skip_existing=is_retry, 

294 clobber_outputs=is_retry, 

295 raise_on_partial_outputs=raise_on_partial_outputs, 

296 ) 

297 qg.build_execution_quanta() 

298 xgraph = qg.quantum_only_xgraph 

299 walker = GraphWalker[uuid.UUID](xgraph.copy()) 

300 for ready in walker: 

301 for quantum_id in ready: 

302 info = xgraph.nodes[quantum_id] 

303 try: 

304 sqe.execute(info["pipeline_node"], info["quantum"], quantum_id) 

305 except AlgorithmError: 

306 walker.fail(quantum_id) 

307 else: 

308 walker.finish(quantum_id) 

309 yield quantum_id 

310 

311 def check_provenance_graph( 

312 self, 

313 pred: PredictedQuantumGraph, 

314 butler: Butler, 

315 expect_failure: bool, 

316 start_time: float, 

317 expect_failures_retried: bool = False, 

318 ) -> ProvenanceQuantumGraph: 

319 """Run a batter of tests on a provenance quantum graph produced by 

320 scanning the graph created by `make_test_repo`. 

321 

322 Parameters 

323 ---------- 

324 pred: `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` 

325 Predicted quantum graph. 

326 prov_reader : \ 

327 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

328 Reader for the provenance quantum graph. 

329 butler : `lsst.daf.butler.Butler` 

330 Client for the data repository. 

331 expect_failure : `bool` 

332 Whether to expect one quantum of 'calibrate' to fail (`True`) or 

333 succeed without writing anything (`False`). 

334 start_time : `float` 

335 A POSIX timestamp that strictly precedes the start time of any 

336 quantum's execution. 

337 expect_failures_retried : `bool`, optional 

338 If `True`, expect an initial attempt with failures prior to the 

339 most recent attempt. 

340 

341 Returns 

342 ------- 

343 prov : `ProvenanceQuantumGraph` 

344 The full provenance quantum graph. 

345 """ 

346 prov: ProvenanceQuantumGraph = butler.get("run_provenance") 

347 self.maxDiff = None 

348 self.assertEqual(prov.header.version, FORMAT_VERSION) 

349 self.assertEqual( 

350 list(butler.collections.get_info(prov.header.output).children), 

351 [prov.header.output_run] 

352 + list(butler.collections.query(prov.header.inputs, flatten_chains=True)), 

353 ) 

354 self.assertEqual(pred.quanta_by_task.keys(), prov.quanta_by_task.keys()) 

355 for task_label in pred.quanta_by_task: 

356 self.assertEqual(pred.quanta_by_task[task_label], prov.quanta_by_task[task_label]) 

357 self.assertEqual(pred.datasets_by_type.keys() - {"packages"}, prov.datasets_by_type.keys()) 

358 for dataset_type_name in prov.datasets_by_type: 

359 self.assertEqual( 

360 pred.datasets_by_type[dataset_type_name], prov.datasets_by_type[dataset_type_name] 

361 ) 

362 self.assertEqual(prov.init_quanta.keys(), pred.quanta_by_task.keys()) 

363 for quantum_id in pred: 

364 # Check consistency between the predicted and provenance quantum 

365 # node attributes. 

366 pred_qinfo: PredictedQuantumInfo = pred.bipartite_xgraph.nodes[quantum_id] 

367 prov_qinfo: ProvenanceQuantumInfo = prov.bipartite_xgraph.nodes[quantum_id] 

368 self.assertEqual(pred_qinfo["task_label"], prov_qinfo["task_label"]) 

369 self.assertEqual(pred_qinfo["data_id"], prov_qinfo["data_id"]) 

370 msg = f"{pred_qinfo['task_label']}@{pred_qinfo['data_id']}" 

371 # Check consistency between the predicted and provenance dataset 

372 # node attributes and edges. Also gather existence information for 

373 # use later. 

374 existence: dict[str, list[bool]] = {} 

375 pipeline_edges: list[Edge] 

376 for dataset_id, _, pipeline_edges in pred.bipartite_xgraph.in_edges( 

377 quantum_id, data="pipeline_edges" 

378 ): 

379 self.assertTrue(prov.bipartite_xgraph.has_predecessor(quantum_id, dataset_id)) 

380 for edge in pipeline_edges: 

381 existence.setdefault(edge.connection_name, []).append( 

382 self.check_dataset(dataset_id, pred, prov, butler) 

383 ) 

384 for _, dataset_id, pipeline_edges in pred.bipartite_xgraph.out_edges( 

385 quantum_id, data="pipeline_edges" 

386 ): 

387 self.assertTrue(prov.bipartite_xgraph.has_successor(quantum_id, dataset_id)) 

388 for edge in pipeline_edges: 

389 existence.setdefault(edge.connection_name, []).append( 

390 self.check_dataset(dataset_id, pred, prov, butler) 

391 ) 

392 # Check quantum status and dataset existence against the known 

393 # structure of the graph and where failures/caveats occur. 

394 match (pred_qinfo["task_label"], dict(pred_qinfo["data_id"].required)): 

395 case "calibrate", {"visit": 2, "detector": 2}: 

396 # This is the quantum that can directly raise. 

397 self._expect_all_exist(existence["input_image"], msg=msg) 

398 self._expect_all_exist(existence["refcat"], msg=msg) 

399 self._expect_none_exist(existence["output_image"], msg=msg) 

400 self._expect_none_exist(existence["output_table"], msg=msg) 

401 if expect_failure: 

402 self._expect_failure(prov_qinfo, existence, msg=msg) 

403 else: 

404 self._expect_successful( 

405 prov_qinfo, 

406 existence, 

407 caveats=( 

408 QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR 

409 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

410 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

411 ), 

412 exception_type="lsst.pipe.base.tests.mocks.MockAlgorithmError", 

413 msg=msg, 

414 ) 

415 if expect_failures_retried: 

416 self.assertEqual(len(prov_qinfo["attempts"]), 2) 

417 self.assertEqual( 

418 prov_qinfo["attempts"][0].exception.type_name, 

419 "lsst.pipe.base.tests.mocks.MockAlgorithmError", 

420 ) 

421 else: 

422 self.assertEqual(len(prov_qinfo["attempts"]), 1) 

423 case "consolidate", {"visit": 2}: 

424 # This quantum will succeed (with one predicted input 

425 # missing) or be blocked. 

426 self._expect_one_missing(existence["input_table"], msg=msg) 

427 if expect_failure: 

428 self._expect_blocked(prov_qinfo, existence, msg=msg) 

429 else: 

430 self._expect_successful(prov_qinfo, existence, msg=msg) 

431 self.assertEqual( 

432 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure 

433 ) 

434 case ( 

435 "resample" | "coadd", 

436 {"tract": 1, "patch": 1} | {"tract": 0, "patch": 5}, 

437 ): 

438 # These quanta will be blocked by an upstream failure or do 

439 # chained caveats, since they won't have enough inputs to 

440 # run. 

441 if expect_failure: 

442 self._expect_blocked(prov_qinfo, existence, msg=msg) 

443 else: 

444 self._expect_successful( 

445 prov_qinfo, 

446 existence, 

447 caveats=( 

448 QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED 

449 | QuantumSuccessCaveats.NO_WORK 

450 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

451 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

452 ), 

453 msg=msg, 

454 ) 

455 self.assertEqual( 

456 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure 

457 ) 

458 case ( 

459 "resample", 

460 {"tract": 0, "patch": 4, "visit": 2} | {"tract": 1, "patch": 0, "visit": 2}, 

461 ): 

462 # This will succeed or be blocked, with one input missing 

463 # regardless. 

464 self._expect_one_missing(existence["input_image"], msg=msg) 

465 if expect_failure: 

466 self._expect_blocked(prov_qinfo, existence, msg=msg) 

467 else: 

468 self._expect_successful(prov_qinfo, existence, msg=msg) 

469 self.assertEqual( 

470 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure 

471 ) 

472 case ( 

473 "coadd", 

474 {"tract": 0, "patch": 4, "band": "r"} | {"tract": 1, "patch": 0, "band": "r"}, 

475 ): 

476 # This will succeed with no inputs missing or be blocked 

477 # with one input missing. 

478 if expect_failure: 

479 self._expect_one_missing(existence["input_image"], msg=msg) 

480 self._expect_blocked(prov_qinfo, existence, msg=msg) 

481 else: 

482 self._expect_all_exist(existence["input_image"], msg=msg) 

483 self._expect_successful(prov_qinfo, existence, msg=msg) 

484 self.assertEqual( 

485 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure 

486 ) 

487 case _: 

488 # All other quanta should succeed and have all inputs 

489 # present. 

490 for connection_name in prov_qinfo["pipeline_node"].inputs.keys(): 

491 self._expect_all_exist(existence[connection_name], msg=msg) 

492 self._expect_successful(prov_qinfo, existence, msg=msg) 

493 self.assertEqual(len(prov_qinfo["attempts"]), 1) 

494 self.check_metadata( 

495 quantum_id, 

496 prov, 

497 butler, 

498 expect_ingested=(prov_qinfo["status"] is QuantumAttemptStatus.SUCCESSFUL), 

499 ) 

500 self.check_log( 

501 quantum_id, 

502 prov, 

503 butler, 

504 expect_ingested=( 

505 prov_qinfo["status"] 

506 in ( 

507 QuantumAttemptStatus.SUCCESSFUL, 

508 QuantumAttemptStatus.FAILED, 

509 ) 

510 ), 

511 ) 

512 self.check_resource_usage_table(prov, expect_failure=expect_failure, start_time=start_time) 

513 self.check_packages(butler) 

514 self.check_configs(butler, prov) 

515 self.check_quantum_table(prov, expect_failure=expect_failure) 

516 self.check_exception_table(prov, expect_failure=expect_failure) 

517 self.check_report(prov, expect_failure=expect_failure) 

518 return prov 

519 

520 def _expect_all_exist(self, existence: list[bool], msg: str) -> None: 

521 self.assertTrue(all(existence), msg=msg) 

522 

523 def _expect_none_exist(self, existence: list[bool], msg: str) -> None: 

524 self.assertFalse(any(existence), msg=msg) 

525 

526 def _expect_one_missing(self, existence: list[bool], msg: str) -> None: 

527 self.assertEqual(existence.count(False), 1, msg=msg) 

528 

529 def _expect_successful( 

530 self, 

531 info: ProvenanceQuantumInfo, 

532 existence: dict[str, list[bool]], 

533 caveats: QuantumSuccessCaveats = QuantumSuccessCaveats.NO_CAVEATS, 

534 exception_type: str | None = None, 

535 *, 

536 msg: str, 

537 ) -> None: 

538 self.assertEqual(info["status"], QuantumAttemptStatus.SUCCESSFUL, msg=msg) 

539 self.assertEqual(info["caveats"], caveats, msg=msg) 

540 if exception_type is None: 

541 self.assertIsNone(info["exception"], msg=msg) 

542 else: 

543 assert info["exception"] is not None 

544 self.assertEqual(info["exception"].type_name, exception_type, msg=msg) 

545 self._expect_all_exist(existence[acc.LOG_OUTPUT_CONNECTION_NAME], msg=msg) 

546 self._expect_all_exist(existence[acc.METADATA_OUTPUT_CONNECTION_NAME], msg=msg) 

547 if not (caveats & QuantumSuccessCaveats.ANY_OUTPUTS_MISSING): 

548 for connection_name in info["pipeline_node"].outputs.keys(): 

549 self._expect_all_exist(existence[connection_name], msg=msg) 

550 if caveats & QuantumSuccessCaveats.ALL_OUTPUTS_MISSING: 

551 for connection_name in info["pipeline_node"].outputs.keys(): 

552 self._expect_none_exist(existence[connection_name], msg=msg) 

553 self.assertIsNotNone(info["resource_usage"], msg=msg) 

554 self.assertGreater(info["resource_usage"].total_time, 0, msg=msg) 

555 self.assertGreater(info["resource_usage"].memory, 0, msg=msg) 

556 

557 def _expect_failure( 

558 self, info: ProvenanceQuantumInfo, existence: dict[str, list[bool]], msg: str 

559 ) -> None: 

560 self.assertEqual(info["status"], QuantumAttemptStatus.FAILED, msg=msg) 

561 self.assertEqual(info["exception"].type_name, "lsst.pipe.base.tests.mocks.MockAlgorithmError") 

562 self._expect_all_exist(existence[acc.LOG_OUTPUT_CONNECTION_NAME], msg=msg) 

563 self._expect_none_exist(existence[acc.METADATA_OUTPUT_CONNECTION_NAME], msg=msg) 

564 for connection_name in info["pipeline_node"].outputs.keys(): 

565 self._expect_none_exist(existence[connection_name], msg=msg) 

566 

567 def _expect_blocked( 

568 self, 

569 info: ProvenanceQuantumInfo, 

570 existence: dict[str, list[bool]], 

571 msg: str, 

572 ) -> None: 

573 self.assertEqual(info["status"], QuantumAttemptStatus.BLOCKED, msg=msg) 

574 self.assertEqual(info["attempts"], []) 

575 self._expect_none_exist(existence[acc.LOG_OUTPUT_CONNECTION_NAME], msg=msg) 

576 self._expect_none_exist(existence[acc.METADATA_OUTPUT_CONNECTION_NAME], msg=msg) 

577 for connection_name in info["pipeline_node"].outputs.keys(): 

578 self._expect_none_exist(existence[connection_name], msg=msg) 

579 

580 def check_dataset( 

581 self, 

582 dataset_id: uuid.UUID, 

583 pred: PredictedQuantumGraph, 

584 prov: ProvenanceQuantumGraph, 

585 butler: Butler, 

586 ) -> bool: 

587 """Check a provenance dataset for consistency with its predicted 

588 counterpart. 

589 

590 Parameters 

591 ---------- 

592 dataset_id : `uuid.UUID` 

593 Unique ID for the dataset. 

594 pred: `lsst.pipe.base.quantum_graph.PredictedQuantumGraph` 

595 Predicted quantum graph. 

596 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph` 

597 Provenance quantum graph. 

598 butler : `lsst.daf.butler.Butler` 

599 Client for the data repository. 

600 

601 Returns 

602 ------- 

603 exists : `bool` 

604 Whether the dataset was marked as existing in the provenance 

605 quantum graph. 

606 """ 

607 pred_info: PredictedDatasetInfo = pred.bipartite_xgraph.nodes[dataset_id] 

608 prov_info: ProvenanceDatasetInfo = prov.bipartite_xgraph.nodes[dataset_id] 

609 self.assertEqual(pred_info["dataset_type_name"], prov_info["dataset_type_name"]) 

610 self.assertEqual(pred_info["data_id"], prov_info["data_id"]) 

611 self.assertEqual(pred_info["run"], prov_info["run"]) 

612 exists = prov_info["produced"] 

613 dataset_type_name = prov_info["dataset_type_name"] 

614 # We can remove this guard when we ingest QG-backed metadata and logs. 

615 if not dataset_type_name.endswith("_metadata") and not dataset_type_name.endswith("_log"): 

616 self.assertEqual( 

617 butler.get_dataset(dataset_id) is not None, 

618 exists, 

619 msg=( 

620 f"Ingest/existence inconsistency for {dataset_type_name}" 

621 f"@{prov_info['data_id']}/{dataset_id}]" 

622 ), 

623 ) 

624 return exists 

625 

626 def check_metadata( 

627 self, quantum_id: uuid.UUID, prov: ProvenanceQuantumGraph, butler: Butler, expect_ingested: bool 

628 ) -> None: 

629 """Check reading a metadata dataset from the butler, and check that the 

630 original metadata file has been deleted. 

631 

632 Parameters 

633 ---------- 

634 quantum_id : `uuid.UUID` 

635 Unique ID for the quantum this metadata belongs to. 

636 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph` 

637 Provenance quantum graph. 

638 butler : `lsst.daf.butler.Butler` 

639 Client for the data repository. 

640 expect_ingested : `bool` 

641 Whether the metadata dataset should have been ingested. 

642 """ 

643 dataset_id = prov.bipartite_xgraph.nodes[quantum_id]["metadata_id"] 

644 ref = butler.get_dataset(dataset_id) 

645 if not expect_ingested: 

646 self.assertIsNone(ref) 

647 return 

648 assert ref is not None 

649 metadata = butler.get(ref) 

650 self.assertIsInstance(metadata, TaskMetadata) 

651 graph_path = butler.getURI("run_provenance") 

652 self.assertEqual(butler.getURI(ref), graph_path) 

653 # We now delete the metadata dataset, in order let us get the original 

654 # location from the butler and check that there's nothing there. Note 

655 # that this doesn't actually delete the file because the butler knows 

656 # it's shared with other datasets. 

657 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True) 

658 original_path = butler.getURI(ref, predict=True) 

659 self.assertTrue(graph_path.exists()) 

660 self.assertNotEqual(graph_path, original_path) 

661 self.assertFalse(original_path.exists()) 

662 

663 def check_log( 

664 self, quantum_id: uuid.UUID, prov: ProvenanceQuantumGraph, butler: Butler, expect_ingested: bool 

665 ) -> None: 

666 """Check reading a log dataset from the butler, and check that the 

667 original log file has been deleted. 

668 

669 Parameters 

670 ---------- 

671 quantum_id : `uuid.UUID` 

672 Unique ID for the quantum this log belongs to. 

673 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph` 

674 Provenance quantum graph. 

675 butler : `lsst.daf.butler.Butler` 

676 Client for the data repository. 

677 expect_ingested : `bool` 

678 Whether the metadata dataset should have been ingested. 

679 """ 

680 dataset_id = prov.bipartite_xgraph.nodes[quantum_id]["log_id"] 

681 ref = butler.get_dataset(dataset_id) 

682 if not expect_ingested: 

683 self.assertIsNone(ref) 

684 return 

685 assert ref is not None 

686 log = butler.get(ref) 

687 self.assertIsInstance(log, ButlerLogRecords) 

688 graph_path = butler.getURI("run_provenance") 

689 self.assertEqual(butler.getURI(ref), graph_path) 

690 # We now delete the log dataset, in order let us get the original 

691 # location from the butler and check that there's nothing there. Note 

692 # that this doesn't actually delete the file because the butler knows 

693 # it's shared with other datasets. 

694 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True) 

695 original_path = butler.getURI(ref, predict=True) 

696 self.assertTrue(graph_path.exists()) 

697 self.assertNotEqual(graph_path, original_path) 

698 self.assertFalse(original_path.exists()) 

699 

700 def check_configs(self, butler: Butler, prov: ProvenanceQuantumGraph) -> None: 

701 for task_node in prov.pipeline_graph.tasks.values(): 

702 config = butler.get(task_node.init.config_output.dataset_type_name) 

703 self.assertIsInstance(config, Config) 

704 

705 def check_packages(self, butler: Butler) -> None: 

706 """Check fetching package versions from the provenance graph. 

707 

708 Parameters 

709 ---------- 

710 butler : `lsst.daf.butler.Butler` 

711 Client for the data repository. 

712 """ 

713 packages = butler.get("run_provenance.packages") 

714 self.assertIsInstance(packages, Packages) 

715 self.assertIn("pipe_base", packages) 

716 

717 def check_resource_usage_table( 

718 self, prov: ProvenanceQuantumGraph, expect_failure: bool, start_time: float 

719 ) -> None: 

720 """Check building a resource usage table from the provenance graph. 

721 

722 Parameters 

723 ---------- 

724 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph` 

725 Reader for the provenance quantum graph. 

726 expect_failure : `bool` 

727 Whether to expect one quantum of 'calibrate' to fail (`True`) or 

728 succeed without writing anything (`False`). 

729 start_time : `float` 

730 A POSIX timestamp that strictly precedes the start time of any 

731 quantum's execution. 

732 """ 

733 tbl = prov.make_task_resource_usage_table("calibrate", include_data_ids=True) 

734 self.assertEqual(len(tbl), prov.header.n_task_quanta["calibrate"]) 

735 self.assertCountEqual( 

736 tbl.colnames, 

737 ["quantum_id"] 

738 + list(prov.pipeline_graph.tasks["calibrate"].dimensions.names) 

739 + list(QuantumResourceUsage.model_fields), 

740 ) 

741 # Check that quantum start times are bounded by the before-execution 

742 # start_time and now. This makes sure we didn't get any timezone 

743 # shenanigans. 

744 end_time = time.time() 

745 for quantum_start_time in tbl["start"]: 

746 self.assertGreater(quantum_start_time, start_time) 

747 self.assertLess(quantum_start_time, end_time) 

748 self.assertTrue(np.all(tbl["init_time"] >= 0.0)) 

749 self.assertTrue(np.all(tbl["prep_time"] > 0.0)) 

750 self.assertTrue(np.all(tbl["run_time"] >= 0.0)) 

751 

752 def check_quantum_table(self, prov: ProvenanceQuantumGraph, expect_failure: bool) -> None: 

753 """Check `ProvenanceQuantumGraph.make_quantum_table`. 

754 

755 Parameters 

756 ---------- 

757 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph` 

758 Reader for the provenance quantum graph. 

759 expect_failure : `bool` 

760 Whether to expect one quantum of 'calibrate' to fail (`True`) or 

761 succeed without writing anything (`False`). 

762 """ 

763 t = prov.make_quantum_table() 

764 self.assertEqual(list(t["Task"]), ["calibrate", "consolidate", "resample", "coadd"]) 

765 self.assertEqual(t["TOTAL"][0], 8) 

766 self.assertEqual(t["EXPECTED"][0], 8) 

767 self.assertEqual(t["Blocked"][0], 0) 

768 self.assertEqual(t["TOTAL"][1], 2) 

769 self.assertEqual(t["EXPECTED"][1], 2) 

770 self.assertEqual(t["TOTAL"][2], 10) 

771 self.assertEqual(t["EXPECTED"][2], 10) 

772 if expect_failure: 

773 # calibrate 

774 self.assertEqual(t["Successful"][0], 7) 

775 self.assertEqual(t["Caveats"][0], "") 

776 self.assertEqual(t["Failed"][0], 1) 

777 # consolidate 

778 self.assertEqual(t["Successful"][1], 1) 

779 self.assertEqual(t["Caveats"][1], "") 

780 self.assertEqual(t["Failed"][1], 0) 

781 self.assertEqual(t["Blocked"][1], 1) 

782 # resample 

783 self.assertEqual(t["Successful"][2], 6) 

784 self.assertEqual(t["Caveats"][2], "") 

785 self.assertEqual(t["Failed"][2], 0) 

786 self.assertEqual(t["Blocked"][2], 4) 

787 else: 

788 # calibrate 

789 self.assertEqual(t["Successful"][0], 8) 

790 self.assertEqual(t["Caveats"][0], "*P(1)") 

791 self.assertEqual(t["Failed"][0], 0) 

792 # consolidate 

793 self.assertEqual(t["Successful"][1], 2)