Coverage for tests/test_aggregator.py: 14%

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 AlgorithmError, QuantumSuccessCaveats, TaskMetadata 

43from lsst.pipe.base import automatic_connection_constants as acc 

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

45from lsst.pipe.base.graph_walker import GraphWalker 

46from lsst.pipe.base.pipeline_graph import Edge 

47from lsst.pipe.base.quantum_graph import ( 

48 PredictedDatasetInfo, 

49 PredictedQuantumGraph, 

50 PredictedQuantumInfo, 

51 ProvenanceDatasetInfo, 

52 ProvenanceQuantumGraph, 

53 ProvenanceQuantumGraphReader, 

54 ProvenanceQuantumInfo, 

55) 

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

57from lsst.pipe.base.quantum_provenance_graph import QuantumRunStatus 

58from lsst.pipe.base.resource_usage import QuantumResourceUsage 

59from lsst.pipe.base.single_quantum_executor import SingleQuantumExecutor 

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

61 DirectButlerRepo, 

62 DynamicConnectionConfig, 

63 DynamicTestPipelineTaskConfig, 

64) 

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

66from lsst.resources import ResourcePath 

67from lsst.utils.packages import Packages 

68 

69 

70@dataclasses.dataclass 

71class PrepInfo: 

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

73 

74 butler: Butler 

75 butler_path: str 

76 predicted: PredictedQuantumGraph 

77 predicted_path: str 

78 config: AggregatorConfig 

79 

80 

81class AggregatorTestCase(unittest.TestCase): 

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

83 

84 @staticmethod 

85 @contextmanager 

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

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

88 

89 Returns 

90 ------- 

91 prep_info : `PrepInfo` 

92 Objects used in aggregator tests. 

93 

94 Notes 

95 ----- 

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

97 

98 ■ calibrate: {detector, visit} 

99 ╭─┤ 

100 ■ │ consolidate: {visit} 

101 │ 

102 ■ resample: {patch, visit} 

103 │ 

104 ■ coadd: {band, patch} 

105 

106 The data can be visualized via:: 

107 

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

109 

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

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

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

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

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

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

116 'coadd' tasks, too. 

117 """ 

118 with patch_deterministic_uuid4(100): 

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

120 calibrate_config = DynamicTestPipelineTaskConfig() 

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

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

123 helper.add_task( 

124 "calibrate", 

125 config=calibrate_config, 

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

127 inputs={ 

128 "input_image": DynamicConnectionConfig( 

129 dataset_type_name="raw", 

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

131 ) 

132 }, 

133 prerequisite_inputs={ 

134 "refcat": DynamicConnectionConfig( 

135 dataset_type_name="references", 

136 dimensions=["htm7"], 

137 multiple=True, 

138 ) 

139 }, 

140 init_outputs={ 

141 "output_schema": DynamicConnectionConfig( 

142 dataset_type_name="source_schema", 

143 ) 

144 }, 

145 outputs={ 

146 "output_image": DynamicConnectionConfig( 

147 dataset_type_name="image", 

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

149 ), 

150 "output_table": DynamicConnectionConfig( 

151 dataset_type_name="source_detector", 

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

153 ), 

154 }, 

155 ) 

156 helper.add_task( 

157 "consolidate", 

158 dimensions=["visit"], 

159 init_inputs={ 

160 "input_schema": DynamicConnectionConfig( 

161 dataset_type_name="source_schema", 

162 ) 

163 }, 

164 inputs={ 

165 "input_table": DynamicConnectionConfig( 

166 dataset_type_name="source_detector", 

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

168 multiple=True, 

169 ) 

170 }, 

171 outputs={ 

172 "output_table": DynamicConnectionConfig( 

173 dataset_type_name="source", 

174 dimensions=["visit"], 

175 ) 

176 }, 

177 ) 

178 helper.add_task( 

179 "resample", 

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

181 inputs={ 

182 "input_image": DynamicConnectionConfig( 

183 dataset_type_name="image", 

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

185 multiple=True, 

186 ) 

187 }, 

188 outputs={ 

189 "output_image": DynamicConnectionConfig( 

190 dataset_type_name="warp", 

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

192 ) 

193 }, 

194 ) 

195 helper.add_task( 

196 "coadd", 

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

198 inputs={ 

199 "input_image": DynamicConnectionConfig( 

200 dataset_type_name="warp", 

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

202 multiple=True, 

203 ) 

204 }, 

205 outputs={ 

206 "output_image": DynamicConnectionConfig( 

207 dataset_type_name="coadd", 

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

209 ), 

210 }, 

211 ) 

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

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

214 # a convenient temporary directory. 

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

216 pqgc.write(predicted_path) 

217 config = AggregatorConfig( 

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

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

220 # affect large graphs. 

221 ingest_batch_size=10, 

222 zstd_dict_size=256, 

223 zstd_dict_n_inputs=16, 

224 ) 

225 yield PrepInfo( 

226 butler=helper.butler, 

227 butler_path=root, 

228 predicted=pqgc.assemble(), 

229 predicted_path=predicted_path, 

230 config=config, 

231 ) 

232 

233 def iter_graph_execution( 

234 self, repo: ResourcePath, qg: PredictedQuantumGraph, raise_on_partial_outputs: bool 

235 ) -> Iterator[uuid.UUID]: 

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

237 

238 Parameters 

239 ---------- 

240 repo : `lsst.resources.ResourcePath` 

241 Butler repository path. 

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

243 Predicted quantum graph. Must have datastore records attached, 

244 since execution uses a quantum-backed butler. 

245 raise_on_partial_outputs : `bool` 

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

247 or treat it as a success with caveats. 

248 

249 Returns 

250 ------- 

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

252 An iterator over successful quantum IDs. Failed and blocked quanta 

253 are not included. 

254 """ 

