lsst-pipe-base 30.0.0rc2__py3-none-any.whl → 30.0.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. lsst/pipe/base/_instrument.py +31 -20
  2. lsst/pipe/base/_quantumContext.py +3 -3
  3. lsst/pipe/base/_status.py +43 -10
  4. lsst/pipe/base/_task_metadata.py +2 -2
  5. lsst/pipe/base/all_dimensions_quantum_graph_builder.py +8 -3
  6. lsst/pipe/base/automatic_connection_constants.py +20 -1
  7. lsst/pipe/base/cli/cmd/__init__.py +18 -2
  8. lsst/pipe/base/cli/cmd/commands.py +149 -4
  9. lsst/pipe/base/connectionTypes.py +72 -160
  10. lsst/pipe/base/connections.py +6 -9
  11. lsst/pipe/base/execution_reports.py +0 -5
  12. lsst/pipe/base/graph/graph.py +11 -10
  13. lsst/pipe/base/graph/quantumNode.py +4 -4
  14. lsst/pipe/base/graph_walker.py +8 -10
  15. lsst/pipe/base/log_capture.py +40 -80
  16. lsst/pipe/base/log_on_close.py +76 -0
  17. lsst/pipe/base/mp_graph_executor.py +51 -15
  18. lsst/pipe/base/pipeline.py +5 -6
  19. lsst/pipe/base/pipelineIR.py +2 -8
  20. lsst/pipe/base/pipelineTask.py +5 -7
  21. lsst/pipe/base/pipeline_graph/_dataset_types.py +2 -2
  22. lsst/pipe/base/pipeline_graph/_edges.py +32 -22
  23. lsst/pipe/base/pipeline_graph/_mapping_views.py +4 -7
  24. lsst/pipe/base/pipeline_graph/_pipeline_graph.py +14 -7
  25. lsst/pipe/base/pipeline_graph/expressions.py +2 -2
  26. lsst/pipe/base/pipeline_graph/io.py +7 -10
  27. lsst/pipe/base/pipeline_graph/visualization/_dot.py +13 -12
  28. lsst/pipe/base/pipeline_graph/visualization/_layout.py +16 -18
  29. lsst/pipe/base/pipeline_graph/visualization/_merge.py +4 -7
  30. lsst/pipe/base/pipeline_graph/visualization/_printer.py +10 -10
  31. lsst/pipe/base/pipeline_graph/visualization/_status_annotator.py +7 -0
  32. lsst/pipe/base/prerequisite_helpers.py +2 -1
  33. lsst/pipe/base/quantum_graph/_common.py +19 -20
  34. lsst/pipe/base/quantum_graph/_multiblock.py +37 -31
  35. lsst/pipe/base/quantum_graph/_predicted.py +113 -15
  36. lsst/pipe/base/quantum_graph/_provenance.py +1136 -45
  37. lsst/pipe/base/quantum_graph/aggregator/__init__.py +0 -1
  38. lsst/pipe/base/quantum_graph/aggregator/_communicators.py +204 -289
  39. lsst/pipe/base/quantum_graph/aggregator/_config.py +87 -9
  40. lsst/pipe/base/quantum_graph/aggregator/_ingester.py +13 -12
  41. lsst/pipe/base/quantum_graph/aggregator/_scanner.py +49 -235
  42. lsst/pipe/base/quantum_graph/aggregator/_structs.py +6 -116
  43. lsst/pipe/base/quantum_graph/aggregator/_supervisor.py +29 -39
  44. lsst/pipe/base/quantum_graph/aggregator/_workers.py +303 -0
  45. lsst/pipe/base/quantum_graph/aggregator/_writer.py +34 -351
  46. lsst/pipe/base/quantum_graph/formatter.py +171 -0
  47. lsst/pipe/base/quantum_graph/ingest_graph.py +413 -0
  48. lsst/pipe/base/quantum_graph/visualization.py +5 -1
  49. lsst/pipe/base/quantum_graph_builder.py +33 -9
  50. lsst/pipe/base/quantum_graph_executor.py +116 -13
  51. lsst/pipe/base/quantum_graph_skeleton.py +31 -35
  52. lsst/pipe/base/quantum_provenance_graph.py +29 -12
  53. lsst/pipe/base/separable_pipeline_executor.py +19 -3
  54. lsst/pipe/base/single_quantum_executor.py +67 -42
  55. lsst/pipe/base/struct.py +4 -0
  56. lsst/pipe/base/testUtils.py +3 -3
  57. lsst/pipe/base/tests/mocks/_storage_class.py +2 -1
  58. lsst/pipe/base/version.py +1 -1
  59. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/METADATA +3 -3
  60. lsst_pipe_base-30.0.1.dist-info/RECORD +129 -0
  61. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/WHEEL +1 -1
  62. lsst_pipe_base-30.0.0rc2.dist-info/RECORD +0 -125
  63. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/entry_points.txt +0 -0
  64. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/licenses/COPYRIGHT +0 -0
  65. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/licenses/LICENSE +0 -0
  66. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/licenses/bsd_license.txt +0 -0
  67. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/licenses/gpl-v3.0.txt +0 -0
  68. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/top_level.txt +0 -0
  69. {lsst_pipe_base-30.0.0rc2.dist-info → lsst_pipe_base-30.0.1.dist-info}/zip-safe +0 -0
@@ -35,37 +35,50 @@ __all__ = (
35
35
  "ProvenanceLogRecordsModel",
36
36
  "ProvenanceQuantumGraph",
37
37
  "ProvenanceQuantumGraphReader",
38
+ "ProvenanceQuantumGraphWriter",
38
39
  "ProvenanceQuantumInfo",
39
40
  "ProvenanceQuantumModel",
41
+ "ProvenanceQuantumReport",
42
+ "ProvenanceQuantumScanData",
43
+ "ProvenanceQuantumScanModels",
44
+ "ProvenanceQuantumScanStatus",
45
+ "ProvenanceReport",
40
46
  "ProvenanceTaskMetadataModel",
41
47
  )
42
48
 
43
-
44
49
  import dataclasses
50
+ import enum
51
+ import itertools
45
52
  import sys
46
53
  import uuid
47
54
  from collections import Counter
48
- from collections.abc import Iterable, Iterator, Mapping
49
- from contextlib import contextmanager
50
- from typing import TYPE_CHECKING, Any, TypedDict, TypeVar
55
+ from collections.abc import Callable, Iterable, Iterator, Mapping
56
+ from contextlib import ExitStack, contextmanager
57
+ from typing import TYPE_CHECKING, Any, TypedDict
51
58
 
52
59
  import astropy.table
53
60
  import networkx
54
61
  import numpy as np
55
62
  import pydantic
56
63
 
57
- from lsst.daf.butler import DataCoordinate
64
+ from lsst.daf.butler import Butler, DataCoordinate
58
65
  from lsst.daf.butler.logging import ButlerLogRecord, ButlerLogRecords
59
- from lsst.resources import ResourcePathExpression
66
+ from lsst.resources import ResourcePath, ResourcePathExpression
67
+ from lsst.utils.iteration import ensure_iterable
68
+ from lsst.utils.logging import LsstLogAdapter, getLogger
60
69
  from lsst.utils.packages import Packages
61
70
 
71
+ from .. import automatic_connection_constants as acc
62
72
  from .._status import ExceptionInfo, QuantumAttemptStatus, QuantumSuccessCaveats
63
73
  from .._task_metadata import TaskMetadata
74
+ from ..log_capture import _ExecutionLogRecordsExtra
75
+ from ..log_on_close import LogOnClose
64
76
  from ..pipeline_graph import PipelineGraph, TaskImportMode, TaskInitNode
65
77
  from ..resource_usage import QuantumResourceUsage
