Coverage for tests/test_aggregator.py: 13%

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 

28import dataclasses 

29import itertools 

30import os 

31import time 

32import unittest.mock 

33import uuid 

34from collections.abc import Iterator 

35from contextlib import contextmanager 

36 

37import numpy as np 

38from click.testing import CliRunner, Result 

39 

40import lsst.utils.tests 

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

42from lsst.pipe.base import ( 

43 AlgorithmError, 

44 QuantumAttemptStatus, 

45 QuantumSuccessCaveats, 

46 TaskMetadata, 

47) 

48from lsst.pipe.base import automatic_connection_constants as acc 

49from lsst.pipe.base.cli.cmd.commands import aggregate_graph as aggregate_graph_cli 

50from lsst.pipe.base.graph_walker import GraphWalker 

51from lsst.pipe.base.pipeline_graph import Edge 

52from lsst.pipe.base.quantum_graph import ( 

53 FORMAT_VERSION, 

54 PredictedDatasetInfo, 

55 PredictedQuantumGraph, 

56 PredictedQuantumInfo, 

57 ProvenanceDatasetInfo, 

58 ProvenanceQuantumGraph, 

59 ProvenanceQuantumGraphReader, 

60 ProvenanceQuantumInfo, 

61) 

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

63from lsst.pipe.base.resource_usage import QuantumResourceUsage 

64from lsst.pipe.base.single_quantum_executor import SingleQuantumExecutor 

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

66 DirectButlerRepo, 

67 DynamicConnectionConfig, 

68 DynamicTestPipelineTaskConfig, 

69) 

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

71from lsst.resources import ResourcePath 

72from lsst.utils.packages import Packages 

73 

74 

75@dataclasses.dataclass 

76class PrepInfo: 

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

78 

79 butler: Butler 

80 butler_path: str 

81 predicted: PredictedQuantumGraph 

82 predicted_path: str 

83 config: AggregatorConfig 

84 

85 

86class AggregatorTestCase(unittest.TestCase): 

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

88 

89 @staticmethod 

90 @contextmanager 

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

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

93 

94 Returns 

95 ------- 

96 prep_info : `PrepInfo` 

97 Objects used in aggregator tests. 

98 

99 Notes 

100 ----- 

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

102 

103 ■ calibrate: {detector, visit} 

104 ╭─┤ 

105 ■ │ consolidate: {visit} 

106 │ 

107 ■ resample: {patch, visit} 

108 │ 

109 ■ coadd: {band, patch} 

110 

111 The data can be visualized via:: 

112 

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

114 

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

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

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

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

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

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

121 'coadd' tasks, too. 