255 qg.init_output_run(qg.make_init_qbb(repo)) 

256 sqe = SingleQuantumExecutor( 

257 limited_butler_factory=lambda quantum: QuantumBackedButler.initialize( 

258 repo, 

259 quantum, 

260 qg.pipeline_graph.universe, 

261 ), 

262 assume_no_existing_outputs=True, 

263 raise_on_partial_outputs=raise_on_partial_outputs, 

264 ) 

265 qg.build_execution_quanta() 

266 xgraph = qg.quantum_only_xgraph 

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

268 for ready in walker: 

269 for quantum_id in ready: 

270 info = xgraph.nodes[quantum_id] 

271 try: 

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

273 except AlgorithmError: 

274 walker.fail(quantum_id) 

275 else: 

276 walker.finish(quantum_id) 

277 yield quantum_id 

278 

279 def check_provenance_graph( 

280 self, 

281 pred: PredictedQuantumGraph, 

282 prov_reader: ProvenanceQuantumGraphReader, 

283 butler: Butler, 

284 expect_failure: bool, 

285 start_time: float, 

286 ) -> None: 

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

288 scanning the graph created by `make_test_repo`. 

289 

290 Parameters 

291 ---------- 

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

293 Predicted quantum graph. 

294 prov_reader : \ 

295 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

296 Reader for the provenance quantum graph. 

297 butler : `lsst.daf.butler.Butler` 

298 Client for the data repository. 

299 expect_failure : `bool` 

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

301 succeed without writing anything (`False`). 

302 start_time : `float` 

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

304 quantum's execution. 

