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.pex.config import Config 

43from lsst.pipe.base import ( 

44 AlgorithmError, 

45 QuantumAttemptStatus, 

46 QuantumSuccessCaveats, 

47 TaskMetadata, 

48) 

49from lsst.pipe.base import automatic_connection_constants as acc 

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

51from lsst.pipe.base.graph_walker import GraphWalker 

52from lsst.pipe.base.pipeline_graph import Edge 

53from lsst.pipe.base.quantum_graph import ( 

54 FORMAT_VERSION, 

55 PredictedDatasetInfo, 

56 PredictedQuantumGraph, 

57 PredictedQuantumInfo, 

58 ProvenanceDatasetInfo, 

59 ProvenanceQuantumGraph, 

60 ProvenanceQuantumInfo, 

61) 

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

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

64from lsst.pipe.base.resource_usage import QuantumResourceUsage 

65from lsst.pipe.base.single_quantum_executor import SingleQuantumExecutor 

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

67 DirectButlerRepo, 

68 DynamicConnectionConfig, 

69 DynamicTestPipelineTaskConfig, 

70) 

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

72from lsst.resources import ResourcePath 

73from lsst.utils.packages import Packages 

74 

75 

76@dataclasses.dataclass 

77class PrepInfo: 

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

79 

80 butler: Butler 

81 butler_path: str 

82 predicted: PredictedQuantumGraph 

83 predicted_path: str 

84 config: AggregatorConfig 

85 

86 

87class AggregatorTestCase(unittest.TestCase): 

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

89 

90 @staticmethod 

91 @contextmanager 

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

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

94 

95 Returns 

96 ------- 

97 prep_info : `PrepInfo` 

98 Objects used in aggregator tests. 

99 

100 Notes 

101 ----- 

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

103 

104 ■ calibrate: {detector, visit} 

105 ╭─┤ 

106 ■ │ consolidate: {visit} 

107 │ 

108 ■ resample: {patch, visit} 

109 │ 

110 ■ coadd: {band, patch} 

111 

112 The data can be visualized via:: 

113 

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

115 

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

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

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

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

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

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

122 'coadd' tasks, too. 

123 """ 

124 with patch_deterministic_uuid4(100): 

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

126 calibrate_config = DynamicTestPipelineTaskConfig() 

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

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

129 helper.add_task( 

130 "calibrate", 

131 config=calibrate_config, 

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

133 inputs={ 

134 "input_image": DynamicConnectionConfig( 

135 dataset_type_name="raw", 

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

137 ) 

138 }, 

139 prerequisite_inputs={ 

140 "refcat": DynamicConnectionConfig( 

141 dataset_type_name="references", 

142 dimensions=["htm7"], 

143 multiple=True, 

144 ) 

145 }, 

146 init_outputs={ 

147 "output_schema": DynamicConnectionConfig( 

148 dataset_type_name="source_schema", 

149 ) 

150 }, 

151 outputs={ 

152 "output_image": DynamicConnectionConfig( 

153 dataset_type_name="image", 

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

155 ), 

156 "output_table": DynamicConnectionConfig( 

157 dataset_type_name="source_detector", 

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

159 ), 

160 }, 

161 ) 

162 helper.add_task( 

163 "consolidate", 

164 dimensions=["visit"], 

165 init_inputs={ 

166 "input_schema": DynamicConnectionConfig( 

167 dataset_type_name="source_schema", 

168 ) 

169 }, 

170 inputs={ 

171 "input_table": DynamicConnectionConfig( 

172 dataset_type_name="source_detector", 

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

174 multiple=True, 

175 ) 

176 }, 

177 outputs={ 

178 "output_table": DynamicConnectionConfig( 

179 dataset_type_name="source", 

180 dimensions=["visit"], 

181 ) 

182 }, 

183 ) 

184 helper.add_task( 

185 "resample", 

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

187 inputs={ 

188 "input_image": DynamicConnectionConfig( 

189 dataset_type_name="image", 

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

191 multiple=True, 

192 ) 

193 }, 

194 outputs={ 

195 "output_image": DynamicConnectionConfig( 

196 dataset_type_name="warp", 

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

198 ) 

199 }, 

200 ) 

201 helper.add_task( 

202 "coadd", 

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

204 inputs={ 

205 "input_image": DynamicConnectionConfig( 

206 dataset_type_name="warp", 

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

208 multiple=True, 

209 ) 

210 }, 

211 outputs={ 

212 "output_image": DynamicConnectionConfig( 

213 dataset_type_name="coadd", 

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

215 ), 

216 }, 

217 ) 

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

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

220 # a convenient temporary directory. 

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

222 pqgc.write(predicted_path) 

223 config = AggregatorConfig( 

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

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

226 # affect large graphs. 

227 ingest_batch_size=10, 

228 zstd_dict_size=256, 

229 zstd_dict_n_inputs=16, 

230 ) 

231 yield PrepInfo( 

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

233 butler_path=root, 

234 predicted=pqgc.assemble(), 

235 predicted_path=predicted_path, 

236 config=config, 

237 ) 

238 

239 def iter_graph_execution( 

240 self, 

241 repo: ResourcePath, 

242 qg: PredictedQuantumGraph, 

243 raise_on_partial_outputs: bool, 

244 is_retry: bool = False, 

245 ) -> Iterator[uuid.UUID]: 

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

247 

248 Parameters 

249 ---------- 

250 repo : `lsst.resources.ResourcePath` 

251 Butler repository path. 

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

253 Predicted quantum graph. Must have datastore records attached, 

254 since execution uses a quantum-backed butler. 

255 raise_on_partial_outputs : `bool` 

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

257 or treat it as a success with caveats. 

258 is_retry : `bool`, optional 

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

260 already be present; skip successes and reprocess failures. 

261 

262 Returns 

263 ------- 

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

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

266 """ 