66
78
  from ._common import (
67
79
  BaseQuantumGraph,
68
80
  BaseQuantumGraphReader,
81
+ BaseQuantumGraphWriter,
69
82
  ConnectionName,
70
83
  DataCoordinateValues,
71
84
  DatasetInfo,
@@ -74,8 +87,24 @@ from ._common import (
74
87
  QuantumInfo,
75
88
  TaskLabel,
76
89
  )
77
- from ._multiblock import MultiblockReader
78
- from ._predicted import PredictedDatasetModel, PredictedQuantumDatasetsModel
90
+ from ._multiblock import Compressor, MultiblockReader, MultiblockWriter
91
+ from ._predicted import (
92
+ PredictedDatasetModel,
93
+ PredictedQuantumDatasetsModel,
94
+ PredictedQuantumGraph,
95
+ PredictedQuantumGraphComponents,
96
+ )
97
+
98
+ # Sphinx needs imports for type annotations of base class members.
99
+ if "sphinx" in sys.modules:
100
+ import zipfile # noqa: F401
101
+
102
+ from ._multiblock import AddressReader, Decompressor # noqa: F401
103
+
104
+
105
+ type LoopWrapper[T] = Callable[[Iterable[T]], Iterable[T]]
106
+
107
+ _LOG = getLogger(__file__)
79
108
 
80
109
  DATASET_ADDRESS_INDEX = 0
81
110
  QUANTUM_ADDRESS_INDEX = 1
@@ -87,7 +116,9 @@ QUANTUM_MB_NAME = "quanta"
87
116
  LOG_MB_NAME = "logs"
88
117
  METADATA_MB_NAME = "metadata"
89
118
 
90
- _I = TypeVar("_I", bound=uuid.UUID | int)
119
+
120
+ def pass_through[T](arg: T) -> T:
121
+ return arg
91
122
 
92
123
 
93
124
  class ProvenanceDatasetInfo(DatasetInfo):
@@ -161,6 +192,12 @@ class ProvenanceQuantumInfo(QuantumInfo):
161
192
  failure.
162
193
  """
163
194
 
195
+ metadata_id: uuid.UUID
196
+ """ID of this quantum's metadata dataset."""
197
+
198
+ log_id: uuid.UUID
199
+ """ID of this quantum's log dataset."""
200
+
164
201
 
165
202
  class ProvenanceInitQuantumInfo(TypedDict):
166
203
  """A typed dictionary that annotates the attributes of the NetworkX graph
@@ -187,6 +224,9 @@ class ProvenanceInitQuantumInfo(TypedDict):
187
224
  pipeline_node: TaskInitNode
188
225
  """Node in the pipeline graph for this task's init-only step."""
189
226
 
227
+ config_id: uuid.UUID
228
+ """ID of this task's config dataset."""
229
+
190
230
 
191
231
  class ProvenanceDatasetModel(PredictedDatasetModel):
192
232
  """Data model for the datasets in a provenance quantum graph file."""
@@ -518,6 +558,131 @@ class ProvenanceTaskMetadataModel(pydantic.BaseModel):
518
558
  return super().model_validate_strings(*args, **kwargs)
519
559
 
520
560
 
561
+ class ProvenanceQuantumReport(pydantic.BaseModel):
562
+ """A Pydantic model that used to report information about a single
563
+ (generally problematic) quantum.
564
+ """
565
+
566
+ quantum_id: uuid.UUID
567
+ data_id: dict[str, int | str]
568
+ attempts: list[ProvenanceQuantumAttemptModel]
569
+
570
+ @classmethod
571
+ def from_info(cls, quantum_id: uuid.UUID, quantum_info: ProvenanceQuantumInfo) -> ProvenanceQuantumReport:
572
+ """Construct from a provenance quantum graph node.
573
+
574
+ Parameters
575
+ ----------
576
+ quantum_id : `uuid.UUID`
577
+ Unique ID for the quantum.
578
+ quantum_info : `ProvenanceQuantumInfo`
579
+ Node attributes for this quantum.
580
+ """
581
+ return cls(
582
+ quantum_id=quantum_id,
583
+ data_id=dict(quantum_info["data_id"].mapping),
584
+ attempts=quantum_info["attempts"],
585
+ )
586
+
587
+ # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
588
+ # when we inherit those docstrings in our public classes.
589
+ if "sphinx" in sys.modules and not TYPE_CHECKING:
590
+
591
+ def copy(self, *args: Any, **kwargs: Any) -> Any:
592
+ """See `pydantic.BaseModel.copy`."""
593
+ return super().copy(*args, **kwargs)
594
+
595
+ def model_dump(self, *args: Any, **kwargs: Any) -> Any:
596
+ """See `pydantic.BaseModel.model_dump`."""
597
+ return super().model_dump(*args, **kwargs)
598
+
599
+ def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
600
+ """See `pydantic.BaseModel.model_dump_json`."""
601
+ return super().model_dump(*args, **kwargs)
602
+
603
+ def model_copy(self, *args: Any, **kwargs: Any) -> Any:
604
+ """See `pydantic.BaseModel.model_copy`."""
605
+ return super().model_copy(*args, **kwargs)
606
+
607
+ @classmethod
608
+ def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
609
+ """See `pydantic.BaseModel.model_construct`."""
610
+ return super().model_construct(*args, **kwargs)
611
+
612
+ @classmethod
613
+ def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
614
+ """See `pydantic.BaseModel.model_json_schema`."""
615
+ return super().model_json_schema(*args, **kwargs)
616
+
617
+ @classmethod
618
+ def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
619
+ """See `pydantic.BaseModel.model_validate`."""
620
+ return super().model_validate(*args, **kwargs)
621
+
622
+ @classmethod
623
+ def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
624
+ """See `pydantic.BaseModel.model_validate_json`."""
625
+ return super().model_validate_json(*args, **kwargs)
626
+
627
+ @classmethod
628
+ def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
629
+ """See `pydantic.BaseModel.model_validate_strings`."""
630
+ return super().model_validate_strings(*args, **kwargs)
631
+
632
+
633
+ class ProvenanceReport(pydantic.RootModel):
634
+ """A Pydantic model that groups quantum information by task label, then
635
+ status (as a string), and then exception type.
636
+ """
637
+
638
+ root: dict[TaskLabel, dict[str, dict[str | None, list[ProvenanceQuantumReport]]]] = {}
639
+
640
+ # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
641
+ # when we inherit those docstrings in our public classes.
642
+ if "sphinx" in sys.modules and not TYPE_CHECKING:
643
+
644
+ def copy(self, *args: Any, **kwargs: Any) -> Any:
645
+ """See `pydantic.BaseModel.copy`."""
646
+ return super().copy(*args, **kwargs)
647
+
648
+ def model_dump(self, *args: Any, **kwargs: Any) -> Any:
649
+ """See `pydantic.BaseModel.model_dump`."""
650
+ return super().model_dump(*args, **kwargs)
651
+
652
+ def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
653
+ """See `pydantic.BaseModel.model_dump_json`."""
654
+ return super().model_dump(*args, **kwargs)
655
+
656
+ def model_copy(self, *args: Any, **kwargs: Any) -> Any:
657
+ """See `pydantic.BaseModel.model_copy`."""
658
+ return super().model_copy(*args, **kwargs)
659
+
660
+ @classmethod
661
+ def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
662
+ """See `pydantic.BaseModel.model_construct`."""
663
+ return super().model_construct(*args, **kwargs)
664
+
665
+ @classmethod
666
+ def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
667
+ """See `pydantic.BaseModel.model_json_schema`."""
668
+ return super().model_json_schema(*args, **kwargs)
669
+
670
+ @classmethod
671
+ def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
672
+ """See `pydantic.BaseModel.model_validate`."""
673
+ return super().model_validate(*args, **kwargs)
674
+
675
+ @classmethod
676
+ def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
677
+ """See `pydantic.BaseModel.model_validate_json`."""
678
+ return super().model_validate_json(*args, **kwargs)
679
+
680
+ @classmethod
681
+ def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
682
+ """See `pydantic.BaseModel.model_validate_strings`."""
683
+ return super().model_validate_strings(*args, **kwargs)
684
+
685
+
521
686
  class ProvenanceQuantumModel(pydantic.BaseModel):
522
687
  """Data model for the quanta in a provenance quantum graph file."""
523
688
 
@@ -621,6 +786,8 @@ class ProvenanceQuantumModel(pydantic.BaseModel):
621
786
  resource_usage=last_attempt.resource_usage,
622
787
  attempts=self.attempts,
623
788
  )
789
+ graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
790
+ graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
624
791
  for connection_name, dataset_ids in self.inputs.items():
625
792
  read_edge = task_node.get_input_edge(connection_name)
626
793
  for dataset_id in dataset_ids:
@@ -630,6 +797,30 @@ class ProvenanceQuantumModel(pydantic.BaseModel):
630
797
  ).append(read_edge)
631
798
  for connection_name, dataset_ids in self.outputs.items():
632
799
  write_edge = task_node.get_output_edge(connection_name)
800
+ if connection_name == acc.METADATA_OUTPUT_CONNECTION_NAME:
801
+ graph._bipartite_xgraph.add_node(
802
+ dataset_ids[0],
803
+ data_id=data_id,
804
+ dataset_type_name=write_edge.dataset_type_name,
805
+ pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
806
+ run=graph.header.output_run,
807
+ produced=last_attempt.status.has_metadata,
808
+ )
809
+ graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
810
+ graph._bipartite_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
811
+ graph._quantum_only_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
812
+ if connection_name == acc.LOG_OUTPUT_CONNECTION_NAME:
813
+ graph._bipartite_xgraph.add_node(
814
+ dataset_ids[0],
815
+ data_id=data_id,
816
+ dataset_type_name=write_edge.dataset_type_name,
817
+ pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
818
+ run=graph.header.output_run,
819
+ produced=last_attempt.status.has_log,
820
+ )
821
+ graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
822
+ graph._bipartite_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
823
+ graph._quantum_only_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
633
824
  for dataset_id in dataset_ids:
634
825
  graph._bipartite_xgraph.add_edge(
635
826
  self.quantum_id,
@@ -638,8 +829,6 @@ class ProvenanceQuantumModel(pydantic.BaseModel):
638
829
  # There can only be one pipeline edge for an output.
639
830
  pipeline_edges=[write_edge],
640
831
  )
641
- graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
642
- graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
643
832
  for dataset_id in graph._bipartite_xgraph.predecessors(self.quantum_id):
644
833
  for upstream_quantum_id in graph._bipartite_xgraph.predecessors(dataset_id):
645
834
  graph._quantum_only_xgraph.add_edge(upstream_quantum_id, self.quantum_id)
@@ -778,6 +967,15 @@ class ProvenanceInitQuantumModel(pydantic.BaseModel):
778
967
  ).append(read_edge)
779
968
  for connection_name, dataset_id in self.outputs.items():
780
969
  write_edge = task_init_node.get_output_edge(connection_name)
970
+ graph._bipartite_xgraph.add_node(
971
+ dataset_id,
972
+ data_id=empty_data_id,
973
+ dataset_type_name=write_edge.dataset_type_name,
974
+ pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
975
+ run=graph.header.output_run,
976
+ produced=True,
977
+ )
978
+ graph._datasets_by_type[write_edge.dataset_type_name][empty_data_id] = dataset_id
781
979
  graph._bipartite_xgraph.add_edge(
782
980
  self.quantum_id,
783
981
  dataset_id,
@@ -785,6 +983,8 @@ class ProvenanceInitQuantumModel(pydantic.BaseModel):
785
983
  # There can only be one pipeline edge for an output.
786
984
  pipeline_edges=[write_edge],
787
985
  )
986
+ if write_edge.connection_name == acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME:
987
+ graph._bipartite_xgraph.nodes[self.quantum_id]["config_id"] = dataset_id
788
988
  graph._init_quanta[self.task_label] = self.quantum_id
789
989
 
790
990
  # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
@@ -929,6 +1129,83 @@ class ProvenanceQuantumGraph(BaseQuantumGraph):
929
1129
  dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
930
1130
  }
931
1131
 
1132
+ @classmethod
1133
+ @contextmanager
1134
+ def from_args(
1135
+ cls,
1136
+ repo_or_filename: str,
1137
+ /,
1138
+ collection: str | None = None,
1139
+ *,
1140
+ quanta: Iterable[uuid.UUID] | None = None,
1141
+ datasets: Iterable[uuid.UUID] | None = None,
1142
+ writeable: bool = False,
1143
+ ) -> Iterator[tuple[ProvenanceQuantumGraph, Butler | None]]:
1144
+ """Construct a `ProvenanceQuantumGraph` fron CLI-friendly arguments for
1145
+ a file or butler-ingested graph dataset.
1146
+
1147
+ Parameters
1148
+ ----------
1149
+ repo_or_filename : `str`
1150
+ Either a provenance quantum graph filename or a butler repository
1151
+ path or alias.
1152
+ collection : `str`, optional
1153
+ Collection to search; presence indicates that the first argument
1154
+ is a butler repository, not a filename.
1155
+ quanta : `~collections.abc.Iterable` [ `str` ] or `None`, optional
1156
+ IDs of the quanta to load, or `None` to load all.
1157
+ datasets : `~collections.abc.Iterable` [ `str` ], optional
1158
+ IDs of the datasets to load, or `None` to load all.
1159
+ writeable : `bool`, optional
1160
+ Whether the butler should be constructed with write support.
1161
+
1162
+ Returns
1163
+ -------
1164
+ context : `contextlib.AbstractContextManager`
1165
+ A context manager that yields a tuple of
1166
+
1167
+ - the `ProvenanceQuantumGraph`
1168
+ - the `Butler` constructed (or `None`)
1169
+
1170
+ when entered.
1171
+ """
1172
+ exit_stack = ExitStack()
1173
+ if collection is not None:
1174
+ try:
1175
+ butler = exit_stack.enter_context(
1176
+ Butler.from_config(repo_or_filename, collections=[collection], writeable=writeable)
1177
+ )
1178
+ except Exception as err:
1179
+ err.add_note(
1180
+ f"Expected {repo_or_filename!r} to be a butler repository path or alias because a "
1181
+ f"collection ({collection}) was provided."
1182
+ )
1183
+ raise
1184
+ with exit_stack:
1185
+ graph = butler.get(
1186
+ acc.PROVENANCE_DATASET_TYPE_NAME, parameters={"quanta": quanta, "datasets": datasets}
1187
+ )
1188
+ yield graph, butler
1189
+ else:
1190
+ try:
1191
+ reader = exit_stack.enter_context(ProvenanceQuantumGraphReader.open(repo_or_filename))
1192
+ except Exception as err:
1193
+ err.add_note(
1194
+ f"Expected a {repo_or_filename} to be a provenance quantum graph filename "
1195
+ f"because no collection was provided."
1196
+ )
1197
+ raise
1198
+ with exit_stack:
1199
+ if quanta is None:
1200
+ reader.read_quanta()
1201
+ elif not quanta:
1202
+ reader.read_quanta(quanta)
1203
+ if datasets is None:
1204
+ reader.read_datasets()
1205
+ elif not datasets:
1206
+ reader.read_datasets(datasets)
1207
+ yield reader.graph, None
1208
+
932
1209
  @property
933
1210
  def init_quanta(self) -> Mapping[TaskLabel, uuid.UUID]:
934
1211
  """A mapping from task label to the ID of the special init quantum for
@@ -969,6 +1246,8 @@ class ProvenanceQuantumGraph(BaseQuantumGraph):
969
1246
  types in the pipeline graph are included, even if none of their
970
1247
  datasets were loaded (i.e. nested mappings may be empty).
971
1248
 
1249
+ Reading a quantum also populates its log and metadata datasets.
1250
+
972
1251
  The returned object may be an internal dictionary; as the type
973
1252
  annotation indicates, it should not be modified in place.
974
1253
  """
@@ -1007,7 +1286,8 @@ class ProvenanceQuantumGraph(BaseQuantumGraph):
1007
1286
  `ProvenanceQuantumGraphReader.read_quanta`) or datasets (via
1008
1287
  `ProvenanceQuantumGraphReader.read_datasets`) will load those nodes
1009
1288
  with full attributes and edges to adjacent nodes with no attributes.
1010
- Loading quanta necessary to populate edge attributes.
1289
+ Loading quanta is necessary to populate edge attributes.
1290
+ Reading a quantum also populates its log and metadata datasets.
1011
1291
 
1012
1292
  Node attributes are described by the
1013
1293
  `ProvenanceQuantumInfo`, `ProvenanceInitQuantumInfo`, and
@@ -1022,10 +1302,16 @@ class ProvenanceQuantumGraph(BaseQuantumGraph):
1022
1302
  """
1023
1303
  return self._bipartite_xgraph.copy(as_view=True)
1024
1304
 
1025
- def make_quantum_table(self) -> astropy.table.Table:
1305
+ def make_quantum_table(self, drop_unused_columns: bool = True) -> astropy.table.Table:
1026
1306
  """Construct an `astropy.table.Table` with a tabular summary of the
1027
1307
  quanta.
1028
1308
 
1309
+ Parameters
1310
+ ----------
1311
+ drop_unused_columns : `bool`, optional
1312
+ Whether to drop columns for rare states that did not actually
1313
+ occur in this run.
1314
+
1029
1315
  Returns
1030
1316
  -------
1031
1317
  table : `astropy.table.Table`
@@ -1061,28 +1347,30 @@ class ProvenanceQuantumGraph(BaseQuantumGraph):
1061
1347
  caveats = f"{code.concise()}({count})" # type: ignore[union-attr]
1062
1348
  else:
1063
1349
  caveats = ""
