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 qbb = qg.make_init_qbb(repo) 

267 self.enterContext(qbb) 

268 qg.init_output_run(qbb) 

269 sqe = SingleQuantumExecutor( 

270 limited_butler_factory=lambda quantum: QuantumBackedButler.initialize( 

271 repo, 

272 quantum, 

273 qg.pipeline_graph.universe, 

274 ), 

275 assume_no_existing_outputs=not is_retry, 

276 skip_existing=is_retry, 

277 clobber_outputs=is_retry, 

278 raise_on_partial_outputs=raise_on_partial_outputs, 

279 ) 

280 qg.build_execution_quanta() 

281 xgraph = qg.quantum_only_xgraph 

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

283 for ready in walker: 

284 for quantum_id in ready: 

285 info = xgraph.nodes[quantum_id] 

286 try: 

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

288 except AlgorithmError: 

289 walker.fail(quantum_id) 

290 else: 

291 walker.finish(quantum_id) 

292 yield quantum_id 

293 

294 def check_provenance_graph( 

295 self, 

296 pred: PredictedQuantumGraph, 

297 prov_reader: ProvenanceQuantumGraphReader, 

298 butler: Butler, 

299 expect_failure: bool, 

300 start_time: float, 

301 expect_failures_retried: bool = False, 

302 ) -> ProvenanceQuantumGraph: 

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

304 scanning the graph created by `make_test_repo`. 

305 

306 Parameters 

307 ---------- 

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

309 Predicted quantum graph. 

310 prov_reader : \ 

311 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

312 Reader for the provenance quantum graph. 

313 butler : `lsst.daf.butler.Butler` 

314 Client for the data repository. 

315 expect_failure : `bool` 

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

317 succeed without writing anything (`False`). 

318 start_time : `float` 

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

320 quantum's execution. 

321 expect_failures_retried : `bool`, optional 

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

323 most recent attempt. 

324 

325 Returns 

326 ------- 

327 prov : `ProvenanceQuantumGraph` 

328 The full provenance quantum graph. 

329 """ 

330 prov_reader.read_full_graph() 

331 prov = prov_reader.graph 

332 checked_some_metadata = False 

333 checked_some_log = False 

334 self.maxDiff = None 

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

336 self.assertEqual( 

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

338 [prov.header.output_run] 

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

340 ) 

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

342 for task_label in pred.quanta_by_task: 

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

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

345 for dataset_type_name in prov.datasets_by_type: 

346 self.assertEqual( 

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

348 ) 

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

350 for quantum_id in pred: 

351 # Check consistency between the predicted and provenance quantum 

352 # node attributes. 

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

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

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

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

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

358 # Check consistency between the predicted and provenance dataset 

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

360 # use later. 

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

362 pipeline_edges: list[Edge] 

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

364 quantum_id, data="pipeline_edges" 

365 ): 

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

367 for edge in pipeline_edges: 

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

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

370 ) 

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

372 quantum_id, data="pipeline_edges" 

373 ): 

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

375 for edge in pipeline_edges: 

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

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

378 ) 

379 # Check quantum status and dataset existence against the known 

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

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

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

383 # This is the quantum that can directly raise. 

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

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

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

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

388 if expect_failure: 

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

390 else: 

391 self._expect_successful( 

392 prov_qinfo, 

393 existence, 

394 caveats=( 

395 QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR 

396 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

397 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

398 ), 

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

400 msg=msg, 

401 ) 

402 if expect_failures_retried: 

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

404 self.assertEqual( 

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

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

407 ) 

408 else: 

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

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

411 # This quantum will succeed (with one predicted input 

412 # missing) or be blocked. 

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

414 if expect_failure: 

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

416 else: 

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

418 self.assertEqual( 

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

420 ) 

421 case ( 

422 "resample" | "coadd", 

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

424 ): 

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

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

427 # run. 

428 if expect_failure: 

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

430 else: 

431 self._expect_successful( 

432 prov_qinfo, 

433 existence, 

434 caveats=( 

435 QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED 

436 | QuantumSuccessCaveats.NO_WORK 

437 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

438 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

439 ), 

440 msg=msg, 

441 ) 

442 self.assertEqual( 

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

444 ) 

445 case ( 

446 "resample", 

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

448 ): 

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

450 # regardless. 

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

452 if expect_failure: 

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

454 else: 

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

456 self.assertEqual( 

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

458 ) 

459 case ( 

460 "coadd", 

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

462 ): 

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

464 # with one input missing. 

465 if expect_failure: 

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

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

468 else: 

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

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

471 self.assertEqual( 

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

473 ) 

474 case _: 

475 # All other quanta should succeed and have all inputs 

476 # present. 

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

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

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

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

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

482 self.check_metadata(quantum_id, prov_reader, butler) 

483 checked_some_metadata = True 

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

485 QuantumAttemptStatus.SUCCESSFUL, 

486 QuantumAttemptStatus.FAILED, 

487 ): 

488 self.check_log(quantum_id, prov_reader, butler) 

489 checked_some_log = True 

490 self.assertTrue(checked_some_metadata) 

491 self.assertTrue(checked_some_log) 

492 self.check_resource_usage_table( 

493 prov_reader.graph, expect_failure=expect_failure, start_time=start_time 

494 ) 

495 self.check_packages(prov_reader, butler) 

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

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

498 return prov 

499 

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

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

502 

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

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

505 

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

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

508 

509 def _expect_successful( 

510 self, 

511 info: ProvenanceQuantumInfo, 

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

513 caveats: QuantumSuccessCaveats = QuantumSuccessCaveats.NO_CAVEATS, 

514 exception_type: str | None = None, 

515 *, 

516 msg: str, 

517 ) -> None: 

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

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

520 if exception_type is None: 

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

522 else: 

523 assert info["exception"] is not None 

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

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

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

527 if not (caveats & QuantumSuccessCaveats.ANY_OUTPUTS_MISSING): 

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

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

530 if caveats & QuantumSuccessCaveats.ALL_OUTPUTS_MISSING: 

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

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

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

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

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

536 

537 def _expect_failure( 

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

539 ) -> None: 

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

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

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

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

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

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

546 

547 def _expect_blocked( 

548 self, 

549 info: ProvenanceQuantumInfo, 

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

551 msg: str, 

552 ) -> None: 

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

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

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

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

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

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

559 

560 def check_dataset( 

561 self, 

562 dataset_id: uuid.UUID, 

563 pred: PredictedQuantumGraph, 

564 prov: ProvenanceQuantumGraph, 

565 butler: Butler, 

566 ) -> bool: 

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

568 counterpart. 

569 

570 Parameters 

571 ---------- 

572 dataset_id : `uuid.UUID` 

573 Unique ID for the dataset. 

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

575 Predicted quantum graph. 

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

577 Provenance quantum graph. 

578 butler : `lsst.daf.butler.Butler` 

579 Client for the data repository. 

580 

581 Returns 

582 ------- 

583 exists : `bool` 

584 Whether the dataset was marked as existing in the provenance 

585 quantum graph. 

586 """ 

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

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

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

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

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