122 """ 

123 with patch_deterministic_uuid4(100): 

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

125 calibrate_config = DynamicTestPipelineTaskConfig() 

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

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

128 helper.add_task( 

129 "calibrate", 

130 config=calibrate_config, 

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

132 inputs={ 

133 "input_image": DynamicConnectionConfig( 

134 dataset_type_name="raw", 

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

136 ) 

137 }, 

138 prerequisite_inputs={ 

139 "refcat": DynamicConnectionConfig( 

140 dataset_type_name="references", 

141 dimensions=["htm7"], 

142 multiple=True, 

143 ) 

144 }, 

145 init_outputs={ 

146 "output_schema": DynamicConnectionConfig( 

147 dataset_type_name="source_schema", 

148 ) 

149 }, 

150 outputs={ 

151 "output_image": DynamicConnectionConfig( 

152 dataset_type_name="image", 

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

154 ), 

155 "output_table": DynamicConnectionConfig( 

156 dataset_type_name="source_detector", 

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

158 ), 

159 }, 

160 ) 

161 helper.add_task( 

162 "consolidate", 

163 dimensions=["visit"], 

164 init_inputs={ 

165 "input_schema": DynamicConnectionConfig( 

166 dataset_type_name="source_schema", 

167 ) 

168 }, 

169 inputs={ 

170 "input_table": DynamicConnectionConfig( 

171 dataset_type_name="source_detector", 

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

173 multiple=True, 

174 ) 

175 }, 

176 outputs={ 

177 "output_table": DynamicConnectionConfig( 

178 dataset_type_name="source", 

179 dimensions=["visit"], 

180 ) 

181 }, 

182 ) 

183 helper.add_task( 

184 "resample", 

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

186 inputs={ 

187 "input_image": DynamicConnectionConfig( 

188 dataset_type_name="image", 

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

190 multiple=True, 

191 ) 

192 }, 

193 outputs={ 

194 "output_image": DynamicConnectionConfig( 

195 dataset_type_name="warp", 

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

197 ) 

198 }, 

199 ) 

200 helper.add_task( 

201 "coadd", 

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

203 inputs={ 

204 "input_image": DynamicConnectionConfig( 

205 dataset_type_name="warp", 

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

207 multiple=True, 

208 ) 

209 }, 

210 outputs={ 

211 "output_image": DynamicConnectionConfig( 

212 dataset_type_name="coadd", 

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

214 ), 

215 }, 

216 ) 

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

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

219 # a convenient temporary directory. 

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

221 pqgc.write(predicted_path) 

222 config = AggregatorConfig( 

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

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

225 # affect large graphs. 

226 ingest_batch_size=10, 

227 zstd_dict_size=256, 

228 zstd_dict_n_inputs=16, 

229 ) 

230 yield PrepInfo( 

231 butler=helper.butler, 

232 butler_path=root, 

233 predicted=pqgc.assemble(), 

234 predicted_path=predicted_path, 

235 config=config, 

236 ) 

237 

238 def iter_graph_execution( 

239 self, 

240 repo: ResourcePath, 

241 qg: PredictedQuantumGraph, 

242 raise_on_partial_outputs: bool, 

243 is_retry: bool = False, 

244 ) -> Iterator[uuid.UUID]: 

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

246 

247 Parameters 

248 ---------- 

249 repo : `lsst.resources.ResourcePath` 

250 Butler repository path. 

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

252 Predicted quantum graph. Must have datastore records attached, 

253 since execution uses a quantum-backed butler. 

254 raise_on_partial_outputs : `bool` 

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

256 or treat it as a success with caveats. 

257 is_retry : `bool`, optional 

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

259 already be present; skip successes and reprocess failures. 

260 

261 Returns 

262 ------- 

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

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

265 """ 

266 qg.init_output_run(qg.make_init_qbb(repo)) 

267 sqe = SingleQuantumExecutor( 

268 limited_butler_factory=lambda quantum: QuantumBackedButler.initialize( 

269 repo, 

270 quantum, 

271 qg.pipeline_graph.universe, 

272 ), 

273 assume_no_existing_outputs=not is_retry, 

274 skip_existing=is_retry, 

275 clobber_outputs=is_retry, 

276 raise_on_partial_outputs=raise_on_partial_outputs, 

277 ) 

278 qg.build_execution_quanta() 

279 xgraph = qg.quantum_only_xgraph 

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

281 for ready in walker: 

282 for quantum_id in ready: 

283 info = xgraph.nodes[quantum_id] 

284 try: 

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

286 except AlgorithmError: 

287 walker.fail(quantum_id) 

288 else: 

289 walker.finish(quantum_id) 

290 yield quantum_id 

291 

292 def check_provenance_graph( 

293 self, 

294 pred: PredictedQuantumGraph, 

295 prov_reader: ProvenanceQuantumGraphReader, 

296 butler: Butler, 

297 expect_failure: bool, 

298 start_time: float, 

299 expect_failures_retried: bool = False, 

300 ) -> ProvenanceQuantumGraph: 

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

302 scanning the graph created by `make_test_repo`. 

303 

304 Parameters 

305 ---------- 

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

307 Predicted quantum graph. 

308 prov_reader : \ 

309 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

310 Reader for the provenance quantum graph. 

311 butler : `lsst.daf.butler.Butler` 

312 Client for the data repository. 

313 expect_failure : `bool` 

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