305 """ 

306 prov_reader.read_full_graph() 

307 prov = prov_reader.graph 

308 checked_some_metadata = False 

309 checked_some_log = False 

310 self.maxDiff = None 

311 self.assertEqual( 

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

313 [prov.header.output_run] 

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

315 ) 

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

317 for task_label in pred.quanta_by_task: 

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

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

320 for dataset_type_name in prov.datasets_by_type: 

321 self.assertEqual( 

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

323 ) 

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

325 for quantum_id in pred: 

326 # Check consistency between the predicted and provenance quantum 

327 # node attributes. 

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

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

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

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

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

333 # Check consistency between the predicted and provenance dataset 

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

335 # use later. 

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

337 pipeline_edges: list[Edge] 

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

339 quantum_id, data="pipeline_edges" 

340 ): 

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

342 for edge in pipeline_edges: 

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

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

345 ) 

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

347 quantum_id, data="pipeline_edges" 

348 ): 

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

350 for edge in pipeline_edges: 

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

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

353 ) 

354 # Check quantum status and dataset existence against the known 

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

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

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

358 # This is the quantum that can directly raise. 

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

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

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

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

363 if expect_failure: 

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

365 else: 

366 self._expect_successful( 

367 prov_qinfo, 

368 existence, 

369 caveats=( 

370 QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR 

371 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

372 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

373 ), 

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

375 msg=msg, 

376 ) 

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

378 # This quantum will succeed (with one predicted input 

379 # missing) or be blocked. 

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

381 if expect_failure: 

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

383 else: 

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

385 case ( 

386 "resample" | "coadd", 

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

388 ): 

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

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

391 # run. 

392 if expect_failure: 

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

394 else: 

395 self._expect_successful( 

396 prov_qinfo, 

397 existence, 

398 caveats=( 

399 QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED 

400 | QuantumSuccessCaveats.NO_WORK 

401 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

402 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

403 ), 

404 msg=msg, 

405 ) 

406 case ( 

407 "resample", 

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

409 ): 

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

411 # regardless. 

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

413 if expect_failure: 

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

415 else: 

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

417 case ( 

418 "coadd", 

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

420 ): 

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

422 # with one input missing. 

423 if expect_failure: 

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

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

426 else: 

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

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

429 case _: 

430 # All other quanta should succeed and have all inputs 

431 # present. 

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

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

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

435 if not checked_some_metadata and prov_qinfo["status"] is QuantumRunStatus.SUCCESSFUL: 

436 self.check_metadata(quantum_id, prov_reader, butler) 

437 checked_some_metadata = True 

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

439 QuantumRunStatus.SUCCESSFUL, 

440 QuantumRunStatus.FAILED, 

441 ): 

442 self.check_log(quantum_id, prov_reader, butler) 

443 checked_some_log = True 

444 self.assertTrue(checked_some_metadata) 

445 self.assertTrue(checked_some_log) 

446 self.check_resource_usage_table( 

447 prov_reader.graph, expect_failure=expect_failure, start_time=start_time 

448 ) 

449 self.check_packages(prov_reader) 

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

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

452 

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

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

455 

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

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

458 

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

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

461 

462 def _expect_successful( 

463 self, 

464 info: ProvenanceQuantumInfo, 

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

466 caveats: QuantumSuccessCaveats = QuantumSuccessCaveats.NO_CAVEATS, 

467 exception_type: str | None = None, 

468 *, 

469 msg: str, 

470 ) -> None: 

471 self.assertEqual(info["status"], QuantumRunStatus.SUCCESSFUL, msg=msg) 

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

473 if exception_type is None: 

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

475 else: 

476 assert info["exception"] is not None 

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

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

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

480 if not (caveats & QuantumSuccessCaveats.ANY_OUTPUTS_MISSING): 

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

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

483 if caveats & QuantumSuccessCaveats.ALL_OUTPUTS_MISSING: 

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

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

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

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

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

489 

490 def _expect_failure( 

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

492 ) -> None: 

493 self.assertEqual(info["status"], QuantumRunStatus.FAILED, msg=msg) 

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

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

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

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

498 

499 def _expect_blocked( 

500 self, 

501 info: ProvenanceQuantumInfo, 

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

503 msg: str, 

504 ) -> None: 

505 self.assertEqual(info["status"], QuantumRunStatus.BLOCKED, msg=msg) 

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

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

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

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

510 

511 def check_dataset( 

512 self, 

513 dataset_id: uuid.UUID, 

514 pred: PredictedQuantumGraph, 

515 prov: ProvenanceQuantumGraph, 

516 butler: Butler, 

517 ) -> bool: 

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

519 counterpart. 

520 

521 Parameters 

522 ---------- 

523 dataset_id : `uuid.UUID` 

524 Unique ID for the dataset. 

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

526 Predicted quantum graph. 

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

528 Provenance quantum graph. 

529 butler : `lsst.daf.butler.Butler` 

530 Client for the data repository. 

531 

532 Returns 

533 ------- 

534 exists : `bool` 

535 Whether the dataset was marked as existing in the provenance 

536 quantum graph. 

537 """ 

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

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

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

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

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

543 exists = prov_info["exists"] 

544 dataset_type_name = prov_info["dataset_type_name"] 

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

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

547 self.assertEqual( 

548 butler.get_dataset(dataset_id) is not None, 

549 exists, 

550 msg=( 

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

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

553 ), 

554 ) 

555 return exists 

556 

557 def check_metadata( 

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

559 ) -> None: 

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

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

562 

563 Parameters 

564 ---------- 

565 quantum_id : `uuid.UUID` 

566 Unique ID for the quantum this metadata belongs to. 

567 provenance_reader : \ 

568 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

569 Reader for the provenance quantum graph. 

570 butler : `lsst.daf.butler.Butler` 

571 Client for the data repository. 

572 """ 

573 # Try reading metadata through the quantum ID. 

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

575 self.assertIsInstance(metadata1, TaskMetadata) 

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

577 quantum_id, data="pipeline_edges" 

578 ): 

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

580 # Also try reading metadata through the dataset ID. 

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

582 break 

583 else: 

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

585 self.assertEqual(metadata1, metadata2) 

586 

587 def check_log( 

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

589 ) -> None: 

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

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

592 

593 Parameters 

594 ---------- 

595 quantum_id : `uuid.UUID` 

596 Unique ID for the quantum this log belongs to. 

597 provenance_reader : \ 

598 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

599 Reader for the provenance quantum graph. 

600 butler : `lsst.daf.butler.Butler` 

601 Client for the data repository. 

602 """ 

603 # Try reading log through the quantum ID. 

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

605 self.assertIsInstance(log1, ButlerLogRecords) 

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

607 quantum_id, data="pipeline_edges" 

608 ): 

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

610 # Also try reading log through the dataset ID. 

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

