Coverage for python / lsst / pipe / base / quantum_graph / _provenance.py: 31%
586 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-22 02:02 -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# (http://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 <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = (
31 "ProvenanceDatasetInfo",
32 "ProvenanceDatasetModel",
33 "ProvenanceInitQuantumInfo",
34 "ProvenanceInitQuantumModel",
35 "ProvenanceLogRecordsModel",
36 "ProvenanceQuantumGraph",
37 "ProvenanceQuantumGraphReader",
38 "ProvenanceQuantumGraphWriter",
39 "ProvenanceQuantumInfo",
40 "ProvenanceQuantumModel",
41 "ProvenanceQuantumScanData",
42 "ProvenanceQuantumScanModels",
43 "ProvenanceQuantumScanStatus",
44 "ProvenanceTaskMetadataModel",
45)
48import dataclasses
49import enum
50import itertools
51import sys
52import uuid
53from collections import Counter
54from collections.abc import Callable, Iterable, Iterator, Mapping
55from contextlib import ExitStack, contextmanager
56from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict, TypeVar
58import astropy.table
59import networkx
60import numpy as np
61import pydantic
63from lsst.daf.butler import DataCoordinate
64from lsst.daf.butler.logging import ButlerLogRecord, ButlerLogRecords
65from lsst.resources import ResourcePathExpression
66from lsst.utils.iteration import ensure_iterable
67from lsst.utils.logging import LsstLogAdapter, getLogger
68from lsst.utils.packages import Packages
70from .. import automatic_connection_constants as acc
71from .._status import ExceptionInfo, QuantumAttemptStatus, QuantumSuccessCaveats
72from .._task_metadata import TaskMetadata
73from ..log_capture import _ExecutionLogRecordsExtra
74from ..log_on_close import LogOnClose
75from ..pipeline_graph import PipelineGraph, TaskImportMode, TaskInitNode
76from ..resource_usage import QuantumResourceUsage
77from ._common import (
78 BaseQuantumGraph,
79 BaseQuantumGraphReader,
80 BaseQuantumGraphWriter,
81 ConnectionName,
82 DataCoordinateValues,
83 DatasetInfo,
84 DatasetTypeName,
85 HeaderModel,
86 QuantumInfo,
87 TaskLabel,
88)
89from ._multiblock import Compressor, MultiblockReader, MultiblockWriter
90from ._predicted import (
91 PredictedDatasetModel,
92 PredictedQuantumDatasetsModel,
93 PredictedQuantumGraph,
94 PredictedQuantumGraphComponents,
95)
97# Sphinx needs imports for type annotations of base class members.
98if "sphinx" in sys.modules:
99 import zipfile # noqa: F401
101 from ._multiblock import AddressReader, Decompressor # noqa: F401
104_T = TypeVar("_T")
106LoopWrapper: TypeAlias = Callable[[Iterable[_T]], Iterable[_T]]
108_LOG = getLogger(__file__)
110DATASET_ADDRESS_INDEX = 0
111QUANTUM_ADDRESS_INDEX = 1
112LOG_ADDRESS_INDEX = 2
113METADATA_ADDRESS_INDEX = 3
115DATASET_MB_NAME = "datasets"
116QUANTUM_MB_NAME = "quanta"
117LOG_MB_NAME = "logs"
118METADATA_MB_NAME = "metadata"
121def pass_through(arg: _T) -> _T:
122 return arg
125class ProvenanceDatasetInfo(DatasetInfo):
126 """A typed dictionary that annotates the attributes of the NetworkX graph
127 node data for a provenance dataset.
129 Since NetworkX types are not generic over their node mapping type, this has
130 to be used explicitly, e.g.::
132 node_data: ProvenanceDatasetInfo = xgraph.nodes[dataset_id]
134 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
135 """
137 dataset_id: uuid.UUID
138 """Unique identifier for the dataset."""
140 produced: bool
141 """Whether this dataset was produced (vs. only predicted).
143 This is always `True` for overall input datasets. It is also `True` for
144 datasets that were produced and then removed before/during transfer back to
145 the central butler repository, so it may not reflect the continued
146 existence of the dataset.
147 """
150class ProvenanceQuantumInfo(QuantumInfo):
151 """A typed dictionary that annotates the attributes of the NetworkX graph
152 node data for a provenance quantum.
154 Since NetworkX types are not generic over their node mapping type, this has
155 to be used explicitly, e.g.::
157 node_data: ProvenanceQuantumInfo = xgraph.nodes[quantum_id]
159 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph` or
160 `ProvenanceQuantumGraph.quantum_only_xgraph`
161 """
163 status: QuantumAttemptStatus
164 """Enumerated status for the quantum.
166 This corresponds to the last attempt to run this quantum, or
167 `QuantumAttemptStatus.BLOCKED` if there were no attempts.
168 """
170 caveats: QuantumSuccessCaveats | None
171 """Flags indicating caveats on successful quanta.
173 This corresponds to the last attempt to run this quantum.
174 """
176 exception: ExceptionInfo | None
177 """Information about an exception raised when the quantum was executing.
179 This corresponds to the last attempt to run this quantum.
180 """
182 resource_usage: QuantumResourceUsage | None
183 """Resource usage information (timing, memory use) for this quantum.
185 This corresponds to the last attempt to run this quantum.
186 """
188 attempts: list[ProvenanceQuantumAttemptModel]
189 """Information about each attempt to run this quantum.
191 An entry is added merely if the quantum *should* have been attempted; an
192 empty `list` is used only for quanta that were blocked by an upstream
193 failure.
194 """
196 metadata_id: uuid.UUID
197 """ID of this quantum's metadata dataset."""
199 log_id: uuid.UUID
200 """ID of this quantum's log dataset."""
203class ProvenanceInitQuantumInfo(TypedDict):
204 """A typed dictionary that annotates the attributes of the NetworkX graph
205 node data for a provenance init quantum.
207 Since NetworkX types are not generic over their node mapping type, this has
208 to be used explicitly, e.g.::
210 node_data: ProvenanceInitQuantumInfo = xgraph.nodes[quantum_id]
212 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
213 """
215 data_id: DataCoordinate
216 """Data ID of the quantum.
218 This is always an empty ID; this key exists to allow init-quanta and
219 regular quanta to be treated more similarly.
220 """
222 task_label: str
223 """Label of the task for this quantum."""
225 pipeline_node: TaskInitNode
226 """Node in the pipeline graph for this task's init-only step."""
228 config_id: uuid.UUID
229 """ID of this task's config dataset."""
232class ProvenanceDatasetModel(PredictedDatasetModel):
233 """Data model for the datasets in a provenance quantum graph file."""
235 produced: bool
236 """Whether this dataset was produced (vs. only predicted).
238 This is always `True` for overall input datasets. It is also `True` for
239 datasets that were produced and then removed before/during transfer back to
240 the central butler repository, so it may not reflect the continued
241 existence of the dataset.
242 """
244 producer: uuid.UUID | None = None
245 """ID of the quantum that produced this dataset.
247 This is `None` for overall inputs to the graph.
248 """
250 consumers: list[uuid.UUID] = pydantic.Field(default_factory=list)
251 """IDs of quanta that were predicted to consume this dataset."""
253 @property
254 def node_id(self) -> uuid.UUID:
255 """Alias for the dataset ID."""
256 return self.dataset_id
258 @classmethod
259 def from_predicted(
260 cls,
261 predicted: PredictedDatasetModel,
262 producer: uuid.UUID | None = None,
263 consumers: Iterable[uuid.UUID] = (),
264 ) -> ProvenanceDatasetModel:
265 """Construct from a predicted dataset model.
267 Parameters
268 ----------
269 predicted : `PredictedDatasetModel`
270 Information about the dataset from the predicted graph.
271 producer : `uuid.UUID` or `None`, optional
272 ID of the quantum that was predicted to produce this dataset.
273 consumers : `~collections.abc.Iterable` [`uuid.UUID`], optional
274 IDs of the quanta that were predicted to consume this dataset.
276 Returns
277 -------
278 provenance : `ProvenanceDatasetModel`
279 Provenance dataset model.
281 Notes
282 -----
283 This initializes `produced` to `True` when ``producer is None`` and
284 `False` otherwise, on the assumption that it will be updated later.
285 """
286 return cls.model_construct(
287 dataset_id=predicted.dataset_id,
288 dataset_type_name=predicted.dataset_type_name,
289 data_coordinate=predicted.data_coordinate,
290 run=predicted.run,
291 produced=(producer is None), # if it's not produced by this QG, it's an overall input
292 producer=producer,
293 consumers=list(consumers),
294 )
296 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
297 """Add this dataset and its edges to quanta to a provenance graph.
299 Parameters
300 ----------
301 graph : `ProvenanceQuantumGraph`
302 Graph to update in place.
304 Notes
305 -----
306 This method adds:
308 - a ``bipartite_xgraph`` dataset node with full attributes;
309 - ``bipartite_xgraph`` edges to adjacent quanta (which adds quantum
310 nodes with no attributes), without populating edge attributes;
311 - ``quantum_only_xgraph`` edges for each pair of quanta in which one
312 produces this dataset and another consumes it (this also adds quantum
313 nodes with no attributes).
314 """
315 dataset_type_node = graph.pipeline_graph.dataset_types[self.dataset_type_name]
316 data_id = DataCoordinate.from_full_values(dataset_type_node.dimensions, tuple(self.data_coordinate))
317 graph._bipartite_xgraph.add_node(
318 self.dataset_id,
319 data_id=data_id,
320 dataset_type_name=self.dataset_type_name,
321 pipeline_node=dataset_type_node,
322 run=self.run,
323 produced=self.produced,
324 )
325 if self.producer is not None:
326 graph._bipartite_xgraph.add_edge(self.producer, self.dataset_id)
327 for consumer_id in self.consumers:
328 graph._bipartite_xgraph.add_edge(self.dataset_id, consumer_id)
329 if self.producer is not None:
330 graph._quantum_only_xgraph.add_edge(self.producer, consumer_id)
331 graph._datasets_by_type[self.dataset_type_name][data_id] = self.dataset_id
333 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
334 # when we inherit those docstrings in our public classes.
335 if "sphinx" in sys.modules and not TYPE_CHECKING:
337 def copy(self, *args: Any, **kwargs: Any) -> Any:
338 """See `pydantic.BaseModel.copy`."""
339 return super().copy(*args, **kwargs)
341 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
342 """See `pydantic.BaseModel.model_dump`."""
343 return super().model_dump(*args, **kwargs)
345 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
346 """See `pydantic.BaseModel.model_dump_json`."""
347 return super().model_dump(*args, **kwargs)
349 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
350 """See `pydantic.BaseModel.model_copy`."""
351 return super().model_copy(*args, **kwargs)
353 @classmethod
354 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
355 """See `pydantic.BaseModel.model_construct`."""
356 return super().model_construct(*args, **kwargs)
358 @classmethod
359 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
360 """See `pydantic.BaseModel.model_json_schema`."""
361 return super().model_json_schema(*args, **kwargs)
363 @classmethod
364 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
365 """See `pydantic.BaseModel.model_validate`."""
366 return super().model_validate(*args, **kwargs)
368 @classmethod
369 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
370 """See `pydantic.BaseModel.model_validate_json`."""
371 return super().model_validate_json(*args, **kwargs)
373 @classmethod
374 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
375 """See `pydantic.BaseModel.model_validate_strings`."""
376 return super().model_validate_strings(*args, **kwargs)
379class ProvenanceQuantumAttemptModel(pydantic.BaseModel):
380 """Data model for a now-superseded attempt to run a quantum in a
381 provenance quantum graph file.
382 """
384 attempt: int = 0
385 """Counter incremented for every attempt to execute this quantum."""
387 status: QuantumAttemptStatus = QuantumAttemptStatus.UNKNOWN
388 """Enumerated status for the quantum."""
390 caveats: QuantumSuccessCaveats | None = None
391 """Flags indicating caveats on successful quanta."""
393 exception: ExceptionInfo | None = None
394 """Information about an exception raised when the quantum was executing."""
396 resource_usage: QuantumResourceUsage | None = None
397 """Resource usage information (timing, memory use) for this quantum."""
399 previous_process_quanta: list[uuid.UUID] = pydantic.Field(default_factory=list)
400 """The IDs of other quanta previously executed in the same process as this
401 one.
402 """
404 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
405 # when we inherit those docstrings in our public classes.
406 if "sphinx" in sys.modules and not TYPE_CHECKING:
408 def copy(self, *args: Any, **kwargs: Any) -> Any:
409 """See `pydantic.BaseModel.copy`."""
410 return super().copy(*args, **kwargs)
412 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
413 """See `pydantic.BaseModel.model_dump`."""
414 return super().model_dump(*args, **kwargs)
416 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
417 """See `pydantic.BaseModel.model_dump_json`."""
418 return super().model_dump(*args, **kwargs)
420 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
421 """See `pydantic.BaseModel.model_copy`."""
422 return super().model_copy(*args, **kwargs)
424 @classmethod
425 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
426 """See `pydantic.BaseModel.model_construct`."""
427 return super().model_construct(*args, **kwargs)
429 @classmethod
430 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
431 """See `pydantic.BaseModel.model_json_schema`."""
432 return super().model_json_schema(*args, **kwargs)
434 @classmethod
435 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
436 """See `pydantic.BaseModel.model_validate`."""
437 return super().model_validate(*args, **kwargs)
439 @classmethod
440 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
441 """See `pydantic.BaseModel.model_validate_json`."""
442 return super().model_validate_json(*args, **kwargs)
444 @classmethod
445 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
446 """See `pydantic.BaseModel.model_validate_strings`."""
447 return super().model_validate_strings(*args, **kwargs)
450class ProvenanceLogRecordsModel(pydantic.BaseModel):
451 """Data model for storing execution logs in a provenance quantum graph
452 file.
453 """
455 attempts: list[list[ButlerLogRecord] | None] = pydantic.Field(default_factory=list)
456 """Logs from attempts to run this task, ordered chronologically from first
457 to last.
458 """
460 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
461 # when we inherit those docstrings in our public classes.
462 if "sphinx" in sys.modules and not TYPE_CHECKING:
464 def copy(self, *args: Any, **kwargs: Any) -> Any:
465 """See `pydantic.BaseModel.copy`."""
466 return super().copy(*args, **kwargs)
468 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
469 """See `pydantic.BaseModel.model_dump`."""
470 return super().model_dump(*args, **kwargs)
472 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
473 """See `pydantic.BaseModel.model_dump_json`."""
474 return super().model_dump(*args, **kwargs)
476 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
477 """See `pydantic.BaseModel.model_copy`."""
478 return super().model_copy(*args, **kwargs)
480 @classmethod
481 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
482 """See `pydantic.BaseModel.model_construct`."""
483 return super().model_construct(*args, **kwargs)
485 @classmethod
486 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
487 """See `pydantic.BaseModel.model_json_schema`."""
488 return super().model_json_schema(*args, **kwargs)
490 @classmethod
491 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
492 """See `pydantic.BaseModel.model_validate`."""
493 return super().model_validate(*args, **kwargs)
495 @classmethod
496 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
497 """See `pydantic.BaseModel.model_validate_json`."""
498 return super().model_validate_json(*args, **kwargs)
500 @classmethod
501 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
502 """See `pydantic.BaseModel.model_validate_strings`."""
503 return super().model_validate_strings(*args, **kwargs)
506class ProvenanceTaskMetadataModel(pydantic.BaseModel):
507 """Data model for storing task metadata in a provenance quantum graph
508 file.
509 """
511 attempts: list[TaskMetadata | None] = pydantic.Field(default_factory=list)
512 """Metadata from attempts to run this task, ordered chronologically from
513 first to last.
514 """
516 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
517 # when we inherit those docstrings in our public classes.
518 if "sphinx" in sys.modules and not TYPE_CHECKING:
520 def copy(self, *args: Any, **kwargs: Any) -> Any:
521 """See `pydantic.BaseModel.copy`."""
522 return super().copy(*args, **kwargs)
524 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
525 """See `pydantic.BaseModel.model_dump`."""
526 return super().model_dump(*args, **kwargs)
528 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
529 """See `pydantic.BaseModel.model_dump_json`."""
530 return super().model_dump(*args, **kwargs)
532 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
533 """See `pydantic.BaseModel.model_copy`."""
534 return super().model_copy(*args, **kwargs)
536 @classmethod
537 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
538 """See `pydantic.BaseModel.model_construct`."""
539 return super().model_construct(*args, **kwargs)
541 @classmethod
542 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
543 """See `pydantic.BaseModel.model_json_schema`."""
544 return super().model_json_schema(*args, **kwargs)
546 @classmethod
547 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
548 """See `pydantic.BaseModel.model_validate`."""
549 return super().model_validate(*args, **kwargs)
551 @classmethod
552 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
553 """See `pydantic.BaseModel.model_validate_json`."""
554 return super().model_validate_json(*args, **kwargs)
556 @classmethod
557 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
558 """See `pydantic.BaseModel.model_validate_strings`."""
559 return super().model_validate_strings(*args, **kwargs)
562class ProvenanceQuantumModel(pydantic.BaseModel):
563 """Data model for the quanta in a provenance quantum graph file."""
565 quantum_id: uuid.UUID
566 """Unique identifier for the quantum."""
568 task_label: TaskLabel
569 """Name of the type of this dataset."""
571 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
572 """The full values (required and implied) of this dataset's data ID."""
574 inputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
575 """IDs of the datasets predicted to be consumed by this quantum, grouped by
576 connection name.
577 """
579 outputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
580 """IDs of the datasets predicted to be produced by this quantum, grouped by
581 connection name.
582 """
584 attempts: list[ProvenanceQuantumAttemptModel] = pydantic.Field(default_factory=list)
585 """Provenance for all attempts to execute this quantum, ordered
586 chronologically from first to last.
588 An entry is added merely if the quantum *should* have been attempted; an
589 empty `list` is used only for quanta that were blocked by an upstream
590 failure.
591 """
593 @property
594 def node_id(self) -> uuid.UUID:
595 """Alias for the quantum ID."""
596 return self.quantum_id
598 @classmethod
599 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceQuantumModel:
600 """Construct from a predicted quantum model.
602 Parameters
603 ----------
604 predicted : `PredictedQuantumDatasetsModel`
605 Information about the quantum from the predicted graph.
607 Returns
608 -------
609 provenance : `ProvenanceQuantumModel`
610 Provenance quantum model.
611 """
612 inputs = {
613 connection_name: [d.dataset_id for d in predicted_inputs]
614 for connection_name, predicted_inputs in predicted.inputs.items()
615 }
616 outputs = {
617 connection_name: [d.dataset_id for d in predicted_outputs]
618 for connection_name, predicted_outputs in predicted.outputs.items()
619 }
620 return cls(
621 quantum_id=predicted.quantum_id,
622 task_label=predicted.task_label,
623 data_coordinate=predicted.data_coordinate,
624 inputs=inputs,
625 outputs=outputs,
626 )
628 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
629 """Add this quantum and its edges to datasets to a provenance graph.
631 Parameters
632 ----------
633 graph : `ProvenanceQuantumGraph`
634 Graph to update in place.
636 Notes
637 -----
638 This method adds:
640 - a ``bipartite_xgraph`` quantum node with full attributes;
641 - a ``quantum_only_xgraph`` quantum node with full attributes;
642 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
643 nodes with no attributes), while populating those edge attributes;
644 - ``quantum_only_xgraph`` edges to any adjacent quantum that has also
645 already been loaded.
646 """
647 task_node = graph.pipeline_graph.tasks[self.task_label]
648 data_id = DataCoordinate.from_full_values(task_node.dimensions, tuple(self.data_coordinate))
649 last_attempt = (
650 self.attempts[-1]
651 if self.attempts
652 else ProvenanceQuantumAttemptModel(status=QuantumAttemptStatus.BLOCKED)
653 )
654 graph._bipartite_xgraph.add_node(
655 self.quantum_id,
656 data_id=data_id,
657 task_label=self.task_label,
658 pipeline_node=task_node,
659 status=last_attempt.status,
660 caveats=last_attempt.caveats,
661 exception=last_attempt.exception,
662 resource_usage=last_attempt.resource_usage,
663 attempts=self.attempts,
664 )
665 graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
666 graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
667 for connection_name, dataset_ids in self.inputs.items():
668 read_edge = task_node.get_input_edge(connection_name)
669 for dataset_id in dataset_ids:
670 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
671 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
672 "pipeline_edges", []
673 ).append(read_edge)
674 for connection_name, dataset_ids in self.outputs.items():
675 write_edge = task_node.get_output_edge(connection_name)
676 if connection_name == acc.METADATA_OUTPUT_CONNECTION_NAME:
677 graph._bipartite_xgraph.add_node(
678 dataset_ids[0],
679 data_id=data_id,
680 dataset_type_name=write_edge.dataset_type_name,
681 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
682 run=graph.header.output_run,
683 produced=last_attempt.status.has_metadata,
684 )
685 graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
686 graph._bipartite_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
687 graph._quantum_only_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
688 if connection_name == acc.LOG_OUTPUT_CONNECTION_NAME:
689 graph._bipartite_xgraph.add_node(
690 dataset_ids[0],
691 data_id=data_id,
692 dataset_type_name=write_edge.dataset_type_name,
693 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
694 run=graph.header.output_run,
695 produced=last_attempt.status.has_log,
696 )
697 graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
698 graph._bipartite_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
699 graph._quantum_only_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
700 for dataset_id in dataset_ids:
701 graph._bipartite_xgraph.add_edge(
702 self.quantum_id,
703 dataset_id,
704 is_read=False,
705 # There can only be one pipeline edge for an output.
706 pipeline_edges=[write_edge],
707 )
708 for dataset_id in graph._bipartite_xgraph.predecessors(self.quantum_id):
709 for upstream_quantum_id in graph._bipartite_xgraph.predecessors(dataset_id):
710 graph._quantum_only_xgraph.add_edge(upstream_quantum_id, self.quantum_id)
711 for dataset_id in graph._bipartite_xgraph.successors(self.quantum_id):
712 for downstream_quantum_id in graph._bipartite_xgraph.successors(dataset_id):
713 graph._quantum_only_xgraph.add_edge(self.quantum_id, downstream_quantum_id)
715 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
716 # when we inherit those docstrings in our public classes.
717 if "sphinx" in sys.modules and not TYPE_CHECKING:
719 def copy(self, *args: Any, **kwargs: Any) -> Any:
720 """See `pydantic.BaseModel.copy`."""
721 return super().copy(*args, **kwargs)
723 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
724 """See `pydantic.BaseModel.model_dump`."""
725 return super().model_dump(*args, **kwargs)
727 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
728 """See `pydantic.BaseModel.model_dump_json`."""
729 return super().model_dump(*args, **kwargs)
731 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
732 """See `pydantic.BaseModel.model_copy`."""
733 return super().model_copy(*args, **kwargs)
735 @classmethod
736 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
737 """See `pydantic.BaseModel.model_construct`."""
738 return super().model_construct(*args, **kwargs)
740 @classmethod
741 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
742 """See `pydantic.BaseModel.model_json_schema`."""
743 return super().model_json_schema(*args, **kwargs)
745 @classmethod
746 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
747 """See `pydantic.BaseModel.model_validate`."""
748 return super().model_validate(*args, **kwargs)
750 @classmethod
751 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
752 """See `pydantic.BaseModel.model_validate_json`."""
753 return super().model_validate_json(*args, **kwargs)
755 @classmethod
756 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
757 """See `pydantic.BaseModel.model_validate_strings`."""
758 return super().model_validate_strings(*args, **kwargs)
761class ProvenanceInitQuantumModel(pydantic.BaseModel):
762 """Data model for the special "init" quanta in a provenance quantum graph
763 file.
764 """
766 quantum_id: uuid.UUID
767 """Unique identifier for the quantum."""
769 task_label: TaskLabel
770 """Name of the type of this dataset.
772 This is always a parent dataset type name, not a component.
774 Note that full dataset type definitions are stored in the pipeline graph.
775 """
777 inputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
778 """IDs of the datasets predicted to be consumed by this quantum, grouped by
779 connection name.
780 """
782 outputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
783 """IDs of the datasets predicted to be produced by this quantum, grouped by
784 connection name.
785 """
787 @classmethod
788 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceInitQuantumModel:
789 """Construct from a predicted quantum model.
791 Parameters
792 ----------
793 predicted : `PredictedQuantumDatasetsModel`
794 Information about the quantum from the predicted graph.
796 Returns
797 -------
798 provenance : `ProvenanceInitQuantumModel`
799 Provenance init quantum model.
800 """
801 inputs = {
802 connection_name: predicted_inputs[0].dataset_id
803 for connection_name, predicted_inputs in predicted.inputs.items()
804 }
805 outputs = {
806 connection_name: predicted_outputs[0].dataset_id
807 for connection_name, predicted_outputs in predicted.outputs.items()
808 }
809 return cls(
810 quantum_id=predicted.quantum_id,
811 task_label=predicted.task_label,
812 inputs=inputs,
813 outputs=outputs,
814 )
816 def _add_to_graph(self, graph: ProvenanceQuantumGraph, empty_data_id: DataCoordinate) -> None:
817 """Add this quantum and its edges to datasets to a provenance graph.
819 Parameters
820 ----------
821 graph : `ProvenanceQuantumGraph`
822 Graph to update in place.
823 empty_data_id : `lsst.daf.butler.DataCoordinate`
824 The empty data ID for the appropriate dimension universe.
826 Notes
827 -----
828 This method adds:
830 - a ``bipartite_xgraph`` quantum node with full attributes;
831 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
832 nodes with no attributes), while populating those edge attributes;
833 """
834 task_init_node = graph.pipeline_graph.tasks[self.task_label].init
835 graph._bipartite_xgraph.add_node(
836 self.quantum_id, data_id=empty_data_id, task_label=self.task_label, pipeline_node=task_init_node
837 )
838 for connection_name, dataset_id in self.inputs.items():
839 read_edge = task_init_node.get_input_edge(connection_name)
840 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
841 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
842 "pipeline_edges", []
843 ).append(read_edge)
844 for connection_name, dataset_id in self.outputs.items():
845 write_edge = task_init_node.get_output_edge(connection_name)
846 graph._bipartite_xgraph.add_node(
847 dataset_id,
848 data_id=empty_data_id,
849 dataset_type_name=write_edge.dataset_type_name,
850 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
851 run=graph.header.output_run,
852 produced=True,
853 )
854 graph._datasets_by_type[write_edge.dataset_type_name][empty_data_id] = dataset_id
855 graph._bipartite_xgraph.add_edge(
856 self.quantum_id,
857 dataset_id,
858 is_read=False,
859 # There can only be one pipeline edge for an output.
860 pipeline_edges=[write_edge],
861 )
862 if write_edge.connection_name == acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME:
863 graph._bipartite_xgraph.nodes[self.quantum_id]["config_id"] = dataset_id
864 graph._init_quanta[self.task_label] = self.quantum_id
866 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
867 # when we inherit those docstrings in our public classes.
868 if "sphinx" in sys.modules and not TYPE_CHECKING:
870 def copy(self, *args: Any, **kwargs: Any) -> Any:
871 """See `pydantic.BaseModel.copy`."""
872 return super().copy(*args, **kwargs)
874 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
875 """See `pydantic.BaseModel.model_dump`."""
876 return super().model_dump(*args, **kwargs)
878 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
879 """See `pydantic.BaseModel.model_dump_json`."""
880 return super().model_dump(*args, **kwargs)
882 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
883 """See `pydantic.BaseModel.model_copy`."""
884 return super().model_copy(*args, **kwargs)
886 @classmethod
887 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
888 """See `pydantic.BaseModel.model_construct`."""
889 return super().model_construct(*args, **kwargs)
891 @classmethod
892 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
893 """See `pydantic.BaseModel.model_json_schema`."""
894 return super().model_json_schema(*args, **kwargs)
896 @classmethod
897 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
898 """See `pydantic.BaseModel.model_validate`."""
899 return super().model_validate(*args, **kwargs)
901 @classmethod
902 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
903 """See `pydantic.BaseModel.model_validate_json`."""
904 return super().model_validate_json(*args, **kwargs)
906 @classmethod
907 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
908 """See `pydantic.BaseModel.model_validate_strings`."""
909 return super().model_validate_strings(*args, **kwargs)
912class ProvenanceInitQuantaModel(pydantic.RootModel):
913 """Data model for the init quanta in a provenance graph."""
915 root: list[ProvenanceInitQuantumModel] = pydantic.Field(default_factory=list)
916 """List of special "init" quanta, one for each task."""
918 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
919 """Add this quantum and its edges to datasets to a provenance graph.
921 Parameters
922 ----------
923 graph : `ProvenanceQuantumGraph`
924 Graph to update in place.
925 """
926 empty_data_id = DataCoordinate.make_empty(graph.pipeline_graph.universe)
927 for init_quantum in self.root:
928 init_quantum._add_to_graph(graph, empty_data_id=empty_data_id)
930 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
931 # when we inherit those docstrings in our public classes.
932 if "sphinx" in sys.modules and not TYPE_CHECKING:
934 def copy(self, *args: Any, **kwargs: Any) -> Any:
935 """See `pydantic.BaseModel.copy`."""
936 return super().copy(*args, **kwargs)
938 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
939 """See `pydantic.BaseModel.model_dump`."""
940 return super().model_dump(*args, **kwargs)
942 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
943 """See `pydantic.BaseModel.model_dump_json`."""
944 return super().model_dump(*args, **kwargs)
946 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
947 """See `pydantic.BaseModel.model_copy`."""
948 return super().model_copy(*args, **kwargs)
950 @classmethod
951 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
952 """See `pydantic.BaseModel.model_construct`."""
953 return super().model_construct(*args, **kwargs)
955 @classmethod
956 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
957 """See `pydantic.BaseModel.model_json_schema`."""
958 return super().model_json_schema(*args, **kwargs)
960 @classmethod
961 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
962 """See `pydantic.BaseModel.model_validate`."""
963 return super().model_validate(*args, **kwargs)
965 @classmethod
966 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
967 """See `pydantic.BaseModel.model_validate_json`."""
968 return super().model_validate_json(*args, **kwargs)
970 @classmethod
971 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
972 """See `pydantic.BaseModel.model_validate_strings`."""
973 return super().model_validate_strings(*args, **kwargs)
976class ProvenanceQuantumGraph(BaseQuantumGraph):
977 """A quantum graph that represents processing that has already been
978 executed.
980 Parameters
981 ----------
982 header : `HeaderModel`
983 General metadata shared with other quantum graph types.
984 pipeline_graph : `.pipeline_graph.PipelineGraph`
985 Graph of tasks and dataset types. May contain a superset of the tasks
986 and dataset types that actually have quanta and datasets in the quantum
987 graph.
989 Notes
990 -----
991 A provenance quantum graph is generally obtained via the
992 `ProvenanceQuantumGraphReader.graph` attribute, which is updated in-place
993 as information is read from disk.
994 """
996 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph) -> None:
997 super().__init__(header, pipeline_graph)
998 self._init_quanta: dict[TaskLabel, uuid.UUID] = {}
999 self._quantum_only_xgraph = networkx.DiGraph()
1000 self._bipartite_xgraph = networkx.DiGraph()
1001 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {
1002 task_label: {} for task_label in self.pipeline_graph.tasks.keys()
1003 }
1004 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {
1005 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
1006 }
1008 @property
1009 def init_quanta(self) -> Mapping[TaskLabel, uuid.UUID]:
1010 """A mapping from task label to the ID of the special init quantum for
1011 that task.
1013 This is populated by the ``init_quanta`` component. Additional
1014 information about each init quantum can be found by using the ID to
1015 look up node attributes in the `bipartite_xgraph`, i.e.::
1017 info: ProvenanceInitQuantumInfo = qg.bipartite_xgraph.nodes[id]
1018 """
1019 return self._init_quanta
1021 @property
1022 def quanta_by_task(self) -> Mapping[TaskLabel, Mapping[DataCoordinate, uuid.UUID]]:
1023 """A nested mapping of all quanta, keyed first by task name and then by
1024 data ID.
1026 Notes
1027 -----
1028 This is populated one quantum at a time as they are read. All tasks in
1029 the pipeline graph are included, even if none of their quanta were
1030 loaded (i.e. nested mappings may be empty).
1032 The returned object may be an internal dictionary; as the type
1033 annotation indicates, it should not be modified in place.
1034 """
1035 return self._quanta_by_task_label
1037 @property
1038 def datasets_by_type(self) -> Mapping[DatasetTypeName, Mapping[DataCoordinate, uuid.UUID]]:
1039 """A nested mapping of all datasets, keyed first by dataset type name
1040 and then by data ID.
1042 Notes
1043 -----
1044 This is populated one dataset at a time as they are read. All dataset
1045 types in the pipeline graph are included, even if none of their
1046 datasets were loaded (i.e. nested mappings may be empty).
1048 Reading a quantum also populates its log and metadata datasets.
1050 The returned object may be an internal dictionary; as the type
1051 annotation indicates, it should not be modified in place.
1052 """
1053 return self._datasets_by_type
1055 @property
1056 def quantum_only_xgraph(self) -> networkx.DiGraph:
1057 """A directed acyclic graph with quanta as nodes (and datasets elided).
1059 Notes
1060 -----
1061 Node keys are quantum UUIDs, and are populated one quantum at a time as
1062 they are loaded. Loading quanta (via
1063 `ProvenanceQuantumGraphReader.read_quanta`) will add the loaded nodes
1064 with full attributes and add edges to adjacent nodes with no
1065 attributes. Loading datasets (via
1066 `ProvenanceQuantumGraphReader.read_datasets`) will also add edges and
1067 nodes with no attributes.
1069 Node attributes are described by the `ProvenanceQuantumInfo` types.
1071 This graph does not include special "init" quanta.
1073 The returned object is a read-only view of an internal one.
1074 """
1075 return self._quantum_only_xgraph.copy(as_view=True)
1077 @property
1078 def bipartite_xgraph(self) -> networkx.DiGraph:
1079 """A directed acyclic graph with quantum and dataset nodes.
1081 Notes
1082 -----
1083 Node keys are quantum or dataset UUIDs, and are populated one quantum
1084 or dataset at a time as they are loaded. Loading quanta (via
1085 `ProvenanceQuantumGraphReader.read_quanta`) or datasets (via
1086 `ProvenanceQuantumGraphReader.read_datasets`) will load those nodes
1087 with full attributes and edges to adjacent nodes with no attributes.
1088 Loading quanta is necessary to populate edge attributes.
1089 Reading a quantum also populates its log and metadata datasets.
1091 Node attributes are described by the
1092 `ProvenanceQuantumInfo`, `ProvenanceInitQuantumInfo`, and
1093 `ProvenanceDatasetInfo` types.
1095 This graph includes init-input and init-output datasets, but it does
1096 *not* reflect the dependency between each task's special "init" quantum
1097 and its runtime quanta (as this would require edges between quanta, and
1098 that would break the "bipartite" property).
1100 The returned object is a read-only view of an internal one.
1101 """
1102 return self._bipartite_xgraph.copy(as_view=True)
1104 def make_quantum_table(self) -> astropy.table.Table:
1105 """Construct an `astropy.table.Table` with a tabular summary of the
1106 quanta.
1108 Returns
1109 -------
1110 table : `astropy.table.Table`
1111 A table view of the quantum information. This only includes
1112 counts of status categories and caveats, not any per-data-ID
1113 detail.
1115 Notes
1116 -----
1117 Success caveats in the table are represented by their
1118 `~QuantumSuccessCaveats.concise` form, so when pretty-printing this
1119 table for users, the `~QuantumSuccessCaveats.legend` should generally
1120 be printed as well.
1121 """
1122 rows = []
1123 for task_label, quanta_for_task in self.quanta_by_task.items():
1124 if not self.header.n_task_quanta[task_label]:
1125 continue
1126 status_counts = Counter[QuantumAttemptStatus](
1127 self._quantum_only_xgraph.nodes[q]["status"] for q in quanta_for_task.values()
1128 )
1129 caveat_counts = Counter[QuantumSuccessCaveats | None](
1130 self._quantum_only_xgraph.nodes[q]["caveats"] for q in quanta_for_task.values()
1131 )
1132 caveat_counts.pop(QuantumSuccessCaveats.NO_CAVEATS, None)
1133 caveat_counts.pop(None, None)
1134 if len(caveat_counts) > 1:
1135 caveats = "(multiple)"
1136 elif len(caveat_counts) == 1:
1137 ((code, count),) = caveat_counts.items()
1138 # MyPy can't tell that the pop(None, None) above makes None
1139 # impossible here.
1140 caveats = f"{code.concise()}({count})" # type: ignore[union-attr]
1141 else:
1142 caveats = ""
1143 rows.append(
1144 {
1145 "Task": task_label,
1146 "Unknown": status_counts.get(QuantumAttemptStatus.UNKNOWN, 0),
1147 "Successful": status_counts.get(QuantumAttemptStatus.SUCCESSFUL, 0),
1148 "Caveats": caveats,
1149 "Blocked": status_counts.get(QuantumAttemptStatus.BLOCKED, 0),
1150 "Failed": status_counts.get(QuantumAttemptStatus.FAILED, 0),
1151 "TOTAL": len(quanta_for_task),
1152 "EXPECTED": self.header.n_task_quanta[task_label],
1153 }
1154 )
1155 return astropy.table.Table(rows)
1157 def make_exception_table(self) -> astropy.table.Table:
1158 """Construct an `astropy.table.Table` with counts for each exception
1159 type raised by each task.
1161 Returns
1162 -------
1163 table : `astropy.table.Table`
1164 A table with columns for task label, exception type, and counts.
1165 """
1166 rows = []
1167 for task_label, quanta_for_task in self.quanta_by_task.items():
1168 counts_by_type = Counter(
1169 exc_info.type_name
1170 for q in quanta_for_task.values()
1171 if (exc_info := self._quantum_only_xgraph.nodes[q]["exception"]) is not None
1172 )
1173 for type_name, count in counts_by_type.items():
1174 rows.append({"Task": task_label, "Exception": type_name, "Count": count})
1175 return astropy.table.Table(rows)
1177 def make_task_resource_usage_table(
1178 self, task_label: TaskLabel, include_data_ids: bool = False
1179 ) -> astropy.table.Table:
1180 """Make a table of resource usage for a single task.
1182 Parameters
1183 ----------
1184 task_label : `str`
1185 Label of the task to extract resource usage for.
1186 include_data_ids : `bool`, optional
1187 Whether to also include data ID columns.
1189 Returns
1190 -------
1191 table : `astropy.table.Table`
1192 A table with columns for quantum ID and all fields in
1193 `QuantumResourceUsage`.
1194 """
1195 quanta_for_task = self.quanta_by_task[task_label]
1196 dtype_terms: list[tuple[str, np.dtype]] = [("quantum_id", np.dtype((np.void, 16)))]
1197 if include_data_ids:
1198 dimensions = self.pipeline_graph.tasks[task_label].dimensions
1199 for dimension_name in dimensions.data_coordinate_keys:
1200 dtype = np.dtype(self.pipeline_graph.universe.dimensions[dimension_name].primary_key.pytype)
1201 dtype_terms.append((dimension_name, dtype))
1202 fields = QuantumResourceUsage.get_numpy_fields()
1203 dtype_terms.extend(fields.items())
1204 row_dtype = np.dtype(dtype_terms)
1205 rows: list[object] = []
1206 for data_id, quantum_id in quanta_for_task.items():
1207 info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1208 if (resource_usage := info["resource_usage"]) is not None:
1209 row: tuple[object, ...] = (quantum_id.bytes,)
1210 if include_data_ids:
1211 row += data_id.full_values
1212 row += resource_usage.get_numpy_row()
1213 rows.append(row)
1214 array = np.array(rows, dtype=row_dtype)
1215 return astropy.table.Table(array, units=QuantumResourceUsage.get_units())
1218@dataclasses.dataclass
1219class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1220 """A helper class for reading provenance quantum graphs.
1222 Notes
1223 -----
1224 The `open` context manager should be used to construct new instances.
1225 Instances cannot be used after the context manager exits, except to access
1226 the `graph` attribute`.
1228 The various ``read_*`` methods in this class update the `graph` attribute
1229 in place.
1230 """
1232 graph: ProvenanceQuantumGraph = dataclasses.field(init=False)
1233 """Loaded provenance graph, populated in place as components are read."""
1235 @classmethod
1236 @contextmanager
1237 def open(
1238 cls,
1239 uri: ResourcePathExpression,
1240 *,
1241 page_size: int | None = None,
1242 import_mode: TaskImportMode = TaskImportMode.DO_NOT_IMPORT,
1243 ) -> Iterator[ProvenanceQuantumGraphReader]:
1244 """Construct a reader from a URI.
1246 Parameters
1247 ----------
1248 uri : convertible to `lsst.resources.ResourcePath`
1249 URI to open. Should have a ``.qg`` extension.
1250 page_size : `int`, optional
1251 Approximate number of bytes to read at once from address files and
1252 multi-block files. Note that this does not set a page size for
1253 *all* reads, but it does affect the smallest, most numerous reads.
1254 Can also be set via the ``LSST_QG_PAGE_SIZE`` environment variable.
1255 import_mode : `.pipeline_graph.TaskImportMode`, optional
1256 How to handle importing the task classes referenced in the pipeline
1257 graph.
1259 Returns
1260 -------
1261 reader : `contextlib.AbstractContextManager` [ \
1262 `ProvenanceQuantumGraphReader` ]
1263 A context manager that returns the reader when entered.
1264 """
1265 with cls._open(
1266 uri,
1267 graph_type="provenance",
1268 address_filename="nodes",
1269 page_size=page_size,
1270 import_mode=import_mode,
1271 n_addresses=4,
1272 ) as self:
1273 yield self
1275 def __post_init__(self) -> None:
1276 self.graph = ProvenanceQuantumGraph(self.header, self.pipeline_graph)
1278 def read_init_quanta(self) -> None:
1279 """Read the thin graph, with all edge information and categorization of
1280 quanta by task label.
1281 """
1282 init_quanta = self._read_single_block("init_quanta", ProvenanceInitQuantaModel)
1283 for init_quantum in init_quanta.root:
1284 self.graph._init_quanta[init_quantum.task_label] = init_quantum.quantum_id
1285 init_quanta._add_to_graph(self.graph)
1287 def read_full_graph(self) -> None:
1288 """Read all bipartite edges and all quantum and dataset node
1289 attributes, fully populating the `graph` attribute.
1291 Notes
1292 -----
1293 This does not read logs, metadata, or packages ; those must always be
1294 fetched explicitly.
1295 """
1296 self.read_init_quanta()
1297 self.read_datasets()
1298 self.read_quanta()
1300 def read_datasets(self, datasets: Iterable[uuid.UUID] | None = None) -> None:
1301 """Read information about the given datasets.
1303 Parameters
1304 ----------
1305 datasets : `~collections.abc.Iterable` [`uuid.UUID`], optional
1306 Iterable of dataset IDs to load. If not provided, all datasets
1307 will be loaded. The UUIDs and indices of quanta will be ignored.
1308 """
1309 self._read_nodes(datasets, DATASET_ADDRESS_INDEX, DATASET_MB_NAME, ProvenanceDatasetModel)
1311 def read_quanta(self, quanta: Iterable[uuid.UUID] | None = None) -> None:
1312 """Read information about the given quanta.
1314 Parameters
1315 ----------
1316 quanta : `~collections.abc.Iterable` [`uuid.UUID`], optional
1317 Iterable of quantum IDs to load. If not provided, all quanta will
1318 be loaded. The UUIDs and indices of datasets and special init
1319 quanta will be ignored.
1320 """
1321 self._read_nodes(quanta, QUANTUM_ADDRESS_INDEX, QUANTUM_MB_NAME, ProvenanceQuantumModel)
1323 def _read_nodes(
1324 self,
1325 nodes: Iterable[uuid.UUID] | None,
1326 address_index: int,
1327 mb_name: str,
1328 model_type: type[ProvenanceDatasetModel] | type[ProvenanceQuantumModel],
1329 ) -> None:
1330 node: ProvenanceDatasetModel | ProvenanceQuantumModel | None
1331 if nodes is None:
1332 self.address_reader.read_all()
1333 nodes = self.address_reader.rows.keys()
1334 for node in MultiblockReader.read_all_models_in_zip(
1335 self.zf,
1336 mb_name,
1337 model_type,
1338 self.decompressor,
1339 int_size=self.header.int_size,
1340 page_size=self.page_size,
1341 ):
1342 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(node.node_id, {}):
1343 # Use the old node to reduce memory usage (since it might
1344 # also have other outstanding reference holders).
1345 continue
1346 node._add_to_graph(self.graph)
1347 else:
1348 with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1349 for node_id_or_index in nodes:
1350 address_row = self.address_reader.find(node_id_or_index)
1351 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1352 # Use the old node to reduce memory usage (since it
1353 # might also have other outstanding reference holders).
1354 continue
1355 node = mb_reader.read_model(
1356 address_row.addresses[address_index], model_type, self.decompressor
1357 )
1358 if node is not None:
1359 node._add_to_graph(self.graph)
1361 def fetch_logs(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[ButlerLogRecords | None]]:
1362 """Fetch log datasets.
1364 Parameters
1365 ----------
1366 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1367 UUIDs of the log datasets themselves or of the quanta they
1368 correspond to.
1370 Returns
1371 -------
1372 logs : `dict` [ `uuid.UUID`, `list` [\
1373 `lsst.daf.butler.ButlerLogRecords` or `None`] ]
1374 Logs for the given IDs. Each value is a list of
1375 `lsst.daf.butler.ButlerLogRecords` instances representing different
1376 execution attempts, ordered chronologically from first to last.
1377 Attempts where logs were missing will have `None` in this list.
1378 """
1379 result: dict[uuid.UUID, list[ButlerLogRecords | None]] = {}
1380 with MultiblockReader.open_in_zip(self.zf, LOG_MB_NAME, int_size=self.header.int_size) as mb_reader:
1381 for node_id_or_index in nodes:
1382 address_row = self.address_reader.find(node_id_or_index)
1383 logs_by_attempt = mb_reader.read_model(
1384 address_row.addresses[LOG_ADDRESS_INDEX], ProvenanceLogRecordsModel, self.decompressor
1385 )
1386 if logs_by_attempt is not None:
1387 result[node_id_or_index] = [
1388 ButlerLogRecords.from_records(attempt_logs) if attempt_logs is not None else None
1389 for attempt_logs in logs_by_attempt.attempts
1390 ]
1391 return result
1393 def fetch_metadata(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[TaskMetadata | None]]:
1394 """Fetch metadata datasets.
1396 Parameters
1397 ----------
1398 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1399 UUIDs of the metadata datasets themselves or of the quanta they
1400 correspond to.
1402 Returns
1403 -------
1404 metadata : `dict` [ `uuid.UUID`, `list` [`.TaskMetadata`] ]
1405 Metadata for the given IDs. Each value is a list of
1406 `.TaskMetadata` instances representing different execution
1407 attempts, ordered chronologically from first to last. Attempts
1408 where metadata was missing (not written even in the fallback extra
1409 provenance in the logs) will have `None` in this list.
1410 """
1411 result: dict[uuid.UUID, list[TaskMetadata | None]] = {}
1412 with MultiblockReader.open_in_zip(
1413 self.zf, METADATA_MB_NAME, int_size=self.header.int_size
1414 ) as mb_reader:
1415 for node_id_or_index in nodes:
1416 address_row = self.address_reader.find(node_id_or_index)
1417 metadata_by_attempt = mb_reader.read_model(
1418 address_row.addresses[METADATA_ADDRESS_INDEX],
1419 ProvenanceTaskMetadataModel,
1420 self.decompressor,
1421 )
1422 if metadata_by_attempt is not None:
1423 result[node_id_or_index] = metadata_by_attempt.attempts
1424 return result
1426 def fetch_packages(self) -> Packages:
1427 """Fetch package version information."""
1428 data = self._read_single_block_raw("packages")
1429 return Packages.fromBytes(data, format="json")
1432class ProvenanceQuantumGraphWriter:
1433 """A struct of low-level writer objects for the main components of a
1434 provenance quantum graph.
1436 Parameters
1437 ----------
1438 output_path : `str`
1439 Path to write the graph to.
1440 exit_stack : `contextlib.ExitStack`
1441 Object that can be used to manage multiple context managers.
1442 log_on_close : `LogOnClose`
1443 Factory for context managers that log when closed.
1444 predicted : `.PredictedQuantumGraphComponents`
1445 Components of the predicted graph.
1446 zstd_level : `int`, optional
1447 Compression level.
1448 cdict_data : `bytes` or `None`, optional
1449 Bytes representation of the compression dictionary used by the
1450 compressor.
1451 loop_wrapper : `~collections.abc.Callable`, optional
1452 A callable that takes an iterable and returns an equivalent one, to be
1453 used in all potentially-large loops. This can be used to add progress
1454 reporting or check for cancelation signals.
1455 log : `LsstLogAdapter`, optional
1456 Logger to use for debug messages.
1457 """
1459 def __init__(
1460 self,
1461 output_path: str,
1462 *,
1463 exit_stack: ExitStack,
1464 log_on_close: LogOnClose,
1465 predicted: PredictedQuantumGraphComponents | PredictedQuantumGraph,
1466 zstd_level: int = 10,
1467 cdict_data: bytes | None = None,
1468 loop_wrapper: LoopWrapper = pass_through,
1469 log: LsstLogAdapter | None = None,
1470 ) -> None:
1471 header = predicted.header.model_copy()
1472 header.graph_type = "provenance"
1473 if log is None:
1474 log = _LOG
1475 self.log = log
1476 self._base_writer = exit_stack.enter_context(
1477 log_on_close.wrap(
1478 BaseQuantumGraphWriter.open(
1479 output_path,
1480 header,
1481 predicted.pipeline_graph,
1482 address_filename="nodes",
1483 zstd_level=zstd_level,
1484 cdict_data=cdict_data,
1485 ),
1486 "Finishing writing provenance quantum graph.",
1487 )
1488 )
1489 self._base_writer.address_writer.addresses = [{}, {}, {}, {}]
1490 self._log_writer = exit_stack.enter_context(
1491 log_on_close.wrap(
1492 MultiblockWriter.open_in_zip(
1493 self._base_writer.zf, LOG_MB_NAME, header.int_size, use_tempfile=True
1494 ),
1495 "Copying logs into zip archive.",
1496 ),
1497 )
1498 self._base_writer.address_writer.addresses[LOG_ADDRESS_INDEX] = self._log_writer.addresses
1499 self._metadata_writer = exit_stack.enter_context(
1500 log_on_close.wrap(
1501 MultiblockWriter.open_in_zip(
1502 self._base_writer.zf, METADATA_MB_NAME, header.int_size, use_tempfile=True
1503 ),
1504 "Copying metadata into zip archive.",
1505 )
1506 )
1507 self._base_writer.address_writer.addresses[METADATA_ADDRESS_INDEX] = self._metadata_writer.addresses
1508 self._dataset_writer = exit_stack.enter_context(
1509 log_on_close.wrap(
1510 MultiblockWriter.open_in_zip(
1511 self._base_writer.zf, DATASET_MB_NAME, header.int_size, use_tempfile=True
1512 ),
1513 "Copying dataset provenance into zip archive.",
1514 )
1515 )
1516 self._base_writer.address_writer.addresses[DATASET_ADDRESS_INDEX] = self._dataset_writer.addresses
1517 self._quantum_writer = exit_stack.enter_context(
1518 log_on_close.wrap(
1519 MultiblockWriter.open_in_zip(
1520 self._base_writer.zf, QUANTUM_MB_NAME, header.int_size, use_tempfile=True
1521 ),
1522 "Copying quantum provenance into zip archive.",
1523 )
1524 )
1525 self._base_writer.address_writer.addresses[QUANTUM_ADDRESS_INDEX] = self._quantum_writer.addresses
1526 self._init_predicted_quanta(predicted)
1527 self._populate_xgraph_and_inputs(loop_wrapper)
1528 self._existing_init_outputs: set[uuid.UUID] = set()
1530 def _init_predicted_quanta(
1531 self, predicted: PredictedQuantumGraph | PredictedQuantumGraphComponents
1532 ) -> None:
1533 self._predicted_init_quanta: list[PredictedQuantumDatasetsModel] = []
1534 self._predicted_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = {}
1535 if isinstance(predicted, PredictedQuantumGraph):
1536 self._predicted_init_quanta.extend(predicted._init_quanta.values())
1537 self._predicted_quanta.update(predicted._quantum_datasets)
1538 else:
1539 self._predicted_init_quanta.extend(predicted.init_quanta.root)
1540 self._predicted_quanta.update(predicted.quantum_datasets)
1541 self._predicted_quanta.update({q.quantum_id: q for q in self._predicted_init_quanta})
1543 def _populate_xgraph_and_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1544 self._xgraph = networkx.DiGraph()
1545 self._overall_inputs: dict[uuid.UUID, PredictedDatasetModel] = {}
1546 output_dataset_ids: set[uuid.UUID] = set()
1547 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1548 if not predicted_quantum.task_label:
1549 # Skip the 'packages' producer quantum.
1550 continue
1551 output_dataset_ids.update(predicted_quantum.iter_output_dataset_ids())
1552 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1553 if not predicted_quantum.task_label:
1554 # Skip the 'packages' producer quantum.
1555 continue
1556 for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
1557 self._xgraph.add_edge(predicted_input.dataset_id, predicted_quantum.quantum_id)
1558 if predicted_input.dataset_id not in output_dataset_ids:
1559 self._overall_inputs.setdefault(predicted_input.dataset_id, predicted_input)
1560 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1561 self._xgraph.add_edge(predicted_quantum.quantum_id, predicted_output.dataset_id)
1563 @property
1564 def compressor(self) -> Compressor:
1565 """Object that should be used to compress all JSON blocks."""
1566 return self._base_writer.compressor
1568 def write_packages(self) -> None:
1569 """Write package version information to the provenance graph."""
1570 packages = Packages.fromSystem(include_all=True)
1571 data = packages.toBytes("json")
1572 self._base_writer.write_single_block("packages", data)
1574 def write_overall_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1575 """Write provenance for overall-input datasets.
1577 Parameters
1578 ----------
1579 loop_wrapper : `~collections.abc.Callable`, optional
1580 A callable that takes an iterable and returns an equivalent one, to
1581 be used in all potentially-large loops. This can be used to add
1582 progress reporting or check for cancelation signals.
1583 """
1584 for predicted_input in loop_wrapper(self._overall_inputs.values()):
1585 if predicted_input.dataset_id not in self._dataset_writer.addresses:
1586 self._dataset_writer.write_model(
1587 predicted_input.dataset_id,
1588 ProvenanceDatasetModel.from_predicted(
1589 predicted_input,
1590 producer=None,
1591 consumers=self._xgraph.successors(predicted_input.dataset_id),
1592 ),
1593 self.compressor,
1594 )
1595 del self._overall_inputs
1597 def write_init_outputs(self, assume_existence: bool = True) -> None:
1598 """Write provenance for init-output datasets and init-quanta.
1600 Parameters
1601 ----------
1602 assume_existence : `bool`, optional
1603 If `True`, just assume all init-outputs exist.
1604 """
1605 init_quanta = ProvenanceInitQuantaModel()
1606 for predicted_init_quantum in self._predicted_init_quanta:
1607 if not predicted_init_quantum.task_label:
1608 # Skip the 'packages' producer quantum.
1609 continue
1610 for predicted_output in itertools.chain.from_iterable(predicted_init_quantum.outputs.values()):
1611 provenance_output = ProvenanceDatasetModel.from_predicted(
1612 predicted_output,
1613 producer=predicted_init_quantum.quantum_id,
1614 consumers=self._xgraph.successors(predicted_output.dataset_id),
1615 )
1616 provenance_output.produced = assume_existence or (
1617 provenance_output.dataset_id in self._existing_init_outputs
1618 )
1619 self._dataset_writer.write_model(
1620 provenance_output.dataset_id, provenance_output, self.compressor
1621 )
1622 init_quanta.root.append(ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum))
1623 self._base_writer.write_single_model("init_quanta", init_quanta)
1625 def write_quantum_provenance(
1626 self, quantum_id: uuid.UUID, metadata: TaskMetadata | None, logs: ButlerLogRecords | None
1627 ) -> None:
1628 """Gather and write provenance for a quantum.
1630 Parameters
1631 ----------
1632 quantum_id : `uuid.UUID`
1633 Unique ID for the quantum.
1634 metadata : `..TaskMetadata` or `None`
1635 Task metadata.
1636 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
1637 Task logs.
1638 """
1639 predicted_quantum = self._predicted_quanta[quantum_id]
1640 provenance_models = ProvenanceQuantumScanModels.from_metadata_and_logs(
1641 predicted_quantum, metadata, logs, incomplete=False
1642 )
1643 scan_data = provenance_models.to_scan_data(predicted_quantum, compressor=self.compressor)
1644 self.write_scan_data(scan_data)
1646 def write_scan_data(self, scan_data: ProvenanceQuantumScanData) -> None:
1647 """Write the output of a quantum provenance scan to disk.
1649 Parameters
1650 ----------
1651 scan_data : `ProvenanceQuantumScanData`
1652 Result of a quantum provenance scan.
1653 """
1654 if scan_data.status is ProvenanceQuantumScanStatus.INIT:
1655 self.log.debug("Handling init-output scan for %s.", scan_data.quantum_id)
1656 self._existing_init_outputs.update(scan_data.existing_outputs)
1657 return
1658 self.log.debug("Handling quantum scan for %s.", scan_data.quantum_id)
1659 # We shouldn't need this predicted quantum after this method runs; pop
1660 # from the dict it in the hopes that'll free up some memory when we're
1661 # done.
1662 predicted_quantum = self._predicted_quanta.pop(scan_data.quantum_id)
1663 outputs: dict[uuid.UUID, bytes] = {}
1664 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1665 provenance_output = ProvenanceDatasetModel.from_predicted(
1666 predicted_output,
1667 producer=predicted_quantum.quantum_id,
1668 consumers=self._xgraph.successors(predicted_output.dataset_id),
1669 )
1670 provenance_output.produced = provenance_output.dataset_id in scan_data.existing_outputs
1671 outputs[provenance_output.dataset_id] = self.compressor.compress(
1672 provenance_output.model_dump_json().encode()
1673 )
1674 if not scan_data.quantum:
1675 scan_data.quantum = (
1676 ProvenanceQuantumModel.from_predicted(predicted_quantum).model_dump_json().encode()
1677 )
1678 if scan_data.is_compressed:
1679 scan_data.quantum = self.compressor.compress(scan_data.quantum)
1680 if not scan_data.is_compressed:
1681 scan_data.quantum = self.compressor.compress(scan_data.quantum)
1682 if scan_data.metadata:
1683 scan_data.metadata = self.compressor.compress(scan_data.metadata)
1684 if scan_data.logs:
1685 scan_data.logs = self.compressor.compress(scan_data.logs)
1686 self.log.debug("Writing quantum %s.", scan_data.quantum_id)
1687 self._quantum_writer.write_bytes(scan_data.quantum_id, scan_data.quantum)
1688 for dataset_id, dataset_data in outputs.items():
1689 self._dataset_writer.write_bytes(dataset_id, dataset_data)
1690 if scan_data.metadata:
1691 (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
1692 address = self._metadata_writer.write_bytes(scan_data.quantum_id, scan_data.metadata)
1693 self._metadata_writer.addresses[metadata_output.dataset_id] = address
1694 if scan_data.logs:
1695 (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
1696 address = self._log_writer.write_bytes(scan_data.quantum_id, scan_data.logs)
1697 self._log_writer.addresses[log_output.dataset_id] = address
1700class ProvenanceQuantumScanStatus(enum.Enum):
1701 """Status enum for quantum scanning.
1703 Note that this records the status for the *scanning* which is distinct
1704 from the status of the quantum's execution.
1705 """
1707 INCOMPLETE = enum.auto()
1708 """The quantum is not necessarily done running, and cannot be scanned
1709 conclusively yet.
1710 """
1712 ABANDONED = enum.auto()
1713 """The quantum's execution appears to have failed but we cannot rule out
1714 the possibility that it could be recovered, but we've also waited long
1715 enough (according to `ScannerTimeConfigDict.retry_timeout`) that it's time
1716 to stop trying for now.
1718 This state means `ProvenanceQuantumScanModels.from_metadata_and_logs` must
1719 be run again with ``incomplete=False``.
1720 """
1722 SUCCESSFUL = enum.auto()
1723 """The quantum was conclusively scanned and was executed successfully,
1724 unblocking scans for downstream quanta.
1725 """
1727 FAILED = enum.auto()
1728 """The quantum was conclusively scanned and failed execution, blocking
1729 scans for downstream quanta.
1730 """
1732 BLOCKED = enum.auto()
1733 """A quantum upstream of this one failed."""
1735 INIT = enum.auto()
1736 """Init quanta need special handling, because they don't have logs and
1737 metadata.
1738 """
1741@dataclasses.dataclass
1742class ProvenanceQuantumScanModels:
1743 """A struct that represents provenance information for a single quantum."""
1745 quantum_id: uuid.UUID
1746 """Unique ID for the quantum."""
1748 status: ProvenanceQuantumScanStatus = ProvenanceQuantumScanStatus.INCOMPLETE
1749 """Combined status for the scan and the execution of the quantum."""
1751 attempts: list[ProvenanceQuantumAttemptModel] = dataclasses.field(default_factory=list)
1752 """Provenance information about each attempt to run the quantum."""
1754 output_existence: dict[uuid.UUID, bool] = dataclasses.field(default_factory=dict)
1755 """Unique IDs of the output datasets mapped to whether they were actually
1756 produced.
1757 """
1759 metadata: ProvenanceTaskMetadataModel = dataclasses.field(default_factory=ProvenanceTaskMetadataModel)
1760 """Task metadata information for each attempt.
1761 """
1763 logs: ProvenanceLogRecordsModel = dataclasses.field(default_factory=ProvenanceLogRecordsModel)
1764 """Log records for each attempt.
1765 """
1767 @classmethod
1768 def from_metadata_and_logs(
1769 cls,
1770 predicted: PredictedQuantumDatasetsModel,
1771 metadata: TaskMetadata | None,
1772 logs: ButlerLogRecords | None,
1773 *,
1774 incomplete: bool = False,
1775 ) -> ProvenanceQuantumScanModels:
1776 """Construct provenance information from task metadata and logs.
1778 Parameters
1779 ----------
1780 predicted : `PredictedQuantumDatasetsModel`
1781 Information about the predicted quantum.
1782 metadata : `..TaskMetadata` or `None`
1783 Task metadata.
1784 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
1785 Task logs.
1786 incomplete : `bool`, optional
1787 If `True`, treat execution failures as possibly-incomplete quanta
1788 and do not fully process them; instead just set the status to
1789 `ProvenanceQuantumScanStatus.ABANDONED` and return.
1791 Returns
1792 -------
1793 scan_models : `ProvenanceQuantumScanModels`
1794 Struct of models that describe quantum provenance.
1796 Notes
1797 -----
1798 This method does not necessarily fully populate the `output_existence`
1799 field; it does what it can given the information in the metadata and
1800 logs, but the caller is responsible for filling in the existence status
1801 for any predicted outputs that are not present at all in that `dict`.
1802 """
1803 self = ProvenanceQuantumScanModels(predicted.quantum_id)
1804 last_attempt = ProvenanceQuantumAttemptModel()
1805 self._process_logs(predicted, logs, last_attempt, incomplete=incomplete)
1806 self._process_metadata(predicted, metadata, last_attempt, incomplete=incomplete)
1807 if self.status is ProvenanceQuantumScanStatus.ABANDONED:
1808 return self
1809 self._reconcile_attempts(last_attempt)
1810 self._extract_output_existence(predicted)
1811 return self
1813 def _process_logs(
1814 self,
1815 predicted: PredictedQuantumDatasetsModel,
1816 logs: ButlerLogRecords | None,
1817 last_attempt: ProvenanceQuantumAttemptModel,
1818 *,
1819 incomplete: bool,
1820 ) -> None:
1821 (predicted_log_dataset,) = predicted.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
1822 if logs is None:
1823 self.output_existence[predicted_log_dataset.dataset_id] = False
1824 if incomplete:
1825 self.status = ProvenanceQuantumScanStatus.ABANDONED
1826 else:
1827 self.status = ProvenanceQuantumScanStatus.FAILED
1828 else:
1829 # Set the attempt's run status to FAILED, since the default is
1830 # UNKNOWN (i.e. logs *and* metadata are missing) and we now know
1831 # the logs exist. This will usually get replaced by SUCCESSFUL
1832 # when we look for metadata next.
1833 last_attempt.status = QuantumAttemptStatus.FAILED
1834 self.output_existence[predicted_log_dataset.dataset_id] = True
1835 if logs.extra:
1836 log_extra = _ExecutionLogRecordsExtra.model_validate(logs.extra)
1837 self._extract_from_log_extra(log_extra, last_attempt=last_attempt)
1838 self.logs.attempts.append(list(logs))
1840 def _extract_from_log_extra(
1841 self,
1842 log_extra: _ExecutionLogRecordsExtra,
1843 last_attempt: ProvenanceQuantumAttemptModel | None,
1844 ) -> None:
1845 for previous_attempt_log_extra in log_extra.previous_attempts:
1846 self._extract_from_log_extra(
1847 previous_attempt_log_extra,
1848 last_attempt=None,
1849 )
1850 quantum_attempt: ProvenanceQuantumAttemptModel
1851 if last_attempt is None:
1852 # This is not the last attempt, so it must be a failure.
1853 quantum_attempt = ProvenanceQuantumAttemptModel(
1854 attempt=len(self.attempts), status=QuantumAttemptStatus.FAILED
1855 )
1856 # We also need to get the logs from this extra provenance, since
1857 # they won't be the main section of the log records.
1858 self.logs.attempts.append(log_extra.logs)
1859 # The special last attempt is only appended after we attempt to
1860 # read metadata later, but we have to append this one now.
1861 self.attempts.append(quantum_attempt)
1862 else:
1863 assert not log_extra.logs, "Logs for the last attempt should not be stored in the extra JSON."
1864 quantum_attempt = last_attempt
1865 if log_extra.exception is not None or log_extra.metadata is not None or last_attempt is None:
1866 # We won't be getting a separate metadata dataset, so anything we
1867 # might get from the metadata has to come from this extra
1868 # provenance in the logs.
1869 quantum_attempt.exception = log_extra.exception
1870 if log_extra.metadata is not None:
1871 quantum_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(log_extra.metadata)
1872 self.metadata.attempts.append(log_extra.metadata)
1873 else:
1874 self.metadata.attempts.append(None)
1875 # Regardless of whether this is the last attempt or not, we can only
1876 # get the previous_process_quanta from the log extra.
1877 quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)
1879 def _process_metadata(
1880 self,
1881 predicted: PredictedQuantumDatasetsModel,
1882 metadata: TaskMetadata | None,
1883 last_attempt: ProvenanceQuantumAttemptModel,
1884 *,
1885 incomplete: bool,
1886 ) -> None:
1887 (predicted_metadata_dataset,) = predicted.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
1888 if metadata is None:
1889 self.output_existence[predicted_metadata_dataset.dataset_id] = False
1890 if incomplete:
1891 self.status = ProvenanceQuantumScanStatus.ABANDONED
1892 else:
1893 self.status = ProvenanceQuantumScanStatus.FAILED
1894 else:
1895 self.status = ProvenanceQuantumScanStatus.SUCCESSFUL
1896 self.output_existence[predicted_metadata_dataset.dataset_id] = True
1897 last_attempt.status = QuantumAttemptStatus.SUCCESSFUL
1898 try:
1899 # Int conversion guards against spurious conversion to
1900 # float that can apparently sometimes happen in
1901 # TaskMetadata.
1902 last_attempt.caveats = QuantumSuccessCaveats(int(metadata["quantum"]["caveats"]))
1903 except LookupError:
1904 pass
1905 try:
1906 last_attempt.exception = ExceptionInfo._from_metadata(
1907 metadata[predicted.task_label]["failure"]
1908 )
1909 except LookupError:
1910 pass
1911 last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)
1912 self.metadata.attempts.append(metadata)
1914 def _reconcile_attempts(self, last_attempt: ProvenanceQuantumAttemptModel) -> None:
1915 last_attempt.attempt = len(self.attempts)
1916 self.attempts.append(last_attempt)
1917 assert self.status is not ProvenanceQuantumScanStatus.INCOMPLETE
1918 assert self.status is not ProvenanceQuantumScanStatus.ABANDONED
1919 if len(self.logs.attempts) < len(self.attempts):
1920 # Logs were not found for this attempt; must have been a hard error
1921 # that kept the `finally` block from running or otherwise
1922 # interrupted the writing of the logs.
1923 self.logs.attempts.append(None)
1924 if self.status is ProvenanceQuantumScanStatus.SUCCESSFUL:
1925 # But we found the metadata! Either that hard error happened
1926 # at a very unlucky time (in between those two writes), or
1927 # something even weirder happened.
1928 self.attempts[-1].status = QuantumAttemptStatus.ABORTED_SUCCESS
1929 else:
1930 self.attempts[-1].status = QuantumAttemptStatus.FAILED
1931 if len(self.metadata.attempts) < len(self.attempts):
1932 # Metadata missing usually just means a failure. In any case, the
1933 # status will already be correct, either because it was set to a
1934 # failure when we read the logs, or left at UNKNOWN if there were
1935 # no logs. Note that scanners never process BLOCKED quanta at all.
1936 self.metadata.attempts.append(None)
1937 assert len(self.logs.attempts) == len(self.attempts) or len(self.metadata.attempts) == len(
1938 self.attempts
1939 ), (
1940 "The only way we can add more than one quantum attempt is by "
1941 "extracting info stored with the logs, and that always appends "
1942 "a log attempt and a metadata attempt, so this must be a bug in "
1943 "this class."
1944 )
1946 def _extract_output_existence(self, predicted: PredictedQuantumDatasetsModel) -> None:
1947 try:
1948 outputs_put = self.metadata.attempts[-1]["quantum"].getArray("outputs") # type: ignore[index]
1949 except (
1950 IndexError, # metadata.attempts is empty
1951 TypeError, # metadata.attempts[-1] is None
1952 LookupError, # no 'quantum' entry in metadata or 'outputs' in that
1953 ):
1954 pass
1955 else:
1956 for id_str in ensure_iterable(outputs_put):
1957 self.output_existence[uuid.UUID(id_str)] = True
1958 # If the metadata told us what it wrote, anything not in that
1959 # list was not written.
1960 for predicted_output in itertools.chain.from_iterable(predicted.outputs.values()):
1961 self.output_existence.setdefault(predicted_output.dataset_id, False)
1963 def to_scan_data(
1964 self: ProvenanceQuantumScanModels,
1965 predicted_quantum: PredictedQuantumDatasetsModel,
1966 compressor: Compressor | None = None,
1967 ) -> ProvenanceQuantumScanData:
1968 """Convert these models to JSON data.
1970 Parameters
1971 ----------
1972 predicted_quantum : `PredictedQuantumDatasetsModel`
1973 Information about the predicted quantum.
1974 compressor : `Compressor`
1975 Object that can compress bytes.
1977 Returns
1978 -------
1979 scan_data : `ProvenanceQuantumScanData`
1980 Scan information ready for serialization.
1981 """
1982 quantum: ProvenanceInitQuantumModel | ProvenanceQuantumModel
1983 if self.status is ProvenanceQuantumScanStatus.INIT:
1984 quantum = ProvenanceInitQuantumModel.from_predicted(predicted_quantum)
1985 else:
1986 quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum)
1987 quantum.attempts = self.attempts
1988 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1989 if predicted_output.dataset_id not in self.output_existence:
1990 raise RuntimeError(
1991 "Logic bug in provenance gathering or execution invariants: "
1992 f"no existence information for output {predicted_output.dataset_id} "
1993 f"({predicted_output.dataset_type_name}@{predicted_output.data_coordinate})."
1994 )
1995 data = ProvenanceQuantumScanData(
1996 self.quantum_id,
1997 self.status,
1998 existing_outputs={
1999 dataset_id for dataset_id, was_produced in self.output_existence.items() if was_produced
2000 },
2001 quantum=quantum.model_dump_json().encode(),
2002 logs=self.logs.model_dump_json().encode() if self.logs.attempts else b"",
2003 metadata=self.metadata.model_dump_json().encode() if self.metadata.attempts else b"",
2004 )
2005 if compressor is not None:
2006 data.compress(compressor)
2007 return data
2010@dataclasses.dataclass
2011class ProvenanceQuantumScanData:
2012 """A struct that represents ready-for-serialization provenance information
2013 for a single quantum.
2014 """
2016 quantum_id: uuid.UUID
2017 """Unique ID for the quantum."""
2019 status: ProvenanceQuantumScanStatus
2020 """Combined status for the scan and the execution of the quantum."""
2022 existing_outputs: set[uuid.UUID] = dataclasses.field(default_factory=set)
2023 """Unique IDs of the output datasets that were actually written."""
2025 quantum: bytes = b""
2026 """Serialized quantum provenance model.
2028 This may be empty for quanta that had no attempts.
2029 """
2031 metadata: bytes = b""
2032 """Serialized task metadata."""
2034 logs: bytes = b""
2035 """Serialized logs."""
2037 is_compressed: bool = False
2038 """Whether the `quantum`, `metadata`, and `log` attributes are
2039 compressed.
2040 """
2042 def compress(self, compressor: Compressor) -> None:
2043 """Compress the data in this struct if it has not been compressed
2044 already.
2046 Parameters
2047 ----------
2048 compressor : `Compressor`
2049 Object with a ``compress`` method that takes and returns `bytes`.
2050 """
2051 if not self.is_compressed:
2052 self.quantum = compressor.compress(self.quantum)
2053 self.logs = compressor.compress(self.logs) if self.logs else b""
2054 self.metadata = compressor.compress(self.metadata) if self.metadata else b""
2055 self.is_compressed = True