Coverage for tests / test_aggregator.py: 12%
503 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-03 01:30 -0800
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/>.
28from __future__ import annotations
30import dataclasses
31import itertools
32import os
33import tempfile
34import time
35import unittest.mock
36import uuid
37from collections.abc import Iterator
38from contextlib import contextmanager
39from typing import Any, cast
41import astropy.table
42import click.testing
43import numpy as np
44import pydantic
45from click.testing import CliRunner, Result
47import lsst.utils.tests
48from lsst.daf.butler import Butler, ButlerLogRecords, QuantumBackedButler
49from lsst.pex.config import Config
50from lsst.pipe.base import (
51 AlgorithmError,
52 QuantumAttemptStatus,
53 QuantumSuccessCaveats,
54 TaskMetadata,
55)
56from lsst.pipe.base import automatic_connection_constants as acc
57from lsst.pipe.base.cli.cmd.commands import (
58 aggregate_graph as aggregate_graph_cli,
59)
60from lsst.pipe.base.cli.cmd.commands import (
61 provenance_report as provenance_report_cli,
62)
63from lsst.pipe.base.graph_walker import GraphWalker
64from lsst.pipe.base.pipeline_graph import Edge
65from lsst.pipe.base.quantum_graph import (
66 FORMAT_VERSION,
67 PredictedDatasetInfo,
68 PredictedQuantumGraph,
69 PredictedQuantumInfo,
70 ProvenanceDatasetInfo,
71 ProvenanceQuantumGraph,
72 ProvenanceQuantumInfo,
73 ProvenanceQuantumReport,
74 ProvenanceReport,
75)
76from lsst.pipe.base.quantum_graph.aggregator import AggregatorConfig, FatalWorkerError, aggregate_graph
77from lsst.pipe.base.quantum_graph.ingest_graph import ingest_graph
78from lsst.pipe.base.resource_usage import QuantumResourceUsage
79from lsst.pipe.base.single_quantum_executor import SingleQuantumExecutor
80from lsst.pipe.base.tests.mocks import (
81 DirectButlerRepo,
82 DynamicConnectionConfig,
83 DynamicTestPipelineTaskConfig,
84)
85from lsst.pipe.base.tests.util import patch_deterministic_uuid4
86from lsst.utils.packages import Packages
89@dataclasses.dataclass
90class PrepInfo:
91 """Struct of objects used in an aggregator test."""
93 butler: Butler
94 butler_path: str
95 predicted: PredictedQuantumGraph
96 predicted_path: str
97 config: AggregatorConfig
100class AggregatorTestCase(unittest.TestCase):
101 """Unit tests for `lsst.pipe.base.quantum_graph.aggregator`."""
103 @staticmethod
104 @contextmanager
105 def make_test_repo() -> Iterator[PrepInfo]:
106 """Make a test data repository and predicted quantum graph.
108 Returns
109 -------
110 prep_info : `PrepInfo`
111 Objects used in aggregator tests.
113 Notes
114 -----
115 The pipeline graph used by this task looks like this:
117 ■ calibrate: {detector, visit}
118 ╭─┤
119 ■ │ consolidate: {visit}
120 │
121 ■ resample: {patch, visit}
122 │
123 ■ coadd: {band, patch}
125 The data can be visualized via::
127 python -m lsst.daf.butler.tests.registry_data.spatial
129 One of the 'calibrate' quanta (visit=2, detector=2) is configured to
130 fail with `lsst.pipe.base.AnnotatedPartialOutputsError`. This lets us
131 test both success-with-caveats and failures, depending on how we
132 configure the executor. This ``{visit: 2, detector: 2}`` data ID is
133 the only one that overlaps ``{tract: 1, patch: 1}`` and
134 ``{tract: 0, patch: 5}``, so it should chain to the 'resample' and
135 'coadd' tasks, too.
136 """
137 with patch_deterministic_uuid4(100):
138 with DirectButlerRepo.make_temporary("base.yaml", "spatial.yaml") as (helper, root):
139 calibrate_config = DynamicTestPipelineTaskConfig()
140 calibrate_config.fail_exception = "lsst.pipe.base.AnnotatedPartialOutputsError"
141 calibrate_config.fail_condition = "visit=2 AND detector=2"
142 helper.add_task(
143 "calibrate",
144 config=calibrate_config,
145 dimensions=["visit", "detector"],
146 inputs={
147 "input_image": DynamicConnectionConfig(
148 dataset_type_name="raw",
149 dimensions=["visit", "detector"],
150 )
151 },
152 prerequisite_inputs={
153 "refcat": DynamicConnectionConfig(
154 dataset_type_name="references",
155 dimensions=["htm7"],
156 multiple=True,
157 )
158 },
159 init_outputs={
160 "output_schema": DynamicConnectionConfig(
161 dataset_type_name="source_schema",
162 )
163 },
164 outputs={
165 "output_image": DynamicConnectionConfig(
166 dataset_type_name="image",
167 dimensions=["visit", "detector"],
168 ),
169 "output_table": DynamicConnectionConfig(
170 dataset_type_name="source_detector",
171 dimensions=["visit", "detector"],
172 ),
173 },
174 )
175 helper.add_task(
176 "consolidate",
177 dimensions=["visit"],
178 init_inputs={
179 "input_schema": DynamicConnectionConfig(
180 dataset_type_name="source_schema",
181 )
182 },
183 inputs={
184 "input_table": DynamicConnectionConfig(
185 dataset_type_name="source_detector",
186 dimensions=["visit", "detector"],
187 multiple=True,
188 )
189 },
190 outputs={
191 "output_table": DynamicConnectionConfig(
192 dataset_type_name="source",
193 dimensions=["visit"],
194 )
195 },
196 )
197 helper.add_task(
198 "resample",
199 dimensions=["patch", "visit"],
200 inputs={
201 "input_image": DynamicConnectionConfig(
202 dataset_type_name="image",
203 dimensions=["visit", "detector"],
204 multiple=True,
205 )
206 },
207 outputs={
208 "output_image": DynamicConnectionConfig(
209 dataset_type_name="warp",
210 dimensions=["patch", "visit"],
211 )
212 },
213 )
214 helper.add_task(
215 "coadd",
216 dimensions=["patch", "band"],
217 inputs={
218 "input_image": DynamicConnectionConfig(
219 dataset_type_name="warp",
220 dimensions=["patch", "visit"],
221 multiple=True,
222 )
223 },
224 outputs={
225 "output_image": DynamicConnectionConfig(
226 dataset_type_name="coadd",
227 dimensions=["patch", "band"],
228 ),
229 },
230 )
231 pqgc = helper.make_quantum_graph_builder().finish(output="out_chain")
232 # We use the butler root for various QG files just because it's
233 # a convenient temporary directory.
234 predicted_path = os.path.join(root, "predicted.qg")
235 pqgc.write(predicted_path)
236 config = AggregatorConfig(
237 output_path=os.path.join(root, "provenance.qg"),
238 # Set these small to see logic paths that otherwise only
239 # affect large graphs.
240 ingest_batch_size=10,
241 zstd_dict_size=256,
242 zstd_dict_n_inputs=16,
243 )
244 yield PrepInfo(
245 butler=helper.butler.clone(collections="out_chain"),
246 butler_path=root,
247 predicted=pqgc.assemble(),
248 predicted_path=predicted_path,
249 config=config,
250 )
252 def iter_graph_execution(
253 self,
254 repo: str,
255 qg: PredictedQuantumGraph,
256 raise_on_partial_outputs: bool,
257 is_retry: bool = False,
258 ) -> Iterator[uuid.UUID]:
259 """Return an iterator that executes and yields quanta one by one.
261 Parameters
262 ----------
263 repo : `str`
264 Butler repository path or alias.
265 qg : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`
266 Predicted quantum graph. Must have datastore records attached,
267 since execution uses a quantum-backed butler.
268 raise_on_partial_outputs : `bool`
269 Whether to raise on `lsst.pipe.base.AnnotatedPartialOutputsError`
270 or treat it as a success with caveats.
271 is_retry : `bool`, optional
272 If `True`, this is a retry attempt and hence some outputs may
273 already be present; skip successes and reprocess failures.
275 Returns
276 -------
277 quanta : `~collections.abc.Iterator` [`uuid.UUID`]
278 An iterator over all executed quantum IDs (not blocked ones).
279 """
280 qbb = qg.make_init_qbb(repo)
281 self.enterContext(qbb)
282 qg.init_output_run(qbb)
283 sqe = SingleQuantumExecutor(
284 limited_butler_factory=lambda quantum: QuantumBackedButler.initialize(
285 repo,
286 quantum,
287 qg.pipeline_graph.universe,
288 ),
289 assume_no_existing_outputs=not is_retry,
290 skip_existing=is_retry,
291 clobber_outputs=is_retry,
292 raise_on_partial_outputs=raise_on_partial_outputs,
293 )
294 qg.build_execution_quanta()
295 xgraph = qg.quantum_only_xgraph
296 walker = GraphWalker[uuid.UUID](xgraph.copy())
297 for ready in walker:
298 for quantum_id in ready:
299 info = xgraph.nodes[quantum_id]
300 try:
301 sqe.execute(info["pipeline_node"], info["quantum"], quantum_id)
302 except AlgorithmError:
303 walker.fail(quantum_id)
304 else:
305 walker.finish(quantum_id)
306 yield quantum_id
308 def check_provenance_graph(
309 self,
310 pred: PredictedQuantumGraph,
311 butler: Butler,
312 expect_failure: bool,
313 start_time: float,
314 expect_failures_retried: bool = False,
315 ) -> ProvenanceQuantumGraph:
316 """Run a batter of tests on a provenance quantum graph produced by
317 scanning the graph created by `make_test_repo`.
319 Parameters
320 ----------
321 pred: `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`
322 Predicted quantum graph.
323 prov_reader : \
324 `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraphReader`
325 Reader for the provenance quantum graph.
326 butler : `lsst.daf.butler.Butler`
327 Client for the data repository.
328 expect_failure : `bool`
329 Whether to expect one quantum of 'calibrate' to fail (`True`) or
330 succeed without writing anything (`False`).
331 start_time : `float`
332 A POSIX timestamp that strictly precedes the start time of any
333 quantum's execution.
334 expect_failures_retried : `bool`, optional
335 If `True`, expect an initial attempt with failures prior to the
336 most recent attempt.
338 Returns
339 -------
340 prov : `ProvenanceQuantumGraph`
341 The full provenance quantum graph.
342 """
343 prov: ProvenanceQuantumGraph = butler.get("run_provenance")
344 self.maxDiff = None
345 self.assertEqual(prov.header.version, FORMAT_VERSION)
346 self.assertEqual(
347 list(butler.collections.get_info(prov.header.output).children),
348 [prov.header.output_run]
349 + list(butler.collections.query(prov.header.inputs, flatten_chains=True)),
350 )
351 self.assertEqual(pred.quanta_by_task.keys(), prov.quanta_by_task.keys())
352 for task_label in pred.quanta_by_task:
353 self.assertEqual(pred.quanta_by_task[task_label], prov.quanta_by_task[task_label])
354 self.assertEqual(pred.datasets_by_type.keys() - {"packages"}, prov.datasets_by_type.keys())
355 for dataset_type_name in prov.datasets_by_type:
356 self.assertEqual(
357 pred.datasets_by_type[dataset_type_name], prov.datasets_by_type[dataset_type_name]
358 )
359 self.assertEqual(prov.init_quanta.keys(), pred.quanta_by_task.keys())
360 for quantum_id in pred:
361 # Check consistency between the predicted and provenance quantum
362 # node attributes.
363 pred_qinfo: PredictedQuantumInfo = pred.bipartite_xgraph.nodes[quantum_id]
364 prov_qinfo: ProvenanceQuantumInfo = prov.bipartite_xgraph.nodes[quantum_id]
365 self.assertEqual(pred_qinfo["task_label"], prov_qinfo["task_label"])
366 self.assertEqual(pred_qinfo["data_id"], prov_qinfo["data_id"])
367 msg = f"{pred_qinfo['task_label']}@{pred_qinfo['data_id']}"
368 # Check consistency between the predicted and provenance dataset
369 # node attributes and edges. Also gather existence information for
370 # use later.
371 existence: dict[str, list[bool]] = {}
372 pipeline_edges: list[Edge]
373 for dataset_id, _, pipeline_edges in pred.bipartite_xgraph.in_edges(
374 quantum_id, data="pipeline_edges"
375 ):
376 self.assertTrue(prov.bipartite_xgraph.has_predecessor(quantum_id, dataset_id))
377 for edge in pipeline_edges:
378 existence.setdefault(edge.connection_name, []).append(
379 self.check_dataset(dataset_id, pred, prov, butler)
380 )
381 for _, dataset_id, pipeline_edges in pred.bipartite_xgraph.out_edges(
382 quantum_id, data="pipeline_edges"
383 ):
384 self.assertTrue(prov.bipartite_xgraph.has_successor(quantum_id, dataset_id))
385 for edge in pipeline_edges:
386 existence.setdefault(edge.connection_name, []).append(
387 self.check_dataset(dataset_id, pred, prov, butler)
388 )
389 # Check quantum status and dataset existence against the known
390 # structure of the graph and where failures/caveats occur.
391 match (pred_qinfo["task_label"], dict(pred_qinfo["data_id"].required)):
392 case "calibrate", {"visit": 2, "detector": 2}:
393 # This is the quantum that can directly raise.
394 self._expect_all_exist(existence["input_image"], msg=msg)
395 self._expect_all_exist(existence["refcat"], msg=msg)
396 self._expect_none_exist(existence["output_image"], msg=msg)
397 self._expect_none_exist(existence["output_table"], msg=msg)
398 if expect_failure:
399 self._expect_failure(prov_qinfo, existence, msg=msg)
400 else:
401 self._expect_successful(
402 prov_qinfo,
403 existence,
404 caveats=(
405 QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR
406 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING
407 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING
408 ),
409 exception_type="lsst.pipe.base.tests.mocks.MockAlgorithmError",
410 msg=msg,
411 )
412 if expect_failures_retried:
413 self.assertEqual(len(prov_qinfo["attempts"]), 2)
414 self.assertEqual(
415 prov_qinfo["attempts"][0].exception.type_name,
416 "lsst.pipe.base.tests.mocks.MockAlgorithmError",
417 )
418 else:
419 self.assertEqual(len(prov_qinfo["attempts"]), 1)
420 case "consolidate", {"visit": 2}:
421 # This quantum will succeed (with one predicted input
422 # missing) or be blocked.
423 self._expect_one_missing(existence["input_table"], msg=msg)
424 if expect_failure:
425 self._expect_blocked(prov_qinfo, existence, msg=msg)
426 else:
427 self._expect_successful(prov_qinfo, existence, msg=msg)
428 self.assertEqual(
429 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure
430 )
431 case (
432 "resample" | "coadd",
433 {"tract": 1, "patch": 1} | {"tract": 0, "patch": 5},
434 ):
435 # These quanta will be blocked by an upstream failure or do
436 # chained caveats, since they won't have enough inputs to
437 # run.
438 if expect_failure:
439 self._expect_blocked(prov_qinfo, existence, msg=msg)
440 else:
441 self._expect_successful(
442 prov_qinfo,
443 existence,
444 caveats=(
445 QuantumSuccessCaveats.ADJUST_QUANTUM_RAISED
446 | QuantumSuccessCaveats.NO_WORK
447 | QuantumSuccessCaveats.ALL_OUTPUTS_MISSING
448 | QuantumSuccessCaveats.ANY_OUTPUTS_MISSING
449 ),
450 msg=msg,
451 )
452 self.assertEqual(
453 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure
454 )
455 case (
456 "resample",
457 {"tract": 0, "patch": 4, "visit": 2} | {"tract": 1, "patch": 0, "visit": 2},
458 ):
459 # This will succeed or be blocked, with one input missing
460 # regardless.
461 self._expect_one_missing(existence["input_image"], msg=msg)
462 if expect_failure:
463 self._expect_blocked(prov_qinfo, existence, msg=msg)
464 else:
465 self._expect_successful(prov_qinfo, existence, msg=msg)
466 self.assertEqual(
467 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure
468 )
469 case (
470 "coadd",
471 {"tract": 0, "patch": 4, "band": "r"} | {"tract": 1, "patch": 0, "band": "r"},
472 ):
473 # This will succeed with no inputs missing or be blocked
474 # with one input missing.
475 if expect_failure:
476 self._expect_one_missing(existence["input_image"], msg=msg)
477 self._expect_blocked(prov_qinfo, existence, msg=msg)
478 else:
479 self._expect_all_exist(existence["input_image"], msg=msg)
480 self._expect_successful(prov_qinfo, existence, msg=msg)
481 self.assertEqual(
482 len(prov_qinfo["attempts"]), expect_failures_retried or not expect_failure
483 )
484 case _:
485 # All other quanta should succeed and have all inputs
486 # present.
487 for connection_name in prov_qinfo["pipeline_node"].inputs.keys():
488 self._expect_all_exist(existence[connection_name], msg=msg)
489 self._expect_successful(prov_qinfo, existence, msg=msg)
490 self.assertEqual(len(prov_qinfo["attempts"]), 1)
491 self.check_metadata(
492 quantum_id,
493 prov,
494 butler,
495 expect_ingested=(prov_qinfo["status"] is QuantumAttemptStatus.SUCCESSFUL),
496 )
497 self.check_log(
498 quantum_id,
499 prov,
500 butler,
501 expect_ingested=(
502 prov_qinfo["status"]
503 in (
504 QuantumAttemptStatus.SUCCESSFUL,
505 QuantumAttemptStatus.FAILED,
506 )
507 ),
508 )
509 self.check_resource_usage_table(prov, expect_failure=expect_failure, start_time=start_time)
510 self.check_packages(butler)
511 self.check_configs(butler, prov)
512 self.check_quantum_table(prov, expect_failure=expect_failure)
513 self.check_exception_table(prov, expect_failure=expect_failure)
514 self.check_report(prov, expect_failure=expect_failure)
515 return prov
517 def _expect_all_exist(self, existence: list[bool], msg: str) -> None:
518 self.assertTrue(all(existence), msg=msg)
520 def _expect_none_exist(self, existence: list[bool], msg: str) -> None:
521 self.assertFalse(any(existence), msg=msg)
523 def _expect_one_missing(self, existence: list[bool], msg: str) -> None:
524 self.assertEqual(existence.count(False), 1, msg=msg)
526 def _expect_successful(
527 self,
528 info: ProvenanceQuantumInfo,
529 existence: dict[str, list[bool]],
530 caveats: QuantumSuccessCaveats = QuantumSuccessCaveats.NO_CAVEATS,
531 exception_type: str | None = None,
532 *,
533 msg: str,
534 ) -> None:
535 self.assertEqual(info["status"], QuantumAttemptStatus.SUCCESSFUL, msg=msg)
536 self.assertEqual(info["caveats"], caveats, msg=msg)
537 if exception_type is None:
538 self.assertIsNone(info["exception"], msg=msg)
539 else:
540 assert info["exception"] is not None
541 self.assertEqual(info["exception"].type_name, exception_type, msg=msg)
542 self._expect_all_exist(existence[acc.LOG_OUTPUT_CONNECTION_NAME], msg=msg)
543 self._expect_all_exist(existence[acc.METADATA_OUTPUT_CONNECTION_NAME], msg=msg)
544 if not (caveats & QuantumSuccessCaveats.ANY_OUTPUTS_MISSING):
545 for connection_name in info["pipeline_node"].outputs.keys():
546 self._expect_all_exist(existence[connection_name], msg=msg)
547 if caveats & QuantumSuccessCaveats.ALL_OUTPUTS_MISSING:
548 for connection_name in info["pipeline_node"].outputs.keys():
549 self._expect_none_exist(existence[connection_name], msg=msg)
550 self.assertIsNotNone(info["resource_usage"], msg=msg)
551 self.assertGreater(info["resource_usage"].total_time, 0, msg=msg)
552 self.assertGreater(info["resource_usage"].memory, 0, msg=msg)
554 def _expect_failure(
555 self, info: ProvenanceQuantumInfo, existence: dict[str, list[bool]], msg: str
556 ) -> None:
557 self.assertEqual(info["status"], QuantumAttemptStatus.FAILED, msg=msg)
558 self.assertEqual(info["exception"].type_name, "lsst.pipe.base.tests.mocks.MockAlgorithmError")
559 self._expect_all_exist(existence[acc.LOG_OUTPUT_CONNECTION_NAME], msg=msg)
560 self._expect_none_exist(existence[acc.METADATA_OUTPUT_CONNECTION_NAME], msg=msg)
561 for connection_name in info["pipeline_node"].outputs.keys():
562 self._expect_none_exist(existence[connection_name], msg=msg)
564 def _expect_blocked(
565 self,
566 info: ProvenanceQuantumInfo,
567 existence: dict[str, list[bool]],
568 msg: str,
569 ) -> None:
570 self.assertEqual(info["status"], QuantumAttemptStatus.BLOCKED, msg=msg)
571 self.assertEqual(info["attempts"], [])
572 self._expect_none_exist(existence[acc.LOG_OUTPUT_CONNECTION_NAME], msg=msg)
573 self._expect_none_exist(existence[acc.METADATA_OUTPUT_CONNECTION_NAME], msg=msg)
574 for connection_name in info["pipeline_node"].outputs.keys():
575 self._expect_none_exist(existence[connection_name], msg=msg)
577 def check_dataset(
578 self,
579 dataset_id: uuid.UUID,
580 pred: PredictedQuantumGraph,
581 prov: ProvenanceQuantumGraph,
582 butler: Butler,
583 ) -> bool:
584 """Check a provenance dataset for consistency with its predicted
585 counterpart.
587 Parameters
588 ----------
589 dataset_id : `uuid.UUID`
590 Unique ID for the dataset.
591 pred: `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`
592 Predicted quantum graph.
593 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph`
594 Provenance quantum graph.
595 butler : `lsst.daf.butler.Butler`
596 Client for the data repository.
598 Returns
599 -------
600 exists : `bool`
601 Whether the dataset was marked as existing in the provenance
602 quantum graph.
603 """
604 pred_info: PredictedDatasetInfo = pred.bipartite_xgraph.nodes[dataset_id]
605 prov_info: ProvenanceDatasetInfo = prov.bipartite_xgraph.nodes[dataset_id]
606 self.assertEqual(pred_info["dataset_type_name"], prov_info["dataset_type_name"])
607 self.assertEqual(pred_info["data_id"], prov_info["data_id"])
608 self.assertEqual(pred_info["run"], prov_info["run"])
609 exists = prov_info["produced"]
610 dataset_type_name = prov_info["dataset_type_name"]
611 # We can remove this guard when we ingest QG-backed metadata and logs.
612 if not dataset_type_name.endswith("_metadata") and not dataset_type_name.endswith("_log"):
613 self.assertEqual(
614 butler.get_dataset(dataset_id) is not None,
615 exists,
616 msg=(
617 f"Ingest/existence inconsistency for {dataset_type_name}"
618 f"@{prov_info['data_id']}/{dataset_id}]"
619 ),
620 )
621 return exists
623 def check_metadata(
624 self, quantum_id: uuid.UUID, prov: ProvenanceQuantumGraph, butler: Butler, expect_ingested: bool
625 ) -> None:
626 """Check reading a metadata dataset from the butler, and check that the
627 original metadata file has been deleted.
629 Parameters
630 ----------
631 quantum_id : `uuid.UUID`
632 Unique ID for the quantum this metadata belongs to.
633 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph`
634 Provenance quantum graph.
635 butler : `lsst.daf.butler.Butler`
636 Client for the data repository.
637 expect_ingested : `bool`
638 Whether the metadata dataset should have been ingested.
639 """
640 dataset_id = prov.bipartite_xgraph.nodes[quantum_id]["metadata_id"]
641 ref = butler.get_dataset(dataset_id)
642 if not expect_ingested:
643 self.assertIsNone(ref)
644 return
645 assert ref is not None
646 metadata = butler.get(ref)
647 self.assertIsInstance(metadata, TaskMetadata)
648 graph_path = butler.getURI("run_provenance")
649 self.assertEqual(butler.getURI(ref), graph_path)
650 # We now delete the metadata dataset, in order let us get the original
651 # location from the butler and check that there's nothing there. Note
652 # that this doesn't actually delete the file because the butler knows
653 # it's shared with other datasets.
654 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
655 original_path = butler.getURI(ref, predict=True)
656 self.assertTrue(graph_path.exists())
657 self.assertNotEqual(graph_path, original_path)
658 self.assertFalse(original_path.exists())
660 def check_log(
661 self, quantum_id: uuid.UUID, prov: ProvenanceQuantumGraph, butler: Butler, expect_ingested: bool
662 ) -> None:
663 """Check reading a log dataset from the butler, and check that the
664 original log file has been deleted.
666 Parameters
667 ----------
668 quantum_id : `uuid.UUID`
669 Unique ID for the quantum this log belongs to.
670 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph`
671 Provenance quantum graph.
672 butler : `lsst.daf.butler.Butler`
673 Client for the data repository.
674 expect_ingested : `bool`
675 Whether the metadata dataset should have been ingested.
676 """
677 dataset_id = prov.bipartite_xgraph.nodes[quantum_id]["log_id"]
678 ref = butler.get_dataset(dataset_id)
679 if not expect_ingested:
680 self.assertIsNone(ref)
681 return
682 assert ref is not None
683 log = butler.get(ref)
684 self.assertIsInstance(log, ButlerLogRecords)
685 graph_path = butler.getURI("run_provenance")
686 self.assertEqual(butler.getURI(ref), graph_path)
687 # We now delete the log dataset, in order let us get the original
688 # location from the butler and check that there's nothing there. Note
689 # that this doesn't actually delete the file because the butler knows
690 # it's shared with other datasets.
691 butler.pruneDatasets([ref], disassociate=True, unstore=True, purge=True)
692 original_path = butler.getURI(ref, predict=True)
693 self.assertTrue(graph_path.exists())
694 self.assertNotEqual(graph_path, original_path)
695 self.assertFalse(original_path.exists())
697 def check_configs(self, butler: Butler, prov: ProvenanceQuantumGraph) -> None:
698 for task_node in prov.pipeline_graph.tasks.values():
699 config = butler.get(task_node.init.config_output.dataset_type_name)
700 self.assertIsInstance(config, Config)
702 def check_packages(self, butler: Butler) -> None:
703 """Check fetching package versions from the provenance graph.
705 Parameters
706 ----------
707 butler : `lsst.daf.butler.Butler`
708 Client for the data repository.
709 """
710 packages = butler.get("run_provenance.packages")
711 self.assertIsInstance(packages, Packages)
712 self.assertIn("pipe_base", packages)
714 def check_resource_usage_table(
715 self, prov: ProvenanceQuantumGraph, expect_failure: bool, start_time: float
716 ) -> None:
717 """Check building a resource usage table from the provenance graph.
719 Parameters
720 ----------
721 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph`
722 Reader for the provenance quantum graph.
723 expect_failure : `bool`
724 Whether to expect one quantum of 'calibrate' to fail (`True`) or
725 succeed without writing anything (`False`).
726 start_time : `float`
727 A POSIX timestamp that strictly precedes the start time of any
728 quantum's execution.
729 """
730 tbl = prov.make_task_resource_usage_table("calibrate", include_data_ids=True)
731 self.assertEqual(len(tbl), prov.header.n_task_quanta["calibrate"])
732 self.assertCountEqual(
733 tbl.colnames,
734 ["quantum_id"]
735 + list(prov.pipeline_graph.tasks["calibrate"].dimensions.names)
736 + list(QuantumResourceUsage.model_fields),
737 )
738 # Check that quantum start times are bounded by the before-execution
739 # start_time and now. This makes sure we didn't get any timezone
740 # shenanigans.
741 end_time = time.time()
742 for quantum_start_time in tbl["start"]:
743 self.assertGreater(quantum_start_time, start_time)
744 self.assertLess(quantum_start_time, end_time)
745 self.assertTrue(np.all(tbl["init_time"] >= 0.0))
746 self.assertTrue(np.all(tbl["prep_time"] > 0.0))
747 self.assertTrue(np.all(tbl["run_time"] >= 0.0))
749 def check_quantum_table(self, prov: ProvenanceQuantumGraph, expect_failure: bool) -> None:
750 """Check `ProvenanceQuantumGraph.make_quantum_table`.
752 Parameters
753 ----------
754 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph`
755 Reader for the provenance quantum graph.
756 expect_failure : `bool`
757 Whether to expect one quantum of 'calibrate' to fail (`True`) or
758 succeed without writing anything (`False`).
759 """
760 t = prov.make_quantum_table()
761 self.assertEqual(list(t["Task"]), ["calibrate", "consolidate", "resample", "coadd"])
762 self.assertEqual(t["TOTAL"][0], 8)
763 self.assertEqual(t["EXPECTED"][0], 8)
764 self.assertEqual(t["Blocked"][0], 0)
765 self.assertEqual(t["TOTAL"][1], 2)
766 self.assertEqual(t["EXPECTED"][1], 2)
767 self.assertEqual(t["TOTAL"][2], 10)
768 self.assertEqual(t["EXPECTED"][2], 10)
769 if expect_failure:
770 # calibrate
771 self.assertEqual(t["Successful"][0], 7)
772 self.assertEqual(t["Caveats"][0], "")
773 self.assertEqual(t["Failed"][0], 1)
774 # consolidate
775 self.assertEqual(t["Successful"][1], 1)
776 self.assertEqual(t["Caveats"][1], "")
777 self.assertEqual(t["Failed"][1], 0)
778 self.assertEqual(t["Blocked"][1], 1)
779 # resample
780 self.assertEqual(t["Successful"][2], 6)
781 self.assertEqual(t["Caveats"][2], "")
782 self.assertEqual(t["Failed"][2], 0)
783 self.assertEqual(t["Blocked"][2], 4)
784 else:
785 # calibrate
786 self.assertEqual(t["Successful"][0], 8)
787 self.assertEqual(t["Caveats"][0], "*P(1)")
788 self.assertEqual(t["Failed"][0], 0)
789 # consolidate
790 self.assertEqual(t["Successful"][1], 2)
791 self.assertEqual(t["Caveats"][1], "")
792 self.assertEqual(t["Failed"][1], 0)
793 self.assertEqual(t["Blocked"][1], 0)
794 # resample
795 self.assertEqual(t["Successful"][2], 10)
796 self.assertEqual(t["Caveats"][2], "*A(2)")
797 self.assertEqual(t["Failed"][2], 0)
798 self.assertEqual(t["Blocked"][2], 0)
800 def check_exception_table(self, prov: ProvenanceQuantumGraph, expect_failure: bool) -> None:
801 """Check `ProvenanceQuantumGraph.make_exception_table`.
803 Parameters
804 ----------
805 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph`
806 Reader for the provenance quantum graph.
807 expect_failure : `bool`
808 Whether to expect one quantum of 'calibrate' to fail (`True`) or
809 succeed without writing anything (`False`).
810 """
811 t = prov.make_exception_table()
812 self.assertEqual(list(t["Task"]), ["calibrate"])
813 self.assertEqual(list(t["Exception"]), ["lsst.pipe.base.tests.mocks.MockAlgorithmError"])
814 self.assertEqual(list(t["Successes"]), [int(not expect_failure)])
815 self.assertEqual(list(t["Failures"]), [int(expect_failure)])
817 def check_report(self, prov: ProvenanceQuantumGraph, expect_failure: bool) -> None:
818 """Check `ProvenanceQuantumGraph.make_report`.
820 Parameters
821 ----------
822 prov : `lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph`
823 Reader for the provenance quantum graph.
824 expect_failure : `bool`
825 Whether to expect one quantum of 'calibrate' to fail (`True`) or
826 succeed without writing anything (`False`).
827 """
828 with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as data_id_table_dir:
829 report = prov.make_status_report(
830 also=QuantumAttemptStatus.SUCCESSFUL, data_id_table_dir=data_id_table_dir
831 )
832 task_label = "calibrate"
833 status_name = "FAILED" if expect_failure else "SUCCESSFUL"
834 exc_type = "lsst.pipe.base.tests.mocks.MockAlgorithmError"
835 self.assertEqual(report.root.keys(), {task_label})
836 self.assertEqual(report.root[task_label].keys(), {status_name})
837 self.assertEqual(report.root[task_label][status_name].keys(), {exc_type})
838 self.assertEqual(len(report.root[task_label][status_name][exc_type]), 1)
839 qr = report.root[task_label][status_name][exc_type][0]
840 self.assertIsInstance(qr, ProvenanceQuantumReport)
841 self.assertEqual(
842 qr.data_id,
843 {
844 "instrument": "Cam1",
845 "visit": 2,
846 "detector": 2,
847 "band": "r",
848 "day_obs": 20210909,
849 "physical_filter": "Cam1-R1",
850 },
851 )
852 tbl = astropy.table.Table.read(
853 os.path.join(data_id_table_dir, task_label, status_name, f"{exc_type}.ecsv")
854 )
855 self.assertCountEqual(tbl.colnames, ["instrument", "visit", "detector"])
856 self.assertEqual(len(tbl), 1)
857 self.assertEqual(list(tbl["instrument"]), ["Cam1"])
858 self.assertEqual(list(tbl["detector"]), [2])
859 self.assertEqual(list(tbl["visit"]), [2])
861 def test_all_successful(self) -> None:
862 """Test running a full graph with no failures, and then scanning the
863 results with incomplete=False.
864 """
865 with self.make_test_repo() as prep:
866 prep.config.incomplete = False
867 start_time = time.time()
868 attempted_quanta = list(
869 self.iter_graph_execution(prep.butler_path, prep.predicted, raise_on_partial_outputs=False)
870 )
871 self.assertCountEqual(attempted_quanta, prep.predicted.quantum_only_xgraph.nodes.keys())
872 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
873 ingest_graph(prep.butler_path, prep.config.output_path, transfer="move", batch_size=10)
874 prov = self.check_provenance_graph(
875 prep.predicted,
876 prep.butler,
877 expect_failure=False,
878 start_time=start_time,
879 )
880 self.check_no_original_dirs(prep.butler_path, prov.header.output_run)
881 for i, quantum_id in enumerate(attempted_quanta):
882 qinfo: ProvenanceQuantumInfo = prov.quantum_only_xgraph.nodes[quantum_id]
883 self.assertEqual(qinfo["attempts"][-1].previous_process_quanta, attempted_quanta[:i])
885 def test_all_successful_two_phase(self) -> None:
886 """Test running some of a graph with no failures, scanning with
887 incomplete=True, then finishing the graph and scanning again.
888 """
889 with self.make_test_repo() as prep:
890 start_time = time.time()
891 execution_iter = self.iter_graph_execution(
892 prep.butler_path, prep.predicted, raise_on_partial_outputs=False
893 )
894 attempted_quanta = list(itertools.islice(execution_iter, 9))
895 self.assertEqual(len(attempted_quanta), 9)
896 # Run the scanner while telling it execution is incomplete, so it
897 # just abandons incomplete quanta and doesn't write the provenance
898 # QG.
899 prep.config.incomplete = True
900 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
901 self.assertFalse(os.path.exists(cast(str, prep.config.output_path)))
902 # Finish executing the quanta.
903 attempted_quanta.extend(execution_iter)
904 # Scan again, and write the provenance QG.
905 prep.config.incomplete = False
906 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
907 ingest_graph(prep.butler_path, prep.config.output_path, transfer="move", batch_size=10)
908 prov = self.check_provenance_graph(
909 prep.predicted,
910 prep.butler,
911 expect_failure=False,
912 start_time=start_time,
913 )
914 for i, quantum_id in enumerate(attempted_quanta):
915 qinfo: ProvenanceQuantumInfo = prov.quantum_only_xgraph.nodes[quantum_id]
916 self.assertEqual(qinfo["attempts"][-1].previous_process_quanta, attempted_quanta[:i])
918 def test_some_failed(self) -> None:
919 """Test running a full graph with some failures, and then scanning the
920 results with incomplete=False.
921 """
922 with self.make_test_repo() as prep:
923 prep.config.incomplete = False
924 start_time = time.time()
925 attempted_quanta = list(
926 self.iter_graph_execution(prep.butler_path, prep.predicted, raise_on_partial_outputs=True)
927 )
928 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
929 ingest_graph(prep.butler_path, prep.config.output_path, transfer="move", batch_size=10)
930 prov = self.check_provenance_graph(
931 prep.predicted,
932 prep.butler,
933 expect_failure=True,
934 start_time=start_time,
935 )
936 self.check_no_original_dirs(prep.butler_path, prov.header.output_run)
937 for i, quantum_id in enumerate(attempted_quanta):
938 qinfo: ProvenanceQuantumInfo = prov.quantum_only_xgraph.nodes[quantum_id]
939 self.assertEqual(qinfo["attempts"][-1].previous_process_quanta, attempted_quanta[:i])
941 def test_some_failed_two_phase(self) -> None:
942 """Test running a full graph with some failures, then scanning the
943 results with incomplete=True, then scanning again with
944 incomplete=False.
945 """
946 with self.make_test_repo() as prep:
947 start_time = time.time()
948 attempted_quanta = list(
949 self.iter_graph_execution(prep.butler_path, prep.predicted, raise_on_partial_outputs=True)
950 )
951 prep.config.incomplete = True
952 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
953 prep.config.incomplete = False
954 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
955 ingest_graph(prep.butler_path, prep.config.output_path, transfer="move", batch_size=10)
956 prov = self.check_provenance_graph(
957 prep.predicted,
958 prep.butler,
959 expect_failure=True,
960 start_time=start_time,
961 )
962 for i, quantum_id in enumerate(attempted_quanta):
963 qinfo: ProvenanceQuantumInfo = prov.quantum_only_xgraph.nodes[quantum_id]
964 self.assertEqual(qinfo["attempts"][-1].previous_process_quanta, attempted_quanta[:i])
966 def test_retry(self) -> None:
967 """Test running a full graph with some failures, rerunning the quanta
968 that failed or were blocked in the first attempt, and then scanning
969 for provenance.
970 """
971 with self.make_test_repo() as prep:
972 start_time = time.time()
973 attempted_quanta_1 = list(
974 self.iter_graph_execution(prep.butler_path, prep.predicted, raise_on_partial_outputs=True)
975 )
976 attempted_quanta_2 = list(
977 self.iter_graph_execution(
978 prep.butler_path, prep.predicted, raise_on_partial_outputs=False, is_retry=True
979 )
980 )
981 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
982 ingest_graph(prep.butler_path, prep.config.output_path, transfer="move", batch_size=10)
983 prov = self.check_provenance_graph(
984 prep.predicted,
985 prep.butler,
986 expect_failure=False,
987 start_time=start_time,
988 expect_failures_retried=True,
989 )
990 for i, quantum_id in enumerate(attempted_quanta_1):
991 qinfo: ProvenanceQuantumInfo = prov.quantum_only_xgraph.nodes[quantum_id]
992 self.assertEqual(qinfo["attempts"][0].previous_process_quanta, attempted_quanta_1[:i])
993 expected: list[uuid.UUID] = []
994 for quantum_id in attempted_quanta_2:
995 qinfo: ProvenanceQuantumInfo = prov.quantum_only_xgraph.nodes[quantum_id]
996 if (
997 quantum_id in attempted_quanta_1
998 and qinfo["attempts"][0].status is QuantumAttemptStatus.SUCCESSFUL
999 ):
1000 # These weren't actually attempted twice, since they
1001 # were already successful in the first round.
1002 self.assertEqual(len(qinfo["attempts"]), 1)
1003 else:
1004 self.assertEqual(qinfo["attempts"][-1].previous_process_quanta, expected)
1005 expected.append(quantum_id)
1007 def test_promise_ingest_graph(self) -> None:
1008 """Test running with promise_ingest_graph=True."""
1009 with self.make_test_repo() as prep:
1010 prep.config.incomplete = False
1011 prep.config.promise_ingest_graph = True
1012 start_time = time.time()
1013 attempted_quanta = list(
1014 self.iter_graph_execution(prep.butler_path, prep.predicted, raise_on_partial_outputs=True)
1015 )
1016 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
1017 self.assertFalse(prep.butler.query_datasets("calibrate_metadata", explain=False))
1018 self.assertFalse(prep.butler.query_datasets("consolidate_log", explain=False))
1019 self.assertFalse(prep.butler.query_datasets("resample_config", explain=False))
1020 ingest_graph(prep.butler_path, prep.config.output_path, transfer="move", batch_size=10)
1021 prov = self.check_provenance_graph(
1022 prep.predicted,
1023 prep.butler,
1024 expect_failure=True,
1025 start_time=start_time,
1026 )
1027 self.check_no_original_dirs(prep.butler_path, prov.header.output_run)
1028 for i, quantum_id in enumerate(attempted_quanta):
1029 qinfo: ProvenanceQuantumInfo = prov.quantum_only_xgraph.nodes[quantum_id]
1030 self.assertEqual(qinfo["attempts"][-1].previous_process_quanta, attempted_quanta[:i])
1032 def test_worker_failures(self) -> None:
1033 """Test that if failures occur on (multiple) workers we shut down
1034 gracefully instead of hanging.
1035 """
1036 with self.make_test_repo() as prep:
1037 with self.assertRaises(FatalWorkerError):
1038 aggregate_graph(prep.predicted_path, "nonexistent", prep.config)
1040 def test_aggregate_graph_cli_overrides(self) -> None:
1041 """Test that command-line options override config attributes as
1042 expected.
1043 """
1045 def mock_run(predicted_path: str, butler_path: str, config: AggregatorConfig) -> None:
1046 print(config.model_dump_json(indent=2))
1048 def check(result: Result, **kwargs: Any) -> None:
1049 self.assertEqual(result.exit_code, 0, msg=result.output)
1050 self.assertEqual(
1051 result.output.strip(), AggregatorConfig(**kwargs).model_dump_json(indent=2).strip()
1052 )
1054 self.maxDiff = None
1055 runner = CliRunner()
1056 with unittest.mock.patch("lsst.pipe.base.quantum_graph.aggregator.aggregate_graph", mock_run):
1057 check(runner.invoke(aggregate_graph_cli, ("pg", "repo")))
1058 check(runner.invoke(aggregate_graph_cli, ("pg", "repo", "--output", "out")), output_path="out")
1059 check(runner.invoke(aggregate_graph_cli, ("pg", "repo", "-j", "4")), n_processes=4)
1060 check(
1061 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--incomplete")),
1062 incomplete=True,
1063 )
1064 check(runner.invoke(aggregate_graph_cli, ("pg", "repo", "--dry-run")), dry_run=True)
1065 check(
1066 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--interactive-status")),
1067 interactive_status=True,
1068 )
1069 check(
1070 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--log-status-interval", "120")),
1071 log_status_interval=120,
1072 )
1073 check(
1074 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--no-register-dataset-types")),
1075 register_dataset_types=False,
1076 )
1077 check(
1078 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--no-update-output-chain")),
1079 update_output_chain=False,
1080 )
1081 check(
1082 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--worker-log-dir", "wlogs")),
1083 worker_log_dir="wlogs",
1084 )
1085 check(
1086 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--worker-log-level", "DEBUG")),
1087 worker_log_level="DEBUG",
1088 )
1089 check(
1090 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--zstd-level", "11")),
1091 zstd_level=11,
1092 )
1093 check(
1094 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--zstd-dict-size", "143")),
1095 zstd_dict_size=143,
1096 )
1097 check(
1098 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--zstd-dict-n-inputs", "2")),
1099 zstd_dict_n_inputs=2,
1100 )
1101 check(
1102 runner.invoke(aggregate_graph_cli, ("pg", "repo", "--mock-storage-classes")),
1103 mock_storage_classes=True,
1104 )
1106 def check_provenance_report(self, result: click.testing.Result, root: str) -> None:
1107 self.maxDiff = None
1108 self.assertEqual(result.exit_code, 0, msg=result.output)
1109 self.assertEqual(
1110 result.output,
1111 " Task Caveats Failed Blocked Successful TOTAL EXPECTED\n"
1112 "----------- ------- ------ ------- ---------- ----- --------\n"
1113 " calibrate 1 0 7 8 8\n"
1114 "consolidate 0 1 1 2 2\n"
1115 " resample 0 4 6 10 10\n"
1116 " coadd 0 4 6 10 10\n"
1117 "\n"
1118 " Task Exception Successes Failures\n"
1119 "--------- --------------------------------------------- --------- --------\n"
1120 "calibrate lsst.pipe.base.tests.mocks.MockAlgorithmError 0 1\n"
1121 "\n",
1122 )
1123 with open(os.path.join(root, "report.json")) as report_file:
1124 report = ProvenanceReport.model_validate_json(report_file.read())
1125 self.assertEqual(report.root.keys(), {"calibrate"})
1126 self.assertEqual(report.root["calibrate"].keys(), {"FAILED"})
1127 self.assertEqual(
1128 report.root["calibrate"]["FAILED"].keys(), {"lsst.pipe.base.tests.mocks.MockAlgorithmError"}
1129 )
1130 self.assertEqual(
1131 len(report.root["calibrate"]["FAILED"]["lsst.pipe.base.tests.mocks.MockAlgorithmError"]), 1
1132 )
1133 self.assertEqual(
1134 list(os.walk(os.path.join(root, "data_ids"))),
1135 [
1136 (os.path.join(root, "data_ids"), ["calibrate"], []),
1137 (os.path.join(root, "data_ids", "calibrate"), ["FAILED"], []),
1138 (
1139 os.path.join(root, "data_ids", "calibrate", "FAILED"),
1140 [],
1141 ["lsst.pipe.base.tests.mocks.MockAlgorithmError.ecsv"],
1142 ),
1143 ],
1144 )
1146 def check_no_original_dirs(self, butler_path: str, output_run: str) -> None:
1147 """Check that there are no config/log/metadata directories in
1148 the butler's directory for this output run.
1149 """
1150 root = os.path.join(butler_path, output_run)
1151 for subdir in os.listdir(root):
1152 if subdir.endswith("_config") or subdir.endswith("_metadata") or subdir.endswith("_log"):
1153 raise AssertionError(f"Directory {os.path.join(root, subdir)} still exists.")
1155 def test_provenance_report_content(self) -> None:
1156 """Test the provenance-report CLI command."""
1157 with self.make_test_repo() as prep:
1158 prep.config.incomplete = False
1159 prep.config.promise_ingest_graph = True
1160 for _ in self.iter_graph_execution(
1161 prep.butler_path, prep.predicted, raise_on_partial_outputs=True
1162 ):
1163 pass
1164 aggregate_graph(prep.predicted_path, prep.butler_path, prep.config)
1165 self.assertFalse(prep.butler.query_datasets("calibrate_metadata", explain=False))
1166 self.assertFalse(prep.butler.query_datasets("consolidate_log", explain=False))
1167 self.assertFalse(prep.butler.query_datasets("resample_config", explain=False))
1169 # First test on a provenance graph file that has not been ingested.
1170 runner = CliRunner()
1171 report_root = os.path.join(prep.butler_path, "uningested")
1172 result = runner.invoke(
1173 provenance_report_cli,
1174 (
1175 cast(str, prep.config.output_path),
1176 "--status-report",
1177 os.path.join(report_root, "report.json"),
1178 "--data-id-table-dir",
1179 os.path.join(report_root, "data_ids"),
1180 ),
1181 )
1182 self.check_provenance_report(result, report_root)
1184 ingest_graph(prep.butler_path, prep.config.output_path, transfer="move", batch_size=10)
1186 runner = CliRunner()
1187 report_root = os.path.join(prep.butler_path, "ingested")
1188 result = runner.invoke(
1189 provenance_report_cli,
1190 (
1191 prep.butler_path,
1192 *prep.butler.collections.defaults,
1193 "--status-report",
1194 os.path.join(report_root, "report.json"),
1195 "--data-id-table-dir",
1196 os.path.join(report_root, "data_ids"),
1197 ),
1198 )
1199 self.check_provenance_report(result, os.path.join(report_root))
1201 def test_provenance_report_cli_overrides(self) -> None:
1202 """Test the provenance-report CLI command with a mocked
1203 implementation.
1204 """
1206 class MakeManyReportsArgs(pydantic.BaseModel):
1207 status_report_file: str | None = None
1208 print_quantum_table: bool = True
1209 print_exception_table: bool = True
1210 states: list[QuantumAttemptStatus] = pydantic.Field(
1211 default_factory=lambda: [
1212 QuantumAttemptStatus.FAILED,
1213 QuantumAttemptStatus.ABORTED,
1214 QuantumAttemptStatus.ABORTED_SUCCESS,
1215 ]
1216 )
1217 with_caveats: QuantumSuccessCaveats | None = None
1218 data_id_table_dir: str | None = None
1220 @pydantic.model_validator(mode="after")
1221 def _sort_states(self) -> MakeManyReportsArgs:
1222 self.states.sort(key=lambda e: e.value)
1223 return self
1225 class MockProvenanceQuantumGraph(pydantic.BaseModel):
1226 repo_or_filename: str
1227 collection: str | None
1228 quanta: list[uuid.UUID] | None = None
1229 datasets: list[uuid.UUID] | None = None
1230 writeable: bool = False
1231 make_many_reports_args: MakeManyReportsArgs | None = None
1233 @classmethod
1234 @contextmanager
1235 def from_args(
1236 cls, repo_or_filename: str, /, **kwargs: Any
1237 ) -> Iterator[tuple[MockProvenanceQuantumGraph, None]]:
1238 yield cls(repo_or_filename=repo_or_filename, **kwargs), None
1240 def make_many_reports(self, **kwargs: Any) -> None:
1241 self.make_many_reports_args = MakeManyReportsArgs(**kwargs)
1242 print(self.model_dump_json(indent=2))
1244 def check(
1245 result: Result, repo_or_filename: str, collection: str | None = None, **kwargs: Any
1246 ) -> None:
1247 self.assertEqual(result.exit_code, 0, msg=result.output)
1248 self.assertEqual(
1249 result.output.strip(),
1250 MockProvenanceQuantumGraph(
1251 repo_or_filename=repo_or_filename,
1252 collection=collection,
1253 datasets=[],
1254 writeable=False,
1255 make_many_reports_args=MakeManyReportsArgs(**kwargs),
1256 ).model_dump_json(indent=2),
1257 )
1259 self.maxDiff = None
1260 runner = CliRunner()
1261 with unittest.mock.patch(
1262 "lsst.pipe.base.quantum_graph.ProvenanceQuantumGraph", MockProvenanceQuantumGraph
1263 ):
1264 check(
1265 runner.invoke(provenance_report_cli, ("repo", "collection1")),
1266 repo_or_filename="repo",
1267 collection="collection1",
1268 )
1269 check(
1270 runner.invoke(provenance_report_cli, ("filename1",)),
1271 repo_or_filename="filename1",
1272 collection=None,
1273 )
1274 check(
1275 runner.invoke(provenance_report_cli, ("repo", "collection1", "--no-quantum-table")),
1276 repo_or_filename="repo",
1277 collection="collection1",
1278 print_quantum_table=False,
1279 )
1280 check(
1281 runner.invoke(provenance_report_cli, ("repo", "collection1", "--no-exception-table")),
1282 repo_or_filename="repo",
1283 collection="collection1",
1284 print_exception_table=False,
1285 )
1286 check(
1287 runner.invoke(provenance_report_cli, ("repo", "collection1", "--status-report", "filename2")),
1288 repo_or_filename="repo",
1289 collection="collection1",
1290 status_report_file="filename2",
1291 )
1292 check(
1293 runner.invoke(provenance_report_cli, ("repo", "collection1", "--state", "SUCCESSFUL")),
1294 repo_or_filename="repo",
1295 collection="collection1",
1296 states=[
1297 QuantumAttemptStatus.SUCCESSFUL,
1298 QuantumAttemptStatus.FAILED,
1299 QuantumAttemptStatus.ABORTED,
1300 QuantumAttemptStatus.ABORTED_SUCCESS,
1301 ],
1302 )
1303 check(
1304 runner.invoke(provenance_report_cli, ("repo", "collection1", "--no-state", "FAILED")),
1305 repo_or_filename="repo",
1306 collection="collection1",
1307 states=[
1308 QuantumAttemptStatus.ABORTED,
1309 QuantumAttemptStatus.ABORTED_SUCCESS,
1310 ],
1311 )
1312 check(
1313 runner.invoke(
1314 provenance_report_cli, ("repo", "collection1", "--caveat", "PARTIAL_OUTPUTS_ERROR")
1315 ),
1316 repo_or_filename="repo",
1317 collection="collection1",
1318 with_caveats=QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR,
1319 states=[
1320 QuantumAttemptStatus.SUCCESSFUL,
1321 QuantumAttemptStatus.FAILED,
1322 QuantumAttemptStatus.ABORTED,
1323 QuantumAttemptStatus.ABORTED_SUCCESS,
1324 ],
1325 )
1326 check(
1327 runner.invoke(provenance_report_cli, ("repo", "collection1", "--data-id-table-dir", "dir1")),
1328 repo_or_filename="repo",
1329 collection="collection1",
1330 data_id_table_dir="dir1",
1331 )
1334if __name__ == "__main__":
1335 lsst.utils.tests.init()
1336 unittest.main()