267 qbb = qg.make_init_qbb(repo) 

268 self.enterContext(qbb) 

269 qg.init_output_run(qbb) 

270 sqe = SingleQuantumExecutor( 

271 limited_butler_factory=lambda quantum: QuantumBackedButler.initialize( 

272 repo, 

273 quantum, 

274 qg.pipeline_graph.universe, 

275 ), 

276 assume_no_existing_outputs=not is_retry, 

277 skip_existing=is_retry, 

278 clobber_outputs=is_retry, 

279 raise_on_partial_outputs=raise_on_partial_outputs, 

280 ) 

281 qg.build_execution_quanta() 

282 xgraph = qg.quantum_only_xgraph 

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

284 for ready in walker: 

285 for quantum_id in ready: 

286 info = xgraph.nodes[quantum_id] 

287 try: 

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

289 except AlgorithmError: 

290 walker.fail(quantum_id) 

291 else: 

292 walker.finish(quantum_id) 

293 yield quantum_id 

294 

295 def check_provenance_graph( 

296 self, 

297 pred: PredictedQuantumGraph, 

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: ProvenanceQuantumGraph = butler.get("run_provenance") 

331 self.maxDiff = None 

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

333 self.assertEqual( 

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

335 [prov.header.output_run] 

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

337 ) 

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

339 for task_label in pred.quanta_by_task: 

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

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

342 for dataset_type_name in prov.datasets_by_type: 

343 self.assertEqual( 

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

345 ) 

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

347 for quantum_id in pred: 

348 # Check consistency between the predicted and provenance quantum 

349 # node attributes. 

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

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

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

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

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

355 # Check consistency between the predicted and provenance dataset 

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

357 # use later. 

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

359 pipeline_edges: list[Edge] 

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

361 quantum_id, data="pipeline_edges" 

362 ): 

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

364 for edge in pipeline_edges: 

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

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

367 ) 

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

369 quantum_id, data="pipeline_edges" 

370 ): 

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

372 for edge in pipeline_edges: 

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

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

375 ) 

376 # Check quantum status and dataset existence against the known 

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

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

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

380 # This is the quantum that can directly raise. 

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

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

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

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

385 if expect_failure: 

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

387 else: 

