Coverage for python / lsst / pipe / base / quantum_graph / _provenance.py: 31%
567 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-15 06:32 -0300
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_T = TypeVar("_T")
99LoopWrapper: TypeAlias = Callable[[Iterable[_T]], Iterable[_T]]
101_LOG = getLogger(__file__)
103DATASET_ADDRESS_INDEX = 0
104QUANTUM_ADDRESS_INDEX = 1
105LOG_ADDRESS_INDEX = 2
106METADATA_ADDRESS_INDEX = 3
108DATASET_MB_NAME = "datasets"
109QUANTUM_MB_NAME = "quanta"
110LOG_MB_NAME = "logs"
111METADATA_MB_NAME = "metadata"
114def pass_through(arg: _T) -> _T:
115 return arg
118class ProvenanceDatasetInfo(DatasetInfo):
119 """A typed dictionary that annotates the attributes of the NetworkX graph
120 node data for a provenance dataset.
122 Since NetworkX types are not generic over their node mapping type, this has
123 to be used explicitly, e.g.::
125 node_data: ProvenanceDatasetInfo = xgraph.nodes[dataset_id]
127 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
128 """
130 dataset_id: uuid.UUID
131 """Unique identifier for the dataset."""
133 produced: bool
134 """Whether this dataset was produced (vs. only predicted).
136 This is always `True` for overall input datasets. It is also `True` for
137 datasets that were produced and then removed before/during transfer back to
138 the central butler repository, so it may not reflect the continued
139 existence of the dataset.
140 """
143class ProvenanceQuantumInfo(QuantumInfo):
144 """A typed dictionary that annotates the attributes of the NetworkX graph
145 node data for a provenance quantum.
147 Since NetworkX types are not generic over their node mapping type, this has
148 to be used explicitly, e.g.::
150 node_data: ProvenanceQuantumInfo = xgraph.nodes[quantum_id]
152 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph` or
153 `ProvenanceQuantumGraph.quantum_only_xgraph`
154 """
156 status: QuantumAttemptStatus
157 """Enumerated status for the quantum.
159 This corresponds to the last attempt to run this quantum, or
160 `QuantumAttemptStatus.BLOCKED` if there were no attempts.
161 """
163 caveats: QuantumSuccessCaveats | None
164 """Flags indicating caveats on successful quanta.
166 This corresponds to the last attempt to run this quantum.
167 """
169 exception: ExceptionInfo | None
170 """Information about an exception raised when the quantum was executing.
172 This corresponds to the last attempt to run this quantum.
173 """
175 resource_usage: QuantumResourceUsage | None
176 """Resource usage information (timing, memory use) for this quantum.
178 This corresponds to the last attempt to run this quantum.
179 """
181 attempts: list[ProvenanceQuantumAttemptModel]
182 """Information about each attempt to run this quantum.
184 An entry is added merely if the quantum *should* have been attempted; an
185 empty `list` is used only for quanta that were blocked by an upstream
186 failure.
187 """
190class ProvenanceInitQuantumInfo(TypedDict):
191 """A typed dictionary that annotates the attributes of the NetworkX graph
192 node data for a provenance init quantum.
194 Since NetworkX types are not generic over their node mapping type, this has
195 to be used explicitly, e.g.::
197 node_data: ProvenanceInitQuantumInfo = xgraph.nodes[quantum_id]
199 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
200 """
202 data_id: DataCoordinate
203 """Data ID of the quantum.
205 This is always an empty ID; this key exists to allow init-quanta and
206 regular quanta to be treated more similarly.
207 """
209 task_label: str
210 """Label of the task for this quantum."""
212 pipeline_node: TaskInitNode
213 """Node in the pipeline graph for this task's init-only step."""
216class ProvenanceDatasetModel(PredictedDatasetModel):
217 """Data model for the datasets in a provenance quantum graph file."""
219 produced: bool
220 """Whether this dataset was produced (vs. only predicted).
222 This is always `True` for overall input datasets. It is also `True` for
223 datasets that were produced and then removed before/during transfer back to
224 the central butler repository, so it may not reflect the continued
225 existence of the dataset.
226 """
228 producer: uuid.UUID | None = None
229 """ID of the quantum that produced this dataset.
231 This is `None` for overall inputs to the graph.
232 """
234 consumers: list[uuid.UUID] = pydantic.Field(default_factory=list)
235 """IDs of quanta that were predicted to consume this dataset."""
237 @property
238 def node_id(self) -> uuid.UUID:
239 """Alias for the dataset ID."""
240 return self.dataset_id
242 @classmethod
243 def from_predicted(
244 cls,
245 predicted: PredictedDatasetModel,
246 producer: uuid.UUID | None = None,
247 consumers: Iterable[uuid.UUID] = (),
248 ) -> ProvenanceDatasetModel:
249 """Construct from a predicted dataset model.
251 Parameters
252 ----------
253 predicted : `PredictedDatasetModel`
254 Information about the dataset from the predicted graph.
255 producer : `uuid.UUID` or `None`, optional
256 ID of the quantum that was predicted to produce this dataset.
257 consumers : `~collections.abc.Iterable` [`uuid.UUID`], optional
258 IDs of the quanta that were predicted to consume this dataset.
260 Returns
261 -------
262 provenance : `ProvenanceDatasetModel`
263 Provenance dataset model.
265 Notes
266 -----
267 This initializes `produced` to `True` when ``producer is None`` and
268 `False` otherwise, on the assumption that it will be updated later.
269 """
270 return cls.model_construct(
271 dataset_id=predicted.dataset_id,
272 dataset_type_name=predicted.dataset_type_name,
273 data_coordinate=predicted.data_coordinate,
274 run=predicted.run,
275 produced=(producer is None), # if it's not produced by this QG, it's an overall input
276 producer=producer,
277 consumers=list(consumers),
278 )
280 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
281 """Add this dataset and its edges to quanta to a provenance graph.
283 Parameters
284 ----------
285 graph : `ProvenanceQuantumGraph`
286 Graph to update in place.
288 Notes
289 -----
290 This method adds:
292 - a ``bipartite_xgraph`` dataset node with full attributes;
293 - ``bipartite_xgraph`` edges to adjacent quanta (which adds quantum
294 nodes with no attributes), without populating edge attributes;
295 - ``quantum_only_xgraph`` edges for each pair of quanta in which one
296 produces this dataset and another consumes it (this also adds quantum
297 nodes with no attributes).
298 """
299 dataset_type_node = graph.pipeline_graph.dataset_types[self.dataset_type_name]
300 data_id = DataCoordinate.from_full_values(dataset_type_node.dimensions, tuple(self.data_coordinate))
301 graph._bipartite_xgraph.add_node(
302 self.dataset_id,
303 data_id=data_id,
304 dataset_type_name=self.dataset_type_name,
305 pipeline_node=dataset_type_node,
306 run=self.run,
307 produced=self.produced,
308 )
309 if self.producer is not None:
310 graph._bipartite_xgraph.add_edge(self.producer, self.dataset_id)
311 for consumer_id in self.consumers:
312 graph._bipartite_xgraph.add_edge(self.dataset_id, consumer_id)
313 if self.producer is not None:
314 graph._quantum_only_xgraph.add_edge(self.producer, consumer_id)
315 graph._datasets_by_type[self.dataset_type_name][data_id] = self.dataset_id
317 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
318 # when we inherit those docstrings in our public classes.
319 if "sphinx" in sys.modules and not TYPE_CHECKING:
321 def copy(self, *args: Any, **kwargs: Any) -> Any:
322 """See `pydantic.BaseModel.copy`."""
323 return super().copy(*args, **kwargs)
325 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
326 """See `pydantic.BaseModel.model_dump`."""
327 return super().model_dump(*args, **kwargs)
329 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
330 """See `pydantic.BaseModel.model_dump_json`."""
331 return super().model_dump(*args, **kwargs)
333 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
334 """See `pydantic.BaseModel.model_copy`."""
335 return super().model_copy(*args, **kwargs)
337 @classmethod
338 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
339 """See `pydantic.BaseModel.model_construct`."""
340 return super().model_construct(*args, **kwargs)
342 @classmethod
343 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
344 """See `pydantic.BaseModel.model_json_schema`."""
345 return super().model_json_schema(*args, **kwargs)
347 @classmethod
348 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
349 """See `pydantic.BaseModel.model_validate`."""
350 return super().model_validate(*args, **kwargs)
352 @classmethod
353 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
354 """See `pydantic.BaseModel.model_validate_json`."""
355 return super().model_validate_json(*args, **kwargs)
357 @classmethod
358 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
359 """See `pydantic.BaseModel.model_validate_strings`."""
360 return super().model_validate_strings(*args, **kwargs)
363class ProvenanceQuantumAttemptModel(pydantic.BaseModel):
364 """Data model for a now-superseded attempt to run a quantum in a
365 provenance quantum graph file.
366 """
368 attempt: int = 0
369 """Counter incremented for every attempt to execute this quantum."""
371 status: QuantumAttemptStatus = QuantumAttemptStatus.UNKNOWN
372 """Enumerated status for the quantum."""
374 caveats: QuantumSuccessCaveats | None = None
375 """Flags indicating caveats on successful quanta."""
377 exception: ExceptionInfo | None = None
378 """Information about an exception raised when the quantum was executing."""
380 resource_usage: QuantumResourceUsage | None = None
381 """Resource usage information (timing, memory use) for this quantum."""
383 previous_process_quanta: list[uuid.UUID] = pydantic.Field(default_factory=list)
384 """The IDs of other quanta previously executed in the same process as this
385 one.
386 """
388 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
389 # when we inherit those docstrings in our public classes.
390 if "sphinx" in sys.modules and not TYPE_CHECKING:
392 def copy(self, *args: Any, **kwargs: Any) -> Any:
393 """See `pydantic.BaseModel.copy`."""
394 return super().copy(*args, **kwargs)
396 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
397 """See `pydantic.BaseModel.model_dump`."""
398 return super().model_dump(*args, **kwargs)
400 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
401 """See `pydantic.BaseModel.model_dump_json`."""
402 return super().model_dump(*args, **kwargs)
404 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
405 """See `pydantic.BaseModel.model_copy`."""
406 return super().model_copy(*args, **kwargs)
408 @classmethod
409 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
410 """See `pydantic.BaseModel.model_construct`."""
411 return super().model_construct(*args, **kwargs)
413 @classmethod
414 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
415 """See `pydantic.BaseModel.model_json_schema`."""
416 return super().model_json_schema(*args, **kwargs)
418 @classmethod
419 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
420 """See `pydantic.BaseModel.model_validate`."""
421 return super().model_validate(*args, **kwargs)
423 @classmethod
424 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
425 """See `pydantic.BaseModel.model_validate_json`."""
426 return super().model_validate_json(*args, **kwargs)
428 @classmethod
429 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
430 """See `pydantic.BaseModel.model_validate_strings`."""
431 return super().model_validate_strings(*args, **kwargs)
434class ProvenanceLogRecordsModel(pydantic.BaseModel):
435 """Data model for storing execution logs in a provenance quantum graph
436 file.
437 """
439 attempts: list[list[ButlerLogRecord] | None] = pydantic.Field(default_factory=list)
440 """Logs from attempts to run this task, ordered chronologically from first
441 to last.
442 """
444 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
445 # when we inherit those docstrings in our public classes.
446 if "sphinx" in sys.modules and not TYPE_CHECKING:
448 def copy(self, *args: Any, **kwargs: Any) -> Any:
449 """See `pydantic.BaseModel.copy`."""
450 return super().copy(*args, **kwargs)
452 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
453 """See `pydantic.BaseModel.model_dump`."""
454 return super().model_dump(*args, **kwargs)
456 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
457 """See `pydantic.BaseModel.model_dump_json`."""
458 return super().model_dump(*args, **kwargs)
460 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
461 """See `pydantic.BaseModel.model_copy`."""
462 return super().model_copy(*args, **kwargs)
464 @classmethod
465 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
466 """See `pydantic.BaseModel.model_construct`."""
467 return super().model_construct(*args, **kwargs)
469 @classmethod
470 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
471 """See `pydantic.BaseModel.model_json_schema`."""
472 return super().model_json_schema(*args, **kwargs)
474 @classmethod
475 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
476 """See `pydantic.BaseModel.model_validate`."""
477 return super().model_validate(*args, **kwargs)
479 @classmethod
480 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
481 """See `pydantic.BaseModel.model_validate_json`."""
482 return super().model_validate_json(*args, **kwargs)
484 @classmethod
485 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
486 """See `pydantic.BaseModel.model_validate_strings`."""
487 return super().model_validate_strings(*args, **kwargs)
490class ProvenanceTaskMetadataModel(pydantic.BaseModel):
491 """Data model for storing task metadata in a provenance quantum graph
492 file.
493 """
495 attempts: list[TaskMetadata | None] = pydantic.Field(default_factory=list)
496 """Metadata from attempts to run this task, ordered chronologically from
497 first to last.
498 """
500 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
501 # when we inherit those docstrings in our public classes.
502 if "sphinx" in sys.modules and not TYPE_CHECKING:
504 def copy(self, *args: Any, **kwargs: Any) -> Any:
505 """See `pydantic.BaseModel.copy`."""
506 return super().copy(*args, **kwargs)
508 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
509 """See `pydantic.BaseModel.model_dump`."""
510 return super().model_dump(*args, **kwargs)
512 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
513 """See `pydantic.BaseModel.model_dump_json`."""
514 return super().model_dump(*args, **kwargs)
516 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
517 """See `pydantic.BaseModel.model_copy`."""
518 return super().model_copy(*args, **kwargs)
520 @classmethod
521 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
522 """See `pydantic.BaseModel.model_construct`."""
523 return super().model_construct(*args, **kwargs)
525 @classmethod
526 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
527 """See `pydantic.BaseModel.model_json_schema`."""
528 return super().model_json_schema(*args, **kwargs)
530 @classmethod
531 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
532 """See `pydantic.BaseModel.model_validate`."""
533 return super().model_validate(*args, **kwargs)
535 @classmethod
536 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
537 """See `pydantic.BaseModel.model_validate_json`."""
538 return super().model_validate_json(*args, **kwargs)
540 @classmethod
541 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
542 """See `pydantic.BaseModel.model_validate_strings`."""
543 return super().model_validate_strings(*args, **kwargs)
546class ProvenanceQuantumModel(pydantic.BaseModel):
547 """Data model for the quanta in a provenance quantum graph file."""
549 quantum_id: uuid.UUID
550 """Unique identifier for the quantum."""
552 task_label: TaskLabel
553 """Name of the type of this dataset."""
555 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
556 """The full values (required and implied) of this dataset's data ID."""
558 inputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
559 """IDs of the datasets predicted to be consumed by this quantum, grouped by
560 connection name.
561 """
563 outputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
564 """IDs of the datasets predicted to be produced by this quantum, grouped by
565 connection name.
566 """
568 attempts: list[ProvenanceQuantumAttemptModel] = pydantic.Field(default_factory=list)
569 """Provenance for all attempts to execute this quantum, ordered
570 chronologically from first to last.
572 An entry is added merely if the quantum *should* have been attempted; an
573 empty `list` is used only for quanta that were blocked by an upstream
574 failure.
575 """
577 @property
578 def node_id(self) -> uuid.UUID:
579 """Alias for the quantum ID."""
580 return self.quantum_id
582 @classmethod
583 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceQuantumModel:
584 """Construct from a predicted quantum model.
586 Parameters
587 ----------
588 predicted : `PredictedQuantumDatasetsModel`
589 Information about the quantum from the predicted graph.
591 Returns
592 -------
593 provenance : `ProvenanceQuantumModel`
594 Provenance quantum model.
595 """
596 inputs = {
597 connection_name: [d.dataset_id for d in predicted_inputs]
598 for connection_name, predicted_inputs in predicted.inputs.items()
599 }
600 outputs = {
601 connection_name: [d.dataset_id for d in predicted_outputs]
602 for connection_name, predicted_outputs in predicted.outputs.items()
603 }
604 return cls(
605 quantum_id=predicted.quantum_id,
606 task_label=predicted.task_label,
607 data_coordinate=predicted.data_coordinate,
608 inputs=inputs,
609 outputs=outputs,
610 )
612 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
613 """Add this quantum and its edges to datasets to a provenance graph.
615 Parameters
616 ----------
617 graph : `ProvenanceQuantumGraph`
618 Graph to update in place.
620 Notes
621 -----
622 This method adds:
624 - a ``bipartite_xgraph`` quantum node with full attributes;
625 - a ``quantum_only_xgraph`` quantum node with full attributes;
626 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
627 nodes with no attributes), while populating those edge attributes;
628 - ``quantum_only_xgraph`` edges to any adjacent quantum that has also
629 already been loaded.
630 """
631 task_node = graph.pipeline_graph.tasks[self.task_label]
632 data_id = DataCoordinate.from_full_values(task_node.dimensions, tuple(self.data_coordinate))
633 last_attempt = (
634 self.attempts[-1]
635 if self.attempts
636 else ProvenanceQuantumAttemptModel(status=QuantumAttemptStatus.BLOCKED)
637 )
638 graph._bipartite_xgraph.add_node(
639 self.quantum_id,
640 data_id=data_id,
641 task_label=self.task_label,
642 pipeline_node=task_node,
643 status=last_attempt.status,
644 caveats=last_attempt.caveats,
645 exception=last_attempt.exception,
646 resource_usage=last_attempt.resource_usage,
647 attempts=self.attempts,
648 )
649 for connection_name, dataset_ids in self.inputs.items():
650 read_edge = task_node.get_input_edge(connection_name)
651 for dataset_id in dataset_ids:
652 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
653 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
654 "pipeline_edges", []
655 ).append(read_edge)
656 for connection_name, dataset_ids in self.outputs.items():
657 write_edge = task_node.get_output_edge(connection_name)
658 for dataset_id in dataset_ids:
659 graph._bipartite_xgraph.add_edge(
660 self.quantum_id,
661 dataset_id,
662 is_read=False,
663 # There can only be one pipeline edge for an output.
664 pipeline_edges=[write_edge],
665 )
666 graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
667 graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
668 for dataset_id in graph._bipartite_xgraph.predecessors(self.quantum_id):
669 for upstream_quantum_id in graph._bipartite_xgraph.predecessors(dataset_id):
670 graph._quantum_only_xgraph.add_edge(upstream_quantum_id, self.quantum_id)
671 for dataset_id in graph._bipartite_xgraph.successors(self.quantum_id):
672 for downstream_quantum_id in graph._bipartite_xgraph.successors(dataset_id):
673 graph._quantum_only_xgraph.add_edge(self.quantum_id, downstream_quantum_id)
675 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
676 # when we inherit those docstrings in our public classes.
677 if "sphinx" in sys.modules and not TYPE_CHECKING:
679 def copy(self, *args: Any, **kwargs: Any) -> Any:
680 """See `pydantic.BaseModel.copy`."""
681 return super().copy(*args, **kwargs)
683 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
684 """See `pydantic.BaseModel.model_dump`."""
685 return super().model_dump(*args, **kwargs)
687 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
688 """See `pydantic.BaseModel.model_dump_json`."""
689 return super().model_dump(*args, **kwargs)
691 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
692 """See `pydantic.BaseModel.model_copy`."""
693 return super().model_copy(*args, **kwargs)
695 @classmethod
696 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
697 """See `pydantic.BaseModel.model_construct`."""
698 return super().model_construct(*args, **kwargs)
700 @classmethod
701 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
702 """See `pydantic.BaseModel.model_json_schema`."""
703 return super().model_json_schema(*args, **kwargs)
705 @classmethod
706 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
707 """See `pydantic.BaseModel.model_validate`."""
708 return super().model_validate(*args, **kwargs)
710 @classmethod
711 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
712 """See `pydantic.BaseModel.model_validate_json`."""
713 return super().model_validate_json(*args, **kwargs)
715 @classmethod
716 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
717 """See `pydantic.BaseModel.model_validate_strings`."""
718 return super().model_validate_strings(*args, **kwargs)
721class ProvenanceInitQuantumModel(pydantic.BaseModel):
722 """Data model for the special "init" quanta in a provenance quantum graph
723 file.
724 """
726 quantum_id: uuid.UUID
727 """Unique identifier for the quantum."""
729 task_label: TaskLabel
730 """Name of the type of this dataset.
732 This is always a parent dataset type name, not a component.
734 Note that full dataset type definitions are stored in the pipeline graph.
735 """
737 inputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
738 """IDs of the datasets predicted to be consumed by this quantum, grouped by
739 connection name.
740 """
742 outputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
743 """IDs of the datasets predicted to be produced by this quantum, grouped by
744 connection name.
745 """
747 @classmethod
748 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceInitQuantumModel:
749 """Construct from a predicted quantum model.
751 Parameters
752 ----------
753 predicted : `PredictedQuantumDatasetsModel`
754 Information about the quantum from the predicted graph.
756 Returns
757 -------
758 provenance : `ProvenanceInitQuantumModel`
759 Provenance init quantum model.
760 """
761 inputs = {
762 connection_name: predicted_inputs[0].dataset_id
763 for connection_name, predicted_inputs in predicted.inputs.items()
764 }
765 outputs = {
766 connection_name: predicted_outputs[0].dataset_id
767 for connection_name, predicted_outputs in predicted.outputs.items()
768 }
769 return cls(
770 quantum_id=predicted.quantum_id,
771 task_label=predicted.task_label,
772 inputs=inputs,
773 outputs=outputs,
774 )
776 def _add_to_graph(self, graph: ProvenanceQuantumGraph, empty_data_id: DataCoordinate) -> None:
777 """Add this quantum and its edges to datasets to a provenance graph.
779 Parameters
780 ----------
781 graph : `ProvenanceQuantumGraph`
782 Graph to update in place.
783 empty_data_id : `lsst.daf.butler.DataCoordinate`
784 The empty data ID for the appropriate dimension universe.
786 Notes
787 -----
788 This method adds:
790 - a ``bipartite_xgraph`` quantum node with full attributes;
791 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
792 nodes with no attributes), while populating those edge attributes;
793 """
794 task_init_node = graph.pipeline_graph.tasks[self.task_label].init
795 graph._bipartite_xgraph.add_node(
796 self.quantum_id, data_id=empty_data_id, task_label=self.task_label, pipeline_node=task_init_node
797 )
798 for connection_name, dataset_id in self.inputs.items():
799 read_edge = task_init_node.get_input_edge(connection_name)
800 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
801 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
802 "pipeline_edges", []
803 ).append(read_edge)
804 for connection_name, dataset_id in self.outputs.items():
805 write_edge = task_init_node.get_output_edge(connection_name)
806 graph._bipartite_xgraph.add_edge(
807 self.quantum_id,
808 dataset_id,
809 is_read=False,
810 # There can only be one pipeline edge for an output.
811 pipeline_edges=[write_edge],
812 )
813 graph._init_quanta[self.task_label] = self.quantum_id
815 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
816 # when we inherit those docstrings in our public classes.
817 if "sphinx" in sys.modules and not TYPE_CHECKING:
819 def copy(self, *args: Any, **kwargs: Any) -> Any:
820 """See `pydantic.BaseModel.copy`."""
821 return super().copy(*args, **kwargs)
823 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
824 """See `pydantic.BaseModel.model_dump`."""
825 return super().model_dump(*args, **kwargs)
827 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
828 """See `pydantic.BaseModel.model_dump_json`."""
829 return super().model_dump(*args, **kwargs)
831 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
832 """See `pydantic.BaseModel.model_copy`."""
833 return super().model_copy(*args, **kwargs)
835 @classmethod
836 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
837 """See `pydantic.BaseModel.model_construct`."""
838 return super().model_construct(*args, **kwargs)
840 @classmethod
841 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
842 """See `pydantic.BaseModel.model_json_schema`."""
843 return super().model_json_schema(*args, **kwargs)
845 @classmethod
846 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
847 """See `pydantic.BaseModel.model_validate`."""
848 return super().model_validate(*args, **kwargs)
850 @classmethod
851 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
852 """See `pydantic.BaseModel.model_validate_json`."""
853 return super().model_validate_json(*args, **kwargs)
855 @classmethod
856 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
857 """See `pydantic.BaseModel.model_validate_strings`."""
858 return super().model_validate_strings(*args, **kwargs)
861class ProvenanceInitQuantaModel(pydantic.RootModel):
862 """Data model for the init quanta in a provenance graph."""
864 root: list[ProvenanceInitQuantumModel] = pydantic.Field(default_factory=list)
865 """List of special "init" quanta, one for each task."""
867 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
868 """Add this quantum and its edges to datasets to a provenance graph.
870 Parameters
871 ----------
872 graph : `ProvenanceQuantumGraph`
873 Graph to update in place.
874 """
875 empty_data_id = DataCoordinate.make_empty(graph.pipeline_graph.universe)
876 for init_quantum in self.root:
877 init_quantum._add_to_graph(graph, empty_data_id=empty_data_id)
879 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
880 # when we inherit those docstrings in our public classes.
881 if "sphinx" in sys.modules and not TYPE_CHECKING:
883 def copy(self, *args: Any, **kwargs: Any) -> Any:
884 """See `pydantic.BaseModel.copy`."""
885 return super().copy(*args, **kwargs)
887 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
888 """See `pydantic.BaseModel.model_dump`."""
889 return super().model_dump(*args, **kwargs)
891 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
892 """See `pydantic.BaseModel.model_dump_json`."""
893 return super().model_dump(*args, **kwargs)
895 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
896 """See `pydantic.BaseModel.model_copy`."""
897 return super().model_copy(*args, **kwargs)
899 @classmethod
900 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
901 """See `pydantic.BaseModel.model_construct`."""
902 return super().model_construct(*args, **kwargs)
904 @classmethod
905 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
906 """See `pydantic.BaseModel.model_json_schema`."""
907 return super().model_json_schema(*args, **kwargs)
909 @classmethod
910 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
911 """See `pydantic.BaseModel.model_validate`."""
912 return super().model_validate(*args, **kwargs)
914 @classmethod
915 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
916 """See `pydantic.BaseModel.model_validate_json`."""
917 return super().model_validate_json(*args, **kwargs)
919 @classmethod
920 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
921 """See `pydantic.BaseModel.model_validate_strings`."""
922 return super().model_validate_strings(*args, **kwargs)
925class ProvenanceQuantumGraph(BaseQuantumGraph):
926 """A quantum graph that represents processing that has already been
927 executed.
929 Parameters
930 ----------
931 header : `HeaderModel`
932 General metadata shared with other quantum graph types.
933 pipeline_graph : `.pipeline_graph.PipelineGraph`
934 Graph of tasks and dataset types. May contain a superset of the tasks
935 and dataset types that actually have quanta and datasets in the quantum
936 graph.
938 Notes
939 -----
940 A provenance quantum graph is generally obtained via the
941 `ProvenanceQuantumGraphReader.graph` attribute, which is updated in-place
942 as information is read from disk.
943 """
945 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph) -> None:
946 super().__init__(header, pipeline_graph)
947 self._init_quanta: dict[TaskLabel, uuid.UUID] = {}
948 self._quantum_only_xgraph = networkx.DiGraph()
949 self._bipartite_xgraph = networkx.DiGraph()
950 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {
951 task_label: {} for task_label in self.pipeline_graph.tasks.keys()
952 }
953 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {
954 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
955 }
957 @property
958 def init_quanta(self) -> Mapping[TaskLabel, uuid.UUID]:
959 """A mapping from task label to the ID of the special init quantum for
960 that task.
962 This is populated by the ``init_quanta`` component. Additional
963 information about each init quantum can be found by using the ID to
964 look up node attributes in the `bipartite_xgraph`, i.e.::
966 info: ProvenanceInitQuantumInfo = qg.bipartite_xgraph.nodes[id]
967 """
968 return self._init_quanta
970 @property
971 def quanta_by_task(self) -> Mapping[TaskLabel, Mapping[DataCoordinate, uuid.UUID]]:
972 """A nested mapping of all quanta, keyed first by task name and then by
973 data ID.
975 Notes
976 -----
977 This is populated one quantum at a time as they are read. All tasks in
978 the pipeline graph are included, even if none of their quanta were
979 loaded (i.e. nested mappings may be empty).
981 The returned object may be an internal dictionary; as the type
982 annotation indicates, it should not be modified in place.
983 """
984 return self._quanta_by_task_label
986 @property
987 def datasets_by_type(self) -> Mapping[DatasetTypeName, Mapping[DataCoordinate, uuid.UUID]]:
988 """A nested mapping of all datasets, keyed first by dataset type name
989 and then by data ID.
991 Notes
992 -----
993 This is populated one dataset at a time as they are read. All dataset
994 types in the pipeline graph are included, even if none of their
995 datasets were loaded (i.e. nested mappings may be empty).
997 The returned object may be an internal dictionary; as the type
998 annotation indicates, it should not be modified in place.
999 """
1000 return self._datasets_by_type
1002 @property
1003 def quantum_only_xgraph(self) -> networkx.DiGraph:
1004 """A directed acyclic graph with quanta as nodes (and datasets elided).
1006 Notes
1007 -----
1008 Node keys are quantum UUIDs, and are populated one quantum at a time as
1009 they are loaded. Loading quanta (via
1010 `ProvenanceQuantumGraphReader.read_quanta`) will add the loaded nodes
1011 with full attributes and add edges to adjacent nodes with no
1012 attributes. Loading datasets (via
1013 `ProvenanceQuantumGraphReader.read_datasets`) will also add edges and
1014 nodes with no attributes.
1016 Node attributes are described by the `ProvenanceQuantumInfo` types.
1018 This graph does not include special "init" quanta.
1020 The returned object is a read-only view of an internal one.
1021 """
1022 return self._quantum_only_xgraph.copy(as_view=True)
1024 @property
1025 def bipartite_xgraph(self) -> networkx.DiGraph:
1026 """A directed acyclic graph with quantum and dataset nodes.
1028 Notes
1029 -----
1030 Node keys are quantum or dataset UUIDs, and are populated one quantum
1031 or dataset at a time as they are loaded. Loading quanta (via
1032 `ProvenanceQuantumGraphReader.read_quanta`) or datasets (via
1033 `ProvenanceQuantumGraphReader.read_datasets`) will load those nodes
1034 with full attributes and edges to adjacent nodes with no attributes.
1035 Loading quanta necessary to populate edge attributes.
1037 Node attributes are described by the
1038 `ProvenanceQuantumInfo`, `ProvenanceInitQuantumInfo`, and
1039 `ProvenanceDatasetInfo` types.
1041 This graph includes init-input and init-output datasets, but it does
1042 *not* reflect the dependency between each task's special "init" quantum
1043 and its runtime quanta (as this would require edges between quanta, and
1044 that would break the "bipartite" property).
1046 The returned object is a read-only view of an internal one.
1047 """
1048 return self._bipartite_xgraph.copy(as_view=True)
1050 def make_quantum_table(self) -> astropy.table.Table:
1051 """Construct an `astropy.table.Table` with a tabular summary of the
1052 quanta.
1054 Returns
1055 -------
1056 table : `astropy.table.Table`
1057 A table view of the quantum information. This only includes
1058 counts of status categories and caveats, not any per-data-ID
1059 detail.
1061 Notes
1062 -----
1063 Success caveats in the table are represented by their
1064 `~QuantumSuccessCaveats.concise` form, so when pretty-printing this
1065 table for users, the `~QuantumSuccessCaveats.legend` should generally
1066 be printed as well.
1067 """
1068 rows = []
1069 for task_label, quanta_for_task in self.quanta_by_task.items():
1070 if not self.header.n_task_quanta[task_label]:
1071 continue
1072 status_counts = Counter[QuantumAttemptStatus](
1073 self._quantum_only_xgraph.nodes[q]["status"] for q in quanta_for_task.values()
1074 )
1075 caveat_counts = Counter[QuantumSuccessCaveats | None](
1076 self._quantum_only_xgraph.nodes[q]["caveats"] for q in quanta_for_task.values()
1077 )
1078 caveat_counts.pop(QuantumSuccessCaveats.NO_CAVEATS, None)
1079 caveat_counts.pop(None, None)
1080 if len(caveat_counts) > 1:
1081 caveats = "(multiple)"
1082 elif len(caveat_counts) == 1:
1083 ((code, count),) = caveat_counts.items()
1084 # MyPy can't tell that the pop(None, None) above makes None
1085 # impossible here.
1086 caveats = f"{code.concise()}({count})" # type: ignore[union-attr]
1087 else:
1088 caveats = ""
1089 rows.append(
1090 {
1091 "Task": task_label,
1092 "Unknown": status_counts.get(QuantumAttemptStatus.UNKNOWN, 0),
1093 "Successful": status_counts.get(QuantumAttemptStatus.SUCCESSFUL, 0),
1094 "Caveats": caveats,
1095 "Blocked": status_counts.get(QuantumAttemptStatus.BLOCKED, 0),
1096 "Failed": status_counts.get(QuantumAttemptStatus.FAILED, 0),
1097 "TOTAL": len(quanta_for_task),
1098 "EXPECTED": self.header.n_task_quanta[task_label],
1099 }
1100 )
1101 return astropy.table.Table(rows)
1103 def make_exception_table(self) -> astropy.table.Table:
1104 """Construct an `astropy.table.Table` with counts for each exception
1105 type raised by each task.
1107 At present this only includes information from partial-outputs-error
1108 successes, since exception information for failures is not tracked.
1109 This may change in the future.
1111 Returns
1112 -------
1113 table : `astropy.table.Table`
1114 A table with columns for task label, exception type, and counts.
1115 """
1116 rows = []
1117 for task_label, quanta_for_task in self.quanta_by_task.items():
1118 counts_by_type = Counter(
1119 exc_info.type_name
1120 for q in quanta_for_task.values()
1121 if (exc_info := self._quantum_only_xgraph.nodes[q]["exception"]) is not None
1122 )
1123 for type_name, count in counts_by_type.items():
1124 rows.append({"Task": task_label, "Exception": type_name, "Count": count})
1125 return astropy.table.Table(rows)
1127 def make_task_resource_usage_table(
1128 self, task_label: TaskLabel, include_data_ids: bool = False
1129 ) -> astropy.table.Table:
1130 """Make a table of resource usage for a single task.
1132 Parameters
1133 ----------
1134 task_label : `str`
1135 Label of the task to extract resource usage for.
1136 include_data_ids : `bool`, optional
1137 Whether to also include data ID columns.
1139 Returns
1140 -------
1141 table : `astropy.table.Table`
1142 A table with columns for quantum ID and all fields in
1143 `QuantumResourceUsage`.
1144 """
1145 quanta_for_task = self.quanta_by_task[task_label]
1146 dtype_terms: list[tuple[str, np.dtype]] = [("quantum_id", np.dtype((np.void, 16)))]
1147 if include_data_ids:
1148 dimensions = self.pipeline_graph.tasks[task_label].dimensions
1149 for dimension_name in dimensions.data_coordinate_keys:
1150 dtype = np.dtype(self.pipeline_graph.universe.dimensions[dimension_name].primary_key.pytype)
1151 dtype_terms.append((dimension_name, dtype))
1152 fields = QuantumResourceUsage.get_numpy_fields()
1153 dtype_terms.extend(fields.items())
1154 row_dtype = np.dtype(dtype_terms)
1155 rows: list[object] = []
1156 for data_id, quantum_id in quanta_for_task.items():
1157 info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1158 if (resource_usage := info["resource_usage"]) is not None:
1159 row: tuple[object, ...] = (quantum_id.bytes,)
1160 if include_data_ids:
1161 row += data_id.full_values
1162 row += resource_usage.get_numpy_row()
1163 rows.append(row)
1164 array = np.array(rows, dtype=row_dtype)
1165 return astropy.table.Table(array, units=QuantumResourceUsage.get_units())
1168@dataclasses.dataclass
1169class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1170 """A helper class for reading provenance quantum graphs.
1172 Notes
1173 -----
1174 The `open` context manager should be used to construct new instances.
1175 Instances cannot be used after the context manager exits, except to access
1176 the `graph` attribute`.
1178 The various ``read_*`` methods in this class update the `graph` attribute
1179 in place.
1180 """
1182 graph: ProvenanceQuantumGraph = dataclasses.field(init=False)
1183 """Loaded provenance graph, populated in place as components are read."""
1185 @classmethod
1186 @contextmanager
1187 def open(
1188 cls,
1189 uri: ResourcePathExpression,
1190 *,
1191 page_size: int | None = None,
1192 import_mode: TaskImportMode = TaskImportMode.DO_NOT_IMPORT,
1193 ) -> Iterator[ProvenanceQuantumGraphReader]:
1194 """Construct a reader from a URI.
1196 Parameters
1197 ----------
1198 uri : convertible to `lsst.resources.ResourcePath`
1199 URI to open. Should have a ``.qg`` extension.
1200 page_size : `int`, optional
1201 Approximate number of bytes to read at once from address files and
1202 multi-block files. Note that this does not set a page size for
1203 *all* reads, but it does affect the smallest, most numerous reads.
1204 Can also be set via the ``LSST_QG_PAGE_SIZE`` environment variable.
1205 import_mode : `.pipeline_graph.TaskImportMode`, optional
1206 How to handle importing the task classes referenced in the pipeline
1207 graph.
1209 Returns
1210 -------
1211 reader : `contextlib.AbstractContextManager` [ \
1212 `ProvenanceQuantumGraphReader` ]
1213 A context manager that returns the reader when entered.
1214 """
1215 with cls._open(
1216 uri,
1217 graph_type="provenance",
1218 address_filename="nodes",
1219 page_size=page_size,
1220 import_mode=import_mode,
1221 n_addresses=4,
1222 ) as self:
1223 yield self
1225 def __post_init__(self) -> None:
1226 self.graph = ProvenanceQuantumGraph(self.header, self.pipeline_graph)
1228 def read_init_quanta(self) -> None:
1229 """Read the thin graph, with all edge information and categorization of
1230 quanta by task label.
1231 """
1232 init_quanta = self._read_single_block("init_quanta", ProvenanceInitQuantaModel)
1233 for init_quantum in init_quanta.root:
1234 self.graph._init_quanta[init_quantum.task_label] = init_quantum.quantum_id
1235 init_quanta._add_to_graph(self.graph)
1237 def read_full_graph(self) -> None:
1238 """Read all bipartite edges and all quantum and dataset node
1239 attributes, fully populating the `graph` attribute.
1241 Notes
1242 -----
1243 This does not read logs, metadata, or packages ; those must always be
1244 fetched explicitly.
1245 """
1246 self.read_init_quanta()
1247 self.read_datasets()
1248 self.read_quanta()
1250 def read_datasets(self, datasets: Iterable[uuid.UUID] | None = None) -> None:
1251 """Read information about the given datasets.
1253 Parameters
1254 ----------
1255 datasets : `~collections.abc.Iterable` [`uuid.UUID`], optional
1256 Iterable of dataset IDs to load. If not provided, all datasets
1257 will be loaded. The UUIDs and indices of quanta will be ignored.
1258 """
1259 self._read_nodes(datasets, DATASET_ADDRESS_INDEX, DATASET_MB_NAME, ProvenanceDatasetModel)
1261 def read_quanta(self, quanta: Iterable[uuid.UUID] | None = None) -> None:
1262 """Read information about the given quanta.
1264 Parameters
1265 ----------
1266 quanta : `~collections.abc.Iterable` [`uuid.UUID`], optional
1267 Iterable of quantum IDs to load. If not provided, all quanta will
1268 be loaded. The UUIDs and indices of datasets and special init
1269 quanta will be ignored.
1270 """
1271 self._read_nodes(quanta, QUANTUM_ADDRESS_INDEX, QUANTUM_MB_NAME, ProvenanceQuantumModel)
1273 def _read_nodes(
1274 self,
1275 nodes: Iterable[uuid.UUID] | None,
1276 address_index: int,
1277 mb_name: str,
1278 model_type: type[ProvenanceDatasetModel] | type[ProvenanceQuantumModel],
1279 ) -> None:
1280 node: ProvenanceDatasetModel | ProvenanceQuantumModel | None
1281 if nodes is None:
1282 self.address_reader.read_all()
1283 nodes = self.address_reader.rows.keys()
1284 for node in MultiblockReader.read_all_models_in_zip(
1285 self.zf,
1286 mb_name,
1287 model_type,
1288 self.decompressor,
1289 int_size=self.header.int_size,
1290 page_size=self.page_size,
1291 ):
1292 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(node.node_id, {}):
1293 # Use the old node to reduce memory usage (since it might
1294 # also have other outstanding reference holders).
1295 continue
1296 node._add_to_graph(self.graph)
1297 return
1298 with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1299 for node_id_or_index in nodes:
1300 address_row = self.address_reader.find(node_id_or_index)
1301 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1302 # Use the old node to reduce memory usage (since it might
1303 # also have other outstanding reference holders).
1304 continue
1305 node = mb_reader.read_model(
1306 address_row.addresses[address_index], model_type, self.decompressor
1307 )
1308 if node is not None:
1309 node._add_to_graph(self.graph)
1311 def fetch_logs(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[ButlerLogRecords | None]]:
1312 """Fetch log datasets.
1314 Parameters
1315 ----------
1316 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1317 UUIDs of the log datasets themselves or of the quanta they
1318 correspond to.
1320 Returns
1321 -------
1322 logs : `dict` [ `uuid.UUID`, `list` [\
1323 `lsst.daf.butler.ButlerLogRecords` or `None`] ]
1324 Logs for the given IDs. Each value is a list of
1325 `lsst.daf.butler.ButlerLogRecords` instances representing different
1326 execution attempts, ordered chronologically from first to last.
1327 Attempts where logs were missing will have `None` in this list.
1328 """
1329 result: dict[uuid.UUID, list[ButlerLogRecords | None]] = {}
1330 with MultiblockReader.open_in_zip(self.zf, LOG_MB_NAME, int_size=self.header.int_size) as mb_reader:
1331 for node_id_or_index in nodes:
1332 address_row = self.address_reader.find(node_id_or_index)
1333 logs_by_attempt = mb_reader.read_model(
1334 address_row.addresses[LOG_ADDRESS_INDEX], ProvenanceLogRecordsModel, self.decompressor
1335 )
1336 if logs_by_attempt is not None:
1337 result[node_id_or_index] = [
1338 ButlerLogRecords.from_records(attempt_logs) if attempt_logs is not None else None
1339 for attempt_logs in logs_by_attempt.attempts
1340 ]
1341 return result
1343 def fetch_metadata(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[TaskMetadata | None]]:
1344 """Fetch metadata datasets.
1346 Parameters
1347 ----------
1348 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1349 UUIDs of the metadata datasets themselves or of the quanta they
1350 correspond to.
1352 Returns
1353 -------
1354 metadata : `dict` [ `uuid.UUID`, `list` [`.TaskMetadata`] ]
1355 Metadata for the given IDs. Each value is a list of
1356 `.TaskMetadata` instances representing different execution
1357 attempts, ordered chronologically from first to last. Attempts
1358 where metadata was missing (not written even in the fallback extra
1359 provenance in the logs) will have `None` in this list.
1360 """
1361 result: dict[uuid.UUID, list[TaskMetadata | None]] = {}
1362 with MultiblockReader.open_in_zip(
1363 self.zf, METADATA_MB_NAME, int_size=self.header.int_size
1364 ) as mb_reader:
1365 for node_id_or_index in nodes:
1366 address_row = self.address_reader.find(node_id_or_index)
1367 metadata_by_attempt = mb_reader.read_model(
1368 address_row.addresses[METADATA_ADDRESS_INDEX],
1369 ProvenanceTaskMetadataModel,
1370 self.decompressor,
1371 )
1372 if metadata_by_attempt is not None:
1373 result[node_id_or_index] = metadata_by_attempt.attempts
1374 return result
1376 def fetch_packages(self) -> Packages:
1377 """Fetch package version information."""
1378 data = self._read_single_block_raw("packages")
1379 return Packages.fromBytes(data, format="json")
1382class ProvenanceQuantumGraphWriter:
1383 """A struct of low-level writer objects for the main components of a
1384 provenance quantum graph.
1386 Parameters
1387 ----------
1388 output_path : `str`
1389 Path to write the graph to.
1390 exit_stack : `contextlib.ExitStack`
1391 Object that can be used to manage multiple context managers.
1392 log_on_close : `LogOnClose`
1393 Factory for context managers that log when closed.
1394 predicted : `.PredictedQuantumGraphComponents`
1395 Components of the predicted graph.
1396 zstd_level : `int`, optional
1397 Compression level.
1398 cdict_data : `bytes` or `None`, optional
1399 Bytes representation of the compression dictionary used by the
1400 compressor.
1401 loop_wrapper : `~collections.abc.Callable`, optional
1402 A callable that takes an iterable and returns an equivalent one, to be
1403 used in all potentially-large loops. This can be used to add progress
1404 reporting or check for cancelation signals.
1405 log : `LsstLogAdapter`, optional
1406 Logger to use for debug messages.
1407 """
1409 def __init__(
1410 self,
1411 output_path: str,
1412 *,
1413 exit_stack: ExitStack,
1414 log_on_close: LogOnClose,
1415 predicted: PredictedQuantumGraphComponents | PredictedQuantumGraph,
1416 zstd_level: int = 10,
1417 cdict_data: bytes | None = None,
1418 loop_wrapper: LoopWrapper = pass_through,
1419 log: LsstLogAdapter | None = None,
1420 ) -> None:
1421 header = predicted.header.model_copy()
1422 header.graph_type = "provenance"
1423 if log is None:
1424 log = _LOG
1425 self.log = log
1426 self._base_writer = exit_stack.enter_context(
1427 log_on_close.wrap(
1428 BaseQuantumGraphWriter.open(
1429 output_path,
1430 header,
1431 predicted.pipeline_graph,
1432 address_filename="nodes",
1433 zstd_level=zstd_level,
1434 cdict_data=cdict_data,
1435 ),
1436 "Finishing writing provenance quantum graph.",
1437 )
1438 )
1439 self._base_writer.address_writer.addresses = [{}, {}, {}, {}]
1440 self._log_writer = exit_stack.enter_context(
1441 log_on_close.wrap(
1442 MultiblockWriter.open_in_zip(
1443 self._base_writer.zf, LOG_MB_NAME, header.int_size, use_tempfile=True
1444 ),
1445 "Copying logs into zip archive.",
1446 ),
1447 )
1448 self._base_writer.address_writer.addresses[LOG_ADDRESS_INDEX] = self._log_writer.addresses
1449 self._metadata_writer = exit_stack.enter_context(
1450 log_on_close.wrap(
1451 MultiblockWriter.open_in_zip(
1452 self._base_writer.zf, METADATA_MB_NAME, header.int_size, use_tempfile=True
1453 ),
1454 "Copying metadata into zip archive.",
1455 )
1456 )
1457 self._base_writer.address_writer.addresses[METADATA_ADDRESS_INDEX] = self._metadata_writer.addresses
1458 self._dataset_writer = exit_stack.enter_context(
1459 log_on_close.wrap(
1460 MultiblockWriter.open_in_zip(
1461 self._base_writer.zf, DATASET_MB_NAME, header.int_size, use_tempfile=True
1462 ),
1463 "Copying dataset provenance into zip archive.",
1464 )
1465 )
1466 self._base_writer.address_writer.addresses[DATASET_ADDRESS_INDEX] = self._dataset_writer.addresses
1467 self._quantum_writer = exit_stack.enter_context(
1468 log_on_close.wrap(
1469 MultiblockWriter.open_in_zip(
1470 self._base_writer.zf, QUANTUM_MB_NAME, header.int_size, use_tempfile=True
1471 ),
1472 "Copying quantum provenance into zip archive.",
1473 )
1474 )
1475 self._base_writer.address_writer.addresses[QUANTUM_ADDRESS_INDEX] = self._quantum_writer.addresses
1476 self._init_predicted_quanta(predicted)
1477 self._populate_xgraph_and_inputs(loop_wrapper)
1478 self._existing_init_outputs: set[uuid.UUID] = set()
1480 def _init_predicted_quanta(
1481 self, predicted: PredictedQuantumGraph | PredictedQuantumGraphComponents
1482 ) -> None:
1483 self._predicted_init_quanta: list[PredictedQuantumDatasetsModel] = []
1484 self._predicted_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = {}
1485 if isinstance(predicted, PredictedQuantumGraph):
1486 self._predicted_init_quanta.extend(predicted._init_quanta.values())
1487 self._predicted_quanta.update(predicted._quantum_datasets)
1488 else:
1489 self._predicted_init_quanta.extend(predicted.init_quanta.root)
1490 self._predicted_quanta.update(predicted.quantum_datasets)
1491 self._predicted_quanta.update({q.quantum_id: q for q in self._predicted_init_quanta})
1493 def _populate_xgraph_and_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1494 self._xgraph = networkx.DiGraph()
1495 self._overall_inputs: dict[uuid.UUID, PredictedDatasetModel] = {}
1496 output_dataset_ids: set[uuid.UUID] = set()
1497 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1498 if not predicted_quantum.task_label:
1499 # Skip the 'packages' producer quantum.
1500 continue
1501 output_dataset_ids.update(predicted_quantum.iter_output_dataset_ids())
1502 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1503 if not predicted_quantum.task_label:
1504 # Skip the 'packages' producer quantum.
1505 continue
1506 for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
1507 self._xgraph.add_edge(predicted_input.dataset_id, predicted_quantum.quantum_id)
1508 if predicted_input.dataset_id not in output_dataset_ids:
1509 self._overall_inputs.setdefault(predicted_input.dataset_id, predicted_input)
1510 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1511 self._xgraph.add_edge(predicted_quantum.quantum_id, predicted_output.dataset_id)
1513 @property
1514 def compressor(self) -> Compressor:
1515 """Object that should be used to compress all JSON blocks."""
1516 return self._base_writer.compressor
1518 def write_packages(self) -> None:
1519 """Write package version information to the provenance graph."""
1520 packages = Packages.fromSystem(include_all=True)
1521 data = packages.toBytes("json")
1522 self._base_writer.write_single_block("packages", data)
1524 def write_overall_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1525 """Write provenance for overall-input datasets.
1527 Parameters
1528 ----------
1529 loop_wrapper : `~collections.abc.Callable`, optional
1530 A callable that takes an iterable and returns an equivalent one, to
1531 be used in all potentially-large loops. This can be used to add
1532 progress reporting or check for cancelation signals.
1533 """
1534 for predicted_input in loop_wrapper(self._overall_inputs.values()):
1535 if predicted_input.dataset_id not in self._dataset_writer.addresses:
1536 self._dataset_writer.write_model(
1537 predicted_input.dataset_id,
1538 ProvenanceDatasetModel.from_predicted(
1539 predicted_input,
1540 producer=None,
1541 consumers=self._xgraph.successors(predicted_input.dataset_id),
1542 ),
1543 self.compressor,
1544 )
1545 del self._overall_inputs
1547 def write_init_outputs(self, assume_existence: bool = True) -> None:
1548 """Write provenance for init-output datasets and init-quanta.
1550 Parameters
1551 ----------
1552 assume_existence : `bool`, optional
1553 If `True`, just assume all init-outputs exist.
1554 """
1555 init_quanta = ProvenanceInitQuantaModel()
1556 for predicted_init_quantum in self._predicted_init_quanta:
1557 if not predicted_init_quantum.task_label:
1558 # Skip the 'packages' producer quantum.
1559 continue
1560 for predicted_output in itertools.chain.from_iterable(predicted_init_quantum.outputs.values()):
1561 provenance_output = ProvenanceDatasetModel.from_predicted(
1562 predicted_output,
1563 producer=predicted_init_quantum.quantum_id,
1564 consumers=self._xgraph.successors(predicted_output.dataset_id),
1565 )
1566 provenance_output.produced = assume_existence or (
1567 provenance_output.dataset_id in self._existing_init_outputs
1568 )
1569 self._dataset_writer.write_model(
1570 provenance_output.dataset_id, provenance_output, self.compressor
1571 )
1572 init_quanta.root.append(ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum))
1573 self._base_writer.write_single_model("init_quanta", init_quanta)
1575 def write_quantum_provenance(
1576 self, quantum_id: uuid.UUID, metadata: TaskMetadata | None, logs: ButlerLogRecords | None
1577 ) -> None:
1578 """Gather and write provenance for a quantum.
1580 Parameters
1581 ----------
1582 quantum_id : `uuid.UUID`
1583 Unique ID for the quantum.
1584 metadata : `..TaskMetadata` or `None`
1585 Task metadata.
1586 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
1587 Task logs.
1588 """
1589 predicted_quantum = self._predicted_quanta[quantum_id]
1590 provenance_models = ProvenanceQuantumScanModels.from_metadata_and_logs(
1591 predicted_quantum, metadata, logs, assume_complete=True
1592 )
1593 scan_data = provenance_models.to_scan_data(predicted_quantum, compressor=self.compressor)
1594 self.write_scan_data(scan_data)
1596 def write_scan_data(self, scan_data: ProvenanceQuantumScanData) -> None:
1597 """Write the output of a quantum provenance scan to disk.
1599 Parameters
1600 ----------
1601 scan_data : `ProvenanceQuantumScanData`
1602 Result of a quantum provenance scan.
1603 """
1604 if scan_data.status is ProvenanceQuantumScanStatus.INIT:
1605 self.log.debug("Handling init-output scan for %s.", scan_data.quantum_id)
1606 self._existing_init_outputs.update(scan_data.existing_outputs)
1607 return
1608 self.log.debug("Handling quantum scan for %s.", scan_data.quantum_id)
1609 # We shouldn't need this predicted quantum after this method runs; pop
1610 # from the dict it in the hopes that'll free up some memory when we're
1611 # done.
1612 predicted_quantum = self._predicted_quanta.pop(scan_data.quantum_id)
1613 outputs: dict[uuid.UUID, bytes] = {}
1614 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1615 provenance_output = ProvenanceDatasetModel.from_predicted(
1616 predicted_output,
1617 producer=predicted_quantum.quantum_id,
1618 consumers=self._xgraph.successors(predicted_output.dataset_id),
1619 )
1620 provenance_output.produced = provenance_output.dataset_id in scan_data.existing_outputs
1621 outputs[provenance_output.dataset_id] = self.compressor.compress(
1622 provenance_output.model_dump_json().encode()
1623 )
1624 if not scan_data.quantum:
1625 scan_data.quantum = (
1626 ProvenanceQuantumModel.from_predicted(predicted_quantum).model_dump_json().encode()
1627 )
1628 if scan_data.is_compressed:
1629 scan_data.quantum = self.compressor.compress(scan_data.quantum)
1630 if not scan_data.is_compressed:
1631 scan_data.quantum = self.compressor.compress(scan_data.quantum)
1632 if scan_data.metadata:
1633 scan_data.metadata = self.compressor.compress(scan_data.metadata)
1634 if scan_data.logs:
1635 scan_data.logs = self.compressor.compress(scan_data.logs)
1636 self.log.debug("Writing quantum %s.", scan_data.quantum_id)
1637 self._quantum_writer.write_bytes(scan_data.quantum_id, scan_data.quantum)
1638 for dataset_id, dataset_data in outputs.items():
1639 self._dataset_writer.write_bytes(dataset_id, dataset_data)
1640 if scan_data.metadata:
1641 (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
1642 address = self._metadata_writer.write_bytes(scan_data.quantum_id, scan_data.metadata)
1643 self._metadata_writer.addresses[metadata_output.dataset_id] = address
1644 if scan_data.logs:
1645 (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
1646 address = self._log_writer.write_bytes(scan_data.quantum_id, scan_data.logs)
1647 self._log_writer.addresses[log_output.dataset_id] = address
1650class ProvenanceQuantumScanStatus(enum.Enum):
1651 """Status enum for quantum scanning.
1653 Note that this records the status for the *scanning* which is distinct
1654 from the status of the quantum's execution.
1655 """
1657 INCOMPLETE = enum.auto()
1658 """The quantum is not necessarily done running, and cannot be scanned
1659 conclusively yet.
1660 """
1662 ABANDONED = enum.auto()
1663 """The quantum's execution appears to have failed but we cannot rule out
1664 the possibility that it could be recovered, but we've also waited long
1665 enough (according to `ScannerTimeConfigDict.retry_timeout`) that it's time
1666 to stop trying for now.
1668 This state means a later run with `ScannerConfig.assume_complete` is
1669 required.
1670 """
1672 SUCCESSFUL = enum.auto()
1673 """The quantum was conclusively scanned and was executed successfully,
1674 unblocking scans for downstream quanta.
1675 """
1677 FAILED = enum.auto()
1678 """The quantum was conclusively scanned and failed execution, blocking
1679 scans for downstream quanta.
1680 """
1682 BLOCKED = enum.auto()
1683 """A quantum upstream of this one failed."""
1685 INIT = enum.auto()
1686 """Init quanta need special handling, because they don't have logs and
1687 metadata.
1688 """
1691@dataclasses.dataclass
1692class ProvenanceQuantumScanModels:
1693 """A struct that represents provenance information for a single quantum."""
1695 quantum_id: uuid.UUID
1696 """Unique ID for the quantum."""
1698 status: ProvenanceQuantumScanStatus = ProvenanceQuantumScanStatus.INCOMPLETE
1699 """Combined status for the scan and the execution of the quantum."""
1701 attempts: list[ProvenanceQuantumAttemptModel] = dataclasses.field(default_factory=list)
1702 """Provenance information about each attempt to run the quantum."""
1704 output_existence: dict[uuid.UUID, bool] = dataclasses.field(default_factory=dict)
1705 """Unique IDs of the output datasets mapped to whether they were actually
1706 produced.
1707 """
1709 metadata: ProvenanceTaskMetadataModel = dataclasses.field(default_factory=ProvenanceTaskMetadataModel)
1710 """Task metadata information for each attempt.
1711 """
1713 logs: ProvenanceLogRecordsModel = dataclasses.field(default_factory=ProvenanceLogRecordsModel)
1714 """Log records for each attempt.
1715 """
1717 @classmethod
1718 def from_metadata_and_logs(
1719 cls,
1720 predicted: PredictedQuantumDatasetsModel,
1721 metadata: TaskMetadata | None,
1722 logs: ButlerLogRecords | None,
1723 *,
1724 assume_complete: bool = True,
1725 ) -> ProvenanceQuantumScanModels:
1726 """Construct provenance information from task metadata and logs.
1728 Parameters
1729 ----------
1730 predicted : `PredictedQuantumDatasetsModel`
1731 Information about the predicted quantum.
1732 metadata : `..TaskMetadata` or `None`
1733 Task metadata.
1734 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
1735 Task logs.
1736 assume_complete : `bool`, optional
1737 If `False`, treat execution failures as possibly-incomplete quanta
1738 and do not fully process them; instead just set the status to
1739 `ProvenanceQuantumScanStatus.ABANDONED` and return.
1741 Returns
1742 -------
1743 scan_models : `ProvenanceQuantumScanModels`
1744 Struct of models that describe quantum provenance.
1746 Notes
1747 -----
1748 This method does not necessarily fully populate the `output_existence`
1749 field; it does what it can given the information in the metadata and
1750 logs, but the caller is responsible for filling in the existence status
1751 for any predicted outputs that are not present at all in that `dict`.
1752 """
1753 self = ProvenanceQuantumScanModels(predicted.quantum_id)
1754 last_attempt = ProvenanceQuantumAttemptModel()
1755 self._process_logs(predicted, logs, last_attempt, assume_complete=assume_complete)
1756 self._process_metadata(predicted, metadata, last_attempt, assume_complete=assume_complete)
1757 if self.status is ProvenanceQuantumScanStatus.ABANDONED:
1758 return self
1759 self._reconcile_attempts(last_attempt)
1760 self._extract_output_existence(predicted)
1761 return self
1763 def _process_logs(
1764 self,
1765 predicted: PredictedQuantumDatasetsModel,
1766 logs: ButlerLogRecords | None,
1767 last_attempt: ProvenanceQuantumAttemptModel,
1768 *,
1769 assume_complete: bool,
1770 ) -> None:
1771 (predicted_log_dataset,) = predicted.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
1772 if logs is None:
1773 self.output_existence[predicted_log_dataset.dataset_id] = False
1774 if assume_complete:
1775 self.status = ProvenanceQuantumScanStatus.FAILED
1776 else:
1777 self.status = ProvenanceQuantumScanStatus.ABANDONED
1778 else:
1779 # Set the attempt's run status to FAILED, since the default is
1780 # UNKNOWN (i.e. logs *and* metadata are missing) and we now know
1781 # the logs exist. This will usually get replaced by SUCCESSFUL
1782 # when we look for metadata next.
1783 last_attempt.status = QuantumAttemptStatus.FAILED
1784 self.output_existence[predicted_log_dataset.dataset_id] = True
1785 if logs.extra:
1786 log_extra = _ExecutionLogRecordsExtra.model_validate(logs.extra)
1787 self._extract_from_log_extra(log_extra, last_attempt=last_attempt)
1788 self.logs.attempts.append(list(logs))
1790 def _extract_from_log_extra(
1791 self,
1792 log_extra: _ExecutionLogRecordsExtra,
1793 last_attempt: ProvenanceQuantumAttemptModel | None,
1794 ) -> None:
1795 for previous_attempt_log_extra in log_extra.previous_attempts:
1796 self._extract_from_log_extra(
1797 previous_attempt_log_extra,
1798 last_attempt=None,
1799 )
1800 quantum_attempt: ProvenanceQuantumAttemptModel
1801 if last_attempt is None:
1802 # This is not the last attempt, so it must be a failure.
1803 quantum_attempt = ProvenanceQuantumAttemptModel(
1804 attempt=len(self.attempts), status=QuantumAttemptStatus.FAILED
1805 )
1806 # We also need to get the logs from this extra provenance, since
1807 # they won't be the main section of the log records.
1808 self.logs.attempts.append(log_extra.logs)
1809 # The special last attempt is only appended after we attempt to
1810 # read metadata later, but we have to append this one now.
1811 self.attempts.append(quantum_attempt)
1812 else:
1813 assert not log_extra.logs, "Logs for the last attempt should not be stored in the extra JSON."
1814 quantum_attempt = last_attempt
1815 if log_extra.exception is not None or log_extra.metadata is not None or last_attempt is None:
1816 # We won't be getting a separate metadata dataset, so anything we
1817 # might get from the metadata has to come from this extra
1818 # provenance in the logs.
1819 quantum_attempt.exception = log_extra.exception
1820 if log_extra.metadata is not None:
1821 quantum_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(log_extra.metadata)
1822 self.metadata.attempts.append(log_extra.metadata)
1823 else:
1824 self.metadata.attempts.append(None)
1825 # Regardless of whether this is the last attempt or not, we can only
1826 # get the previous_process_quanta from the log extra.
1827 quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)
1829 def _process_metadata(
1830 self,
1831 predicted: PredictedQuantumDatasetsModel,
1832 metadata: TaskMetadata | None,
1833 last_attempt: ProvenanceQuantumAttemptModel,
1834 *,
1835 assume_complete: bool,
1836 ) -> None:
1837 (predicted_metadata_dataset,) = predicted.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
1838 if metadata is None:
1839 self.output_existence[predicted_metadata_dataset.dataset_id] = False
1840 if assume_complete:
1841 self.status = ProvenanceQuantumScanStatus.FAILED
1842 else:
1843 self.status = ProvenanceQuantumScanStatus.ABANDONED
1844 else:
1845 self.status = ProvenanceQuantumScanStatus.SUCCESSFUL
1846 self.output_existence[predicted_metadata_dataset.dataset_id] = True
1847 last_attempt.status = QuantumAttemptStatus.SUCCESSFUL
1848 try:
1849 # Int conversion guards against spurious conversion to
1850 # float that can apparently sometimes happen in
1851 # TaskMetadata.
1852 last_attempt.caveats = QuantumSuccessCaveats(int(metadata["quantum"]["caveats"]))
1853 except LookupError:
1854 pass
1855 try:
1856 last_attempt.exception = ExceptionInfo._from_metadata(
1857 metadata[predicted.task_label]["failure"]
1858 )
1859 except LookupError:
1860 pass
1861 last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)
1862 self.metadata.attempts.append(metadata)
1864 def _reconcile_attempts(self, last_attempt: ProvenanceQuantumAttemptModel) -> None:
1865 last_attempt.attempt = len(self.attempts)
1866 self.attempts.append(last_attempt)
1867 assert self.status is not ProvenanceQuantumScanStatus.INCOMPLETE
1868 assert self.status is not ProvenanceQuantumScanStatus.ABANDONED
1869 if len(self.logs.attempts) < len(self.attempts):
1870 # Logs were not found for this attempt; must have been a hard error
1871 # that kept the `finally` block from running or otherwise
1872 # interrupted the writing of the logs.
1873 self.logs.attempts.append(None)
1874 if self.status is ProvenanceQuantumScanStatus.SUCCESSFUL:
1875 # But we found the metadata! Either that hard error happened
1876 # at a very unlucky time (in between those two writes), or
1877 # something even weirder happened.
1878 self.attempts[-1].status = QuantumAttemptStatus.LOGS_MISSING
1879 else:
1880 self.attempts[-1].status = QuantumAttemptStatus.FAILED
1881 if len(self.metadata.attempts) < len(self.attempts):
1882 # Metadata missing usually just means a failure. In any case, the
1883 # status will already be correct, either because it was set to a
1884 # failure when we read the logs, or left at UNKNOWN if there were
1885 # no logs. Note that scanners never process BLOCKED quanta at all.
1886 self.metadata.attempts.append(None)
1887 assert len(self.logs.attempts) == len(self.attempts) or len(self.metadata.attempts) == len(
1888 self.attempts
1889 ), (
1890 "The only way we can add more than one quantum attempt is by "
1891 "extracting info stored with the logs, and that always appends "
1892 "a log attempt and a metadata attempt, so this must be a bug in "
1893 "this class."
1894 )
1896 def _extract_output_existence(self, predicted: PredictedQuantumDatasetsModel) -> None:
1897 try:
1898 outputs_put = self.metadata.attempts[-1]["quantum"].getArray("outputs") # type: ignore[index]
1899 except (
1900 IndexError, # metadata.attempts is empty
1901 TypeError, # metadata.attempts[-1] is None
1902 LookupError, # no 'quantum' entry in metadata or 'outputs' in that
1903 ):
1904 pass
1905 else:
1906 for id_str in ensure_iterable(outputs_put):
1907 self.output_existence[uuid.UUID(id_str)] = True
1908 # If the metadata told us what it wrote, anything not in that
1909 # list was not written.
1910 for predicted_output in itertools.chain.from_iterable(predicted.outputs.values()):
1911 self.output_existence.setdefault(predicted_output.dataset_id, False)
1913 def to_scan_data(
1914 self: ProvenanceQuantumScanModels,
1915 predicted_quantum: PredictedQuantumDatasetsModel,
1916 compressor: Compressor | None = None,
1917 ) -> ProvenanceQuantumScanData:
1918 """Convert these models to JSON data.
1920 Parameters
1921 ----------
1922 predicted_quantum : `PredictedQuantumDatasetsModel`
1923 Information about the predicted quantum.
1924 compressor : `Compressor`
1925 Object that can compress bytes.
1927 Returns
1928 -------
1929 scan_data : `ProvenanceQuantumScanData`
1930 Scan information ready for serialization.
1931 """
1932 quantum: ProvenanceInitQuantumModel | ProvenanceQuantumModel
1933 if self.status is ProvenanceQuantumScanStatus.INIT:
1934 quantum = ProvenanceInitQuantumModel.from_predicted(predicted_quantum)
1935 else:
1936 quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum)
1937 quantum.attempts = self.attempts
1938 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1939 if predicted_output.dataset_id not in self.output_existence:
1940 raise RuntimeError(
1941 "Logic bug in provenance gathering or execution invariants: "
1942 f"no existence information for output {predicted_output.dataset_id} "
1943 f"({predicted_output.dataset_type_name}@{predicted_output.data_coordinate})."
1944 )
1945 data = ProvenanceQuantumScanData(
1946 self.quantum_id,
1947 self.status,
1948 existing_outputs={
1949 dataset_id for dataset_id, was_produced in self.output_existence.items() if was_produced
1950 },
1951 quantum=quantum.model_dump_json().encode(),
1952 logs=self.logs.model_dump_json().encode() if self.logs.attempts else b"",
1953 metadata=self.metadata.model_dump_json().encode() if self.metadata.attempts else b"",
1954 )
1955 if compressor is not None:
1956 data.compress(compressor)
1957 return data
1960@dataclasses.dataclass
1961class ProvenanceQuantumScanData:
1962 """A struct that represents ready-for-serialization provenance information
1963 for a single quantum.
1964 """
1966 quantum_id: uuid.UUID
1967 """Unique ID for the quantum."""
1969 status: ProvenanceQuantumScanStatus
1970 """Combined status for the scan and the execution of the quantum."""
1972 existing_outputs: set[uuid.UUID] = dataclasses.field(default_factory=set)
1973 """Unique IDs of the output datasets that were actually written."""
1975 quantum: bytes = b""
1976 """Serialized quantum provenance model.
1978 This may be empty for quanta that had no attempts.
1979 """
1981 metadata: bytes = b""
1982 """Serialized task metadata."""
1984 logs: bytes = b""
1985 """Serialized logs."""
1987 is_compressed: bool = False
1988 """Whether the `quantum`, `metadata`, and `log` attributes are
1989 compressed.
1990 """
1992 def compress(self, compressor: Compressor) -> None:
1993 """Compress the data in this struct if it has not been compressed
1994 already.
1996 Parameters
1997 ----------
1998 compressor : `Compressor`
1999 Object with a ``compress`` method that takes and returns `bytes`.
2000 """
2001 if not self.is_compressed:
2002 self.quantum = compressor.compress(self.quantum)
2003 self.logs = compressor.compress(self.logs) if self.logs else b""
2004 self.metadata = compressor.compress(self.metadata) if self.metadata else b""
2005 self.is_compressed = True