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 PredictedDatasetInfo, 

54 PredictedQuantumGraph, 

55 PredictedQuantumInfo, 

56 ProvenanceDatasetInfo, 

57 ProvenanceQuantumGraph, 

58 ProvenanceQuantumGraphReader, 

59 ProvenanceQuantumInfo, 

60) 

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

62from lsst.pipe.base.resource_usage import QuantumResourceUsage 

63from lsst.pipe.base.single_quantum_executor import SingleQuantumExecutor 

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

65 DirectButlerRepo, 

66 DynamicConnectionConfig, 

67 DynamicTestPipelineTaskConfig, 

68) 

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

70from lsst.resources import ResourcePath 

71from lsst.utils.packages import Packages 

72 

73 

74@dataclasses.dataclass 

75class PrepInfo: 

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

77 

78 butler: Butler 

79 butler_path: str 

80 predicted: PredictedQuantumGraph 

81 predicted_path: str 

82 config: AggregatorConfig 

83 

84 

85class AggregatorTestCase(unittest.TestCase): 

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

87 

88 @staticmethod 

89 @contextmanager 

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

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

92 

93 Returns 

94 ------- 

95 prep_info : `PrepInfo` 

96 Objects used in aggregator tests. 

97 

98 Notes 

99 ----- 

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

101 

102 ■ calibrate: {detector, visit} 

103 ╭─┤ 

104 ■ │ consolidate: {visit} 

105 │ 

106 ■ resample: {patch, visit} 

107 │ 

108 ■ coadd: {band, patch} 

109 

110 The data can be visualized via:: 

111 

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

113 

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

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

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

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

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

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

120 'coadd' tasks, too. 

121 """ 

122 with patch_deterministic_uuid4(100): 

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

124 calibrate_config = DynamicTestPipelineTaskConfig() 

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

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

127 helper.add_task( 

128 "calibrate", 

129 config=calibrate_config, 

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

131 inputs={ 

132 "input_image": DynamicConnectionConfig( 

133 dataset_type_name="raw", 

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

135 ) 

136 }, 

137 prerequisite_inputs={ 

138 "refcat": DynamicConnectionConfig( 

139 dataset_type_name="references", 

140 dimensions=["htm7"], 

141 multiple=True, 

142 ) 

143 }, 

144 init_outputs={ 

145 "output_schema": DynamicConnectionConfig( 

146 dataset_type_name="source_schema", 

147 ) 

148 }, 

149 outputs={ 

150 "output_image": DynamicConnectionConfig( 

151 dataset_type_name="image", 

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

153 ), 

154 "output_table": DynamicConnectionConfig( 

155 dataset_type_name="source_detector", 

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

157 ), 

158 }, 

159 ) 

160 helper.add_task( 

161 "consolidate", 

162 dimensions=["visit"], 

163 init_inputs={ 

164 "input_schema": DynamicConnectionConfig( 

165 dataset_type_name="source_schema", 

166 ) 

167 }, 

168 inputs={ 

169 "input_table": DynamicConnectionConfig( 

170 dataset_type_name="source_detector", 

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

172 multiple=True, 

173 ) 

174 }, 

175 outputs={ 

176 "output_table": DynamicConnectionConfig( 

177 dataset_type_name="source", 

178 dimensions=["visit"], 

179 ) 

180 }, 

181 ) 

182 helper.add_task( 

183 "resample", 

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

185 inputs={ 

186 "input_image": DynamicConnectionConfig( 

187 dataset_type_name="image", 

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

189 multiple=True, 

190 ) 

191 }, 

192 outputs={ 

193 "output_image": DynamicConnectionConfig( 

194 dataset_type_name="warp", 

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

196 ) 

197 }, 

198 ) 

199 helper.add_task( 

200 "coadd", 

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

202 inputs={ 

203 "input_image": DynamicConnectionConfig( 

204 dataset_type_name="warp", 

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

206 multiple=True, 

207 ) 

208 }, 

209 outputs={ 

210 "output_image": DynamicConnectionConfig( 

211 dataset_type_name="coadd", 

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

213 ), 

214 }, 

215 ) 

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

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

218 # a convenient temporary directory. 

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

220 pqgc.write(predicted_path) 

221 config = AggregatorConfig( 

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

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

224 # affect large graphs. 

225 ingest_batch_size=10, 

226 zstd_dict_size=256, 

227 zstd_dict_n_inputs=16, 

228 ) 

229 yield PrepInfo( 

230 butler=helper.butler, 

231 butler_path=root, 

232 predicted=pqgc.assemble(), 

233 predicted_path=predicted_path, 

234 config=config, 

235 ) 

236 

237 def iter_graph_execution( 

238 self, 

239 repo: ResourcePath, 

240 qg: PredictedQuantumGraph, 

241 raise_on_partial_outputs: bool, 

242 is_retry: bool = False, 

243 ) -> Iterator[uuid.UUID]: 

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

245 

246 Parameters 

247 ---------- 

248 repo : `lsst.resources.ResourcePath` 

249 Butler repository path. 

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

251 Predicted quantum graph. Must have datastore records attached, 

252 since execution uses a quantum-backed butler. 

253 raise_on_partial_outputs : `bool` 

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

255 or treat it as a success with caveats. 

256 is_retry : `bool`, optional 

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

258 already be present; skip successes and reprocess failures. 

259 

260 Returns 

261 ------- 

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

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

264 """ 