315 succeed without writing anything (`False`). 

316 start_time : `float` 

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

318 quantum's execution. 

319 expect_failures_retried : `bool`, optional 

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

321 most recent attempt. 

322 

323 Returns 

324 ------- 

325 prov : `ProvenanceQuantumGraph` 

326 The full provenance quantum graph. 

327 """ 

328 prov_reader.read_full_graph() 

329 prov = prov_reader.graph 

330 checked_some_metadata = False 

331 checked_some_log = False 

332 self.maxDiff = None 

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

334 self.assertEqual( 

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

336 [prov.header.output_run] 

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

338 ) 

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

340 for task_label in pred.quanta_by_task: 

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

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

343 for dataset_type_name in prov.datasets_by_type: 

344 self.assertEqual( 

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

346 ) 

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

348 for quantum_id in pred: 

349 # Check consistency between the predicted and provenance quantum 

350 # node attributes. 

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

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

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

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

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

356 # Check consistency between the predicted and provenance dataset 

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

358 # use later. 

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

360 pipeline_edges: list[Edge] 

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

362 quantum_id, data="pipeline_edges" 

363 ): 

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

365 for edge in pipeline_edges: 

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

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

368 ) 

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

370 quantum_id, data="pipeline_edges" 

371 ): 

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

373 for edge in pipeline_edges: 

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

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

376 ) 

377 # Check quantum status and dataset existence against the known 

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

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

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

381 # This is the quantum that can directly raise. 

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

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

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

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

386 if expect_failure: 

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

388 else: 

389 self._expect_successful( 

390 prov_qinfo, 

391 existence, 

392 caveats=( 

393 QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR 

394 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

395 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

396 ), 

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

398 msg=msg, 

399 ) 

400 if expect_failures_retried: 

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

402 self.assertEqual( 

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

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

405 ) 

406 else: 

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

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

409 # This quantum will succeed (with one predicted input 

410 # missing) or be blocked. 

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

412 if expect_failure: 

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

414 else: 

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

416 self.assertEqual( 

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

418 ) 

419 case ( 

420 "resample" | "coadd", 

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

422 ): 

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

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

425 # run. 

426 if expect_failure: 

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

428 else: 

429 self._expect_successful( 

430 prov_qinfo, 

431 existence, 

432 caveats=( 

433 QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED 

434 | QuantumSuccessCaveats.NO_WORK 

435 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

436 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

437 ), 

438 msg=msg, 

439 ) 

440 self.assertEqual( 

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

442 ) 

443 case ( 

444 "resample", 

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

446 ): 

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

448 # regardless. 

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

450 if expect_failure: 

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

452 else: 

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

454 self.assertEqual( 

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

456 ) 

457 case ( 

458 "coadd", 

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

460 ): 

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

462 # with one input missing. 

463 if expect_failure: 

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

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

466 else: 

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

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 # All other quanta should succeed and have all inputs 

474 # present. 

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

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

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

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

479 if not checked_some_metadata and prov_qinfo["status"] is QuantumAttemptStatus.SUCCESSFUL: 

480 self.check_metadata(quantum_id, prov_reader, butler) 

481 checked_some_metadata = True 

482 if not checked_some_log and prov_qinfo["status"] in ( 

483 QuantumAttemptStatus.SUCCESSFUL, 

484 QuantumAttemptStatus.FAILED, 

485 ): 

486 self.check_log(quantum_id, prov_reader, butler) 

487 checked_some_log = True 

488 self.assertTrue(checked_some_metadata) 

489 self.assertTrue(checked_some_log) 

490 self.check_resource_usage_table( 

491 prov_reader.graph, expect_failure=expect_failure, start_time=start_time 

492 ) 

493 self.check_packages(prov_reader, butler) 

494 self.check_quantum_table(prov_reader.graph, expect_failure=expect_failure) 

495 self.check_exception_table(prov_reader.graph, expect_failure=expect_failure) 

496 return prov 

497 

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

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

500 

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

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

503 

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

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

506 

507 def _expect_successful( 

508 self, 

509 info: ProvenanceQuantumInfo, 

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

511 caveats: QuantumSuccessCaveats = QuantumSuccessCaveats.NO_CAVEATS, 

512 exception_type: str | None = None, 

513 *, 

514 msg: str, 

515 ) -> None: 

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

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

518 if exception_type is None: 

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

520 else: 

521 assert info["exception"] is not None 

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

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

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

525 if not (caveats & QuantumSuccessCaveats.ANY_OUTPUTS_MISSING): 

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

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

528 if caveats & QuantumSuccessCaveats.ALL_OUTPUTS_MISSING: 

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

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

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

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

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

534 

535 def _expect_failure( 

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

537 ) -> None: 

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

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

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

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

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

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

544 

545 def _expect_blocked( 

546 self, 

547 info: ProvenanceQuantumInfo, 

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

549 msg: str, 

550 ) -> None: 

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

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

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

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

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

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

557 

558 def check_dataset( 

559 self, 

560 dataset_id: uuid.UUID, 

561 pred: PredictedQuantumGraph, 

562 prov: ProvenanceQuantumGraph, 

563 butler: Butler, 

564 ) -> bool: 

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

566 counterpart. 

567 

568 Parameters 

569 ---------- 

570 dataset_id : `uuid.UUID` 

571 Unique ID for the dataset. 

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

573 Predicted quantum graph. 

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

575 Provenance quantum graph. 

576 butler : `lsst.daf.butler.Butler` 

577 Client for the data repository. 

578 

579 Returns 

580 ------- 

581 exists : `bool` 

582 Whether the dataset was marked as existing in the provenance 

583 quantum graph. 

584 """ 

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

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

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

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

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

