Coverage for python / lsst / pipe / base / quantum_graph / _provenance.py: 27%
693 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-08 16:40 -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 "ProvenanceQuantumReport",
42 "ProvenanceQuantumScanData",
43 "ProvenanceQuantumScanModels",
44 "ProvenanceQuantumScanStatus",
45 "ProvenanceReport",
46 "ProvenanceTaskMetadataModel",
47)
49import dataclasses
50import enum
51import itertools
52import sys
53import uuid
54from collections import Counter
55from collections.abc import Callable, Iterable, Iterator, Mapping
56from contextlib import ExitStack, contextmanager
57from typing import TYPE_CHECKING, Any, TypedDict
59import astropy.table
60import networkx
61import numpy as np
62import pydantic
64from lsst.daf.butler import Butler, DataCoordinate
65from lsst.daf.butler.logging import ButlerLogRecord, ButlerLogRecords
66from lsst.resources import ResourcePath, ResourcePathExpression
67from lsst.utils.iteration import ensure_iterable
68from lsst.utils.logging import LsstLogAdapter, getLogger
69from lsst.utils.packages import Packages
71from .. import automatic_connection_constants as acc
72from .._status import ExceptionInfo, QuantumAttemptStatus, QuantumSuccessCaveats
73from .._task_metadata import TaskMetadata
74from ..log_capture import _ExecutionLogRecordsExtra
75from ..log_on_close import LogOnClose
76from ..pipeline_graph import PipelineGraph, TaskImportMode, TaskInitNode
77from ..resource_usage import QuantumResourceUsage
78from ._common import (
79 BaseQuantumGraph,
80 BaseQuantumGraphReader,
81 BaseQuantumGraphWriter,
82 ConnectionName,
83 DataCoordinateValues,
84 DatasetInfo,
85 DatasetTypeName,
86 HeaderModel,
87 QuantumInfo,
88 TaskLabel,
89)
90from ._multiblock import Compressor, MultiblockReader, MultiblockWriter
91from ._predicted import (
92 PredictedDatasetModel,
93 PredictedQuantumDatasetsModel,
94 PredictedQuantumGraph,
95 PredictedQuantumGraphComponents,
96)
98# Sphinx needs imports for type annotations of base class members.
99if "sphinx" in sys.modules:
100 import zipfile # noqa: F401
102 from ._multiblock import AddressReader, Decompressor # noqa: F401
105type LoopWrapper[T] = Callable[[Iterable[T]], Iterable[T]]
107_LOG = getLogger(__file__)
109DATASET_ADDRESS_INDEX = 0
110QUANTUM_ADDRESS_INDEX = 1
111LOG_ADDRESS_INDEX = 2
112METADATA_ADDRESS_INDEX = 3
114DATASET_MB_NAME = "datasets"
115QUANTUM_MB_NAME = "quanta"
116LOG_MB_NAME = "logs"
117METADATA_MB_NAME = "metadata"
120def pass_through[T](arg: T) -> T:
121 return arg
124class ProvenanceDatasetInfo(DatasetInfo):
125 """A typed dictionary that annotates the attributes of the NetworkX graph
126 node data for a provenance dataset.
128 Since NetworkX types are not generic over their node mapping type, this has
129 to be used explicitly, e.g.::
131 node_data: ProvenanceDatasetInfo = xgraph.nodes[dataset_id]
133 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
134 """
136 dataset_id: uuid.UUID
137 """Unique identifier for the dataset."""
139 produced: bool
140 """Whether this dataset was produced (vs. only predicted).
142 This is always `True` for overall input datasets. It is also `True` for
143 datasets that were produced and then removed before/during transfer back to
144 the central butler repository, so it may not reflect the continued
145 existence of the dataset.
146 """
149class ProvenanceQuantumInfo(QuantumInfo):
150 """A typed dictionary that annotates the attributes of the NetworkX graph
151 node data for a provenance quantum.
153 Since NetworkX types are not generic over their node mapping type, this has
154 to be used explicitly, e.g.::
156 node_data: ProvenanceQuantumInfo = xgraph.nodes[quantum_id]
158 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph` or
159 `ProvenanceQuantumGraph.quantum_only_xgraph`
160 """
162 status: QuantumAttemptStatus
163 """Enumerated status for the quantum.
165 This corresponds to the last attempt to run this quantum, or
166 `QuantumAttemptStatus.BLOCKED` if there were no attempts.
167 """
169 caveats: QuantumSuccessCaveats | None
170 """Flags indicating caveats on successful quanta.
172 This corresponds to the last attempt to run this quantum.
173 """
175 exception: ExceptionInfo | None
176 """Information about an exception raised when the quantum was executing.
178 This corresponds to the last attempt to run this quantum.
179 """
181 resource_usage: QuantumResourceUsage | None
182 """Resource usage information (timing, memory use) for this quantum.
184 This corresponds to the last attempt to run this quantum.
185 """
187 attempts: list[ProvenanceQuantumAttemptModel]
188 """Information about each attempt to run this quantum.
190 An entry is added merely if the quantum *should* have been attempted; an
191 empty `list` is used only for quanta that were blocked by an upstream
192 failure.
193 """
195 metadata_id: uuid.UUID
196 """ID of this quantum's metadata dataset."""
198 log_id: uuid.UUID
199 """ID of this quantum's log dataset."""
202class ProvenanceInitQuantumInfo(TypedDict):
203 """A typed dictionary that annotates the attributes of the NetworkX graph
204 node data for a provenance init quantum.
206 Since NetworkX types are not generic over their node mapping type, this has
207 to be used explicitly, e.g.::
209 node_data: ProvenanceInitQuantumInfo = xgraph.nodes[quantum_id]
211 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
212 """
214 data_id: DataCoordinate
215 """Data ID of the quantum.
217 This is always an empty ID; this key exists to allow init-quanta and
218 regular quanta to be treated more similarly.
219 """
221 task_label: str
222 """Label of the task for this quantum."""
224 pipeline_node: TaskInitNode
225 """Node in the pipeline graph for this task's init-only step."""
227 config_id: uuid.UUID
228 """ID of this task's config dataset."""
231class ProvenanceDatasetModel(PredictedDatasetModel):
232 """Data model for the datasets in a provenance quantum graph file."""
234 produced: bool
235 """Whether this dataset was produced (vs. only predicted).
237 This is always `True` for overall input datasets. It is also `True` for
238 datasets that were produced and then removed before/during transfer back to
239 the central butler repository, so it may not reflect the continued
240 existence of the dataset.
241 """
243 producer: uuid.UUID | None = None
244 """ID of the quantum that produced this dataset.
246 This is `None` for overall inputs to the graph.
247 """
249 consumers: list[uuid.UUID] = pydantic.Field(default_factory=list)
250 """IDs of quanta that were predicted to consume this dataset."""
252 @property
253 def node_id(self) -> uuid.UUID:
254 """Alias for the dataset ID."""
255 return self.dataset_id
257 @classmethod
258 def from_predicted(
259 cls,
260 predicted: PredictedDatasetModel,
261 producer: uuid.UUID | None = None,
262 consumers: Iterable[uuid.UUID] = (),
263 ) -> ProvenanceDatasetModel:
264 """Construct from a predicted dataset model.
266 Parameters
267 ----------
268 predicted : `PredictedDatasetModel`
269 Information about the dataset from the predicted graph.
270 producer : `uuid.UUID` or `None`, optional
271 ID of the quantum that was predicted to produce this dataset.
272 consumers : `~collections.abc.Iterable` [`uuid.UUID`], optional
273 IDs of the quanta that were predicted to consume this dataset.
275 Returns
276 -------
277 provenance : `ProvenanceDatasetModel`
278 Provenance dataset model.
280 Notes
281 -----
282 This initializes `produced` to `True` when ``producer is None`` and
283 `False` otherwise, on the assumption that it will be updated later.
284 """
285 return cls.model_construct(
286 dataset_id=predicted.dataset_id,
287 dataset_type_name=predicted.dataset_type_name,
288 data_coordinate=predicted.data_coordinate,
289 run=predicted.run,
290 produced=(producer is None), # if it's not produced by this QG, it's an overall input
291 producer=producer,
292 consumers=list(consumers),
293 )
295 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
296 """Add this dataset and its edges to quanta to a provenance graph.
298 Parameters
299 ----------
300 graph : `ProvenanceQuantumGraph`
301 Graph to update in place.
303 Notes
304 -----
305 This method adds:
307 - a ``bipartite_xgraph`` dataset node with full attributes;
308 - ``bipartite_xgraph`` edges to adjacent quanta (which adds quantum
309 nodes with no attributes), without populating edge attributes;
310 - ``quantum_only_xgraph`` edges for each pair of quanta in which one
311 produces this dataset and another consumes it (this also adds quantum
312 nodes with no attributes).
313 """
314 dataset_type_node = graph.pipeline_graph.dataset_types[self.dataset_type_name]
315 data_id = DataCoordinate.from_full_values(dataset_type_node.dimensions, tuple(self.data_coordinate))
316 graph._bipartite_xgraph.add_node(
317 self.dataset_id,
318 data_id=data_id,
319 dataset_type_name=self.dataset_type_name,
320 pipeline_node=dataset_type_node,
321 run=self.run,
322 produced=self.produced,
323 )
324 if self.producer is not None:
325 graph._bipartite_xgraph.add_edge(self.producer, self.dataset_id)
326 for consumer_id in self.consumers:
327 graph._bipartite_xgraph.add_edge(self.dataset_id, consumer_id)
328 if self.producer is not None:
329 graph._quantum_only_xgraph.add_edge(self.producer, consumer_id)
330 graph._datasets_by_type[self.dataset_type_name][data_id] = self.dataset_id
332 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
333 # when we inherit those docstrings in our public classes.
334 if "sphinx" in sys.modules and not TYPE_CHECKING:
336 def copy(self, *args: Any, **kwargs: Any) -> Any:
337 """See `pydantic.BaseModel.copy`."""
338 return super().copy(*args, **kwargs)
340 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
341 """See `pydantic.BaseModel.model_dump`."""
342 return super().model_dump(*args, **kwargs)
344 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
345 """See `pydantic.BaseModel.model_dump_json`."""
346 return super().model_dump(*args, **kwargs)
348 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
349 """See `pydantic.BaseModel.model_copy`."""
350 return super().model_copy(*args, **kwargs)
352 @classmethod
353 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
354 """See `pydantic.BaseModel.model_construct`."""
355 return super().model_construct(*args, **kwargs)
357 @classmethod
358 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
359 """See `pydantic.BaseModel.model_json_schema`."""
360 return super().model_json_schema(*args, **kwargs)
362 @classmethod
363 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
364 """See `pydantic.BaseModel.model_validate`."""
365 return super().model_validate(*args, **kwargs)
367 @classmethod
368 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
369 """See `pydantic.BaseModel.model_validate_json`."""
370 return super().model_validate_json(*args, **kwargs)
372 @classmethod
373 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
374 """See `pydantic.BaseModel.model_validate_strings`."""
375 return super().model_validate_strings(*args, **kwargs)
378class ProvenanceQuantumAttemptModel(pydantic.BaseModel):
379 """Data model for a now-superseded attempt to run a quantum in a
380 provenance quantum graph file.
381 """
383 attempt: int = 0
384 """Counter incremented for every attempt to execute this quantum."""
386 status: QuantumAttemptStatus = QuantumAttemptStatus.UNKNOWN
387 """Enumerated status for the quantum."""
389 caveats: QuantumSuccessCaveats | None = None
390 """Flags indicating caveats on successful quanta."""
392 exception: ExceptionInfo | None = None
393 """Information about an exception raised when the quantum was executing."""
395 resource_usage: QuantumResourceUsage | None = None
396 """Resource usage information (timing, memory use) for this quantum."""
398 previous_process_quanta: list[uuid.UUID] = pydantic.Field(default_factory=list)
399 """The IDs of other quanta previously executed in the same process as this
400 one.
401 """
403 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
404 # when we inherit those docstrings in our public classes.
405 if "sphinx" in sys.modules and not TYPE_CHECKING:
407 def copy(self, *args: Any, **kwargs: Any) -> Any:
408 """See `pydantic.BaseModel.copy`."""
409 return super().copy(*args, **kwargs)
411 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
412 """See `pydantic.BaseModel.model_dump`."""
413 return super().model_dump(*args, **kwargs)
415 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
416 """See `pydantic.BaseModel.model_dump_json`."""
417 return super().model_dump(*args, **kwargs)
419 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
420 """See `pydantic.BaseModel.model_copy`."""
421 return super().model_copy(*args, **kwargs)
423 @classmethod
424 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
425 """See `pydantic.BaseModel.model_construct`."""
426 return super().model_construct(*args, **kwargs)
428 @classmethod
429 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
430 """See `pydantic.BaseModel.model_json_schema`."""
431 return super().model_json_schema(*args, **kwargs)
433 @classmethod
434 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
435 """See `pydantic.BaseModel.model_validate`."""
436 return super().model_validate(*args, **kwargs)
438 @classmethod
439 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
440 """See `pydantic.BaseModel.model_validate_json`."""
441 return super().model_validate_json(*args, **kwargs)
443 @classmethod
444 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
445 """See `pydantic.BaseModel.model_validate_strings`."""
446 return super().model_validate_strings(*args, **kwargs)
449class ProvenanceLogRecordsModel(pydantic.BaseModel):
450 """Data model for storing execution logs in a provenance quantum graph
451 file.
452 """
454 attempts: list[list[ButlerLogRecord] | None] = pydantic.Field(default_factory=list)
455 """Logs from attempts to run this task, ordered chronologically from first
456 to last.
457 """
459 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
460 # when we inherit those docstrings in our public classes.
461 if "sphinx" in sys.modules and not TYPE_CHECKING:
463 def copy(self, *args: Any, **kwargs: Any) -> Any:
464 """See `pydantic.BaseModel.copy`."""
465 return super().copy(*args, **kwargs)
467 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
468 """See `pydantic.BaseModel.model_dump`."""
469 return super().model_dump(*args, **kwargs)
471 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
472 """See `pydantic.BaseModel.model_dump_json`."""
473 return super().model_dump(*args, **kwargs)
475 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
476 """See `pydantic.BaseModel.model_copy`."""
477 return super().model_copy(*args, **kwargs)
479 @classmethod
480 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
481 """See `pydantic.BaseModel.model_construct`."""
482 return super().model_construct(*args, **kwargs)
484 @classmethod
485 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
486 """See `pydantic.BaseModel.model_json_schema`."""
487 return super().model_json_schema(*args, **kwargs)
489 @classmethod
490 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
491 """See `pydantic.BaseModel.model_validate`."""
492 return super().model_validate(*args, **kwargs)
494 @classmethod
495 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
496 """See `pydantic.BaseModel.model_validate_json`."""
497 return super().model_validate_json(*args, **kwargs)
499 @classmethod
500 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
501 """See `pydantic.BaseModel.model_validate_strings`."""
502 return super().model_validate_strings(*args, **kwargs)
505class ProvenanceTaskMetadataModel(pydantic.BaseModel):
506 """Data model for storing task metadata in a provenance quantum graph
507 file.
508 """
510 # We want to convert infs and nans to constants, not null. Unfortunately
511 # the fact that TaskMetadata _also_ sets this is ignored when that model
512 # is nested here.
513 model_config = pydantic.ConfigDict(ser_json_inf_nan="constants")
515 attempts: list[TaskMetadata | None] = pydantic.Field(default_factory=list)
516 """Metadata from attempts to run this task, ordered chronologically from
517 first to last.
518 """
520 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
521 # when we inherit those docstrings in our public classes.
522 if "sphinx" in sys.modules and not TYPE_CHECKING:
524 def copy(self, *args: Any, **kwargs: Any) -> Any:
525 """See `pydantic.BaseModel.copy`."""
526 return super().copy(*args, **kwargs)
528 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
529 """See `pydantic.BaseModel.model_dump`."""
530 return super().model_dump(*args, **kwargs)
532 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
533 """See `pydantic.BaseModel.model_dump_json`."""
534 return super().model_dump(*args, **kwargs)
536 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
537 """See `pydantic.BaseModel.model_copy`."""
538 return super().model_copy(*args, **kwargs)
540 @classmethod
541 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
542 """See `pydantic.BaseModel.model_construct`."""
543 return super().model_construct(*args, **kwargs)
545 @classmethod
546 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
547 """See `pydantic.BaseModel.model_json_schema`."""
548 return super().model_json_schema(*args, **kwargs)
550 @classmethod
551 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
552 """See `pydantic.BaseModel.model_validate`."""
553 return super().model_validate(*args, **kwargs)
555 @classmethod
556 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
557 """See `pydantic.BaseModel.model_validate_json`."""
558 return super().model_validate_json(*args, **kwargs)
560 @classmethod
561 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
562 """See `pydantic.BaseModel.model_validate_strings`."""
563 return super().model_validate_strings(*args, **kwargs)
566class ProvenanceQuantumReport(pydantic.BaseModel):
567 """A Pydantic model that used to report information about a single
568 (generally problematic) quantum.
569 """
571 quantum_id: uuid.UUID
572 data_id: dict[str, int | str]
573 attempts: list[ProvenanceQuantumAttemptModel]
575 @classmethod
576 def from_info(cls, quantum_id: uuid.UUID, quantum_info: ProvenanceQuantumInfo) -> ProvenanceQuantumReport:
577 """Construct from a provenance quantum graph node.
579 Parameters
580 ----------
581 quantum_id : `uuid.UUID`
582 Unique ID for the quantum.
583 quantum_info : `ProvenanceQuantumInfo`
584 Node attributes for this quantum.
585 """
586 return cls(
587 quantum_id=quantum_id,
588 data_id=dict(quantum_info["data_id"].mapping),
589 attempts=quantum_info["attempts"],
590 )
592 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
593 # when we inherit those docstrings in our public classes.
594 if "sphinx" in sys.modules and not TYPE_CHECKING:
596 def copy(self, *args: Any, **kwargs: Any) -> Any:
597 """See `pydantic.BaseModel.copy`."""
598 return super().copy(*args, **kwargs)
600 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
601 """See `pydantic.BaseModel.model_dump`."""
602 return super().model_dump(*args, **kwargs)
604 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
605 """See `pydantic.BaseModel.model_dump_json`."""
606 return super().model_dump(*args, **kwargs)
608 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
609 """See `pydantic.BaseModel.model_copy`."""
610 return super().model_copy(*args, **kwargs)
612 @classmethod
613 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
614 """See `pydantic.BaseModel.model_construct`."""
615 return super().model_construct(*args, **kwargs)
617 @classmethod
618 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
619 """See `pydantic.BaseModel.model_json_schema`."""
620 return super().model_json_schema(*args, **kwargs)
622 @classmethod
623 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
624 """See `pydantic.BaseModel.model_validate`."""
625 return super().model_validate(*args, **kwargs)
627 @classmethod
628 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
629 """See `pydantic.BaseModel.model_validate_json`."""
630 return super().model_validate_json(*args, **kwargs)
632 @classmethod
633 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
634 """See `pydantic.BaseModel.model_validate_strings`."""
635 return super().model_validate_strings(*args, **kwargs)
638class ProvenanceReport(pydantic.RootModel):
639 """A Pydantic model that groups quantum information by task label, then
640 status (as a string), and then exception type.
641 """
643 root: dict[TaskLabel, dict[str, dict[str | None, list[ProvenanceQuantumReport]]]] = {}
645 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
646 # when we inherit those docstrings in our public classes.
647 if "sphinx" in sys.modules and not TYPE_CHECKING:
649 def copy(self, *args: Any, **kwargs: Any) -> Any:
650 """See `pydantic.BaseModel.copy`."""
651 return super().copy(*args, **kwargs)
653 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
654 """See `pydantic.BaseModel.model_dump`."""
655 return super().model_dump(*args, **kwargs)
657 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
658 """See `pydantic.BaseModel.model_dump_json`."""
659 return super().model_dump(*args, **kwargs)
661 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
662 """See `pydantic.BaseModel.model_copy`."""
663 return super().model_copy(*args, **kwargs)
665 @classmethod
666 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
667 """See `pydantic.BaseModel.model_construct`."""
668 return super().model_construct(*args, **kwargs)
670 @classmethod
671 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
672 """See `pydantic.BaseModel.model_json_schema`."""
673 return super().model_json_schema(*args, **kwargs)
675 @classmethod
676 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
677 """See `pydantic.BaseModel.model_validate`."""
678 return super().model_validate(*args, **kwargs)
680 @classmethod
681 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
682 """See `pydantic.BaseModel.model_validate_json`."""
683 return super().model_validate_json(*args, **kwargs)
685 @classmethod
686 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
687 """See `pydantic.BaseModel.model_validate_strings`."""
688 return super().model_validate_strings(*args, **kwargs)
691class ProvenanceQuantumModel(pydantic.BaseModel):
692 """Data model for the quanta in a provenance quantum graph file."""
694 quantum_id: uuid.UUID
695 """Unique identifier for the quantum."""
697 task_label: TaskLabel
698 """Name of the type of this dataset."""
700 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
701 """The full values (required and implied) of this dataset's data ID."""
703 inputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
704 """IDs of the datasets predicted to be consumed by this quantum, grouped by
705 connection name.
706 """
708 outputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
709 """IDs of the datasets predicted to be produced by this quantum, grouped by
710 connection name.
711 """
713 attempts: list[ProvenanceQuantumAttemptModel] = pydantic.Field(default_factory=list)
714 """Provenance for all attempts to execute this quantum, ordered
715 chronologically from first to last.
717 An entry is added merely if the quantum *should* have been attempted; an
718 empty `list` is used only for quanta that were blocked by an upstream
719 failure.
720 """
722 @property
723 def node_id(self) -> uuid.UUID:
724 """Alias for the quantum ID."""
725 return self.quantum_id
727 @classmethod
728 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceQuantumModel:
729 """Construct from a predicted quantum model.
731 Parameters
732 ----------
733 predicted : `PredictedQuantumDatasetsModel`
734 Information about the quantum from the predicted graph.
736 Returns
737 -------
738 provenance : `ProvenanceQuantumModel`
739 Provenance quantum model.
740 """
741 inputs = {
742 connection_name: [d.dataset_id for d in predicted_inputs]
743 for connection_name, predicted_inputs in predicted.inputs.items()
744 }
745 outputs = {
746 connection_name: [d.dataset_id for d in predicted_outputs]
747 for connection_name, predicted_outputs in predicted.outputs.items()
748 }
749 return cls(
750 quantum_id=predicted.quantum_id,
751 task_label=predicted.task_label,
752 data_coordinate=predicted.data_coordinate,
753 inputs=inputs,
754 outputs=outputs,
755 )
757 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
758 """Add this quantum and its edges to datasets to a provenance graph.
760 Parameters
761 ----------
762 graph : `ProvenanceQuantumGraph`
763 Graph to update in place.
765 Notes
766 -----
767 This method adds:
769 - a ``bipartite_xgraph`` quantum node with full attributes;
770 - a ``quantum_only_xgraph`` quantum node with full attributes;
771 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
772 nodes with no attributes), while populating those edge attributes;
773 - ``quantum_only_xgraph`` edges to any adjacent quantum that has also
774 already been loaded.
775 """
776 task_node = graph.pipeline_graph.tasks[self.task_label]
777 data_id = DataCoordinate.from_full_values(task_node.dimensions, tuple(self.data_coordinate))
778 last_attempt = (
779 self.attempts[-1]
780 if self.attempts
781 else ProvenanceQuantumAttemptModel(status=QuantumAttemptStatus.BLOCKED)
782 )
783 graph._bipartite_xgraph.add_node(
784 self.quantum_id,
785 data_id=data_id,
786 task_label=self.task_label,
787 pipeline_node=task_node,
788 status=last_attempt.status,
789 caveats=last_attempt.caveats,
790 exception=last_attempt.exception,
791 resource_usage=last_attempt.resource_usage,
792 attempts=self.attempts,
793 )
794 graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
795 graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
796 for connection_name, dataset_ids in self.inputs.items():
797 read_edge = task_node.get_input_edge(connection_name)
798 for dataset_id in dataset_ids:
799 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
800 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
801 "pipeline_edges", []
802 ).append(read_edge)
803 for connection_name, dataset_ids in self.outputs.items():
804 write_edge = task_node.get_output_edge(connection_name)
805 if connection_name == acc.METADATA_OUTPUT_CONNECTION_NAME:
806 graph._bipartite_xgraph.add_node(
807 dataset_ids[0],
808 data_id=data_id,
809 dataset_type_name=write_edge.dataset_type_name,
810 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
811 run=graph.header.output_run,
812 produced=last_attempt.status.has_metadata,
813 )
814 graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
815 graph._bipartite_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
816 graph._quantum_only_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
817 if connection_name == acc.LOG_OUTPUT_CONNECTION_NAME:
818 graph._bipartite_xgraph.add_node(
819 dataset_ids[0],
820 data_id=data_id,
821 dataset_type_name=write_edge.dataset_type_name,
822 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
823 run=graph.header.output_run,
824 produced=last_attempt.status.has_log,
825 )
826 graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
827 graph._bipartite_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
828 graph._quantum_only_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
829 for dataset_id in dataset_ids:
830 graph._bipartite_xgraph.add_edge(
831 self.quantum_id,
832 dataset_id,
833 is_read=False,
834 # There can only be one pipeline edge for an output.
835 pipeline_edges=[write_edge],
836 )
837 for dataset_id in graph._bipartite_xgraph.predecessors(self.quantum_id):
838 for upstream_quantum_id in graph._bipartite_xgraph.predecessors(dataset_id):
839 graph._quantum_only_xgraph.add_edge(upstream_quantum_id, self.quantum_id)
840 for dataset_id in graph._bipartite_xgraph.successors(self.quantum_id):
841 for downstream_quantum_id in graph._bipartite_xgraph.successors(dataset_id):
842 graph._quantum_only_xgraph.add_edge(self.quantum_id, downstream_quantum_id)
844 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
845 # when we inherit those docstrings in our public classes.
846 if "sphinx" in sys.modules and not TYPE_CHECKING:
848 def copy(self, *args: Any, **kwargs: Any) -> Any:
849 """See `pydantic.BaseModel.copy`."""
850 return super().copy(*args, **kwargs)
852 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
853 """See `pydantic.BaseModel.model_dump`."""
854 return super().model_dump(*args, **kwargs)
856 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
857 """See `pydantic.BaseModel.model_dump_json`."""
858 return super().model_dump(*args, **kwargs)
860 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
861 """See `pydantic.BaseModel.model_copy`."""
862 return super().model_copy(*args, **kwargs)
864 @classmethod
865 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
866 """See `pydantic.BaseModel.model_construct`."""
867 return super().model_construct(*args, **kwargs)
869 @classmethod
870 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
871 """See `pydantic.BaseModel.model_json_schema`."""
872 return super().model_json_schema(*args, **kwargs)
874 @classmethod
875 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
876 """See `pydantic.BaseModel.model_validate`."""
877 return super().model_validate(*args, **kwargs)
879 @classmethod
880 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
881 """See `pydantic.BaseModel.model_validate_json`."""
882 return super().model_validate_json(*args, **kwargs)
884 @classmethod
885 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
886 """See `pydantic.BaseModel.model_validate_strings`."""
887 return super().model_validate_strings(*args, **kwargs)
890class ProvenanceInitQuantumModel(pydantic.BaseModel):
891 """Data model for the special "init" quanta in a provenance quantum graph
892 file.
893 """
895 quantum_id: uuid.UUID
896 """Unique identifier for the quantum."""
898 task_label: TaskLabel
899 """Name of the type of this dataset.
901 This is always a parent dataset type name, not a component.
903 Note that full dataset type definitions are stored in the pipeline graph.
904 """
906 inputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
907 """IDs of the datasets predicted to be consumed by this quantum, grouped by
908 connection name.
909 """
911 outputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
912 """IDs of the datasets predicted to be produced by this quantum, grouped by
913 connection name.
914 """
916 @classmethod
917 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceInitQuantumModel:
918 """Construct from a predicted quantum model.
920 Parameters
921 ----------
922 predicted : `PredictedQuantumDatasetsModel`
923 Information about the quantum from the predicted graph.
925 Returns
926 -------
927 provenance : `ProvenanceInitQuantumModel`
928 Provenance init quantum model.
929 """
930 inputs = {
931 connection_name: predicted_inputs[0].dataset_id
932 for connection_name, predicted_inputs in predicted.inputs.items()
933 }
934 outputs = {
935 connection_name: predicted_outputs[0].dataset_id
936 for connection_name, predicted_outputs in predicted.outputs.items()
937 }
938 return cls(
939 quantum_id=predicted.quantum_id,
940 task_label=predicted.task_label,
941 inputs=inputs,
942 outputs=outputs,
943 )
945 def _add_to_graph(self, graph: ProvenanceQuantumGraph, empty_data_id: DataCoordinate) -> None:
946 """Add this quantum and its edges to datasets to a provenance graph.
948 Parameters
949 ----------
950 graph : `ProvenanceQuantumGraph`
951 Graph to update in place.
952 empty_data_id : `lsst.daf.butler.DataCoordinate`
953 The empty data ID for the appropriate dimension universe.
955 Notes
956 -----
957 This method adds:
959 - a ``bipartite_xgraph`` quantum node with full attributes;
960 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
961 nodes with no attributes), while populating those edge attributes;
962 """
963 task_init_node = graph.pipeline_graph.tasks[self.task_label].init
964 graph._bipartite_xgraph.add_node(
965 self.quantum_id, data_id=empty_data_id, task_label=self.task_label, pipeline_node=task_init_node
966 )
967 for connection_name, dataset_id in self.inputs.items():
968 read_edge = task_init_node.get_input_edge(connection_name)
969 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
970 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
971 "pipeline_edges", []
972 ).append(read_edge)
973 for connection_name, dataset_id in self.outputs.items():
974 write_edge = task_init_node.get_output_edge(connection_name)
975 graph._bipartite_xgraph.add_node(
976 dataset_id,
977 data_id=empty_data_id,
978 dataset_type_name=write_edge.dataset_type_name,
979 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
980 run=graph.header.output_run,
981 produced=True,
982 )
983 graph._datasets_by_type[write_edge.dataset_type_name][empty_data_id] = dataset_id
984 graph._bipartite_xgraph.add_edge(
985 self.quantum_id,
986 dataset_id,
987 is_read=False,
988 # There can only be one pipeline edge for an output.
989 pipeline_edges=[write_edge],
990 )
991 if write_edge.connection_name == acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME:
992 graph._bipartite_xgraph.nodes[self.quantum_id]["config_id"] = dataset_id
993 graph._init_quanta[self.task_label] = self.quantum_id
995 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
996 # when we inherit those docstrings in our public classes.
997 if "sphinx" in sys.modules and not TYPE_CHECKING:
999 def copy(self, *args: Any, **kwargs: Any) -> Any:
1000 """See `pydantic.BaseModel.copy`."""
1001 return super().copy(*args, **kwargs)
1003 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
1004 """See `pydantic.BaseModel.model_dump`."""
1005 return super().model_dump(*args, **kwargs)
1007 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
1008 """See `pydantic.BaseModel.model_dump_json`."""
1009 return super().model_dump(*args, **kwargs)
1011 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
1012 """See `pydantic.BaseModel.model_copy`."""
1013 return super().model_copy(*args, **kwargs)
1015 @classmethod
1016 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
1017 """See `pydantic.BaseModel.model_construct`."""
1018 return super().model_construct(*args, **kwargs)
1020 @classmethod
1021 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
1022 """See `pydantic.BaseModel.model_json_schema`."""
1023 return super().model_json_schema(*args, **kwargs)
1025 @classmethod
1026 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
1027 """See `pydantic.BaseModel.model_validate`."""
1028 return super().model_validate(*args, **kwargs)
1030 @classmethod
1031 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
1032 """See `pydantic.BaseModel.model_validate_json`."""
1033 return super().model_validate_json(*args, **kwargs)
1035 @classmethod
1036 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
1037 """See `pydantic.BaseModel.model_validate_strings`."""
1038 return super().model_validate_strings(*args, **kwargs)
1041class ProvenanceInitQuantaModel(pydantic.RootModel):
1042 """Data model for the init quanta in a provenance graph."""
1044 root: list[ProvenanceInitQuantumModel] = pydantic.Field(default_factory=list)
1045 """List of special "init" quanta, one for each task."""
1047 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
1048 """Add this quantum and its edges to datasets to a provenance graph.
1050 Parameters
1051 ----------
1052 graph : `ProvenanceQuantumGraph`
1053 Graph to update in place.
1054 """
1055 empty_data_id = DataCoordinate.make_empty(graph.pipeline_graph.universe)
1056 for init_quantum in self.root:
1057 init_quantum._add_to_graph(graph, empty_data_id=empty_data_id)
1059 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
1060 # when we inherit those docstrings in our public classes.
1061 if "sphinx" in sys.modules and not TYPE_CHECKING:
1063 def copy(self, *args: Any, **kwargs: Any) -> Any:
1064 """See `pydantic.BaseModel.copy`."""
1065 return super().copy(*args, **kwargs)
1067 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
1068 """See `pydantic.BaseModel.model_dump`."""
1069 return super().model_dump(*args, **kwargs)
1071 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
1072 """See `pydantic.BaseModel.model_dump_json`."""
1073 return super().model_dump(*args, **kwargs)
1075 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
1076 """See `pydantic.BaseModel.model_copy`."""
1077 return super().model_copy(*args, **kwargs)
1079 @classmethod
1080 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
1081 """See `pydantic.BaseModel.model_construct`."""
1082 return super().model_construct(*args, **kwargs)
1084 @classmethod
1085 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
1086 """See `pydantic.BaseModel.model_json_schema`."""
1087 return super().model_json_schema(*args, **kwargs)
1089 @classmethod
1090 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
1091 """See `pydantic.BaseModel.model_validate`."""
1092 return super().model_validate(*args, **kwargs)
1094 @classmethod
1095 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
1096 """See `pydantic.BaseModel.model_validate_json`."""
1097 return super().model_validate_json(*args, **kwargs)
1099 @classmethod
1100 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
1101 """See `pydantic.BaseModel.model_validate_strings`."""
1102 return super().model_validate_strings(*args, **kwargs)
1105class ProvenanceQuantumGraph(BaseQuantumGraph):
1106 """A quantum graph that represents processing that has already been
1107 executed.
1109 Parameters
1110 ----------
1111 header : `HeaderModel`
1112 General metadata shared with other quantum graph types.
1113 pipeline_graph : `.pipeline_graph.PipelineGraph`
1114 Graph of tasks and dataset types. May contain a superset of the tasks
1115 and dataset types that actually have quanta and datasets in the quantum
1116 graph.
1118 Notes
1119 -----
1120 A provenance quantum graph is generally obtained via the
1121 `ProvenanceQuantumGraphReader.graph` attribute, which is updated in-place
1122 as information is read from disk.
1123 """
1125 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph) -> None:
1126 super().__init__(header, pipeline_graph)
1127 self._init_quanta: dict[TaskLabel, uuid.UUID] = {}
1128 self._quantum_only_xgraph = networkx.DiGraph()
1129 self._bipartite_xgraph = networkx.DiGraph()
1130 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {
1131 task_label: {} for task_label in self.pipeline_graph.tasks.keys()
1132 }
1133 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {
1134 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
1135 }
1137 @classmethod
1138 @contextmanager
1139 def from_args(
1140 cls,
1141 repo_or_filename: str,
1142 /,
1143 collection: str | None = None,
1144 *,
1145 quanta: Iterable[uuid.UUID] | None = None,
1146 datasets: Iterable[uuid.UUID] | None = None,
1147 writeable: bool = False,
1148 ) -> Iterator[tuple[ProvenanceQuantumGraph, Butler | None]]:
1149 """Construct a `ProvenanceQuantumGraph` fron CLI-friendly arguments for
1150 a file or butler-ingested graph dataset.
1152 Parameters
1153 ----------
1154 repo_or_filename : `str`
1155 Either a provenance quantum graph filename or a butler repository
1156 path or alias.
1157 collection : `str`, optional
1158 Collection to search; presence indicates that the first argument
1159 is a butler repository, not a filename.
1160 quanta : `~collections.abc.Iterable` [ `str` ] or `None`, optional
1161 IDs of the quanta to load, or `None` to load all.
1162 datasets : `~collections.abc.Iterable` [ `str` ], optional
1163 IDs of the datasets to load, or `None` to load all.
1164 writeable : `bool`, optional
1165 Whether the butler should be constructed with write support.
1167 Returns
1168 -------
1169 context : `contextlib.AbstractContextManager`
1170 A context manager that yields a tuple of
1172 - the `ProvenanceQuantumGraph`
1173 - the `Butler` constructed (or `None`)
1175 when entered.
1176 """
1177 exit_stack = ExitStack()
1178 if collection is not None:
1179 try:
1180 butler = exit_stack.enter_context(
1181 Butler.from_config(repo_or_filename, collections=[collection], writeable=writeable)
1182 )
1183 except Exception as err:
1184 err.add_note(
1185 f"Expected {repo_or_filename!r} to be a butler repository path or alias because a "
1186 f"collection ({collection}) was provided."
1187 )
1188 raise
1189 with exit_stack:
1190 graph = butler.get(
1191 acc.PROVENANCE_DATASET_TYPE_NAME, parameters={"quanta": quanta, "datasets": datasets}
1192 )
1193 yield graph, butler
1194 else:
1195 try:
1196 reader = exit_stack.enter_context(ProvenanceQuantumGraphReader.open(repo_or_filename))
1197 except Exception as err:
1198 err.add_note(
1199 f"Expected a {repo_or_filename} to be a provenance quantum graph filename "
1200 f"because no collection was provided."
1201 )
1202 raise
1203 with exit_stack:
1204 if quanta is None:
1205 reader.read_quanta()
1206 elif not quanta:
1207 reader.read_quanta(quanta)
1208 if datasets is None:
1209 reader.read_datasets()
1210 elif not datasets:
1211 reader.read_datasets(datasets)
1212 yield reader.graph, None
1214 @property
1215 def init_quanta(self) -> Mapping[TaskLabel, uuid.UUID]:
1216 """A mapping from task label to the ID of the special init quantum for
1217 that task.
1219 This is populated by the ``init_quanta`` component. Additional
1220 information about each init quantum can be found by using the ID to
1221 look up node attributes in the `bipartite_xgraph`, i.e.::
1223 info: ProvenanceInitQuantumInfo = qg.bipartite_xgraph.nodes[id]
1224 """
1225 return self._init_quanta
1227 @property
1228 def quanta_by_task(self) -> Mapping[TaskLabel, Mapping[DataCoordinate, uuid.UUID]]:
1229 """A nested mapping of all quanta, keyed first by task name and then by
1230 data ID.
1232 Notes
1233 -----
1234 This is populated one quantum at a time as they are read. All tasks in
1235 the pipeline graph are included, even if none of their quanta were
1236 loaded (i.e. nested mappings may be empty).
1238 The returned object may be an internal dictionary; as the type
1239 annotation indicates, it should not be modified in place.
1240 """
1241 return self._quanta_by_task_label
1243 @property
1244 def datasets_by_type(self) -> Mapping[DatasetTypeName, Mapping[DataCoordinate, uuid.UUID]]:
1245 """A nested mapping of all datasets, keyed first by dataset type name
1246 and then by data ID.
1248 Notes
1249 -----
1250 This is populated one dataset at a time as they are read. All dataset
1251 types in the pipeline graph are included, even if none of their
1252 datasets were loaded (i.e. nested mappings may be empty).
1254 Reading a quantum also populates its log and metadata datasets.
1256 The returned object may be an internal dictionary; as the type
1257 annotation indicates, it should not be modified in place.
1258 """
1259 return self._datasets_by_type
1261 @property
1262 def quantum_only_xgraph(self) -> networkx.DiGraph:
1263 """A directed acyclic graph with quanta as nodes (and datasets elided).
1265 Notes
1266 -----
1267 Node keys are quantum UUIDs, and are populated one quantum at a time as
1268 they are loaded. Loading quanta (via
1269 `ProvenanceQuantumGraphReader.read_quanta`) will add the loaded nodes
1270 with full attributes and add edges to adjacent nodes with no
1271 attributes. Loading datasets (via
1272 `ProvenanceQuantumGraphReader.read_datasets`) will also add edges and
1273 nodes with no attributes.
1275 Node attributes are described by the `ProvenanceQuantumInfo` types.
1277 This graph does not include special "init" quanta.
1279 The returned object is a read-only view of an internal one.
1280 """
1281 return self._quantum_only_xgraph.copy(as_view=True)
1283 @property
1284 def bipartite_xgraph(self) -> networkx.DiGraph:
1285 """A directed acyclic graph with quantum and dataset nodes.
1287 Notes
1288 -----
1289 Node keys are quantum or dataset UUIDs, and are populated one quantum
1290 or dataset at a time as they are loaded. Loading quanta (via
1291 `ProvenanceQuantumGraphReader.read_quanta`) or datasets (via
1292 `ProvenanceQuantumGraphReader.read_datasets`) will load those nodes
1293 with full attributes and edges to adjacent nodes with no attributes.
1294 Loading quanta is necessary to populate edge attributes.
1295 Reading a quantum also populates its log and metadata datasets.
1297 Node attributes are described by the
1298 `ProvenanceQuantumInfo`, `ProvenanceInitQuantumInfo`, and
1299 `ProvenanceDatasetInfo` types.
1301 This graph includes init-input and init-output datasets, but it does
1302 *not* reflect the dependency between each task's special "init" quantum
1303 and its runtime quanta (as this would require edges between quanta, and
1304 that would break the "bipartite" property).
1306 The returned object is a read-only view of an internal one.
1307 """
1308 return self._bipartite_xgraph.copy(as_view=True)
1310 def make_quantum_table(self, drop_unused_columns: bool = True) -> astropy.table.Table:
1311 """Construct an `astropy.table.Table` with a tabular summary of the
1312 quanta.
1314 Parameters
1315 ----------
1316 drop_unused_columns : `bool`, optional
1317 Whether to drop columns for rare states that did not actually
1318 occur in this run.
1320 Returns
1321 -------
1322 table : `astropy.table.Table`
1323 A table view of the quantum information. This only includes
1324 counts of status categories and caveats, not any per-data-ID
1325 detail.
1327 Notes
1328 -----
1329 Success caveats in the table are represented by their
1330 `~QuantumSuccessCaveats.concise` form, so when pretty-printing this
1331 table for users, the `~QuantumSuccessCaveats.legend` should generally
1332 be printed as well.
1333 """
1334 rows = []
1335 for task_label, quanta_for_task in self.quanta_by_task.items():
1336 if not self.header.n_task_quanta[task_label]:
1337 continue
1338 status_counts = Counter[QuantumAttemptStatus](
1339 self._quantum_only_xgraph.nodes[q]["status"] for q in quanta_for_task.values()
1340 )
1341 caveat_counts = Counter[QuantumSuccessCaveats | None](
1342 self._quantum_only_xgraph.nodes[q]["caveats"] for q in quanta_for_task.values()
1343 )
1344 caveat_counts.pop(QuantumSuccessCaveats.NO_CAVEATS, None)
1345 caveat_counts.pop(None, None)
1346 if len(caveat_counts) > 1:
1347 caveats = "(multiple)"
1348 elif len(caveat_counts) == 1:
1349 ((code, count),) = caveat_counts.items()
1350 # MyPy can't tell that the pop(None, None) above makes None
1351 # impossible here.
1352 caveats = f"{code.concise()}({count})" # type: ignore[union-attr]
1353 else:
1354 caveats = ""
1355 row: dict[str, Any] = {
1356 "Task": task_label,
1357 "Caveats": caveats,
1358 }
1359 for status in QuantumAttemptStatus:
1360 row[status.title] = status_counts.get(status, 0)
1361 row.update(
1362 {
1363 "TOTAL": len(quanta_for_task),
1364 "EXPECTED": self.header.n_task_quanta[task_label],
1365 }
1366 )
1367 rows.append(row)
1368 table = astropy.table.Table(rows)
1369 if drop_unused_columns:
1370 for status in QuantumAttemptStatus:
1371 if status.is_rare and not table[status.title].any():
1372 del table[status.title]
1373 return table
1375 def make_exception_table(self) -> astropy.table.Table:
1376 """Construct an `astropy.table.Table` with counts for each exception
1377 type raised by each task.
1379 Returns
1380 -------
1381 table : `astropy.table.Table`
1382 A table with columns for task label, exception type, and counts.
1383 """
1384 rows = []
1385 for task_label, quanta_for_task in self.quanta_by_task.items():
1386 success_counts = Counter[str]()
1387 failed_counts = Counter[str]()
1388 for quantum_id in quanta_for_task.values():
1389 quantum_info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1390 exc_info = quantum_info["exception"]
1391 if exc_info is not None:
1392 if quantum_info["status"] is QuantumAttemptStatus.SUCCESSFUL:
1393 success_counts[exc_info.type_name] += 1
1394 else:
1395 failed_counts[exc_info.type_name] += 1
1396 for type_name in sorted(success_counts.keys() | failed_counts.keys()):
1397 rows.append(
1398 {
1399 "Task": task_label,
1400 "Exception": type_name,
1401 "Successes": success_counts.get(type_name, 0),
1402 "Failures": failed_counts.get(type_name, 0),
1403 }
1404 )
1405 return astropy.table.Table(rows)
1407 def make_task_resource_usage_table(
1408 self, task_label: TaskLabel, include_data_ids: bool = False
1409 ) -> astropy.table.Table:
1410 """Make a table of resource usage for a single task.
1412 Parameters
1413 ----------
1414 task_label : `str`
1415 Label of the task to extract resource usage for.
1416 include_data_ids : `bool`, optional
1417 Whether to also include data ID columns.
1419 Returns
1420 -------
1421 table : `astropy.table.Table`
1422 A table with columns for quantum ID and all fields in
1423 `QuantumResourceUsage`.
1424 """
1425 quanta_for_task = self.quanta_by_task[task_label]
1426 dtype_terms: list[tuple[str, np.dtype]] = [("quantum_id", np.dtype((np.void, 16)))]
1427 if include_data_ids:
1428 dimensions = self.pipeline_graph.tasks[task_label].dimensions
1429 for dimension_name in dimensions.data_coordinate_keys:
1430 dtype = np.dtype(self.pipeline_graph.universe.dimensions[dimension_name].primary_key.pytype)
1431 dtype_terms.append((dimension_name, dtype))
1432 fields = QuantumResourceUsage.get_numpy_fields()
1433 dtype_terms.extend(fields.items())
1434 row_dtype = np.dtype(dtype_terms)
1435 rows: list[object] = []
1436 for data_id, quantum_id in quanta_for_task.items():
1437 info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1438 if (resource_usage := info["resource_usage"]) is not None:
1439 row: tuple[object, ...] = (quantum_id.bytes,)
1440 if include_data_ids:
1441 row += data_id.full_values
1442 row += resource_usage.get_numpy_row()
1443 rows.append(row)
1444 array = np.array(rows, dtype=row_dtype)
1445 return astropy.table.Table(array, units=QuantumResourceUsage.get_units())
1447 def make_status_report(
1448 self,
1449 states: Iterable[QuantumAttemptStatus] = (
1450 QuantumAttemptStatus.FAILED,
1451 QuantumAttemptStatus.ABORTED,
1452 QuantumAttemptStatus.ABORTED_SUCCESS,
1453 ),
1454 *,
1455 also: QuantumAttemptStatus | Iterable[QuantumAttemptStatus] = (),
1456 with_caveats: QuantumSuccessCaveats | None = QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR,
1457 data_id_table_dir: ResourcePathExpression | None = None,
1458 ) -> ProvenanceReport:
1459 """Make a JSON- or YAML-friendly report of all quanta with the given
1460 states.
1462 Parameters
1463 ----------
1464 states : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1465 `..QuantumAttemptStatus`, optional
1466 A quantum is included if it has any of these states. Defaults to
1467 states that clearly represent problems.
1468 also : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1469 `..QuantumAttemptStatus`, optional
1470 Additional states to consider; unioned with ``states``. This is
1471 provided so users can easily request additional states while also
1472 getting the defaults.
1473 with_caveats : `..QuantumSuccessCaveats` or `None`, optional
1474 If `..QuantumAttemptStatus.SUCCESSFUL` is in ``states``, only
1475 include quanta with these caveat flags. May be set to `None`
1476 to report on all successful quanta.
1477 data_id_table_dir : convertible to `~lsst.resources.ResourcePath`, \
1478 optional
1479 If provided, a directory to write data ID tables (in ECSV format)
1480 with all of the data IDs with the given states, for use with the
1481 ``--data-id-tables`` argument to the quantum graph builder.
1482 Subdirectories for each task and status will created within this
1483 directory, with one file for each exception type (or ``UNKNOWN``
1484 when there is no exception).
1486 Returns
1487 -------
1488 report : `ProvenanceModel`
1489 A Pydantic model that groups quanta by task label and exception
1490 type.
1491 """
1492 states = set(ensure_iterable(states))
1493 states.update(ensure_iterable(also))
1494 result = ProvenanceReport(root={})
1495 if data_id_table_dir is not None:
1496 data_id_table_dir = ResourcePath(data_id_table_dir)
1497 for task_label, quanta_for_task in self.quanta_by_task.items():
1498 reports_for_task: dict[str, dict[str | None, list[ProvenanceQuantumReport]]] = {}
1499 table_rows_for_task: dict[str, dict[str | None, list[tuple[int | str, ...]]]] = {}
1500 for quantum_id in quanta_for_task.values():
1501 quantum_info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1502 quantum_status = quantum_info["status"]
1503 if quantum_status not in states:
1504 continue
1505 if (
1506 quantum_status is QuantumAttemptStatus.SUCCESSFUL
1507 and with_caveats is not None
1508 and (quantum_info["caveats"] is None or not (quantum_info["caveats"] & with_caveats))
1509 ):
1510 continue
1511 key1 = quantum_status.name
1512 exc_info = quantum_info["exception"]
1513 key2 = exc_info.type_name if exc_info is not None else None
1514 reports_for_task.setdefault(key1, {}).setdefault(key2, []).append(
1515 ProvenanceQuantumReport.from_info(quantum_id, quantum_info)
1516 )
1517 if data_id_table_dir:
1518 table_rows_for_task.setdefault(key1, {}).setdefault(key2, []).append(
1519 quantum_info["data_id"].required_values
1520 )
1521 if reports_for_task:
1522 result.root[task_label] = reports_for_task
1523 if table_rows_for_task:
1524 assert data_id_table_dir is not None, "table_rows_for_task should be empty"
1525 for status_name, table_rows_for_status in table_rows_for_task.items():
1526 dir_for_task_and_status = data_id_table_dir.join(task_label, forceDirectory=True).join(
1527 status_name, forceDirectory=True
1528 )
1529 if dir_for_task_and_status.isLocal:
1530 dir_for_task_and_status.mkdir()
1531 for exc_name, data_id_rows in table_rows_for_status.items():
1532 table = astropy.table.Table(
1533 rows=data_id_rows,
1534 names=list(self.pipeline_graph.tasks[task_label].dimensions.required),
1535 )
1536 filename = f"{exc_name}.ecsv" if exc_name is not None else "UNKNOWN.ecsv"
1537 with dir_for_task_and_status.join(filename).open("w") as stream:
1538 table.write(stream, format="ecsv")
1539 return result
1541 def make_many_reports(
1542 self,
1543 states: Iterable[QuantumAttemptStatus] = (
1544 QuantumAttemptStatus.FAILED,
1545 QuantumAttemptStatus.ABORTED,
1546 QuantumAttemptStatus.ABORTED_SUCCESS,
1547 ),
1548 *,
1549 status_report_file: ResourcePathExpression | None = None,
1550 print_quantum_table: bool = False,
1551 print_exception_table: bool = False,
1552 also: QuantumAttemptStatus | Iterable[QuantumAttemptStatus] = (),
1553 with_caveats: QuantumSuccessCaveats | None = None,
1554 data_id_table_dir: ResourcePathExpression | None = None,
1555 ) -> None:
1556 """Write multiple reports.
1558 Parameters
1559 ----------
1560 states : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1561 `..QuantumAttemptStatus`, optional
1562 A quantum is included in the status report and data ID tables if it
1563 has any of these states. Defaults to states that clearly represent
1564 problems.
1565 status_report_file : convertible to `~lsst.resources.ResourcePath`,
1566 optional
1567 Filename for the JSON status report (see `make_status_report`).
1568 print_quantum_table : `bool`, optional
1569 If `True`, print a quantum summary table (counts only) to STDOUT.
1570 print_exception_table : `bool`, optional
1571 If `True`, print an exception-type summary table (counts only) to
1572 STDOUT.
1573 also : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1574 `..QuantumAttemptStatus`, optional
1575 Additional states to consider in the status report and data ID
1576 tables; unioned with ``states``. This is provided so users can
1577 easily request additional states while also getting the defaults.
1578 with_caveats : `..QuantumSuccessCaveats` or `None`, optional
1579 Only include quanta with these caveat flags in the status report
1580 and data ID tables. May be set to `None` to report on all
1581 successful quanta (an empty sequence reports on only quanta with no
1582 caveats). If provided, `QuantumAttemptStatus.SUCCESSFUL` is
1583 automatically included in ``states``.
1584 data_id_table_dir : convertible to `~lsst.resources.ResourcePath`, \
1585 optional
1586 If provided, a directory to write data ID tables (in ECSV format)
1587 with all of the data IDs with the given states, for use with the
1588 ``--data-id-tables`` argument to the quantum graph builder.
1589 Subdirectories for each task and status will created within this
1590 directory, with one file for each exception type (or ``UNKNOWN``
1591 when there is no exception).
1592 """
1593 if status_report_file is not None or data_id_table_dir is not None:
1594 status_report = self.make_status_report(
1595 states, also=also, with_caveats=with_caveats, data_id_table_dir=data_id_table_dir
1596 )
1597 if status_report_file is not None:
1598 status_report_file = ResourcePath(status_report_file)
1599 if status_report_file.isLocal:
1600 status_report_file.dirname().mkdir()
1601 with ResourcePath(status_report_file).open("w") as stream:
1602 stream.write(status_report.model_dump_json(indent=2))
1603 if print_quantum_table:
1604 quantum_table = self.make_quantum_table()
1605 quantum_table.pprint_all()
1606 print("")
1607 if print_exception_table:
1608 exception_table = self.make_exception_table()
1609 exception_table.pprint_all()
1610 print("")
1613@dataclasses.dataclass
1614class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1615 """A helper class for reading provenance quantum graphs.
1617 Notes
1618 -----
1619 The `open` context manager should be used to construct new instances.
1620 Instances cannot be used after the context manager exits, except to access
1621 the `graph` attribute`.
1623 The various ``read_*`` methods in this class update the `graph` attribute
1624 in place.
1625 """
1627 graph: ProvenanceQuantumGraph = dataclasses.field(init=False)
1628 """Loaded provenance graph, populated in place as components are read."""
1630 @classmethod
1631 @contextmanager
1632 def open(
1633 cls,
1634 uri: ResourcePathExpression,
1635 *,
1636 page_size: int | None = None,
1637 import_mode: TaskImportMode = TaskImportMode.DO_NOT_IMPORT,
1638 ) -> Iterator[ProvenanceQuantumGraphReader]:
1639 """Construct a reader from a URI.
1641 Parameters
1642 ----------
1643 uri : convertible to `lsst.resources.ResourcePath`
1644 URI to open. Should have a ``.qg`` extension.
1645 page_size : `int`, optional
1646 Approximate number of bytes to read at once from address files and
1647 multi-block files. Note that this does not set a page size for
1648 *all* reads, but it does affect the smallest, most numerous reads.
1649 Can also be set via the ``LSST_QG_PAGE_SIZE`` environment variable.
1650 import_mode : `.pipeline_graph.TaskImportMode`, optional
1651 How to handle importing the task classes referenced in the pipeline
1652 graph.
1654 Returns
1655 -------
1656 reader : `contextlib.AbstractContextManager` [ \
1657 `ProvenanceQuantumGraphReader` ]
1658 A context manager that returns the reader when entered.
1659 """
1660 with cls._open(
1661 uri,
1662 graph_type="provenance",
1663 address_filename="nodes",
1664 page_size=page_size,
1665 import_mode=import_mode,
1666 n_addresses=4,
1667 ) as self:
1668 yield self
1670 def __post_init__(self) -> None:
1671 self.graph = ProvenanceQuantumGraph(self.header, self.pipeline_graph)
1673 def read_init_quanta(self) -> None:
1674 """Read the thin graph, with all edge information and categorization of
1675 quanta by task label.
1676 """
1677 init_quanta = self._read_single_block("init_quanta", ProvenanceInitQuantaModel)
1678 for init_quantum in init_quanta.root:
1679 self.graph._init_quanta[init_quantum.task_label] = init_quantum.quantum_id
1680 init_quanta._add_to_graph(self.graph)
1682 def read_full_graph(self) -> None:
1683 """Read all bipartite edges and all quantum and dataset node
1684 attributes, fully populating the `graph` attribute.
1686 Notes
1687 -----
1688 This does not read logs, metadata, or packages ; those must always be
1689 fetched explicitly.
1690 """
1691 self.read_init_quanta()
1692 self.read_datasets()
1693 self.read_quanta()
1695 def read_datasets(self, datasets: Iterable[uuid.UUID] | None = None) -> None:
1696 """Read information about the given datasets.
1698 Parameters
1699 ----------
1700 datasets : `~collections.abc.Iterable` [`uuid.UUID`], optional
1701 Iterable of dataset IDs to load. If not provided, all datasets
1702 will be loaded. The UUIDs and indices of quanta will be ignored.
1703 """
1704 self._read_nodes(datasets, DATASET_ADDRESS_INDEX, DATASET_MB_NAME, ProvenanceDatasetModel)
1706 def read_quanta(self, quanta: Iterable[uuid.UUID] | None = None) -> None:
1707 """Read information about the given quanta.
1709 Parameters
1710 ----------
1711 quanta : `~collections.abc.Iterable` [`uuid.UUID`], optional
1712 Iterable of quantum IDs to load. If not provided, all quanta will
1713 be loaded. The UUIDs and indices of datasets and special init
1714 quanta will be ignored.
1715 """
1716 self._read_nodes(quanta, QUANTUM_ADDRESS_INDEX, QUANTUM_MB_NAME, ProvenanceQuantumModel)
1718 def _read_nodes(
1719 self,
1720 nodes: Iterable[uuid.UUID] | None,
1721 address_index: int,
1722 mb_name: str,
1723 model_type: type[ProvenanceDatasetModel] | type[ProvenanceQuantumModel],
1724 ) -> None:
1725 node: ProvenanceDatasetModel | ProvenanceQuantumModel | None
1726 if nodes is None:
1727 self.address_reader.read_all()
1728 nodes = self.address_reader.rows.keys()
1729 for node in MultiblockReader.read_all_models_in_zip(
1730 self.zf,
1731 mb_name,
1732 model_type,
1733 self.decompressor,
1734 int_size=self.header.int_size,
1735 page_size=self.page_size,
1736 ):
1737 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(node.node_id, {}):
1738 # Use the old node to reduce memory usage (since it might
1739 # also have other outstanding reference holders).
1740 continue
1741 node._add_to_graph(self.graph)
1742 else:
1743 with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1744 for node_id_or_index in nodes:
1745 address_row = self.address_reader.find(node_id_or_index)
1746 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1747 # Use the old node to reduce memory usage (since it
1748 # might also have other outstanding reference holders).
1749 continue
1750 node = mb_reader.read_model(
1751 address_row.addresses[address_index], model_type, self.decompressor
1752 )
1753 if node is not None:
1754 node._add_to_graph(self.graph)
1756 def fetch_logs(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[ButlerLogRecords | None]]:
1757 """Fetch log datasets.
1759 Parameters
1760 ----------
1761 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1762 UUIDs of the log datasets themselves or of the quanta they
1763 correspond to.
1765 Returns
1766 -------
1767 logs : `dict` [ `uuid.UUID`, `list` [\
1768 `lsst.daf.butler.ButlerLogRecords` or `None`] ]
1769 Logs for the given IDs. Each value is a list of
1770 `lsst.daf.butler.ButlerLogRecords` instances representing different
1771 execution attempts, ordered chronologically from first to last.
1772 Attempts where logs were missing will have `None` in this list.
1773 """
1774 result: dict[uuid.UUID, list[ButlerLogRecords | None]] = {}
1775 with MultiblockReader.open_in_zip(self.zf, LOG_MB_NAME, int_size=self.header.int_size) as mb_reader:
1776 for node_id_or_index in nodes:
1777 address_row = self.address_reader.find(node_id_or_index)
1778 logs_by_attempt = mb_reader.read_model(
1779 address_row.addresses[LOG_ADDRESS_INDEX], ProvenanceLogRecordsModel, self.decompressor
1780 )
1781 if logs_by_attempt is not None:
1782 result[node_id_or_index] = [
1783 ButlerLogRecords.from_records(attempt_logs) if attempt_logs is not None else None
1784 for attempt_logs in logs_by_attempt.attempts
1785 ]
1786 return result
1788 def fetch_metadata(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[TaskMetadata | None]]:
1789 """Fetch metadata datasets.
1791 Parameters
1792 ----------
1793 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1794 UUIDs of the metadata datasets themselves or of the quanta they
1795 correspond to.
1797 Returns
1798 -------
1799 metadata : `dict` [ `uuid.UUID`, `list` [`.TaskMetadata`] ]
1800 Metadata for the given IDs. Each value is a list of
1801 `.TaskMetadata` instances representing different execution
1802 attempts, ordered chronologically from first to last. Attempts
1803 where metadata was missing (not written even in the fallback extra
1804 provenance in the logs) will have `None` in this list.
1805 """
1806 result: dict[uuid.UUID, list[TaskMetadata | None]] = {}
1807 with MultiblockReader.open_in_zip(
1808 self.zf, METADATA_MB_NAME, int_size=self.header.int_size
1809 ) as mb_reader:
1810 for node_id_or_index in nodes:
1811 address_row = self.address_reader.find(node_id_or_index)
1812 metadata_by_attempt = mb_reader.read_model(
1813 address_row.addresses[METADATA_ADDRESS_INDEX],
1814 ProvenanceTaskMetadataModel,
1815 self.decompressor,
1816 )
1817 if metadata_by_attempt is not None:
1818 result[node_id_or_index] = metadata_by_attempt.attempts
1819 return result
1821 def fetch_packages(self) -> Packages:
1822 """Fetch package version information."""
1823 data = self._read_single_block_raw("packages")
1824 return Packages.fromBytes(data, format="json")
1827class ProvenanceQuantumGraphWriter:
1828 """A struct of low-level writer objects for the main components of a
1829 provenance quantum graph.
1831 Parameters
1832 ----------
1833 output_path : `str`
1834 Path to write the graph to.
1835 exit_stack : `contextlib.ExitStack`
1836 Object that can be used to manage multiple context managers.
1837 log_on_close : `LogOnClose`
1838 Factory for context managers that log when closed.
1839 predicted : `.PredictedQuantumGraphComponents`
1840 Components of the predicted graph.
1841 zstd_level : `int`, optional
1842 Compression level.
1843 cdict_data : `bytes` or `None`, optional
1844 Bytes representation of the compression dictionary used by the
1845 compressor.
1846 loop_wrapper : `~collections.abc.Callable`, optional
1847 A callable that takes an iterable and returns an equivalent one, to be
1848 used in all potentially-large loops. This can be used to add progress
1849 reporting or check for cancelation signals.
1850 log : `LsstLogAdapter`, optional
1851 Logger to use for debug messages.
1852 """
1854 def __init__(
1855 self,
1856 output_path: str,
1857 *,
1858 exit_stack: ExitStack,
1859 log_on_close: LogOnClose,
1860 predicted: PredictedQuantumGraphComponents | PredictedQuantumGraph,
1861 zstd_level: int = 10,
1862 cdict_data: bytes | None = None,
1863 loop_wrapper: LoopWrapper = pass_through,
1864 log: LsstLogAdapter | None = None,
1865 ) -> None:
1866 header = predicted.header.model_copy()
1867 header.graph_type = "provenance"
1868 if log is None:
1869 log = _LOG
1870 self.log = log
1871 self._base_writer = exit_stack.enter_context(
1872 log_on_close.wrap(
1873 BaseQuantumGraphWriter.open(
1874 output_path,
1875 header,
1876 predicted.pipeline_graph,
1877 address_filename="nodes",
1878 zstd_level=zstd_level,
1879 cdict_data=cdict_data,
1880 ),
1881 "Finishing writing provenance quantum graph.",
1882 )
1883 )
1884 self._base_writer.address_writer.addresses = [{}, {}, {}, {}]
1885 self._log_writer = exit_stack.enter_context(
1886 log_on_close.wrap(
1887 MultiblockWriter.open_in_zip(
1888 self._base_writer.zf, LOG_MB_NAME, header.int_size, use_tempfile=True
1889 ),
1890 "Copying logs into zip archive.",
1891 ),
1892 )
1893 self._base_writer.address_writer.addresses[LOG_ADDRESS_INDEX] = self._log_writer.addresses
1894 self._metadata_writer = exit_stack.enter_context(
1895 log_on_close.wrap(
1896 MultiblockWriter.open_in_zip(
1897 self._base_writer.zf, METADATA_MB_NAME, header.int_size, use_tempfile=True
1898 ),
1899 "Copying metadata into zip archive.",
1900 )
1901 )
1902 self._base_writer.address_writer.addresses[METADATA_ADDRESS_INDEX] = self._metadata_writer.addresses
1903 self._dataset_writer = exit_stack.enter_context(
1904 log_on_close.wrap(
1905 MultiblockWriter.open_in_zip(
1906 self._base_writer.zf, DATASET_MB_NAME, header.int_size, use_tempfile=True
1907 ),
1908 "Copying dataset provenance into zip archive.",
1909 )
1910 )
1911 self._base_writer.address_writer.addresses[DATASET_ADDRESS_INDEX] = self._dataset_writer.addresses
1912 self._quantum_writer = exit_stack.enter_context(
1913 log_on_close.wrap(
1914 MultiblockWriter.open_in_zip(
1915 self._base_writer.zf, QUANTUM_MB_NAME, header.int_size, use_tempfile=True
1916 ),
1917 "Copying quantum provenance into zip archive.",
1918 )
1919 )
1920 self._base_writer.address_writer.addresses[QUANTUM_ADDRESS_INDEX] = self._quantum_writer.addresses
1921 self._init_predicted_quanta(predicted)
1922 self._populate_xgraph_and_inputs(loop_wrapper)
1923 self._existing_init_outputs: set[uuid.UUID] = set()
1925 def _init_predicted_quanta(
1926 self, predicted: PredictedQuantumGraph | PredictedQuantumGraphComponents
1927 ) -> None:
1928 self._predicted_init_quanta: list[PredictedQuantumDatasetsModel] = []
1929 self._predicted_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = {}
1930 if isinstance(predicted, PredictedQuantumGraph):
1931 self._predicted_init_quanta.extend(predicted._init_quanta.values())
1932 self._predicted_quanta.update(predicted._quantum_datasets)
1933 else:
1934 self._predicted_init_quanta.extend(predicted.init_quanta.root)
1935 self._predicted_quanta.update(predicted.quantum_datasets)
1936 self._predicted_quanta.update({q.quantum_id: q for q in self._predicted_init_quanta})
1938 def _populate_xgraph_and_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1939 self._xgraph = networkx.DiGraph()
1940 self._overall_inputs: dict[uuid.UUID, PredictedDatasetModel] = {}
1941 output_dataset_ids: set[uuid.UUID] = set()
1942 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1943 if not predicted_quantum.task_label:
1944 # Skip the 'packages' producer quantum.
1945 continue
1946 output_dataset_ids.update(predicted_quantum.iter_output_dataset_ids())
1947 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1948 if not predicted_quantum.task_label:
1949 # Skip the 'packages' producer quantum.
1950 continue
1951 for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
1952 self._xgraph.add_edge(predicted_input.dataset_id, predicted_quantum.quantum_id)
1953 if predicted_input.dataset_id not in output_dataset_ids:
1954 self._overall_inputs.setdefault(predicted_input.dataset_id, predicted_input)
1955 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1956 self._xgraph.add_edge(predicted_quantum.quantum_id, predicted_output.dataset_id)
1958 @property
1959 def compressor(self) -> Compressor:
1960 """Object that should be used to compress all JSON blocks."""
1961 return self._base_writer.compressor
1963 def write_packages(self) -> None:
1964 """Write package version information to the provenance graph."""
1965 packages = Packages.fromSystem(include_all=True)
1966 data = packages.toBytes("json")
1967 self._base_writer.write_single_block("packages", data)
1969 def write_overall_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1970 """Write provenance for overall-input datasets.
1972 Parameters
1973 ----------
1974 loop_wrapper : `~collections.abc.Callable`, optional
1975 A callable that takes an iterable and returns an equivalent one, to
1976 be used in all potentially-large loops. This can be used to add
1977 progress reporting or check for cancelation signals.
1978 """
1979 for predicted_input in loop_wrapper(self._overall_inputs.values()):
1980 if predicted_input.dataset_id not in self._dataset_writer.addresses:
1981 self._dataset_writer.write_model(
1982 predicted_input.dataset_id,
1983 ProvenanceDatasetModel.from_predicted(
1984 predicted_input,
1985 producer=None,
1986 consumers=self._xgraph.successors(predicted_input.dataset_id),
1987 ),
1988 self.compressor,
1989 )
1990 del self._overall_inputs
1992 def write_init_outputs(self, assume_existence: bool = True) -> None:
1993 """Write provenance for init-output datasets and init-quanta.
1995 Parameters
1996 ----------
1997 assume_existence : `bool`, optional
1998 If `True`, just assume all init-outputs exist.
1999 """
2000 init_quanta = ProvenanceInitQuantaModel()
2001 for predicted_init_quantum in self._predicted_init_quanta:
2002 if not predicted_init_quantum.task_label:
2003 # Skip the 'packages' producer quantum.
2004 continue
2005 for predicted_output in itertools.chain.from_iterable(predicted_init_quantum.outputs.values()):
2006 provenance_output = ProvenanceDatasetModel.from_predicted(
2007 predicted_output,
2008 producer=predicted_init_quantum.quantum_id,
2009 consumers=self._xgraph.successors(predicted_output.dataset_id),
2010 )
2011 provenance_output.produced = assume_existence or (
2012 provenance_output.dataset_id in self._existing_init_outputs
2013 )
2014 self._dataset_writer.write_model(
2015 provenance_output.dataset_id, provenance_output, self.compressor
2016 )
2017 init_quanta.root.append(ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum))
2018 self._base_writer.write_single_model("init_quanta", init_quanta)
2020 def write_quantum_provenance(
2021 self, quantum_id: uuid.UUID, metadata: TaskMetadata | None, logs: ButlerLogRecords | None
2022 ) -> None:
2023 """Gather and write provenance for a quantum.
2025 Parameters
2026 ----------
2027 quantum_id : `uuid.UUID`
2028 Unique ID for the quantum.
2029 metadata : `..TaskMetadata` or `None`
2030 Task metadata.
2031 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
2032 Task logs.
2033 """
2034 predicted_quantum = self._predicted_quanta[quantum_id]
2035 provenance_models = ProvenanceQuantumScanModels.from_metadata_and_logs(
2036 predicted_quantum, metadata, logs, incomplete=False
2037 )
2038 scan_data = provenance_models.to_scan_data(predicted_quantum, compressor=self.compressor)
2039 self.write_scan_data(scan_data)
2041 def write_scan_data(self, scan_data: ProvenanceQuantumScanData) -> None:
2042 """Write the output of a quantum provenance scan to disk.
2044 Parameters
2045 ----------
2046 scan_data : `ProvenanceQuantumScanData`
2047 Result of a quantum provenance scan.
2048 """
2049 if scan_data.status is ProvenanceQuantumScanStatus.INIT:
2050 self.log.debug("Handling init-output scan for %s.", scan_data.quantum_id)
2051 self._existing_init_outputs.update(scan_data.existing_outputs)
2052 return
2053 self.log.debug("Handling quantum scan for %s.", scan_data.quantum_id)
2054 # We shouldn't need this predicted quantum after this method runs; pop
2055 # from the dict it in the hopes that'll free up some memory when we're
2056 # done.
2057 predicted_quantum = self._predicted_quanta.pop(scan_data.quantum_id)
2058 outputs: dict[uuid.UUID, bytes] = {}
2059 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
2060 provenance_output = ProvenanceDatasetModel.from_predicted(
2061 predicted_output,
2062 producer=predicted_quantum.quantum_id,
2063 consumers=self._xgraph.successors(predicted_output.dataset_id),
2064 )
2065 provenance_output.produced = provenance_output.dataset_id in scan_data.existing_outputs
2066 outputs[provenance_output.dataset_id] = self.compressor.compress(
2067 provenance_output.model_dump_json().encode()
2068 )
2069 if not scan_data.quantum:
2070 scan_data.quantum = (
2071 ProvenanceQuantumModel.from_predicted(predicted_quantum).model_dump_json().encode()
2072 )
2073 if scan_data.is_compressed:
2074 scan_data.quantum = self.compressor.compress(scan_data.quantum)
2075 if not scan_data.is_compressed:
2076 scan_data.quantum = self.compressor.compress(scan_data.quantum)
2077 if scan_data.metadata:
2078 scan_data.metadata = self.compressor.compress(scan_data.metadata)
2079 if scan_data.logs:
2080 scan_data.logs = self.compressor.compress(scan_data.logs)
2081 self.log.debug("Writing quantum %s.", scan_data.quantum_id)
2082 self._quantum_writer.write_bytes(scan_data.quantum_id, scan_data.quantum)
2083 for dataset_id, dataset_data in outputs.items():
2084 self._dataset_writer.write_bytes(dataset_id, dataset_data)
2085 if scan_data.metadata:
2086 (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
2087 address = self._metadata_writer.write_bytes(scan_data.quantum_id, scan_data.metadata)
2088 self._metadata_writer.addresses[metadata_output.dataset_id] = address
2089 if scan_data.logs:
2090 (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
2091 address = self._log_writer.write_bytes(scan_data.quantum_id, scan_data.logs)
2092 self._log_writer.addresses[log_output.dataset_id] = address
2095class ProvenanceQuantumScanStatus(enum.Enum):
2096 """Status enum for quantum scanning.
2098 Note that this records the status for the *scanning* which is distinct
2099 from the status of the quantum's execution.
2100 """
2102 INCOMPLETE = enum.auto()
2103 """The quantum is not necessarily done running, and cannot be scanned
2104 conclusively yet.
2105 """
2107 ABANDONED = enum.auto()
2108 """The quantum's execution appears to have failed but we cannot rule out
2109 the possibility that it could be recovered, but we've also waited long
2110 enough (according to `ScannerTimeConfigDict.retry_timeout`) that it's time
2111 to stop trying for now.
2113 This state means `ProvenanceQuantumScanModels.from_metadata_and_logs` must
2114 be run again with ``incomplete=False``.
2115 """
2117 SUCCESSFUL = enum.auto()
2118 """The quantum was conclusively scanned and was executed successfully,
2119 unblocking scans for downstream quanta.
2120 """
2122 FAILED = enum.auto()
2123 """The quantum was conclusively scanned and failed execution, blocking
2124 scans for downstream quanta.
2125 """
2127 BLOCKED = enum.auto()
2128 """A quantum upstream of this one failed."""
2130 INIT = enum.auto()
2131 """Init quanta need special handling, because they don't have logs and
2132 metadata.
2133 """
2136@dataclasses.dataclass
2137class ProvenanceQuantumScanModels:
2138 """A struct that represents provenance information for a single quantum."""
2140 quantum_id: uuid.UUID
2141 """Unique ID for the quantum."""
2143 status: ProvenanceQuantumScanStatus = ProvenanceQuantumScanStatus.INCOMPLETE
2144 """Combined status for the scan and the execution of the quantum."""
2146 attempts: list[ProvenanceQuantumAttemptModel] = dataclasses.field(default_factory=list)
2147 """Provenance information about each attempt to run the quantum."""
2149 output_existence: dict[uuid.UUID, bool] = dataclasses.field(default_factory=dict)
2150 """Unique IDs of the output datasets mapped to whether they were actually
2151 produced.
2152 """
2154 metadata: ProvenanceTaskMetadataModel = dataclasses.field(default_factory=ProvenanceTaskMetadataModel)
2155 """Task metadata information for each attempt.
2156 """
2158 logs: ProvenanceLogRecordsModel = dataclasses.field(default_factory=ProvenanceLogRecordsModel)
2159 """Log records for each attempt.
2160 """
2162 @classmethod
2163 def from_metadata_and_logs(
2164 cls,
2165 predicted: PredictedQuantumDatasetsModel,
2166 metadata: TaskMetadata | None,
2167 logs: ButlerLogRecords | None,
2168 *,
2169 incomplete: bool = False,
2170 ) -> ProvenanceQuantumScanModels:
2171 """Construct provenance information from task metadata and logs.
2173 Parameters
2174 ----------
2175 predicted : `PredictedQuantumDatasetsModel`
2176 Information about the predicted quantum.
2177 metadata : `..TaskMetadata` or `None`
2178 Task metadata.
2179 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
2180 Task logs.
2181 incomplete : `bool`, optional
2182 If `True`, treat execution failures as possibly-incomplete quanta
2183 and do not fully process them; instead just set the status to
2184 `ProvenanceQuantumScanStatus.ABANDONED` and return.
2186 Returns
2187 -------
2188 scan_models : `ProvenanceQuantumScanModels`
2189 Struct of models that describe quantum provenance.
2191 Notes
2192 -----
2193 This method does not necessarily fully populate the `output_existence`
2194 field; it does what it can given the information in the metadata and
2195 logs, but the caller is responsible for filling in the existence status
2196 for any predicted outputs that are not present at all in that `dict`.
2197 """
2198 self = ProvenanceQuantumScanModels(predicted.quantum_id)
2199 last_attempt = ProvenanceQuantumAttemptModel()
2200 self._process_logs(predicted, logs, last_attempt, incomplete=incomplete)
2201 self._process_metadata(predicted, metadata, last_attempt, incomplete=incomplete)
2202 if self.status is ProvenanceQuantumScanStatus.ABANDONED:
2203 return self
2204 self._reconcile_attempts(last_attempt)
2205 self._extract_output_existence(predicted)
2206 return self
2208 def _process_logs(
2209 self,
2210 predicted: PredictedQuantumDatasetsModel,
2211 logs: ButlerLogRecords | None,
2212 last_attempt: ProvenanceQuantumAttemptModel,
2213 *,
2214 incomplete: bool,
2215 ) -> None:
2216 (predicted_log_dataset,) = predicted.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
2217 if logs is None:
2218 self.output_existence[predicted_log_dataset.dataset_id] = False
2219 if incomplete:
2220 self.status = ProvenanceQuantumScanStatus.ABANDONED
2221 else:
2222 self.status = ProvenanceQuantumScanStatus.FAILED
2223 else:
2224 # Set the attempt's run status to FAILED, since the default is
2225 # UNKNOWN (i.e. logs *and* metadata are missing) and we now know
2226 # the logs exist. This will usually get replaced by SUCCESSFUL
2227 # when we look for metadata next.
2228 last_attempt.status = QuantumAttemptStatus.FAILED
2229 self.output_existence[predicted_log_dataset.dataset_id] = True
2230 if logs.extra:
2231 log_extra = _ExecutionLogRecordsExtra.model_validate(logs.extra)
2232 self._extract_from_log_extra(log_extra, last_attempt=last_attempt)
2233 self.logs.attempts.append(list(logs))
2235 def _extract_from_log_extra(
2236 self,
2237 log_extra: _ExecutionLogRecordsExtra,
2238 last_attempt: ProvenanceQuantumAttemptModel | None,
2239 ) -> None:
2240 for previous_attempt_log_extra in log_extra.previous_attempts:
2241 self._extract_from_log_extra(
2242 previous_attempt_log_extra,
2243 last_attempt=None,
2244 )
2245 quantum_attempt: ProvenanceQuantumAttemptModel
2246 if last_attempt is None:
2247 # This is not the last attempt, so it must be a failure.
2248 quantum_attempt = ProvenanceQuantumAttemptModel(
2249 attempt=len(self.attempts), status=QuantumAttemptStatus.FAILED
2250 )
2251 # We also need to get the logs from this extra provenance, since
2252 # they won't be the main section of the log records.
2253 self.logs.attempts.append(log_extra.logs)
2254 # The special last attempt is only appended after we attempt to
2255 # read metadata later, but we have to append this one now.
2256 self.attempts.append(quantum_attempt)
2257 else:
2258 assert not log_extra.logs, "Logs for the last attempt should not be stored in the extra JSON."
2259 quantum_attempt = last_attempt
2260 if log_extra.exception is not None or log_extra.metadata is not None or last_attempt is None:
2261 # We won't be getting a separate metadata dataset, so anything we
2262 # might get from the metadata has to come from this extra
2263 # provenance in the logs.
2264 quantum_attempt.exception = log_extra.exception
2265 if log_extra.metadata is not None:
2266 quantum_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(log_extra.metadata)
2267 self.metadata.attempts.append(log_extra.metadata)
2268 else:
2269 self.metadata.attempts.append(None)
2270 # Regardless of whether this is the last attempt or not, we can only
2271 # get the previous_process_quanta from the log extra.
2272 quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)
2274 def _process_metadata(
2275 self,
2276 predicted: PredictedQuantumDatasetsModel,
2277 metadata: TaskMetadata | None,
2278 last_attempt: ProvenanceQuantumAttemptModel,
2279 *,
2280 incomplete: bool,
2281 ) -> None:
2282 (predicted_metadata_dataset,) = predicted.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
2283 if metadata is None:
2284 self.output_existence[predicted_metadata_dataset.dataset_id] = False
2285 if incomplete:
2286 self.status = ProvenanceQuantumScanStatus.ABANDONED
2287 else:
2288 self.status = ProvenanceQuantumScanStatus.FAILED
2289 else:
2290 self.status = ProvenanceQuantumScanStatus.SUCCESSFUL
2291 self.output_existence[predicted_metadata_dataset.dataset_id] = True
2292 last_attempt.status = QuantumAttemptStatus.SUCCESSFUL
2293 try:
2294 # Int conversion guards against spurious conversion to
2295 # float that can apparently sometimes happen in
2296 # TaskMetadata.
2297 last_attempt.caveats = QuantumSuccessCaveats(int(metadata["quantum"]["caveats"]))
2298 except LookupError:
2299 pass
2300 try:
2301 last_attempt.exception = ExceptionInfo._from_metadata(
2302 metadata[predicted.task_label]["failure"]
2303 )
2304 except LookupError:
2305 pass
2306 last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)
2307 self.metadata.attempts.append(metadata)
2309 def _reconcile_attempts(self, last_attempt: ProvenanceQuantumAttemptModel) -> None:
2310 last_attempt.attempt = len(self.attempts)
2311 self.attempts.append(last_attempt)
2312 assert self.status is not ProvenanceQuantumScanStatus.INCOMPLETE
2313 assert self.status is not ProvenanceQuantumScanStatus.ABANDONED
2314 if len(self.logs.attempts) < len(self.attempts):
2315 # Logs were not found for this attempt; must have been a hard error
2316 # that kept the `finally` block from running or otherwise
2317 # interrupted the writing of the logs.
2318 self.logs.attempts.append(None)
2319 if self.status is ProvenanceQuantumScanStatus.SUCCESSFUL:
2320 # But we found the metadata! Either that hard error happened
2321 # at a very unlucky time (in between those two writes), or
2322 # something even weirder happened.
2323 self.attempts[-1].status = QuantumAttemptStatus.ABORTED_SUCCESS
2324 else:
2325 self.attempts[-1].status = QuantumAttemptStatus.FAILED
2326 if len(self.metadata.attempts) < len(self.attempts):
2327 # Metadata missing usually just means a failure. In any case, the
2328 # status will already be correct, either because it was set to a
2329 # failure when we read the logs, or left at UNKNOWN if there were
2330 # no logs. Note that scanners never process BLOCKED quanta at all.
2331 self.metadata.attempts.append(None)
2332 assert len(self.logs.attempts) == len(self.attempts) or len(self.metadata.attempts) == len(
2333 self.attempts
2334 ), (
2335 "The only way we can add more than one quantum attempt is by "
2336 "extracting info stored with the logs, and that always appends "
2337 "a log attempt and a metadata attempt, so this must be a bug in "
2338 "this class."
2339 )
2341 def _extract_output_existence(self, predicted: PredictedQuantumDatasetsModel) -> None:
2342 try:
2343 outputs_put = self.metadata.attempts[-1]["quantum"].getArray("outputs") # type: ignore[index]
2344 except (
2345 IndexError, # metadata.attempts is empty
2346 TypeError, # metadata.attempts[-1] is None
2347 LookupError, # no 'quantum' entry in metadata or 'outputs' in that
2348 ):
2349 pass
2350 else:
2351 for id_str in ensure_iterable(outputs_put):
2352 self.output_existence[uuid.UUID(id_str)] = True
2353 # If the metadata told us what it wrote, anything not in that
2354 # list was not written.
2355 for predicted_output in itertools.chain.from_iterable(predicted.outputs.values()):
2356 self.output_existence.setdefault(predicted_output.dataset_id, False)
2358 def to_scan_data(
2359 self: ProvenanceQuantumScanModels,
2360 predicted_quantum: PredictedQuantumDatasetsModel,
2361 compressor: Compressor | None = None,
2362 ) -> ProvenanceQuantumScanData:
2363 """Convert these models to JSON data.
2365 Parameters
2366 ----------
2367 predicted_quantum : `PredictedQuantumDatasetsModel`
2368 Information about the predicted quantum.
2369 compressor : `Compressor`
2370 Object that can compress bytes.
2372 Returns
2373 -------
2374 scan_data : `ProvenanceQuantumScanData`
2375 Scan information ready for serialization.
2376 """
2377 quantum: ProvenanceInitQuantumModel | ProvenanceQuantumModel
2378 if self.status is ProvenanceQuantumScanStatus.INIT:
2379 quantum = ProvenanceInitQuantumModel.from_predicted(predicted_quantum)
2380 else:
2381 quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum)
2382 quantum.attempts = self.attempts
2383 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
2384 if predicted_output.dataset_id not in self.output_existence:
2385 raise RuntimeError(
2386 "Logic bug in provenance gathering or execution invariants: "
2387 f"no existence information for output {predicted_output.dataset_id} "
2388 f"({predicted_output.dataset_type_name}@{predicted_output.data_coordinate})."
2389 )
2390 data = ProvenanceQuantumScanData(
2391 self.quantum_id,
2392 self.status,
2393 existing_outputs={
2394 dataset_id for dataset_id, was_produced in self.output_existence.items() if was_produced
2395 },
2396 quantum=quantum.model_dump_json().encode(),
2397 logs=self.logs.model_dump_json().encode() if self.logs.attempts else b"",
2398 metadata=self.metadata.model_dump_json().encode() if self.metadata.attempts else b"",
2399 )
2400 if compressor is not None:
2401 data.compress(compressor)
2402 return data
2405@dataclasses.dataclass
2406class ProvenanceQuantumScanData:
2407 """A struct that represents ready-for-serialization provenance information
2408 for a single quantum.
2409 """
2411 quantum_id: uuid.UUID
2412 """Unique ID for the quantum."""
2414 status: ProvenanceQuantumScanStatus
2415 """Combined status for the scan and the execution of the quantum."""
2417 existing_outputs: set[uuid.UUID] = dataclasses.field(default_factory=set)
2418 """Unique IDs of the output datasets that were actually written."""
2420 quantum: bytes = b""
2421 """Serialized quantum provenance model.
2423 This may be empty for quanta that had no attempts.
2424 """
2426 metadata: bytes = b""
2427 """Serialized task metadata."""
2429 logs: bytes = b""
2430 """Serialized logs."""
2432 is_compressed: bool = False
2433 """Whether the ``quantum``, ``metadata``, and ``log`` attributes are
2434 compressed.
2435 """
2437 def compress(self, compressor: Compressor) -> None:
2438 """Compress the data in this struct if it has not been compressed
2439 already.
2441 Parameters
2442 ----------
2443 compressor : `Compressor`
2444 Object with a ``compress`` method that takes and returns `bytes`.
2445 """
2446 if not self.is_compressed:
2447 self.quantum = compressor.compress(self.quantum)
2448 self.logs = compressor.compress(self.logs) if self.logs else b""
2449 self.metadata = compressor.compress(self.metadata) if self.metadata else b""
2450 self.is_compressed = True