592 exists = prov_info["produced"] 

593 dataset_type_name = prov_info["dataset_type_name"] 

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

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

596 self.assertEqual( 

597 butler.get_dataset(dataset_id) is not None, 

598 exists, 

599 msg=( 

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

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

602 ), 

603 ) 

604 return exists 

605 

606 def check_metadata( 

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

608 ) -> None: 

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

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

611 

612 Parameters 

613 ---------- 

614 quantum_id : `uuid.UUID` 

615 Unique ID for the quantum this metadata belongs to. 

616 provenance_reader : \ 

617 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

618 Reader for the provenance quantum graph. 

619 butler : `lsst.daf.butler.Butler` 

620 Client for the data repository. 

621 """ 

622 # Try reading metadata through the quantum ID. 

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

624 self.assertIsInstance(metadata1, TaskMetadata) 

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

626 quantum_id, data="pipeline_edges" 

627 ): 

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

629 # Also try reading metadata through the dataset ID. 

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

631 break 

632 else: 

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

634 self.assertEqual(metadata1, metadata2) 

635 # Also get the metadata from the butler. 

636 ref = butler.get_dataset(dataset_id) 

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

638 

639 def check_log( 

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

641 ) -> None: 

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

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

644 

645 Parameters 

646 ---------- 

647 quantum_id : `uuid.UUID` 

648 Unique ID for the quantum this log belongs to. 

649 provenance_reader : \ 

650 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

651 Reader for the provenance quantum graph. 

652 butler : `lsst.daf.butler.Butler` 

653 Client for the data repository. 

654 """ 

655 # Try reading log through the quantum ID. 

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

657 self.assertIsInstance(log1, ButlerLogRecords) 

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

659 quantum_id, data="pipeline_edges" 

660 ): 

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

662 # Also try reading log through the dataset ID. 

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

664 break 

665 else: 

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

667 self.assertEqual(log1, log2) 

668 # Also get the logs from the butler. 

669 ref = butler.get_dataset(dataset_id) 

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

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

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

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

674 

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

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

677 

678 Parameters 

679 ---------- 

680 provenance_reader : \ 

681 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

682 Reader for the provenance quantum graph. 

683 butler : `lsst.daf.butler.Butler` 

684 Client for the data repository. 

685 """ 

686 packages = provenance_reader.fetch_packages() 

687 self.assertIsInstance(packages, Packages) 

688 self.assertIn("pipe_base", packages) 

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

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

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

692 # EUPS tag information to the versions. 

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

694 self.assertFalse( 

695 [ 

696 name 

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

698 if not ver1.startswith(ver2) 

699 ] 

700 ) 

701 

702 def check_resource_usage_table( 

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

704 ) -> None: 

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

706 

707 Parameters 

708 ---------- 

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

710 Reader for the provenance quantum graph. 

711 expect_failure : `bool` 

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

713 succeed without writing anything (`False`). 

714 start_time : `float` 

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

716 quantum's execution. 

717 """ 

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

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

720 self.assertCountEqual( 

721 tbl.colnames, 

722 ["quantum_id"] 

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

724 + list(QuantumResourceUsage.model_fields), 

725 ) 

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

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

728 # shenanigans. 

729 end_time = time.time() 

730 for quantum_start_time in tbl["start"]: 

731 self.assertGreater(quantum_start_time, start_time) 

732 self.assertLess(quantum_start_time, end_time) 

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

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

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

736 

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

738 """Check `ProvenanceQuantumGraph.make_quantum_table`. 

739 

740 Parameters 

741 ---------- 

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

743 Reader for the provenance quantum graph. 

744 expect_failure : `bool` 

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

746 succeed without writing anything (`False`). 

747 """ 

748 t = prov.make_quantum_table() 

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

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

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

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

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

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

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

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

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

758 if expect_failure: 

759 # calibrate 

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

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

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

763 # consolidate 

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

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

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

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

768 # resample 

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

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

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

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

773 else: 

774 # calibrate 

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

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

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

778 # consolidate 

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

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

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

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

783 # resample 

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

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

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

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