590 exists = prov_info["produced"] 

591 dataset_type_name = prov_info["dataset_type_name"] 

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

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

594 self.assertEqual( 

595 butler.get_dataset(dataset_id) is not None, 

596 exists, 

597 msg=( 

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

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

600 ), 

601 ) 

602 return exists 

603 

604 def check_metadata( 

605 self, quantum_id: uuid.UUID, provenance_reader: ProvenanceQuantumGraphReader, butler: Butler 

606 ) -> None: 

607 """Check reading a metadata dataset from the provenance reader, 

608 and check that the original metadata file has been deleted. 

609 

610 Parameters 

611 ---------- 

612 quantum_id : `uuid.UUID` 

613 Unique ID for the quantum this metadata belongs to. 

614 provenance_reader : \ 

615 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

616 Reader for the provenance quantum graph. 

617 butler : `lsst.daf.butler.Butler` 

618 Client for the data repository. 

619 """ 

620 # Try reading metadata through the quantum ID. 

621 ((metadata1,),) = provenance_reader.fetch_metadata([quantum_id]).values() 

622 self.assertIsInstance(metadata1, TaskMetadata) 

623 for _, dataset_id, pipeline_edges in provenance_reader.graph.bipartite_xgraph.out_edges( 

624 quantum_id, data="pipeline_edges" 

625 ): 

626 if pipeline_edges[0].connection_name == acc.METADATA_OUTPUT_CONNECTION_NAME: 

627 # Also try reading metadata through the dataset ID. 

628 ((metadata2,),) = provenance_reader.fetch_metadata([dataset_id]).values() 

629 break 

630 else: 

631 raise AssertionError("No metadata connection found.") 

632 self.assertEqual(metadata1, metadata2) 

633 # Also get the metadata from the butler. 

634 ref = butler.get_dataset(dataset_id) 

635 self.assertEqual(butler.get(ref), metadata1) 

636 

637 def check_log( 

638 self, quantum_id: uuid.UUID, provenance_reader: ProvenanceQuantumGraphReader, butler: Butler 

639 ) -> None: 