388 self._expect_successful( 

389 prov_qinfo, 

390 existence, 

391 caveats=( 

392 QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR 

393 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

394 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

395 ), 

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

397 msg=msg, 

398 ) 

399 if expect_failures_retried: 

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

401 self.assertEqual( 

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

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

404 ) 

405 else: 

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

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

408 # This quantum will succeed (with one predicted input 

409 # missing) or be blocked. 

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

411 if expect_failure: 

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

413 else: 

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

415 self.assertEqual( 

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

417 ) 

418 case ( 

419 "resample" | "coadd", 

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

421 ): 

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

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

424 # run. 

425 if expect_failure: 

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

427 else: 

428 self._expect_successful( 

429 prov_qinfo, 

430 existence, 

431 caveats=( 

432 QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED 

433 | QuantumSuccessCaveats.NO_WORK 

434 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

435 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

436 ), 

437 msg=msg, 

438 ) 

439 self.assertEqual( 

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

441 ) 

442 case ( 

443 "resample", 

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

445 ): 

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

447 # regardless. 

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

449 if expect_failure: 

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

451 else: 

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

453 self.assertEqual( 

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

455 ) 

456 case ( 

457 "coadd", 

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

459 ): 

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

461 # with one input missing. 

462 if expect_failure: 

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

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

465 else: 

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

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

468 self.assertEqual( 

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

470 ) 

471 case _: 

472 # All other quanta should succeed and have all inputs 

473 # present. 

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

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

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

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

478 self.check_metadata( 

479 quantum_id, 

480 prov, 

481 butler, 

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

483 ) 

484 self.check_log( 

485 quantum_id, 

486 prov, 

487 butler, 

488 expect_ingested=( 

489 prov_qinfo["status"] 

490 in ( 

491 QuantumAttemptStatus.SUCCESSFUL, 

492 QuantumAttemptStatus.FAILED, 

493 ) 

494 ), 

495 ) 

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

497 self.check_packages(butler) 

498 self.check_configs(butler, prov) 

499 self.check_quantum_table(prov, expect_failure=expect_failure) 

500 self.check_exception_table(prov, expect_failure=expect_failure) 

501 return prov 

502 

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

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

505 

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

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

508 

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

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

511 

512 def _expect_successful( 

513 self, 

514 info: ProvenanceQuantumInfo, 

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

516 caveats: QuantumSuccessCaveats = QuantumSuccessCaveats.NO_CAVEATS, 

517 exception_type: str | None = None, 

518 *, 

519 msg: str, 

520 ) -> None: 

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

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

523 if exception_type is None: 

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

525 else: 

526 assert info["exception"] is not None 

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

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

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

530 if not (caveats & QuantumSuccessCaveats.ANY_OUTPUTS_MISSING): 

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

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

533 if caveats & QuantumSuccessCaveats.ALL_OUTPUTS_MISSING: 

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

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

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

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

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

539 

540 def _expect_failure( 

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

542 ) -> None: 

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

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

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

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

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

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

549 

550 def _expect_blocked( 

551 self, 

552 info: ProvenanceQuantumInfo, 

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

554 msg: str, 

555 ) -> None: 

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

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

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

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

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

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

562 

563 def check_dataset( 

564 self, 

565 dataset_id: uuid.UUID, 

566 pred: PredictedQuantumGraph, 

567 prov: ProvenanceQuantumGraph, 

568 butler: Butler, 

569 ) -> bool: 

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

571 counterpart. 

572 

573 Parameters 

574 ---------- 

575 dataset_id : `uuid.UUID` 

576 Unique ID for the dataset. 

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

578 Predicted quantum graph. 

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

580 Provenance quantum graph. 

581 butler : `lsst.daf.butler.Butler` 

582 Client for the data repository. 

583 

584 Returns 

585 ------- 

586 exists : `bool` 

587 Whether the dataset was marked as existing in the provenance 

588 quantum graph. 