265 qg.init_output_run(qg.make_init_qbb(repo)) 

266 sqe = SingleQuantumExecutor( 

267 limited_butler_factory=lambda quantum: QuantumBackedButler.initialize( 

268 repo, 

269 quantum, 

270 qg.pipeline_graph.universe, 

271 ), 

272 assume_no_existing_outputs=not is_retry, 

273 skip_existing=is_retry, 

274 clobber_outputs=is_retry, 

275 raise_on_partial_outputs=raise_on_partial_outputs, 

276 ) 

277 qg.build_execution_quanta() 

278 xgraph = qg.quantum_only_xgraph 

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

280 for ready in walker: 

281 for quantum_id in ready: 

282 info = xgraph.nodes[quantum_id] 

283 try: 

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

285 except AlgorithmError: 

286 walker.fail(quantum_id) 

287 else: 

288 walker.finish(quantum_id) 

289 yield quantum_id 

290 

291 def check_provenance_graph( 

292 self, 

293 pred: PredictedQuantumGraph, 

294 prov_reader: ProvenanceQuantumGraphReader, 

295 butler: Butler, 

296 expect_failure: bool, 

297 start_time: float, 

298 expect_failures_retried: bool = False, 

299 ) -> ProvenanceQuantumGraph: 

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

301 scanning the graph created by `make_test_repo`. 

302 

303 Parameters 

304 ---------- 

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

306 Predicted quantum graph. 

307 prov_reader : \ 

308 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

309 Reader for the provenance quantum graph. 

310 butler : `lsst.daf.butler.Butler` 

311 Client for the data repository. 

312 expect_failure : `bool` 

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

314 succeed without writing anything (`False`). 

315 start_time : `float` 

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

317 quantum's execution. 

318 expect_failures_retried : `bool`, optional 

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

320 most recent attempt. 

321 

322 Returns 

323 ------- 

324 prov : `ProvenanceQuantumGraph` 

325 The full provenance quantum graph. 