1064
- rows.append(
1350
+ row: dict[str, Any] = {
1351
+ "Task": task_label,
1352
+ "Caveats": caveats,
1353
+ }
1354
+ for status in QuantumAttemptStatus:
1355
+ row[status.title] = status_counts.get(status, 0)
1356
+ row.update(
1065
1357
  {
1066
- "Task": task_label,
1067
- "Unknown": status_counts.get(QuantumAttemptStatus.UNKNOWN, 0),
1068
- "Successful": status_counts.get(QuantumAttemptStatus.SUCCESSFUL, 0),
1069
- "Caveats": caveats,
1070
- "Blocked": status_counts.get(QuantumAttemptStatus.BLOCKED, 0),
1071
- "Failed": status_counts.get(QuantumAttemptStatus.FAILED, 0),
1072
1358
  "TOTAL": len(quanta_for_task),
1073
1359
  "EXPECTED": self.header.n_task_quanta[task_label],
1074
1360
  }
1075
1361
  )
1076
- return astropy.table.Table(rows)
1362
+ rows.append(row)
1363
+ table = astropy.table.Table(rows)
1364
+ if drop_unused_columns:
1365
+ for status in QuantumAttemptStatus:
1366
+ if status.is_rare and not table[status.title].any():
1367
+ del table[status.title]
1368
+ return table
1077
1369
 
1078
1370
  def make_exception_table(self) -> astropy.table.Table:
1079
1371
  """Construct an `astropy.table.Table` with counts for each exception
1080
1372
  type raised by each task.
1081
1373
 
1082
- At present this only includes information from partial-outputs-error
1083
- successes, since exception information for failures is not tracked.
1084
- This may change in the future.
1085
-
1086
1374
  Returns
1087
1375
  -------
1088
1376
  table : `astropy.table.Table`
@@ -1090,13 +1378,25 @@ class ProvenanceQuantumGraph(BaseQuantumGraph):
1090
1378
  """
1091
1379
  rows = []
1092
1380
  for task_label, quanta_for_task in self.quanta_by_task.items():
1093
- counts_by_type = Counter(
1094
- exc_info.type_name
1095
- for q in quanta_for_task.values()
1096
- if (exc_info := self._quantum_only_xgraph.nodes[q]["exception"]) is not None
1097
- )
1098
- for type_name, count in counts_by_type.items():
1099
- rows.append({"Task": task_label, "Exception": type_name, "Count": count})
1381
+ success_counts = Counter[str]()
1382
+ failed_counts = Counter[str]()
1383
+ for quantum_id in quanta_for_task.values():
1384
+ quantum_info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1385
+ exc_info = quantum_info["exception"]
1386
+ if exc_info is not None:
1387
+ if quantum_info["status"] is QuantumAttemptStatus.SUCCESSFUL:
1388
+ success_counts[exc_info.type_name] += 1
1389
+ else:
1390
+ failed_counts[exc_info.type_name] += 1
1391
+ for type_name in sorted(success_counts.keys() | failed_counts.keys()):
1392
+ rows.append(
1393
+ {
1394
+ "Task": task_label,
1395
+ "Exception": type_name,
1396
+ "Successes": success_counts.get(type_name, 0),
1397
+ "Failures": failed_counts.get(type_name, 0),
1398
+ }
1399
+ )
1100
1400
  return astropy.table.Table(rows)
1101
1401
 
1102
1402
  def make_task_resource_usage_table(
@@ -1139,6 +1439,171 @@ class ProvenanceQuantumGraph(BaseQuantumGraph):
1139
1439
  array = np.array(rows, dtype=row_dtype)
1140
1440
  return astropy.table.Table(array, units=QuantumResourceUsage.get_units())
1141
1441
 
1442
+ def make_status_report(
1443
+ self,
1444
+ states: Iterable[QuantumAttemptStatus] = (
1445
+ QuantumAttemptStatus.FAILED,
1446
+ QuantumAttemptStatus.ABORTED,
1447
+ QuantumAttemptStatus.ABORTED_SUCCESS,
1448
+ ),
1449
+ *,
1450
+ also: QuantumAttemptStatus | Iterable[QuantumAttemptStatus] = (),
1451
+ with_caveats: QuantumSuccessCaveats | None = QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR,
1452
+ data_id_table_dir: ResourcePathExpression | None = None,
1453
+ ) -> ProvenanceReport:
1454
+ """Make a JSON- or YAML-friendly report of all quanta with the given
1455
+ states.
1456
+
1457
+ Parameters
1458
+ ----------
1459
+ states : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1460
+ `..QuantumAttemptStatus`, optional
1461
+ A quantum is included if it has any of these states. Defaults to
1462
+ states that clearly represent problems.
1463
+ also : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1464
+ `..QuantumAttemptStatus`, optional
1465
+ Additional states to consider; unioned with ``states``. This is
1466
+ provided so users can easily request additional states while also
1467
+ getting the defaults.
1468
+ with_caveats : `..QuantumSuccessCaveats` or `None`, optional
1469
+ If `..QuantumAttemptStatus.SUCCESSFUL` is in ``states``, only
1470
+ include quanta with these caveat flags. May be set to `None`
1471
+ to report on all successful quanta.
1472
+ data_id_table_dir : convertible to `~lsst.resources.ResourcePath`, \
1473
+ optional
1474
+ If provided, a directory to write data ID tables (in ECSV format)
1475
+ with all of the data IDs with the given states, for use with the
1476
+ ``--data-id-tables`` argument to the quantum graph builder.
1477
+ Subdirectories for each task and status will created within this
1478
+ directory, with one file for each exception type (or ``UNKNOWN``
1479
+ when there is no exception).
1480
+
1481
+ Returns
1482
+ -------
1483
+ report : `ProvenanceModel`
1484
+ A Pydantic model that groups quanta by task label and exception
1485
+ type.
1486
+ """
1487
+ states = set(ensure_iterable(states))
1488
+ states.update(ensure_iterable(also))
1489
+ result = ProvenanceReport(root={})
1490
+ if data_id_table_dir is not None:
1491
+ data_id_table_dir = ResourcePath(data_id_table_dir)
1492
+ for task_label, quanta_for_task in self.quanta_by_task.items():
1493
+ reports_for_task: dict[str, dict[str | None, list[ProvenanceQuantumReport]]] = {}
1494
+ table_rows_for_task: dict[str, dict[str | None, list[tuple[int | str, ...]]]] = {}
1495
+ for quantum_id in quanta_for_task.values():
1496
+ quantum_info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1497
+ quantum_status = quantum_info["status"]
1498
+ if quantum_status not in states:
1499
+ continue
1500
+ if (
1501
+ quantum_status is QuantumAttemptStatus.SUCCESSFUL
1502
+ and with_caveats is not None
1503
+ and (quantum_info["caveats"] is None or not (quantum_info["caveats"] & with_caveats))
1504
+ ):
1505
+ continue
1506
+ key1 = quantum_status.name
1507
+ exc_info = quantum_info["exception"]
1508
+ key2 = exc_info.type_name if exc_info is not None else None
1509
+ reports_for_task.setdefault(key1, {}).setdefault(key2, []).append(
1510
+ ProvenanceQuantumReport.from_info(quantum_id, quantum_info)
1511
+ )
1512
+ if data_id_table_dir:
1513
+ table_rows_for_task.setdefault(key1, {}).setdefault(key2, []).append(
1514
+ quantum_info["data_id"].required_values
1515
+ )
1516
+ if reports_for_task:
1517
+ result.root[task_label] = reports_for_task
1518
+ if table_rows_for_task:
1519
+ assert data_id_table_dir is not None, "table_rows_for_task should be empty"
1520
+ for status_name, table_rows_for_status in table_rows_for_task.items():
1521
+ dir_for_task_and_status = data_id_table_dir.join(task_label, forceDirectory=True).join(
1522
+ status_name, forceDirectory=True
1523
+ )
1524
+ if dir_for_task_and_status.isLocal:
1525
+ dir_for_task_and_status.mkdir()
1526
+ for exc_name, data_id_rows in table_rows_for_status.items():
1527
+ table = astropy.table.Table(
1528
+ rows=data_id_rows,
1529
+ names=list(self.pipeline_graph.tasks[task_label].dimensions.required),
1530
+ )
1531
+ filename = f"{exc_name}.ecsv" if exc_name is not None else "UNKNOWN.ecsv"
1532
+ with dir_for_task_and_status.join(filename).open("w") as stream:
1533
+ table.write(stream, format="ecsv")
1534
+ return result
1535
+
1536
+ def make_many_reports(
1537
+ self,
1538
+ states: Iterable[QuantumAttemptStatus] = (
1539
+ QuantumAttemptStatus.FAILED,
1540
+ QuantumAttemptStatus.ABORTED,
1541
+ QuantumAttemptStatus.ABORTED_SUCCESS,
1542
+ ),
1543
+ *,
1544
+ status_report_file: ResourcePathExpression | None = None,
1545
+ print_quantum_table: bool = False,
1546
+ print_exception_table: bool = False,
1547
+ also: QuantumAttemptStatus | Iterable[QuantumAttemptStatus] = (),
1548
+ with_caveats: QuantumSuccessCaveats | None = None,
1549
+ data_id_table_dir: ResourcePathExpression | None = None,
1550
+ ) -> None:
1551
+ """Write multiple reports.
1552
+
1553
+ Parameters
1554
+ ----------
1555
+ states : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1556
+ `..QuantumAttemptStatus`, optional
1557
+ A quantum is included in the status report and data ID tables if it
1558
+ has any of these states. Defaults to states that clearly represent
1559
+ problems.
1560
+ status_report_file : convertible to `~lsst.resources.ResourcePath`,
1561
+ optional
1562
+ Filename for the JSON status report (see `make_status_report`).
1563
+ print_quantum_table : `bool`, optional
1564
+ If `True`, print a quantum summary table (counts only) to STDOUT.
1565
+ print_exception_table : `bool`, optional
1566
+ If `True`, print an exception-type summary table (counts only) to
1567
+ STDOUT.
1568
+ also : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1569
+ `..QuantumAttemptStatus`, optional
1570
+ Additional states to consider in the status report and data ID
1571
+ tables; unioned with ``states``. This is provided so users can
1572
+ easily request additional states while also getting the defaults.
1573
+ with_caveats : `..QuantumSuccessCaveats` or `None`, optional
1574
+ Only include quanta with these caveat flags in the status report
1575
+ and data ID tables. May be set to `None` to report on all
1576
+ successful quanta (an empty sequence reports on only quanta with no
1577
+ caveats). If provided, `QuantumAttemptStatus.SUCCESSFUL` is
1578
+ automatically included in ``states``.
1579
+ data_id_table_dir : convertible to `~lsst.resources.ResourcePath`, \
1580
+ optional
1581
+ If provided, a directory to write data ID tables (in ECSV format)
1582
+ with all of the data IDs with the given states, for use with the
1583
+ ``--data-id-tables`` argument to the quantum graph builder.
1584
+ Subdirectories for each task and status will created within this
1585
+ directory, with one file for each exception type (or ``UNKNOWN``
1586
+ when there is no exception).
1587
+ """
1588
+ if status_report_file is not None or data_id_table_dir is not None:
1589
+ status_report = self.make_status_report(
1590
+ states, also=also, with_caveats=with_caveats, data_id_table_dir=data_id_table_dir
1591
+ )
1592
+ if status_report_file is not None:
1593
+ status_report_file = ResourcePath(status_report_file)
1594
+ if status_report_file.isLocal:
1595
+ status_report_file.dirname().mkdir()
1596
+ with ResourcePath(status_report_file).open("w") as stream:
1597
+ stream.write(status_report.model_dump_json(indent=2))
1598
+ if print_quantum_table:
1599
+ quantum_table = self.make_quantum_table()
1600
+ quantum_table.pprint_all()
1601
+ print("")
1602
+ if print_exception_table:
1603
+ exception_table = self.make_exception_table()
1604
+ exception_table.pprint_all()
1605
+ print("")
1606
+
1142
1607
 
1143
1608
  @dataclasses.dataclass
1144
1609
  class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
@@ -1269,19 +1734,19 @@ class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1269
1734
  # also have other outstanding reference holders).
1270
1735
  continue
1271
1736
  node._add_to_graph(self.graph)
1272
- return
1273
- with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1274
- for node_id_or_index in nodes:
1275
- address_row = self.address_reader.find(node_id_or_index)
1276
- if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1277
- # Use the old node to reduce memory usage (since it might
1278
- # also have other outstanding reference holders).
1279
- continue
1280
- node = mb_reader.read_model(
1281
- address_row.addresses[address_index], model_type, self.decompressor
1282
- )
1283
- if node is not None:
1284
- node._add_to_graph(self.graph)
1737
+ else:
1738
+ with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1739
+ for node_id_or_index in nodes:
1740
+ address_row = self.address_reader.find(node_id_or_index)
1741
+ if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1742
+ # Use the old node to reduce memory usage (since it
1743
+ # might also have other outstanding reference holders).
1744
+ continue
1745
+ node = mb_reader.read_model(
1746
+ address_row.addresses[address_index], model_type, self.decompressor
1747
+ )
1748
+ if node is not None:
1749
+ node._add_to_graph(self.graph)
1285
1750
 
1286
1751
  def fetch_logs(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[ButlerLogRecords | None]]:
1287
1752
  """Fetch log datasets.
@@ -1352,3 +1817,629 @@ class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1352
1817
  """Fetch package version information."""
1353
1818
  data = self._read_single_block_raw("packages")
1354
1819
  return Packages.fromBytes(data, format="json")
1820
+
1821
+
1822
+ class ProvenanceQuantumGraphWriter:
1823
+ """A struct of low-level writer objects for the main components of a
1824
+ provenance quantum graph.
1825
+
1826
+ Parameters
1827
+ ----------
1828
+ output_path : `str`
1829
+ Path to write the graph to.
1830
+ exit_stack : `contextlib.ExitStack`
1831
+ Object that can be used to manage multiple context managers.
1832
+ log_on_close : `LogOnClose`
1833
+ Factory for context managers that log when closed.
1834
+ predicted : `.PredictedQuantumGraphComponents`
1835
+ Components of the predicted graph.
1836
+ zstd_level : `int`, optional
1837
+ Compression level.
1838
+ cdict_data : `bytes` or `None`, optional
1839
+ Bytes representation of the compression dictionary used by the
1840
+ compressor.
1841
+ loop_wrapper : `~collections.abc.Callable`, optional
1842
+ A callable that takes an iterable and returns an equivalent one, to be
1843
+ used in all potentially-large loops. This can be used to add progress
1844
+ reporting or check for cancelation signals.
1845
+ log : `LsstLogAdapter`, optional
1846
+ Logger to use for debug messages.
1847
+ """
1848
+
1849
+ def __init__(
1850
+ self,
1851
+ output_path: str,
1852
+ *,
1853
+ exit_stack: ExitStack,
1854
+ log_on_close: LogOnClose,
1855
+ predicted: PredictedQuantumGraphComponents | PredictedQuantumGraph,
1856
+ zstd_level: int = 10,
1857
+ cdict_data: bytes | None = None,
1858
+ loop_wrapper: LoopWrapper = pass_through,
1859
+ log: LsstLogAdapter | None = None,
1860
+ ) -> None:
1861
+ header = predicted.header.model_copy()
1862
+ header.graph_type = "provenance"
1863
+ if log is None:
1864
+ log = _LOG
1865
+ self.log = log
1866
+ self._base_writer = exit_stack.enter_context(
1867
+ log_on_close.wrap(
1868
+ BaseQuantumGraphWriter.open(
1869
+ output_path,
1870
+ header,
1871
+ predicted.pipeline_graph,
1872
+ address_filename="nodes",
1873
+ zstd_level=zstd_level,
1874
+ cdict_data=cdict_data,
1875
+ ),
1876
+ "Finishing writing provenance quantum graph.",
1877
+ )
1878
+ )
1879
+ self._base_writer.address_writer.addresses = [{}, {}, {}, {}]
1880
+ self._log_writer = exit_stack.enter_context(
1881
+ log_on_close.wrap(
1882
+ MultiblockWriter.open_in_zip(
1883
+ self._base_writer.zf, LOG_MB_NAME, header.int_size, use_tempfile=True
1884
+ ),
1885
+ "Copying logs into zip archive.",
1886
+ ),
1887
+ )
1888
+ self._base_writer.address_writer.addresses[LOG_ADDRESS_INDEX] = self._log_writer.addresses
1889
+ self._metadata_writer = exit_stack.enter_context(
1890
+ log_on_close.wrap(
1891
+ MultiblockWriter.open_in_zip(
1892
+ self._base_writer.zf, METADATA_MB_NAME, header.int_size, use_tempfile=True
1893
+ ),
1894
+ "Copying metadata into zip archive.",
1895
+ )
1896
+ )
1897
+ self._base_writer.address_writer.addresses[METADATA_ADDRESS_INDEX] = self._metadata_writer.addresses
1898
+ self._dataset_writer = exit_stack.enter_context(
1899
+ log_on_close.wrap(
1900
+ MultiblockWriter.open_in_zip(
1901
+ self._base_writer.zf, DATASET_MB_NAME, header.int_size, use_tempfile=True
1902
+ ),
1903
+ "Copying dataset provenance into zip archive.",
1904
+ )
1905
+ )
1906
+ self._base_writer.address_writer.addresses[DATASET_ADDRESS_INDEX] = self._dataset_writer.addresses
1907
+ self._quantum_writer = exit_stack.enter_context(
1908
+ log_on_close.wrap(
1909
+ MultiblockWriter.open_in_zip(
1910
+ self._base_writer.zf, QUANTUM_MB_NAME, header.int_size, use_tempfile=True
1911
+ ),
1912
+ "Copying quantum provenance into zip archive.",
1913
+ )
1914
+ )
1915
+ self._base_writer.address_writer.addresses[QUANTUM_ADDRESS_INDEX] = self._quantum_writer.addresses
1916
+ self._init_predicted_quanta(predicted)
1917
+ self._populate_xgraph_and_inputs(loop_wrapper)
1918
+ self._existing_init_outputs: set[uuid.UUID] = set()
1919
+
1920
+ def _init_predicted_quanta(
1921
+ self, predicted: PredictedQuantumGraph | PredictedQuantumGraphComponents
1922
+ ) -> None:
1923
+ self._predicted_init_quanta: list[PredictedQuantumDatasetsModel] = []
1924
+ self._predicted_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = {}
1925
+ if isinstance(predicted, PredictedQuantumGraph):
1926
+ self._predicted_init_quanta.extend(predicted._init_quanta.values())
1927
+ self._predicted_quanta.update(predicted._quantum_datasets)
1928
+ else:
1929
+ self._predicted_init_quanta.extend(predicted.init_quanta.root)
1930
+ self._predicted_quanta.update(predicted.quantum_datasets)
1931
+ self._predicted_quanta.update({q.quantum_id: q for q in self._predicted_init_quanta})
1932
+
1933
+ def _populate_xgraph_and_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1934
+ self._xgraph = networkx.DiGraph()
1935
+ self._overall_inputs: dict[uuid.UUID, PredictedDatasetModel] = {}
1936
+ output_dataset_ids: set[uuid.UUID] = set()
1937
+ for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1938
+ if not predicted_quantum.task_label:
1939
+ # Skip the 'packages' producer quantum.
1940
+ continue
1941
+ output_dataset_ids.update(predicted_quantum.iter_output_dataset_ids())
1942
+ for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1943
+ if not predicted_quantum.task_label:
1944
+ # Skip the 'packages' producer quantum.
1945
+ continue
1946
+ for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
1947
+ self._xgraph.add_edge(predicted_input.dataset_id, predicted_quantum.quantum_id)
1948
+ if predicted_input.dataset_id not in output_dataset_ids:
1949
+ self._overall_inputs.setdefault(predicted_input.dataset_id, predicted_input)
1950
+ for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1951
+ self._xgraph.add_edge(predicted_quantum.quantum_id, predicted_output.dataset_id)
1952
+
1953
+ @property
1954
+ def compressor(self) -> Compressor:
1955
+ """Object that should be used to compress all JSON blocks."""
1956
+ return self._base_writer.compressor
1957
+
1958
+ def write_packages(self) -> None:
1959
+ """Write package version information to the provenance graph."""
1960
+ packages = Packages.fromSystem(include_all=True)
1961
+ data = packages.toBytes("json")
1962
+ self._base_writer.write_single_block("packages", data)
1963
+
1964
+ def write_overall_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1965
+ """Write provenance for overall-input datasets.
1966
+
1967
+ Parameters
1968
+ ----------
1969
+ loop_wrapper : `~collections.abc.Callable`, optional
1970
+ A callable that takes an iterable and returns an equivalent one, to
1971
+ be used in all potentially-large loops. This can be used to add
1972
+ progress reporting or check for cancelation signals.
1973
+ """
1974
+ for predicted_input in loop_wrapper(self._overall_inputs.values()):
1975
+ if predicted_input.dataset_id not in self._dataset_writer.addresses:
1976
+ self._dataset_writer.write_model(
1977
+ predicted_input.dataset_id,
1978
+ ProvenanceDatasetModel.from_predicted(
1979
+ predicted_input,
1980
+ producer=None,
1981
+ consumers=self._xgraph.successors(predicted_input.dataset_id),
1982
+ ),
1983
+ self.compressor,
1984
+ )
1985
+ del self._overall_inputs
1986
+
1987
+ def write_init_outputs(self, assume_existence: bool = True) -> None:
1988
+ """Write provenance for init-output datasets and init-quanta.
1989
+
1990
+ Parameters
1991
+ ----------
1992
+ assume_existence : `bool`, optional
1993
+ If `True`, just assume all init-outputs exist.
1994
+ """
1995
+ init_quanta = ProvenanceInitQuantaModel()
1996
+ for predicted_init_quantum in self._predicted_init_quanta:
1997
+ if not predicted_init_quantum.task_label:
1998
+ # Skip the 'packages' producer quantum.
1999
+ continue
2000
+ for predicted_output in itertools.chain.from_iterable(predicted_init_quantum.outputs.values()):
2001
+ provenance_output = ProvenanceDatasetModel.from_predicted(
2002
+ predicted_output,
2003
+ producer=predicted_init_quantum.quantum_id,
2004
+ consumers=self._xgraph.successors(predicted_output.dataset_id),
2005
+ )
2006
+ provenance_output.produced = assume_existence or (
2007
+ provenance_output.dataset_id in self._existing_init_outputs
2008
+ )
2009
+ self._dataset_writer.write_model(
2010
+ provenance_output.dataset_id, provenance_output, self.compressor
2011
+ )
2012
+ init_quanta.root.append(ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum))
2013
+ self._base_writer.write_single_model("init_quanta", init_quanta)
2014
+
2015
+ def write_quantum_provenance(
2016
+ self, quantum_id: uuid.UUID, metadata: TaskMetadata | None, logs: ButlerLogRecords | None
2017
+ ) -> None:
2018
+ """Gather and write provenance for a quantum.
2019
+
2020
+ Parameters
2021
+ ----------
2022
+ quantum_id : `uuid.UUID`
2023
+ Unique ID for the quantum.
2024
+ metadata : `..TaskMetadata` or `None`
2025
+ Task metadata.
2026
+ logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
2027
+ Task logs.
2028
+ """
2029
+ predicted_quantum = self._predicted_quanta[quantum_id]
2030
+ provenance_models = ProvenanceQuantumScanModels.from_metadata_and_logs(
2031
+ predicted_quantum, metadata, logs, incomplete=False
2032
+ )
2033
+ scan_data = provenance_models.to_scan_data(predicted_quantum, compressor=self.compressor)
2034
+ self.write_scan_data(scan_data)
2035
+
2036
+ def write_scan_data(self, scan_data: ProvenanceQuantumScanData) -> None:
2037
+ """Write the output of a quantum provenance scan to disk.
2038
+
2039
+ Parameters
2040
+ ----------
2041
+ scan_data : `ProvenanceQuantumScanData`
2042
+ Result of a quantum provenance scan.
2043
+ """
2044
+ if scan_data.status is ProvenanceQuantumScanStatus.INIT:
2045
+ self.log.debug("Handling init-output scan for %s.", scan_data.quantum_id)
2046
+ self._existing_init_outputs.update(scan_data.existing_outputs)
2047
+ return
2048
+ self.log.debug("Handling quantum scan for %s.", scan_data.quantum_id)
2049
+ # We shouldn't need this predicted quantum after this method runs; pop
2050
+ # from the dict it in the hopes that'll free up some memory when we're
2051
+ # done.
2052
+ predicted_quantum = self._predicted_quanta.pop(scan_data.quantum_id)
2053
+ outputs: dict[uuid.UUID, bytes] = {}
2054
+ for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
2055
+ provenance_output = ProvenanceDatasetModel.from_predicted(
2056
+ predicted_output,
2057
+ producer=predicted_quantum.quantum_id,
2058
+ consumers=self._xgraph.successors(predicted_output.dataset_id),
2059
+ )
2060
+ provenance_output.produced = provenance_output.dataset_id in scan_data.existing_outputs
2061
+ outputs[provenance_output.dataset_id] = self.compressor.compress(
2062
+ provenance_output.model_dump_json().encode()
2063
+ )
2064
+ if not scan_data.quantum:
2065
+ scan_data.quantum = (
2066
+ ProvenanceQuantumModel.from_predicted(predicted_quantum).model_dump_json().encode()
2067
+ )
2068
+ if scan_data.is_compressed:
2069
+ scan_data.quantum = self.compressor.compress(scan_data.quantum)
2070
+ if not scan_data.is_compressed:
2071
+ scan_data.quantum = self.compressor.compress(scan_data.quantum)
2072
+ if scan_data.metadata:
2073
+ scan_data.metadata = self.compressor.compress(scan_data.metadata)
2074
+ if scan_data.logs:
2075
+ scan_data.logs = self.compressor.compress(scan_data.logs)
2076
+ self.log.debug("Writing quantum %s.", scan_data.quantum_id)
2077
+ self._quantum_writer.write_bytes(scan_data.quantum_id, scan_data.quantum)
2078
+ for dataset_id, dataset_data in outputs.items():
2079
+ self._dataset_writer.write_bytes(dataset_id, dataset_data)
2080
+ if scan_data.metadata:
2081
+ (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
2082
+ address = self._metadata_writer.write_bytes(scan_data.quantum_id, scan_data.metadata)
2083
+ self._metadata_writer.addresses[metadata_output.dataset_id] = address
2084
+ if scan_data.logs:
2085
+ (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
2086
+ address = self._log_writer.write_bytes(scan_data.quantum_id, scan_data.logs)
2087
+ self._log_writer.addresses[log_output.dataset_id] = address
2088
+
2089
+
2090
+ class ProvenanceQuantumScanStatus(enum.Enum):
2091
+ """Status enum for quantum scanning.
2092
+
2093
+ Note that this records the status for the *scanning* which is distinct
2094
+ from the status of the quantum's execution.
2095
+ """
2096
+
2097
+ INCOMPLETE = enum.auto()
2098
+ """The quantum is not necessarily done running, and cannot be scanned
2099
+ conclusively yet.
2100
+ """
2101
+
2102
+ ABANDONED = enum.auto()
2103
+ """The quantum's execution appears to have failed but we cannot rule out
2104
+ the possibility that it could be recovered, but we've also waited long
2105
+ enough (according to `ScannerTimeConfigDict.retry_timeout`) that it's time
2106
+ to stop trying for now.
2107
+
2108
+ This state means `ProvenanceQuantumScanModels.from_metadata_and_logs` must
2109
+ be run again with ``incomplete=False``.
2110
+ """
2111
+
2112
+ SUCCESSFUL = enum.auto()
2113
+ """The quantum was conclusively scanned and was executed successfully,
2114
+ unblocking scans for downstream quanta.
2115
+ """
2116
+
2117
+ FAILED = enum.auto()
2118
+ """The quantum was conclusively scanned and failed execution, blocking
2119
+ scans for downstream quanta.
2120
+ """
2121
+
2122
+ BLOCKED = enum.auto()
2123
+ """A quantum upstream of this one failed."""
2124
+
2125
+ INIT = enum.auto()
2126
+ """Init quanta need special handling, because they don't have logs and
2127
+ metadata.
2128
+ """
2129
+
2130
+
2131
+ @dataclasses.dataclass
2132
+ class ProvenanceQuantumScanModels:
2133
+ """A struct that represents provenance information for a single quantum."""
2134
+
2135
+ quantum_id: uuid.UUID
2136
+ """Unique ID for the quantum."""
2137
+
2138
+ status: ProvenanceQuantumScanStatus = ProvenanceQuantumScanStatus.INCOMPLETE
2139
+ """Combined status for the scan and the execution of the quantum."""
2140
+
2141
+ attempts: list[ProvenanceQuantumAttemptModel] = dataclasses.field(default_factory=list)
2142
+ """Provenance information about each attempt to run the quantum."""
2143
+
2144
+ output_existence: dict[uuid.UUID, bool] = dataclasses.field(default_factory=dict)
2145
+ """Unique IDs of the output datasets mapped to whether they were actually
2146
+ produced.
2147
+ """
2148
+
2149
+ metadata: ProvenanceTaskMetadataModel = dataclasses.field(default_factory=ProvenanceTaskMetadataModel)
2150
+ """Task metadata information for each attempt.
2151
+ """
2152
+
2153
+ logs: ProvenanceLogRecordsModel = dataclasses.field(default_factory=ProvenanceLogRecordsModel)
2154
+ """Log records for each attempt.
2155
+ """
2156
+
2157
+ @classmethod
2158
+ def from_metadata_and_logs(
2159
+ cls,
2160
+ predicted: PredictedQuantumDatasetsModel,
2161
+ metadata: TaskMetadata | None,
2162
+ logs: ButlerLogRecords | None,
2163
+ *,
2164
+ incomplete: bool = False,
2165
+ ) -> ProvenanceQuantumScanModels:
2166
+ """Construct provenance information from task metadata and logs.
2167
+
2168
+ Parameters
2169
+ ----------
2170
+ predicted : `PredictedQuantumDatasetsModel`
2171
+ Information about the predicted quantum.
2172
+ metadata : `..TaskMetadata` or `None`
2173
+ Task metadata.
2174
+ logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
2175
+ Task logs.
2176
+ incomplete : `bool`, optional
2177
+ If `True`, treat execution failures as possibly-incomplete quanta
2178
+ and do not fully process them; instead just set the status to
2179
+ `ProvenanceQuantumScanStatus.ABANDONED` and return.
2180
+
2181
+ Returns
2182
+ -------
2183
+ scan_models : `ProvenanceQuantumScanModels`
2184
+ Struct of models that describe quantum provenance.
2185
+
2186
+ Notes
2187
+ -----
2188
+ This method does not necessarily fully populate the `output_existence`
2189
+ field; it does what it can given the information in the metadata and
2190
+ logs, but the caller is responsible for filling in the existence status
2191
+ for any predicted outputs that are not present at all in that `dict`.
2192
+ """
2193
+ self = ProvenanceQuantumScanModels(predicted.quantum_id)
2194
+ last_attempt = ProvenanceQuantumAttemptModel()
2195
+ self._process_logs(predicted, logs, last_attempt, incomplete=incomplete)
2196
+ self._process_metadata(predicted, metadata, last_attempt, incomplete=incomplete)
2197
+ if self.status is ProvenanceQuantumScanStatus.ABANDONED:
2198
+ return self
2199
+ self._reconcile_attempts(last_attempt)
2200
+ self._extract_output_existence(predicted)
2201
+ return self
2202
+
2203
+ def _process_logs(
2204
+ self,
2205
+ predicted: PredictedQuantumDatasetsModel,
2206
+ logs: ButlerLogRecords | None,
2207
+ last_attempt: ProvenanceQuantumAttemptModel,
2208
+ *,
2209
+ incomplete: bool,
2210
+ ) -> None:
2211
+ (predicted_log_dataset,) = predicted.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
2212
+ if logs is None:
2213
+ self.output_existence[predicted_log_dataset.dataset_id] = False
2214
+ if incomplete:
2215
+ self.status = ProvenanceQuantumScanStatus.ABANDONED
2216
+ else:
2217
+ self.status = ProvenanceQuantumScanStatus.FAILED
2218
+ else:
2219
+ # Set the attempt's run status to FAILED, since the default is
2220
+ # UNKNOWN (i.e. logs *and* metadata are missing) and we now know
2221
+ # the logs exist. This will usually get replaced by SUCCESSFUL
2222
+ # when we look for metadata next.
2223
+ last_attempt.status = QuantumAttemptStatus.FAILED
2224
+ self.output_existence[predicted_log_dataset.dataset_id] = True
2225
+ if logs.extra:
2226
+ log_extra = _ExecutionLogRecordsExtra.model_validate(logs.extra)
2227
+ self._extract_from_log_extra(log_extra, last_attempt=last_attempt)
2228
+ self.logs.attempts.append(list(logs))
2229
+
2230
+ def _extract_from_log_extra(
2231
+ self,
2232
+ log_extra: _ExecutionLogRecordsExtra,
2233
+ last_attempt: ProvenanceQuantumAttemptModel | None,
2234
+ ) -> None:
2235
+ for previous_attempt_log_extra in log_extra.previous_attempts:
2236
+ self._extract_from_log_extra(
2237
+ previous_attempt_log_extra,
2238
+ last_attempt=None,
2239
+ )
2240
+ quantum_attempt: ProvenanceQuantumAttemptModel
2241
+ if last_attempt is None:
2242
+ # This is not the last attempt, so it must be a failure.
2243
+ quantum_attempt = ProvenanceQuantumAttemptModel(
2244
+ attempt=len(self.attempts), status=QuantumAttemptStatus.FAILED
2245
+ )
2246
+ # We also need to get the logs from this extra provenance, since
2247
+ # they won't be the main section of the log records.
2248
+ self.logs.attempts.append(log_extra.logs)
2249
+ # The special last attempt is only appended after we attempt to
2250
+ # read metadata later, but we have to append this one now.
2251
+ self.attempts.append(quantum_attempt)
2252
+ else:
2253
+ assert not log_extra.logs, "Logs for the last attempt should not be stored in the extra JSON."
2254
+ quantum_attempt = last_attempt
2255
+ if log_extra.exception is not None or log_extra.metadata is not None or last_attempt is None:
2256
+ # We won't be getting a separate metadata dataset, so anything we
2257
+ # might get from the metadata has to come from this extra
2258
+ # provenance in the logs.
2259
+ quantum_attempt.exception = log_extra.exception
2260
+ if log_extra.metadata is not None:
2261
+ quantum_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(log_extra.metadata)
2262
+ self.metadata.attempts.append(log_extra.metadata)
2263
+ else:
2264
+ self.metadata.attempts.append(None)
2265
+ # Regardless of whether this is the last attempt or not, we can only
2266
+ # get the previous_process_quanta from the log extra.
2267
+ quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)
2268
+
2269
+ def _process_metadata(
2270
+ self,
2271
+ predicted: PredictedQuantumDatasetsModel,
2272
+ metadata: TaskMetadata | None,
2273
+ last_attempt: ProvenanceQuantumAttemptModel,
2274
+ *,
2275
+ incomplete: bool,
2276
+ ) -> None:
2277
+ (predicted_metadata_dataset,) = predicted.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
2278
+ if metadata is None:
2279
+ self.output_existence[predicted_metadata_dataset.dataset_id] = False
2280
+ if incomplete:
2281
+ self.status = ProvenanceQuantumScanStatus.ABANDONED
2282
+ else:
2283
+ self.status = ProvenanceQuantumScanStatus.FAILED
2284
+ else:
2285
+ self.status = ProvenanceQuantumScanStatus.SUCCESSFUL
2286
+ self.output_existence[predicted_metadata_dataset.dataset_id] = True
2287
+ last_attempt.status = QuantumAttemptStatus.SUCCESSFUL
2288
+ try:
2289
+ # Int conversion guards against spurious conversion to
2290
+ # float that can apparently sometimes happen in
2291
+ # TaskMetadata.
2292
+ last_attempt.caveats = QuantumSuccessCaveats(int(metadata["quantum"]["caveats"]))
2293
+ except LookupError:
2294
+ pass
2295
+ try:
2296
+ last_attempt.exception = ExceptionInfo._from_metadata(
2297
+ metadata[predicted.task_label]["failure"]
2298
+ )
2299
+ except LookupError:
2300
+ pass
2301
+ last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)
2302
+ self.metadata.attempts.append(metadata)
2303
+
2304
+ def _reconcile_attempts(self, last_attempt: ProvenanceQuantumAttemptModel) -> None:
2305
+ last_attempt.attempt = len(self.attempts)
2306
+ self.attempts.append(last_attempt)
2307
+ assert self.status is not ProvenanceQuantumScanStatus.INCOMPLETE
2308
+ assert self.status is not ProvenanceQuantumScanStatus.ABANDONED
2309
+ if len(self.logs.attempts) < len(self.attempts):
2310
+ # Logs were not found for this attempt; must have been a hard error
2311
+ # that kept the `finally` block from running or otherwise
2312
+ # interrupted the writing of the logs.
2313
+ self.logs.attempts.append(None)
2314
+ if self.status is ProvenanceQuantumScanStatus.SUCCESSFUL:
2315
+ # But we found the metadata! Either that hard error happened
2316
+ # at a very unlucky time (in between those two writes), or
2317
+ # something even weirder happened.
2318
+ self.attempts[-1].status = QuantumAttemptStatus.ABORTED_SUCCESS
2319
+ else:
2320
+ self.attempts[-1].status = QuantumAttemptStatus.FAILED
2321
+ if len(self.metadata.attempts) < len(self.attempts):
2322
+ # Metadata missing usually just means a failure. In any case, the
2323
+ # status will already be correct, either because it was set to a
2324
+ # failure when we read the logs, or left at UNKNOWN if there were
2325
+ # no logs. Note that scanners never process BLOCKED quanta at all.
2326
+ self.metadata.attempts.append(None)
2327
+ assert len(self.logs.attempts) == len(self.attempts) or len(self.metadata.attempts) == len(
2328
+ self.attempts
2329
+ ), (
2330
+ "The only way we can add more than one quantum attempt is by "
2331
+ "extracting info stored with the logs, and that always appends "
2332
+ "a log attempt and a metadata attempt, so this must be a bug in "
2333
+ "this class."
2334
+ )
2335
+
2336
+ def _extract_output_existence(self, predicted: PredictedQuantumDatasetsModel) -> None:
2337
+ try:
2338
+ outputs_put = self.metadata.attempts[-1]["quantum"].getArray("outputs") # type: ignore[index]
2339
+ except (
2340
+ IndexError, # metadata.attempts is empty
2341
+ TypeError, # metadata.attempts[-1] is None
2342
+ LookupError, # no 'quantum' entry in metadata or 'outputs' in that
2343
+ ):
2344
+ pass
2345
+ else:
2346
+ for id_str in ensure_iterable(outputs_put):
2347
+ self.output_existence[uuid.UUID(id_str)] = True
2348
+ # If the metadata told us what it wrote, anything not in that
2349
+ # list was not written.
2350
+ for predicted_output in itertools.chain.from_iterable(predicted.outputs.values()):
2351
+ self.output_existence.setdefault(predicted_output.dataset_id, False)
2352
+
2353
+ def to_scan_data(
2354
+ self: ProvenanceQuantumScanModels,
2355
+ predicted_quantum: PredictedQuantumDatasetsModel,
2356
+ compressor: Compressor | None = None,
2357
+ ) -> ProvenanceQuantumScanData:
2358
+ """Convert these models to JSON data.
2359
+
2360
+ Parameters
2361
+ ----------
2362
+ predicted_quantum : `PredictedQuantumDatasetsModel`
2363
+ Information about the predicted quantum.
2364
+ compressor : `Compressor`
2365
+ Object that can compress bytes.
2366
+
2367
+ Returns
2368
+ -------
2369
+ scan_data : `ProvenanceQuantumScanData`
2370
+ Scan information ready for serialization.
2371
+ """
2372
+ quantum: ProvenanceInitQuantumModel | ProvenanceQuantumModel
2373
+ if self.status is ProvenanceQuantumScanStatus.INIT:
2374
+ quantum = ProvenanceInitQuantumModel.from_predicted(predicted_quantum)
2375
+ else:
2376
+ quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum)
2377
+ quantum.attempts = self.attempts
2378
+ for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
2379
+ if predicted_output.dataset_id not in self.output_existence:
2380
+ raise RuntimeError(
2381
+ "Logic bug in provenance gathering or execution invariants: "
2382
+ f"no existence information for output {predicted_output.dataset_id} "
2383
+ f"({predicted_output.dataset_type_name}@{predicted_output.data_coordinate})."
2384
+ )
2385
+ data = ProvenanceQuantumScanData(
2386
+ self.quantum_id,
2387
+ self.status,
2388
+ existing_outputs={
2389
+ dataset_id for dataset_id, was_produced in self.output_existence.items() if was_produced
2390
+ },
2391
+ quantum=quantum.model_dump_json().encode(),
2392
+ logs=self.logs.model_dump_json().encode() if self.logs.attempts else b"",
2393
+ metadata=self.metadata.model_dump_json().encode() if self.metadata.attempts else b"",
2394
+ )
2395
+ if compressor is not None:
2396
+ data.compress(compressor)
2397
+ return data
2398
+
2399
+
2400
+ @dataclasses.dataclass
2401
+ class ProvenanceQuantumScanData:
2402
+ """A struct that represents ready-for-serialization provenance information
2403
+ for a single quantum.
2404
+ """
2405
+
2406
+ quantum_id: uuid.UUID
2407
+ """Unique ID for the quantum."""
2408
+
2409
+ status: ProvenanceQuantumScanStatus
2410
+ """Combined status for the scan and the execution of the quantum."""
2411
+
2412
+ existing_outputs: set[uuid.UUID] = dataclasses.field(default_factory=set)
2413
+ """Unique IDs of the output datasets that were actually written."""
2414
+
2415
+ quantum: bytes = b""
2416
+ """Serialized quantum provenance model.
2417
+
2418
+ This may be empty for quanta that had no attempts.
2419
+ """
2420
+
2421
+ metadata: bytes = b""
2422
+ """Serialized task metadata."""
2423
+
2424
+ logs: bytes = b""
2425
+ """Serialized logs."""
2426
+
2427
+ is_compressed: bool = False
2428
+ """Whether the ``quantum``, ``metadata``, and ``log`` attributes are
2429
+ compressed.
2430
+ """
2431
+
2432
+ def compress(self, compressor: Compressor) -> None:
2433
+ """Compress the data in this struct if it has not been compressed
2434
+ already.
2435
+
2436
+ Parameters
2437
+ ----------
2438
+ compressor : `Compressor`
2439
+ Object with a ``compress`` method that takes and returns `bytes`.
2440
+ """
2441
+ if not self.is_compressed:
2442
+ self.quantum = compressor.compress(self.quantum)
2443
+ self.logs = compressor.compress(self.logs) if self.logs else b""
2444
+ self.metadata = compressor.compress(self.metadata) if self.metadata else b""
2445
+ self.is_compressed = True