589 """ 

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

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

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

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

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

595 exists = prov_info["produced"] 

596 dataset_type_name = prov_info["dataset_type_name"] 

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

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

599 self.assertEqual( 

600 butler.get_dataset(dataset_id) is not None, 

601 exists, 

602 msg=( 

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

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

605 ), 

606 ) 

607 return exists 

608 

609 def check_metadata( 

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

611 ) -> None: 

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

613 original metadata file has been deleted. 

614 

615 Parameters 

616 ---------- 

617 quantum_id : `uuid.UUID` 

618 Unique ID for the quantum this metadata belongs to. 

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

620 Provenance quantum graph. 

621 butler : `lsst.daf.butler.Butler` 

622 Client for the data repository. 

623 expect_ingested : `bool` 

624 Whether the metadata dataset should have been ingested. 

625 """ 

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

627 ref = butler.get_dataset(dataset_id) 

628 if not expect_ingested: 

629 self.assertIsNone(ref) 

630 return 

631 metadata = butler.get(ref) 

632 self.assertIsInstance(metadata, TaskMetadata) 

633 graph_path = butler.getURI("run_provenance") 

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

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

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

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

638 # it's shared with other datasets. 

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

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

641 self.assertTrue(graph_path.exists()) 

642 self.assertNotEqual(graph_path, original_path) 

643 self.assertFalse(original_path.exists()) 

644 

645 def check_log( 

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

647 ) -> None: 

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

649 original log file has been deleted. 

650 

651 Parameters 

652 ---------- 

653 quantum_id : `uuid.UUID` 

654 Unique ID for the quantum this log belongs to. 

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

656 Provenance quantum graph. 

657 butler : `lsst.daf.butler.Butler` 

658 Client for the data repository. 

659 expect_ingested : `bool` 

660 Whether the metadata dataset should have been ingested. 

661 """ 

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

663 ref = butler.get_dataset(dataset_id) 

664 if not expect_ingested: 

665 self.assertIsNone(ref) 

666 return 

667 log = butler.get(ref) 

668 self.assertIsInstance(log, ButlerLogRecords) 

669 graph_path = butler.getURI("run_provenance") 

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

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

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

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

674 # it's shared with other datasets. 

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

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

677 self.assertTrue(graph_path.exists()) 

678 self.assertNotEqual(graph_path, original_path) 

679 self.assertFalse(original_path.exists()) 

680 

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

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

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

684 self.assertIsInstance(config, Config) 

685 

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

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

688 

689 Parameters 

690 ---------- 

691 butler : `lsst.daf.butler.Butler` 

692 Client for the data repository. 

693 """ 

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

695 self.assertIsInstance(packages, Packages) 

696 self.assertIn("pipe_base", packages) 

697 

698 def check_resource_usage_table( 

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

700 ) -> None: 

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

702 

703 Parameters 

704 ---------- 

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

706 Reader for the provenance quantum graph. 

707 expect_failure : `bool` 

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

709 succeed without writing anything (`False`). 

710 start_time : `float` 

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

712 quantum's execution. 

713 """ 

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

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

716 self.assertCountEqual( 

717 tbl.colnames, 

718 ["quantum_id"] 

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

720 + list(QuantumResourceUsage.model_fields), 

721 ) 

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

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

724 # shenanigans. 

725 end_time = time.time() 

726 for quantum_start_time in tbl["start"]: 

727 self.assertGreater(quantum_start_time, start_time) 

728 self.assertLess(quantum_start_time, end_time) 

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

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

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

732 

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

734 """Check `ProvenanceQuantumGraph.make_quantum_table`. 

735 

736 Parameters 

737 ---------- 

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

739 Reader for the provenance quantum graph. 

740 expect_failure : `bool` 

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

742 succeed without writing anything (`False`). 

743 """ 

744 t = prov.make_quantum_table() 

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

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

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

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

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

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

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

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

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

754 if expect_failure: 

755 # calibrate 

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

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

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

759 # consolidate 

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

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

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

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

764 # resample 

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

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

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

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

769 else: 

770 # calibrate 

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

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

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

774 # consolidate 

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

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

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

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

779 # resample 

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

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

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

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