326 """ 

327 prov_reader.read_full_graph() 

328 prov = prov_reader.graph 

329 checked_some_metadata = False 

330 checked_some_log = False 

331 self.maxDiff = None 

332 self.assertEqual( 

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

334 [prov.header.output_run] 

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

336 ) 

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

338 for task_label in pred.quanta_by_task: 

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

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

341 for dataset_type_name in prov.datasets_by_type: 

342 self.assertEqual( 

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

344 ) 

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

346 for quantum_id in pred: 

347 # Check consistency between the predicted and provenance quantum 

348 # node attributes. 

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

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

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

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

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

354 # Check consistency between the predicted and provenance dataset 

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

356 # use later. 

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

358 pipeline_edges: list[Edge] 

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

360 quantum_id, data="pipeline_edges" 

361 ): 

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

363 for edge in pipeline_edges: 

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

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

366 ) 

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

368 quantum_id, data="pipeline_edges" 

369 ): 

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

371 for edge in pipeline_edges: 

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

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

374 ) 

375 # Check quantum status and dataset existence against the known 

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

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

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

379 # This is the quantum that can directly raise. 

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

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

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

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

384 if expect_failure: 

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

386 else: 

387 self._expect_successful( 

388 prov_qinfo, 

389 existence, 

390 caveats=( 

391 QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR 

392 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

393 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

394 ), 

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

396 msg=msg, 

397 ) 

398 if expect_failures_retried: 

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

400 self.assertEqual( 

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

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

403 ) 

404 else: 

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

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

407 # This quantum will succeed (with one predicted input 

408 # missing) or be blocked. 

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

410 if expect_failure: 

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

412 else: 

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

414 self.assertEqual( 

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

416 ) 

417 case ( 

418 "resample" | "coadd", 

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

420 ): 

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

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

423 # run. 

424 if expect_failure: 

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

426 else: 

427 self._expect_successful( 

428 prov_qinfo, 

429 existence, 

430 caveats=( 

431 QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED 

432 | QuantumSuccessCaveats.NO_WORK 

433 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING 

434 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING 

435 ), 

436 msg=msg, 

437 ) 

438 self.assertEqual( 

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

440 ) 

441 case ( 

442 "resample", 

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

444 ): 

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

446 # regardless. 

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

448 if expect_failure: 

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

450 else: 

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

452 self.assertEqual( 

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

454 ) 

455 case ( 

456 "coadd", 

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

458 ): 

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

460 # with one input missing. 

461 if expect_failure: 

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

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

464 else: 

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

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

467 self.assertEqual( 

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

469 ) 

470 case _: 

471 # All other quanta should succeed and have all inputs 

472 # present. 

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

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

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

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

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

478 self.check_metadata(quantum_id, prov_reader, butler) 

479 checked_some_metadata = True 

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

481 QuantumAttemptStatus.SUCCESSFUL, 

482 QuantumAttemptStatus.FAILED, 

483 ): 

484 self.check_log(quantum_id, prov_reader, butler) 

485 checked_some_log = True 

486 self.assertTrue(checked_some_metadata) 

487 self.assertTrue(checked_some_log) 

488 self.check_resource_usage_table( 

489 prov_reader.graph, expect_failure=expect_failure, start_time=start_time 

490 ) 

491 self.check_packages(prov_reader) 

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

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

494 return prov 

495 

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

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

498 

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

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

501 

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

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

504 

505 def _expect_successful( 

506 self, 

507 info: ProvenanceQuantumInfo, 

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

509 caveats: QuantumSuccessCaveats = QuantumSuccessCaveats.NO_CAVEATS, 

510 exception_type: str | None = None, 

511 *, 

512 msg: str, 

513 ) -> None: 

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

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

516 if exception_type is None: 

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

518 else: 

519 assert info["exception"] is not None 

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

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

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

523 if not (caveats & QuantumSuccessCaveats.ANY_OUTPUTS_MISSING): 

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

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

526 if caveats & QuantumSuccessCaveats.ALL_OUTPUTS_MISSING: 

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

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

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

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

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

532 

533 def _expect_failure( 

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

535 ) -> None: 

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

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

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

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

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

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

542 

543 def _expect_blocked( 

544 self, 

545 info: ProvenanceQuantumInfo, 

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

547 msg: str, 

548 ) -> None: 

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

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

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

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

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

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

555 

556 def check_dataset( 

557 self, 

558 dataset_id: uuid.UUID, 

559 pred: PredictedQuantumGraph, 

560 prov: ProvenanceQuantumGraph, 

561 butler: Butler, 

562 ) -> bool: 

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

564 counterpart. 

565 

566 Parameters 

567 ---------- 

568 dataset_id : `uuid.UUID` 

569 Unique ID for the dataset. 

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

571 Predicted quantum graph. 

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

573 Provenance quantum graph. 

574 butler : `lsst.daf.butler.Butler` 

575 Client for the data repository. 

576 

577 Returns 

578 ------- 

579 exists : `bool` 

580 Whether the dataset was marked as existing in the provenance 

581 quantum graph. 

582 """ 

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

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

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

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

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