612 break 

613 else: 

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

615 self.assertEqual(log1, log2) 

616 

617 def check_packages(self, provenance_reader: ProvenanceQuantumGraphReader) -> None: 

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

619 

620 Parameters 

621 ---------- 

622 provenance_reader : \ 

623 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

624 Reader for the provenance quantum graph. 

625 """ 

626 packages = provenance_reader.fetch_packages() 

627 self.assertIsInstance(packages, Packages) 

628 self.assertIn("pipe_base", packages) 

629 

630 def check_resource_usage_table( 

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

632 ) -> None: 

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

634 

635 Parameters 

636 ---------- 

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

638 Reader for the provenance quantum graph. 

639 expect_failure : `bool` 

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

641 succeed without writing anything (`False`). 

642 start_time : `float` 

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

644 quantum's execution. 

645 """ 

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

647 if expect_failure: 

648 self.assertEqual(len(tbl), prov.header.n_task_quanta["calibrate"] - 1) 

649 else: 

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

651 self.assertCountEqual( 

652 tbl.colnames, 

653 ["quantum_id"] 

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

655 + list(QuantumResourceUsage.model_fields), 

656 ) 

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

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

659 # shenanigans. 

660 end_time = time.time() 

661 for quantum_start_time in tbl["start"]: 

662 self.assertGreater(quantum_start_time, start_time) 

663 self.assertLess(quantum_start_time, end_time) 

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

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

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

667 

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

669 """Check `ProvenanceQuantumGraph.make_quantum_table`. 

670 

671 Parameters 

672 ---------- 

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

674 Reader for the provenance quantum graph. 

675 expect_failure : `bool` 

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

677 succeed without writing anything (`False`). 

678 """ 

679 t = prov.make_quantum_table() 

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

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

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

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

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

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

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

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

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

689 if expect_failure: 

690 # calibrate 

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

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

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

694 # consolidate 

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

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

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

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

699 # resample 

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

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

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

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

704 else: 

705 # calibrate 

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

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

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

709 # consolidate 

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

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

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

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

714 # resample 

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

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

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

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

719 

720 def check_exception_table(self, prov: ProvenanceQuantumGraph, expect_failure: bool) -> None: 

721 """Check `ProvenanceQuantumGraph.make_exception_table`. 

722 

723 Parameters 

724 ---------- 

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

726 Reader for the provenance quantum graph. 

727 expect_failure : `bool` 

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

729 succeed without writing anything (`False`). 

730 """ 

731 t = prov.make_exception_table() 

732 if expect_failure: 

733 self.assertEqual(len(t), 0) 

734 else: 

735 self.assertEqual(list(t["Task"]), ["calibrate"]) 

736 self.assertEqual(list(t["Exception"]), ["lsst.pipe.base.tests.mocks.MockAlgorithmError"]) 

737 self.assertEqual(list(t["Count"]), [1]) 

738 

739 def test_all_successful(self) -> None: 

740 """Test running a full graph with no failures, and then scanning the 

741 results with assume_complete=False. 

742 

743 Because there are no failures, assume_complete should not be necessary. 

744 """ 

745 with self.make_test_repo() as prep: 

746 prep.config.assume_complete = False 

747 start_time = time.time() 

748 executed_quanta = list( 

749 self.iter_graph_execution(prep.butler_path, prep.predicted, raise_on_partial_outputs=False) 

750 ) 

751 self.assertCountEqual(executed_quanta, prep.predicted.quantum_only_xgraph.nodes.keys()) 

752 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config) 

753 with ProvenanceQuantumGraphReader.open(prep.config.output_path) as reader: 

754 self.check_provenance_graph( 

755 prep.predicted, reader, prep.butler, expect_failure=False, start_time=start_time 

756 ) 

757 

758 def test_all_successful_two_phase(self) -> None: 

759 """Test running some of a graph with no failures, scanning with 

760 assume_complete=False, then finishing the graph and scanning again. 

761 """ 

762 with self.make_test_repo() as prep: 

763 start_time = time.time() 

764 execution_iter = self.iter_graph_execution( 

765 prep.butler_path, prep.predicted, raise_on_partial_outputs=False 

766 ) 

767 executed_quanta = list(itertools.islice(execution_iter, 9)) 

768 self.assertEqual(len(executed_quanta), 9) 

769 # Run the scanner while telling it to assume failures might change, 

770 # so it just waits for incomplete quanta to finish (and then times 

771 # out). 

772 prep.config.assume_complete = False 

773 with self.assertRaises(RuntimeError): 

774 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config) 

775 # Finish executing the quanta. 

776 executed_quanta.extend(execution_iter) 

777 # Scan again, and write the provenance QG. 

778 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config) 

779 # Run the scanner again.