640 """Check reading a log dataset from the provenance reader, 

641 and check that the original log file has been deleted. 

642 

643 Parameters 

644 ---------- 

645 quantum_id : `uuid.UUID` 

646 Unique ID for the quantum this log belongs to. 

647 provenance_reader : \ 

648 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

649 Reader for the provenance quantum graph. 

650 butler : `lsst.daf.butler.Butler` 

651 Client for the data repository. 

652 """ 

653 # Try reading log through the quantum ID. 

654 ((log1,),) = provenance_reader.fetch_logs([quantum_id]).values() 

655 self.assertIsInstance(log1, ButlerLogRecords) 

656 for _, dataset_id, pipeline_edges in provenance_reader.graph.bipartite_xgraph.out_edges( 

657 quantum_id, data="pipeline_edges" 

658 ): 

659 if pipeline_edges[0].connection_name == acc.LOG_OUTPUT_CONNECTION_NAME: 

660 # Also try reading log through the dataset ID. 

661 ((log2,),) = provenance_reader.fetch_logs([dataset_id]).values() 

662 break 

663 else: 

664 raise AssertionError("No log connection found.") 

665 self.assertEqual(log1, log2) 

666 # Also get the logs from the butler. 

667 ref = butler.get_dataset(dataset_id) 

668 # We have to compare the logs are lists because the 'extra' data is 

669 # different, for now; when the logs are backed by the provenance QG 

670 # this will change (but then we might stop ingesting logs at all). 

671 self.assertEqual(list(butler.get(ref)), list(log1)) 

672 

673 def check_packages(self, provenance_reader: ProvenanceQuantumGraphReader, butler: Butler) -> None: 

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

675 

676 Parameters 

677 ---------- 

678 provenance_reader : \ 

679 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

680 Reader for the provenance quantum graph. 

681 butler : `lsst.daf.butler.Butler` 

682 Client for the data repository. 

683 """ 

684 packages = provenance_reader.fetch_packages() 

685 self.assertIsInstance(packages, Packages) 

686 self.assertIn("pipe_base", packages) 

687 # Also check that the packages are available from the butler. They are 

688 # unfortunately not easy to compare because one passes include_all=True 

689 # and the other passes include_all=False, and latter adds non-portable 

690 # EUPS tag information to the versions. 

691 butler_packages = butler.get("packages", collections=[provenance_reader.header.output_run]) 

692 self.assertFalse( 

693 [ 

694 name 

695 for name, (ver1, ver2) in packages.difference(butler_packages).items() 

696 if not ver1.startswith(ver2) 

697 ] 

698 ) 

699 

700 def check_resource_usage_table( 

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

702 ) -> None: 

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

704 

705 Parameters 

706 ---------- 

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

708 Reader for the provenance quantum graph. 

709 expect_failure : `bool` 

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

711 succeed without writing anything (`False`). 

712 start_time : `float` 

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

714 quantum's execution. 

715 """ 

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

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

718 self.assertCountEqual( 

719 tbl.colnames, 

720 ["quantum_id"] 

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

722 + list(QuantumResourceUsage.model_fields), 

723 ) 

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

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

726 # shenanigans. 

727 end_time = time.time() 

728 for quantum_start_time in tbl["start"]: 

729 self.assertGreater(quantum_start_time, start_time) 

730 self.assertLess(quantum_start_time, end_time) 

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

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

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

734 

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

736 """Check `ProvenanceQuantumGraph.make_quantum_table`. 

737 

738 Parameters 

739 ---------- 

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

741 Reader for the provenance quantum graph. 

742 expect_failure : `bool` 

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

744 succeed without writing anything (`False`). 

745 """ 

746 t = prov.make_quantum_table() 

747 self.assertTrue(np.all(t["Unknown"] == 0)) 

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

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

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

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

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

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

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

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

756 if expect_failure: 

757 # calibrate 

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

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

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

761 # consolidate 

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

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

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

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

766 # resample 

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

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

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

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

771 else: 

772 # calibrate 

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

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

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

776 # consolidate 

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

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

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

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

781 # resample 

782 self.assertEqual(t["Successful"][2], 10) 

783 self.assertEqual(t["Caveats"][2], "*A(2)") 

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

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

786