588 exists = prov_info["produced"] 

589 dataset_type_name = prov_info["dataset_type_name"] 

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

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

592 self.assertEqual( 

593 butler.get_dataset(dataset_id) is not None, 

594 exists, 

595 msg=( 

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

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

598 ), 

599 ) 

600 return exists 

601 

602 def check_metadata( 

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

604 ) -> None: 

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

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

607 

608 Parameters 

609 ---------- 

610 quantum_id : `uuid.UUID` 

611 Unique ID for the quantum this metadata belongs to. 

612 provenance_reader : \ 

613 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

614 Reader for the provenance quantum graph. 

615 butler : `lsst.daf.butler.Butler` 

616 Client for the data repository. 

617 """ 

618 # Try reading metadata through the quantum ID. 

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

620 self.assertIsInstance(metadata1, TaskMetadata) 

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

622 quantum_id, data="pipeline_edges" 

623 ): 

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

625 # Also try reading metadata through the dataset ID. 

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

627 break 

628 else: 

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

630 self.assertEqual(metadata1, metadata2) 

631 

632 def check_log( 

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

634 ) -> None: 

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

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

637 

638 Parameters 

639 ---------- 

640 quantum_id : `uuid.UUID` 

641 Unique ID for the quantum this log belongs to. 

642 provenance_reader : \ 

643 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

644 Reader for the provenance quantum graph. 

645 butler : `lsst.daf.butler.Butler` 

646 Client for the data repository. 

647 """ 

648 # Try reading log through the quantum ID. 

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

650 self.assertIsInstance(log1, ButlerLogRecords) 

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

652 quantum_id, data="pipeline_edges" 

653 ): 

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

655 # Also try reading log through the dataset ID. 

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

657 break 

658 else: 

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

660 self.assertEqual(log1, log2) 

661 

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

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

664 

665 Parameters 

666 ---------- 

667 provenance_reader : \ 

668 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader` 

669 Reader for the provenance quantum graph. 

670 """ 

671 packages = provenance_reader.fetch_packages() 

672 self.assertIsInstance(packages, Packages) 

673 self.assertIn("pipe_base", packages) 

674 

675 def check_resource_usage_table( 

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

677 ) -> None: 

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

679 

680 Parameters 

681 ---------- 

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

683 Reader for the provenance quantum graph. 

684 expect_failure : `bool` 

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

686 succeed without writing anything (`False`). 

687 start_time : `float` 

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

689 quantum's execution. 

690 """ 

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

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

693 self.assertCountEqual( 

694 tbl.colnames, 

695 ["quantum_id"] 

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

697 + list(QuantumResourceUsage.model_fields), 

698 ) 

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

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

701 # shenanigans. 

702 end_time = time.time() 

703 for quantum_start_time in tbl["start"]: 

704 self.assertGreater(quantum_start_time, start_time) 

705 self.assertLess(quantum_start_time, end_time) 

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

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

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

709 

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

711 """Check `ProvenanceQuantumGraph.make_quantum_table`. 

712 

713 Parameters 

714 ---------- 

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

716 Reader for the provenance quantum graph. 

717 expect_failure : `bool` 

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

719 succeed without writing anything (`False`). 

720 """ 

721 t = prov.make_quantum_table() 

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

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

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

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

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

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

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

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

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

731 if expect_failure: 

732 # calibrate 

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

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

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

736 # consolidate 

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

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

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

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

741 # resample 

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

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

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

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

746 else: 

747 # calibrate 

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

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

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

751 # consolidate 

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

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

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

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

756 # resample 

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

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

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

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

761 

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

763 """Check `ProvenanceQuantumGraph.make_exception_table`. 

764 

765 Parameters 

766 ---------- 

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

768 Reader for the provenance quantum graph. 

769 expect_failure : `bool` 

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

771 succeed without writing anything (`False`). 

772 """ 

773 t = prov.make_exception_table() 

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

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

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

777 

778 def test_all_successful(self) -> None: 

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

780 results with assume_complete=False. 

781 

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

783 """ 

784 with self.make_test_repo() as prep: 

785 prep.config.assume_complete = False 

786 start_time = time.time() 

787 attempted_quanta = list( 

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