Coverage for python / lsst / pipe / base / quantum_graph / _provenance.py: 27%
692 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-04 02:47 -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 attempts: list[TaskMetadata | None] = pydantic.Field(default_factory=list)
511 """Metadata from attempts to run this task, ordered chronologically from
512 first to last.
513 """
515 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
516 # when we inherit those docstrings in our public classes.
517 if "sphinx" in sys.modules and not TYPE_CHECKING:
519 def copy(self, *args: Any, **kwargs: Any) -> Any:
520 """See `pydantic.BaseModel.copy`."""
521 return super().copy(*args, **kwargs)
523 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
524 """See `pydantic.BaseModel.model_dump`."""
525 return super().model_dump(*args, **kwargs)
527 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
528 """See `pydantic.BaseModel.model_dump_json`."""
529 return super().model_dump(*args, **kwargs)
531 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
532 """See `pydantic.BaseModel.model_copy`."""
533 return super().model_copy(*args, **kwargs)
535 @classmethod
536 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
537 """See `pydantic.BaseModel.model_construct`."""
538 return super().model_construct(*args, **kwargs)
540 @classmethod
541 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
542 """See `pydantic.BaseModel.model_json_schema`."""
543 return super().model_json_schema(*args, **kwargs)
545 @classmethod
546 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
547 """See `pydantic.BaseModel.model_validate`."""
548 return super().model_validate(*args, **kwargs)
550 @classmethod
551 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
552 """See `pydantic.BaseModel.model_validate_json`."""
553 return super().model_validate_json(*args, **kwargs)
555 @classmethod
556 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
557 """See `pydantic.BaseModel.model_validate_strings`."""
558 return super().model_validate_strings(*args, **kwargs)
561class ProvenanceQuantumReport(pydantic.BaseModel):
562 """A Pydantic model that used to report information about a single
563 (generally problematic) quantum.
564 """
566 quantum_id: uuid.UUID
567 data_id: dict[str, int | str]
568 attempts: list[ProvenanceQuantumAttemptModel]
570 @classmethod
571 def from_info(cls, quantum_id: uuid.UUID, quantum_info: ProvenanceQuantumInfo) -> ProvenanceQuantumReport:
572 """Construct from a provenance quantum graph node.
574 Parameters
575 ----------
576 quantum_id : `uuid.UUID`
577 Unique ID for the quantum.
578 quantum_info : `ProvenanceQuantumInfo`
579 Node attributes for this quantum.
580 """
581 return cls(
582 quantum_id=quantum_id,
583 data_id=dict(quantum_info["data_id"].mapping),
584 attempts=quantum_info["attempts"],
585 )
587 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
588 # when we inherit those docstrings in our public classes.
589 if "sphinx" in sys.modules and not TYPE_CHECKING:
591 def copy(self, *args: Any, **kwargs: Any) -> Any:
592 """See `pydantic.BaseModel.copy`."""
593 return super().copy(*args, **kwargs)
595 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
596 """See `pydantic.BaseModel.model_dump`."""
597 return super().model_dump(*args, **kwargs)
599 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
600 """See `pydantic.BaseModel.model_dump_json`."""
601 return super().model_dump(*args, **kwargs)
603 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
604 """See `pydantic.BaseModel.model_copy`."""
605 return super().model_copy(*args, **kwargs)
607 @classmethod
608 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
609 """See `pydantic.BaseModel.model_construct`."""
610 return super().model_construct(*args, **kwargs)
612 @classmethod
613 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
614 """See `pydantic.BaseModel.model_json_schema`."""
615 return super().model_json_schema(*args, **kwargs)
617 @classmethod
618 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
619 """See `pydantic.BaseModel.model_validate`."""
620 return super().model_validate(*args, **kwargs)
622 @classmethod
623 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
624 """See `pydantic.BaseModel.model_validate_json`."""
625 return super().model_validate_json(*args, **kwargs)
627 @classmethod
628 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
629 """See `pydantic.BaseModel.model_validate_strings`."""
630 return super().model_validate_strings(*args, **kwargs)
633class ProvenanceReport(pydantic.RootModel):
634 """A Pydantic model that groups quantum information by task label, then
635 status (as a string), and then exception type.
636 """
638 root: dict[TaskLabel, dict[str, dict[str | None, list[ProvenanceQuantumReport]]]] = {}
640 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
641 # when we inherit those docstrings in our public classes.
642 if "sphinx" in sys.modules and not TYPE_CHECKING:
644 def copy(self, *args: Any, **kwargs: Any) -> Any:
645 """See `pydantic.BaseModel.copy`."""
646 return super().copy(*args, **kwargs)
648 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
649 """See `pydantic.BaseModel.model_dump`."""
650 return super().model_dump(*args, **kwargs)
652 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
653 """See `pydantic.BaseModel.model_dump_json`."""
654 return super().model_dump(*args, **kwargs)
656 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
657 """See `pydantic.BaseModel.model_copy`."""
658 return super().model_copy(*args, **kwargs)
660 @classmethod
661 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
662 """See `pydantic.BaseModel.model_construct`."""
663 return super().model_construct(*args, **kwargs)
665 @classmethod
666 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
667 """See `pydantic.BaseModel.model_json_schema`."""
668 return super().model_json_schema(*args, **kwargs)
670 @classmethod
671 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
672 """See `pydantic.BaseModel.model_validate`."""
673 return super().model_validate(*args, **kwargs)
675 @classmethod
676 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
677 """See `pydantic.BaseModel.model_validate_json`."""
678 return super().model_validate_json(*args, **kwargs)
680 @classmethod
681 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
682 """See `pydantic.BaseModel.model_validate_strings`."""
683 return super().model_validate_strings(*args, **kwargs)
686class ProvenanceQuantumModel(pydantic.BaseModel):
687 """Data model for the quanta in a provenance quantum graph file."""
689 quantum_id: uuid.UUID
690 """Unique identifier for the quantum."""
692 task_label: TaskLabel
693 """Name of the type of this dataset."""
695 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
696 """The full values (required and implied) of this dataset's data ID."""
698 inputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
699 """IDs of the datasets predicted to be consumed by this quantum, grouped by
700 connection name.
701 """
703 outputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
704 """IDs of the datasets predicted to be produced by this quantum, grouped by
705 connection name.
706 """
708 attempts: list[ProvenanceQuantumAttemptModel] = pydantic.Field(default_factory=list)
709 """Provenance for all attempts to execute this quantum, ordered
710 chronologically from first to last.
712 An entry is added merely if the quantum *should* have been attempted; an
713 empty `list` is used only for quanta that were blocked by an upstream
714 failure.
715 """
717 @property
718 def node_id(self) -> uuid.UUID:
719 """Alias for the quantum ID."""
720 return self.quantum_id
722 @classmethod
723 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceQuantumModel:
724 """Construct from a predicted quantum model.
726 Parameters
727 ----------
728 predicted : `PredictedQuantumDatasetsModel`
729 Information about the quantum from the predicted graph.
731 Returns
732 -------
733 provenance : `ProvenanceQuantumModel`
734 Provenance quantum model.
735 """
736 inputs = {
737 connection_name: [d.dataset_id for d in predicted_inputs]
738 for connection_name, predicted_inputs in predicted.inputs.items()
739 }
740 outputs = {
741 connection_name: [d.dataset_id for d in predicted_outputs]
742 for connection_name, predicted_outputs in predicted.outputs.items()
743 }
744 return cls(
745 quantum_id=predicted.quantum_id,
746 task_label=predicted.task_label,
747 data_coordinate=predicted.data_coordinate,
748 inputs=inputs,
749 outputs=outputs,
750 )
752 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
753 """Add this quantum and its edges to datasets to a provenance graph.
755 Parameters
756 ----------
757 graph : `ProvenanceQuantumGraph`
758 Graph to update in place.
760 Notes
761 -----
762 This method adds:
764 - a ``bipartite_xgraph`` quantum node with full attributes;
765 - a ``quantum_only_xgraph`` quantum node with full attributes;
766 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
767 nodes with no attributes), while populating those edge attributes;
768 - ``quantum_only_xgraph`` edges to any adjacent quantum that has also
769 already been loaded.
770 """
771 task_node = graph.pipeline_graph.tasks[self.task_label]
772 data_id = DataCoordinate.from_full_values(task_node.dimensions, tuple(self.data_coordinate))
773 last_attempt = (
774 self.attempts[-1]
775 if self.attempts
776 else ProvenanceQuantumAttemptModel(status=QuantumAttemptStatus.BLOCKED)
777 )
778 graph._bipartite_xgraph.add_node(
779 self.quantum_id,
780 data_id=data_id,
781 task_label=self.task_label,
782 pipeline_node=task_node,
783 status=last_attempt.status,
784 caveats=last_attempt.caveats,
785 exception=last_attempt.exception,
786 resource_usage=last_attempt.resource_usage,
787 attempts=self.attempts,
788 )
789 graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
790 graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
791 for connection_name, dataset_ids in self.inputs.items():
792 read_edge = task_node.get_input_edge(connection_name)
793 for dataset_id in dataset_ids:
794 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
795 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
796 "pipeline_edges", []
797 ).append(read_edge)
798 for connection_name, dataset_ids in self.outputs.items():
799 write_edge = task_node.get_output_edge(connection_name)
800 if connection_name == acc.METADATA_OUTPUT_CONNECTION_NAME:
801 graph._bipartite_xgraph.add_node(
802 dataset_ids[0],
803 data_id=data_id,
804 dataset_type_name=write_edge.dataset_type_name,
805 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
806 run=graph.header.output_run,
807 produced=last_attempt.status.has_metadata,
808 )
809 graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
810 graph._bipartite_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
811 graph._quantum_only_xgraph.nodes[self.quantum_id]["metadata_id"] = dataset_ids[0]
812 if connection_name == acc.LOG_OUTPUT_CONNECTION_NAME:
813 graph._bipartite_xgraph.add_node(
814 dataset_ids[0],
815 data_id=data_id,
816 dataset_type_name=write_edge.dataset_type_name,
817 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
818 run=graph.header.output_run,
819 produced=last_attempt.status.has_log,
820 )
821 graph._datasets_by_type[write_edge.dataset_type_name][data_id] = dataset_ids[0]
822 graph._bipartite_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
823 graph._quantum_only_xgraph.nodes[self.quantum_id]["log_id"] = dataset_ids[0]
824 for dataset_id in dataset_ids:
825 graph._bipartite_xgraph.add_edge(
826 self.quantum_id,
827 dataset_id,
828 is_read=False,
829 # There can only be one pipeline edge for an output.
830 pipeline_edges=[write_edge],
831 )
832 for dataset_id in graph._bipartite_xgraph.predecessors(self.quantum_id):
833 for upstream_quantum_id in graph._bipartite_xgraph.predecessors(dataset_id):
834 graph._quantum_only_xgraph.add_edge(upstream_quantum_id, self.quantum_id)
835 for dataset_id in graph._bipartite_xgraph.successors(self.quantum_id):
836 for downstream_quantum_id in graph._bipartite_xgraph.successors(dataset_id):
837 graph._quantum_only_xgraph.add_edge(self.quantum_id, downstream_quantum_id)
839 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
840 # when we inherit those docstrings in our public classes.
841 if "sphinx" in sys.modules and not TYPE_CHECKING:
843 def copy(self, *args: Any, **kwargs: Any) -> Any:
844 """See `pydantic.BaseModel.copy`."""
845 return super().copy(*args, **kwargs)
847 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
848 """See `pydantic.BaseModel.model_dump`."""
849 return super().model_dump(*args, **kwargs)
851 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
852 """See `pydantic.BaseModel.model_dump_json`."""
853 return super().model_dump(*args, **kwargs)
855 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
856 """See `pydantic.BaseModel.model_copy`."""
857 return super().model_copy(*args, **kwargs)
859 @classmethod
860 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
861 """See `pydantic.BaseModel.model_construct`."""
862 return super().model_construct(*args, **kwargs)
864 @classmethod
865 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
866 """See `pydantic.BaseModel.model_json_schema`."""
867 return super().model_json_schema(*args, **kwargs)
869 @classmethod
870 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
871 """See `pydantic.BaseModel.model_validate`."""
872 return super().model_validate(*args, **kwargs)
874 @classmethod
875 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
876 """See `pydantic.BaseModel.model_validate_json`."""
877 return super().model_validate_json(*args, **kwargs)
879 @classmethod
880 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
881 """See `pydantic.BaseModel.model_validate_strings`."""
882 return super().model_validate_strings(*args, **kwargs)
885class ProvenanceInitQuantumModel(pydantic.BaseModel):
886 """Data model for the special "init" quanta in a provenance quantum graph
887 file.
888 """
890 quantum_id: uuid.UUID
891 """Unique identifier for the quantum."""
893 task_label: TaskLabel
894 """Name of the type of this dataset.
896 This is always a parent dataset type name, not a component.
898 Note that full dataset type definitions are stored in the pipeline graph.
899 """
901 inputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
902 """IDs of the datasets predicted to be consumed by this quantum, grouped by
903 connection name.
904 """
906 outputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
907 """IDs of the datasets predicted to be produced by this quantum, grouped by
908 connection name.
909 """
911 @classmethod
912 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceInitQuantumModel:
913 """Construct from a predicted quantum model.
915 Parameters
916 ----------
917 predicted : `PredictedQuantumDatasetsModel`
918 Information about the quantum from the predicted graph.
920 Returns
921 -------
922 provenance : `ProvenanceInitQuantumModel`
923 Provenance init quantum model.
924 """
925 inputs = {
926 connection_name: predicted_inputs[0].dataset_id
927 for connection_name, predicted_inputs in predicted.inputs.items()
928 }
929 outputs = {
930 connection_name: predicted_outputs[0].dataset_id
931 for connection_name, predicted_outputs in predicted.outputs.items()
932 }
933 return cls(
934 quantum_id=predicted.quantum_id,
935 task_label=predicted.task_label,
936 inputs=inputs,
937 outputs=outputs,
938 )
940 def _add_to_graph(self, graph: ProvenanceQuantumGraph, empty_data_id: DataCoordinate) -> None:
941 """Add this quantum and its edges to datasets to a provenance graph.
943 Parameters
944 ----------
945 graph : `ProvenanceQuantumGraph`
946 Graph to update in place.
947 empty_data_id : `lsst.daf.butler.DataCoordinate`
948 The empty data ID for the appropriate dimension universe.
950 Notes
951 -----
952 This method adds:
954 - a ``bipartite_xgraph`` quantum node with full attributes;
955 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
956 nodes with no attributes), while populating those edge attributes;
957 """
958 task_init_node = graph.pipeline_graph.tasks[self.task_label].init
959 graph._bipartite_xgraph.add_node(
960 self.quantum_id, data_id=empty_data_id, task_label=self.task_label, pipeline_node=task_init_node
961 )
962 for connection_name, dataset_id in self.inputs.items():
963 read_edge = task_init_node.get_input_edge(connection_name)
964 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
965 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
966 "pipeline_edges", []
967 ).append(read_edge)
968 for connection_name, dataset_id in self.outputs.items():
969 write_edge = task_init_node.get_output_edge(connection_name)
970 graph._bipartite_xgraph.add_node(
971 dataset_id,
972 data_id=empty_data_id,
973 dataset_type_name=write_edge.dataset_type_name,
974 pipeline_node=graph.pipeline_graph.dataset_types[write_edge.dataset_type_name],
975 run=graph.header.output_run,
976 produced=True,
977 )
978 graph._datasets_by_type[write_edge.dataset_type_name][empty_data_id] = dataset_id
979 graph._bipartite_xgraph.add_edge(
980 self.quantum_id,
981 dataset_id,
982 is_read=False,
983 # There can only be one pipeline edge for an output.
984 pipeline_edges=[write_edge],
985 )
986 if write_edge.connection_name == acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME:
987 graph._bipartite_xgraph.nodes[self.quantum_id]["config_id"] = dataset_id
988 graph._init_quanta[self.task_label] = self.quantum_id
990 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
991 # when we inherit those docstrings in our public classes.
992 if "sphinx" in sys.modules and not TYPE_CHECKING:
994 def copy(self, *args: Any, **kwargs: Any) -> Any:
995 """See `pydantic.BaseModel.copy`."""
996 return super().copy(*args, **kwargs)
998 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
999 """See `pydantic.BaseModel.model_dump`."""
1000 return super().model_dump(*args, **kwargs)
1002 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
1003 """See `pydantic.BaseModel.model_dump_json`."""
1004 return super().model_dump(*args, **kwargs)
1006 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
1007 """See `pydantic.BaseModel.model_copy`."""
1008 return super().model_copy(*args, **kwargs)
1010 @classmethod
1011 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
1012 """See `pydantic.BaseModel.model_construct`."""
1013 return super().model_construct(*args, **kwargs)
1015 @classmethod
1016 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
1017 """See `pydantic.BaseModel.model_json_schema`."""
1018 return super().model_json_schema(*args, **kwargs)
1020 @classmethod
1021 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
1022 """See `pydantic.BaseModel.model_validate`."""
1023 return super().model_validate(*args, **kwargs)
1025 @classmethod
1026 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
1027 """See `pydantic.BaseModel.model_validate_json`."""
1028 return super().model_validate_json(*args, **kwargs)
1030 @classmethod
1031 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
1032 """See `pydantic.BaseModel.model_validate_strings`."""
1033 return super().model_validate_strings(*args, **kwargs)
1036class ProvenanceInitQuantaModel(pydantic.RootModel):
1037 """Data model for the init quanta in a provenance graph."""
1039 root: list[ProvenanceInitQuantumModel] = pydantic.Field(default_factory=list)
1040 """List of special "init" quanta, one for each task."""
1042 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
1043 """Add this quantum and its edges to datasets to a provenance graph.
1045 Parameters
1046 ----------
1047 graph : `ProvenanceQuantumGraph`
1048 Graph to update in place.
1049 """
1050 empty_data_id = DataCoordinate.make_empty(graph.pipeline_graph.universe)
1051 for init_quantum in self.root:
1052 init_quantum._add_to_graph(graph, empty_data_id=empty_data_id)
1054 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
1055 # when we inherit those docstrings in our public classes.
1056 if "sphinx" in sys.modules and not TYPE_CHECKING:
1058 def copy(self, *args: Any, **kwargs: Any) -> Any:
1059 """See `pydantic.BaseModel.copy`."""
1060 return super().copy(*args, **kwargs)
1062 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
1063 """See `pydantic.BaseModel.model_dump`."""
1064 return super().model_dump(*args, **kwargs)
1066 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
1067 """See `pydantic.BaseModel.model_dump_json`."""
1068 return super().model_dump(*args, **kwargs)
1070 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
1071 """See `pydantic.BaseModel.model_copy`."""
1072 return super().model_copy(*args, **kwargs)
1074 @classmethod
1075 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
1076 """See `pydantic.BaseModel.model_construct`."""
1077 return super().model_construct(*args, **kwargs)
1079 @classmethod
1080 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
1081 """See `pydantic.BaseModel.model_json_schema`."""
1082 return super().model_json_schema(*args, **kwargs)
1084 @classmethod
1085 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
1086 """See `pydantic.BaseModel.model_validate`."""
1087 return super().model_validate(*args, **kwargs)
1089 @classmethod
1090 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
1091 """See `pydantic.BaseModel.model_validate_json`."""
1092 return super().model_validate_json(*args, **kwargs)
1094 @classmethod
1095 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
1096 """See `pydantic.BaseModel.model_validate_strings`."""
1097 return super().model_validate_strings(*args, **kwargs)
1100class ProvenanceQuantumGraph(BaseQuantumGraph):
1101 """A quantum graph that represents processing that has already been
1102 executed.
1104 Parameters
1105 ----------
1106 header : `HeaderModel`
1107 General metadata shared with other quantum graph types.
1108 pipeline_graph : `.pipeline_graph.PipelineGraph`
1109 Graph of tasks and dataset types. May contain a superset of the tasks
1110 and dataset types that actually have quanta and datasets in the quantum
1111 graph.
1113 Notes
1114 -----
1115 A provenance quantum graph is generally obtained via the
1116 `ProvenanceQuantumGraphReader.graph` attribute, which is updated in-place
1117 as information is read from disk.
1118 """
1120 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph) -> None:
1121 super().__init__(header, pipeline_graph)
1122 self._init_quanta: dict[TaskLabel, uuid.UUID] = {}
1123 self._quantum_only_xgraph = networkx.DiGraph()
1124 self._bipartite_xgraph = networkx.DiGraph()
1125 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {
1126 task_label: {} for task_label in self.pipeline_graph.tasks.keys()
1127 }
1128 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {
1129 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
1130 }
1132 @classmethod
1133 @contextmanager
1134 def from_args(
1135 cls,
1136 repo_or_filename: str,
1137 /,
1138 collection: str | None = None,
1139 *,
1140 quanta: Iterable[uuid.UUID] | None = None,
1141 datasets: Iterable[uuid.UUID] | None = None,
1142 writeable: bool = False,
1143 ) -> Iterator[tuple[ProvenanceQuantumGraph, Butler | None]]:
1144 """Construct a `ProvenanceQuantumGraph` fron CLI-friendly arguments for
1145 a file or butler-ingested graph dataset.
1147 Parameters
1148 ----------
1149 repo_or_filename : `str`
1150 Either a provenance quantum graph filename or a butler repository
1151 path or alias.
1152 collection : `str`, optional
1153 Collection to search; presence indicates that the first argument
1154 is a butler repository, not a filename.
1155 quanta : `~collections.abc.Iterable` [ `str` ] or `None`, optional
1156 IDs of the quanta to load, or `None` to load all.
1157 datasets : `~collections.abc.Iterable` [ `str` ], optional
1158 IDs of the datasets to load, or `None` to load all.
1159 writeable : `bool`, optional
1160 Whether the butler should be constructed with write support.
1162 Returns
1163 -------
1164 context : `contextlib.AbstractContextManager`
1165 A context manager that yields a tuple of
1167 - the `ProvenanceQuantumGraph`
1168 - the `Butler` constructed (or `None`)
1170 when entered.
1171 """
1172 exit_stack = ExitStack()
1173 if collection is not None:
1174 try:
1175 butler = exit_stack.enter_context(
1176 Butler.from_config(repo_or_filename, collections=[collection], writeable=writeable)
1177 )
1178 except Exception as err:
1179 err.add_note(
1180 f"Expected {repo_or_filename!r} to be a butler repository path or alias because a "
1181 f"collection ({collection}) was provided."
1182 )
1183 raise
1184 with exit_stack:
1185 graph = butler.get(
1186 acc.PROVENANCE_DATASET_TYPE_NAME, parameters={"quanta": quanta, "datasets": datasets}
1187 )
1188 yield graph, butler
1189 else:
1190 try:
1191 reader = exit_stack.enter_context(ProvenanceQuantumGraphReader.open(repo_or_filename))
1192 except Exception as err:
1193 err.add_note(
1194 f"Expected a {repo_or_filename} to be a provenance quantum graph filename "
1195 f"because no collection was provided."
1196 )
1197 raise
1198 with exit_stack:
1199 if quanta is None:
1200 reader.read_quanta()
1201 elif not quanta:
1202 reader.read_quanta(quanta)
1203 if datasets is None:
1204 reader.read_datasets()
1205 elif not datasets:
1206 reader.read_datasets(datasets)
1207 yield reader.graph, None
1209 @property
1210 def init_quanta(self) -> Mapping[TaskLabel, uuid.UUID]:
1211 """A mapping from task label to the ID of the special init quantum for
1212 that task.
1214 This is populated by the ``init_quanta`` component. Additional
1215 information about each init quantum can be found by using the ID to
1216 look up node attributes in the `bipartite_xgraph`, i.e.::
1218 info: ProvenanceInitQuantumInfo = qg.bipartite_xgraph.nodes[id]
1219 """
1220 return self._init_quanta
1222 @property
1223 def quanta_by_task(self) -> Mapping[TaskLabel, Mapping[DataCoordinate, uuid.UUID]]:
1224 """A nested mapping of all quanta, keyed first by task name and then by
1225 data ID.
1227 Notes
1228 -----
1229 This is populated one quantum at a time as they are read. All tasks in
1230 the pipeline graph are included, even if none of their quanta were
1231 loaded (i.e. nested mappings may be empty).
1233 The returned object may be an internal dictionary; as the type
1234 annotation indicates, it should not be modified in place.
1235 """
1236 return self._quanta_by_task_label
1238 @property
1239 def datasets_by_type(self) -> Mapping[DatasetTypeName, Mapping[DataCoordinate, uuid.UUID]]:
1240 """A nested mapping of all datasets, keyed first by dataset type name
1241 and then by data ID.
1243 Notes
1244 -----
1245 This is populated one dataset at a time as they are read. All dataset
1246 types in the pipeline graph are included, even if none of their
1247 datasets were loaded (i.e. nested mappings may be empty).
1249 Reading a quantum also populates its log and metadata datasets.
1251 The returned object may be an internal dictionary; as the type
1252 annotation indicates, it should not be modified in place.
1253 """
1254 return self._datasets_by_type
1256 @property
1257 def quantum_only_xgraph(self) -> networkx.DiGraph:
1258 """A directed acyclic graph with quanta as nodes (and datasets elided).
1260 Notes
1261 -----
1262 Node keys are quantum UUIDs, and are populated one quantum at a time as
1263 they are loaded. Loading quanta (via
1264 `ProvenanceQuantumGraphReader.read_quanta`) will add the loaded nodes
1265 with full attributes and add edges to adjacent nodes with no
1266 attributes. Loading datasets (via
1267 `ProvenanceQuantumGraphReader.read_datasets`) will also add edges and
1268 nodes with no attributes.
1270 Node attributes are described by the `ProvenanceQuantumInfo` types.
1272 This graph does not include special "init" quanta.
1274 The returned object is a read-only view of an internal one.
1275 """
1276 return self._quantum_only_xgraph.copy(as_view=True)
1278 @property
1279 def bipartite_xgraph(self) -> networkx.DiGraph:
1280 """A directed acyclic graph with quantum and dataset nodes.
1282 Notes
1283 -----
1284 Node keys are quantum or dataset UUIDs, and are populated one quantum
1285 or dataset at a time as they are loaded. Loading quanta (via
1286 `ProvenanceQuantumGraphReader.read_quanta`) or datasets (via
1287 `ProvenanceQuantumGraphReader.read_datasets`) will load those nodes
1288 with full attributes and edges to adjacent nodes with no attributes.
1289 Loading quanta is necessary to populate edge attributes.
1290 Reading a quantum also populates its log and metadata datasets.
1292 Node attributes are described by the
1293 `ProvenanceQuantumInfo`, `ProvenanceInitQuantumInfo`, and
1294 `ProvenanceDatasetInfo` types.
1296 This graph includes init-input and init-output datasets, but it does
1297 *not* reflect the dependency between each task's special "init" quantum
1298 and its runtime quanta (as this would require edges between quanta, and
1299 that would break the "bipartite" property).
1301 The returned object is a read-only view of an internal one.
1302 """
1303 return self._bipartite_xgraph.copy(as_view=True)
1305 def make_quantum_table(self, drop_unused_columns: bool = True) -> astropy.table.Table:
1306 """Construct an `astropy.table.Table` with a tabular summary of the
1307 quanta.
1309 Parameters
1310 ----------
1311 drop_unused_columns : `bool`, optional
1312 Whether to drop columns for rare states that did not actually
1313 occur in this run.
1315 Returns
1316 -------
1317 table : `astropy.table.Table`
1318 A table view of the quantum information. This only includes
1319 counts of status categories and caveats, not any per-data-ID
1320 detail.
1322 Notes
1323 -----
1324 Success caveats in the table are represented by their
1325 `~QuantumSuccessCaveats.concise` form, so when pretty-printing this
1326 table for users, the `~QuantumSuccessCaveats.legend` should generally
1327 be printed as well.
1328 """
1329 rows = []
1330 for task_label, quanta_for_task in self.quanta_by_task.items():
1331 if not self.header.n_task_quanta[task_label]:
1332 continue
1333 status_counts = Counter[QuantumAttemptStatus](
1334 self._quantum_only_xgraph.nodes[q]["status"] for q in quanta_for_task.values()
1335 )
1336 caveat_counts = Counter[QuantumSuccessCaveats | None](
1337 self._quantum_only_xgraph.nodes[q]["caveats"] for q in quanta_for_task.values()
1338 )
1339 caveat_counts.pop(QuantumSuccessCaveats.NO_CAVEATS, None)
1340 caveat_counts.pop(None, None)
1341 if len(caveat_counts) > 1:
1342 caveats = "(multiple)"
1343 elif len(caveat_counts) == 1:
1344 ((code, count),) = caveat_counts.items()
1345 # MyPy can't tell that the pop(None, None) above makes None
1346 # impossible here.
1347 caveats = f"{code.concise()}({count})" # type: ignore[union-attr]
1348 else:
1349 caveats = ""
1350 row: dict[str, Any] = {
1351 "Task": task_label,
1352 "Caveats": caveats,
1353 }
1354 for status in QuantumAttemptStatus:
1355 row[status.title] = status_counts.get(status, 0)
1356 row.update(
1357 {
1358 "TOTAL": len(quanta_for_task),
1359 "EXPECTED": self.header.n_task_quanta[task_label],
1360 }
1361 )
1362 rows.append(row)
1363 table = astropy.table.Table(rows)
1364 if drop_unused_columns:
1365 for status in QuantumAttemptStatus:
1366 if status.is_rare and not table[status.title].any():
1367 del table[status.title]
1368 return table
1370 def make_exception_table(self) -> astropy.table.Table:
1371 """Construct an `astropy.table.Table` with counts for each exception
1372 type raised by each task.
1374 Returns
1375 -------
1376 table : `astropy.table.Table`
1377 A table with columns for task label, exception type, and counts.
1378 """
1379 rows = []
1380 for task_label, quanta_for_task in self.quanta_by_task.items():
1381 success_counts = Counter[str]()
1382 failed_counts = Counter[str]()
1383 for quantum_id in quanta_for_task.values():
1384 quantum_info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1385 exc_info = quantum_info["exception"]
1386 if exc_info is not None:
1387 if quantum_info["status"] is QuantumAttemptStatus.SUCCESSFUL:
1388 success_counts[exc_info.type_name] += 1
1389 else:
1390 failed_counts[exc_info.type_name] += 1
1391 for type_name in sorted(success_counts.keys() | failed_counts.keys()):
1392 rows.append(
1393 {
1394 "Task": task_label,
1395 "Exception": type_name,
1396 "Successes": success_counts.get(type_name, 0),
1397 "Failures": failed_counts.get(type_name, 0),
1398 }
1399 )
1400 return astropy.table.Table(rows)
1402 def make_task_resource_usage_table(
1403 self, task_label: TaskLabel, include_data_ids: bool = False
1404 ) -> astropy.table.Table:
1405 """Make a table of resource usage for a single task.
1407 Parameters
1408 ----------
1409 task_label : `str`
1410 Label of the task to extract resource usage for.
1411 include_data_ids : `bool`, optional
1412 Whether to also include data ID columns.
1414 Returns
1415 -------
1416 table : `astropy.table.Table`
1417 A table with columns for quantum ID and all fields in
1418 `QuantumResourceUsage`.
1419 """
1420 quanta_for_task = self.quanta_by_task[task_label]
1421 dtype_terms: list[tuple[str, np.dtype]] = [("quantum_id", np.dtype((np.void, 16)))]
1422 if include_data_ids:
1423 dimensions = self.pipeline_graph.tasks[task_label].dimensions
1424 for dimension_name in dimensions.data_coordinate_keys:
1425 dtype = np.dtype(self.pipeline_graph.universe.dimensions[dimension_name].primary_key.pytype)
1426 dtype_terms.append((dimension_name, dtype))
1427 fields = QuantumResourceUsage.get_numpy_fields()
1428 dtype_terms.extend(fields.items())
1429 row_dtype = np.dtype(dtype_terms)
1430 rows: list[object] = []
1431 for data_id, quantum_id in quanta_for_task.items():
1432 info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1433 if (resource_usage := info["resource_usage"]) is not None:
1434 row: tuple[object, ...] = (quantum_id.bytes,)
1435 if include_data_ids:
1436 row += data_id.full_values
1437 row += resource_usage.get_numpy_row()
1438 rows.append(row)
1439 array = np.array(rows, dtype=row_dtype)
1440 return astropy.table.Table(array, units=QuantumResourceUsage.get_units())
1442 def make_status_report(
1443 self,
1444 states: Iterable[QuantumAttemptStatus] = (
1445 QuantumAttemptStatus.FAILED,
1446 QuantumAttemptStatus.ABORTED,
1447 QuantumAttemptStatus.ABORTED_SUCCESS,
1448 ),
1449 *,
1450 also: QuantumAttemptStatus | Iterable[QuantumAttemptStatus] = (),
1451 with_caveats: QuantumSuccessCaveats | None = QuantumSuccessCaveats.PARTIAL_OUTPUTS_ERROR,
1452 data_id_table_dir: ResourcePathExpression | None = None,
1453 ) -> ProvenanceReport:
1454 """Make a JSON- or YAML-friendly report of all quanta with the given
1455 states.
1457 Parameters
1458 ----------
1459 states : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1460 `..QuantumAttemptStatus`, optional
1461 A quantum is included if it has any of these states. Defaults to
1462 states that clearly represent problems.
1463 also : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1464 `..QuantumAttemptStatus`, optional
1465 Additional states to consider; unioned with ``states``. This is
1466 provided so users can easily request additional states while also
1467 getting the defaults.
1468 with_caveats : `..QuantumSuccessCaveats` or `None`, optional
1469 If `..QuantumAttemptStatus.SUCCESSFUL` is in ``states``, only
1470 include quanta with these caveat flags. May be set to `None`
1471 to report on all successful quanta.
1472 data_id_table_dir : convertible to `~lsst.resources.ResourcePath`, \
1473 optional
1474 If provided, a directory to write data ID tables (in ECSV format)
1475 with all of the data IDs with the given states, for use with the
1476 ``--data-id-tables`` argument to the quantum graph builder.
1477 Subdirectories for each task and status will created within this
1478 directory, with one file for each exception type (or ``UNKNOWN``
1479 when there is no exception).
1481 Returns
1482 -------
1483 report : `ProvenanceModel`
1484 A Pydantic model that groups quanta by task label and exception
1485 type.
1486 """
1487 states = set(ensure_iterable(states))
1488 states.update(ensure_iterable(also))
1489 result = ProvenanceReport(root={})
1490 if data_id_table_dir is not None:
1491 data_id_table_dir = ResourcePath(data_id_table_dir)
1492 for task_label, quanta_for_task in self.quanta_by_task.items():
1493 reports_for_task: dict[str, dict[str | None, list[ProvenanceQuantumReport]]] = {}
1494 table_rows_for_task: dict[str, dict[str | None, list[tuple[int | str, ...]]]] = {}
1495 for quantum_id in quanta_for_task.values():
1496 quantum_info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1497 quantum_status = quantum_info["status"]
1498 if quantum_status not in states:
1499 continue
1500 if (
1501 quantum_status is QuantumAttemptStatus.SUCCESSFUL
1502 and with_caveats is not None
1503 and (quantum_info["caveats"] is None or not (quantum_info["caveats"] & with_caveats))
1504 ):
1505 continue
1506 key1 = quantum_status.name
1507 exc_info = quantum_info["exception"]
1508 key2 = exc_info.type_name if exc_info is not None else None
1509 reports_for_task.setdefault(key1, {}).setdefault(key2, []).append(
1510 ProvenanceQuantumReport.from_info(quantum_id, quantum_info)
1511 )
1512 if data_id_table_dir:
1513 table_rows_for_task.setdefault(key1, {}).setdefault(key2, []).append(
1514 quantum_info["data_id"].required_values
1515 )
1516 if reports_for_task:
1517 result.root[task_label] = reports_for_task
1518 if table_rows_for_task:
1519 assert data_id_table_dir is not None, "table_rows_for_task should be empty"
1520 for status_name, table_rows_for_status in table_rows_for_task.items():
1521 dir_for_task_and_status = data_id_table_dir.join(task_label, forceDirectory=True).join(
1522 status_name, forceDirectory=True
1523 )
1524 if dir_for_task_and_status.isLocal:
1525 dir_for_task_and_status.mkdir()
1526 for exc_name, data_id_rows in table_rows_for_status.items():
1527 table = astropy.table.Table(
1528 rows=data_id_rows,
1529 names=list(self.pipeline_graph.tasks[task_label].dimensions.required),
1530 )
1531 filename = f"{exc_name}.ecsv" if exc_name is not None else "UNKNOWN.ecsv"
1532 with dir_for_task_and_status.join(filename).open("w") as stream:
1533 table.write(stream, format="ecsv")
1534 return result
1536 def make_many_reports(
1537 self,
1538 states: Iterable[QuantumAttemptStatus] = (
1539 QuantumAttemptStatus.FAILED,
1540 QuantumAttemptStatus.ABORTED,
1541 QuantumAttemptStatus.ABORTED_SUCCESS,
1542 ),
1543 *,
1544 status_report_file: ResourcePathExpression | None = None,
1545 print_quantum_table: bool = False,
1546 print_exception_table: bool = False,
1547 also: QuantumAttemptStatus | Iterable[QuantumAttemptStatus] = (),
1548 with_caveats: QuantumSuccessCaveats | None = None,
1549 data_id_table_dir: ResourcePathExpression | None = None,
1550 ) -> None:
1551 """Write multiple reports.
1553 Parameters
1554 ----------
1555 states : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1556 `..QuantumAttemptStatus`, optional
1557 A quantum is included in the status report and data ID tables if it
1558 has any of these states. Defaults to states that clearly represent
1559 problems.
1560 status_report_file : convertible to `~lsst.resources.ResourcePath`,
1561 optional
1562 Filename for the JSON status report (see `make_status_report`).
1563 print_quantum_table : `bool`, optional
1564 If `True`, print a quantum summary table (counts only) to STDOUT.
1565 print_exception_table : `bool`, optional
1566 If `True`, print an exception-type summary table (counts only) to
1567 STDOUT.
1568 also : `~collections.abc.Iterable` [`..QuantumAttemptStatus`] or \
1569 `..QuantumAttemptStatus`, optional
1570 Additional states to consider in the status report and data ID
1571 tables; unioned with ``states``. This is provided so users can
1572 easily request additional states while also getting the defaults.
1573 with_caveats : `..QuantumSuccessCaveats` or `None`, optional
1574 Only include quanta with these caveat flags in the status report
1575 and data ID tables. May be set to `None` to report on all
1576 successful quanta (an empty sequence reports on only quanta with no
1577 caveats). If provided, `QuantumAttemptStatus.SUCCESSFUL` is
1578 automatically included in ``states``.
1579 data_id_table_dir : convertible to `~lsst.resources.ResourcePath`, \
1580 optional
1581 If provided, a directory to write data ID tables (in ECSV format)
1582 with all of the data IDs with the given states, for use with the
1583 ``--data-id-tables`` argument to the quantum graph builder.
1584 Subdirectories for each task and status will created within this
1585 directory, with one file for each exception type (or ``UNKNOWN``
1586 when there is no exception).
1587 """
1588 if status_report_file is not None or data_id_table_dir is not None:
1589 status_report = self.make_status_report(
1590 states, also=also, with_caveats=with_caveats, data_id_table_dir=data_id_table_dir
1591 )
1592 if status_report_file is not None:
1593 status_report_file = ResourcePath(status_report_file)
1594 if status_report_file.isLocal:
1595 status_report_file.dirname().mkdir()
1596 with ResourcePath(status_report_file).open("w") as stream:
1597 stream.write(status_report.model_dump_json(indent=2))
1598 if print_quantum_table:
1599 quantum_table = self.make_quantum_table()
1600 quantum_table.pprint_all()
1601 print("")
1602 if print_exception_table:
1603 exception_table = self.make_exception_table()
1604 exception_table.pprint_all()
1605 print("")
1608@dataclasses.dataclass
1609class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1610 """A helper class for reading provenance quantum graphs.
1612 Notes
1613 -----
1614 The `open` context manager should be used to construct new instances.
1615 Instances cannot be used after the context manager exits, except to access
1616 the `graph` attribute`.
1618 The various ``read_*`` methods in this class update the `graph` attribute
1619 in place.
1620 """
1622 graph: ProvenanceQuantumGraph = dataclasses.field(init=False)
1623 """Loaded provenance graph, populated in place as components are read."""
1625 @classmethod
1626 @contextmanager
1627 def open(
1628 cls,
1629 uri: ResourcePathExpression,
1630 *,
1631 page_size: int | None = None,
1632 import_mode: TaskImportMode = TaskImportMode.DO_NOT_IMPORT,
1633 ) -> Iterator[ProvenanceQuantumGraphReader]:
1634 """Construct a reader from a URI.
1636 Parameters
1637 ----------
1638 uri : convertible to `lsst.resources.ResourcePath`
1639 URI to open. Should have a ``.qg`` extension.
1640 page_size : `int`, optional
1641 Approximate number of bytes to read at once from address files and
1642 multi-block files. Note that this does not set a page size for
1643 *all* reads, but it does affect the smallest, most numerous reads.
1644 Can also be set via the ``LSST_QG_PAGE_SIZE`` environment variable.
1645 import_mode : `.pipeline_graph.TaskImportMode`, optional
1646 How to handle importing the task classes referenced in the pipeline
1647 graph.
1649 Returns
1650 -------
1651 reader : `contextlib.AbstractContextManager` [ \
1652 `ProvenanceQuantumGraphReader` ]
1653 A context manager that returns the reader when entered.
1654 """
1655 with cls._open(
1656 uri,
1657 graph_type="provenance",
1658 address_filename="nodes",
1659 page_size=page_size,
1660 import_mode=import_mode,
1661 n_addresses=4,
1662 ) as self:
1663 yield self
1665 def __post_init__(self) -> None:
1666 self.graph = ProvenanceQuantumGraph(self.header, self.pipeline_graph)
1668 def read_init_quanta(self) -> None:
1669 """Read the thin graph, with all edge information and categorization of
1670 quanta by task label.
1671 """
1672 init_quanta = self._read_single_block("init_quanta", ProvenanceInitQuantaModel)
1673 for init_quantum in init_quanta.root:
1674 self.graph._init_quanta[init_quantum.task_label] = init_quantum.quantum_id
1675 init_quanta._add_to_graph(self.graph)
1677 def read_full_graph(self) -> None:
1678 """Read all bipartite edges and all quantum and dataset node
1679 attributes, fully populating the `graph` attribute.
1681 Notes
1682 -----
1683 This does not read logs, metadata, or packages ; those must always be
1684 fetched explicitly.
1685 """
1686 self.read_init_quanta()
1687 self.read_datasets()
1688 self.read_quanta()
1690 def read_datasets(self, datasets: Iterable[uuid.UUID] | None = None) -> None:
1691 """Read information about the given datasets.
1693 Parameters
1694 ----------
1695 datasets : `~collections.abc.Iterable` [`uuid.UUID`], optional
1696 Iterable of dataset IDs to load. If not provided, all datasets
1697 will be loaded. The UUIDs and indices of quanta will be ignored.
1698 """
1699 self._read_nodes(datasets, DATASET_ADDRESS_INDEX, DATASET_MB_NAME, ProvenanceDatasetModel)
1701 def read_quanta(self, quanta: Iterable[uuid.UUID] | None = None) -> None:
1702 """Read information about the given quanta.
1704 Parameters
1705 ----------
1706 quanta : `~collections.abc.Iterable` [`uuid.UUID`], optional
1707 Iterable of quantum IDs to load. If not provided, all quanta will
1708 be loaded. The UUIDs and indices of datasets and special init
1709 quanta will be ignored.
1710 """
1711 self._read_nodes(quanta, QUANTUM_ADDRESS_INDEX, QUANTUM_MB_NAME, ProvenanceQuantumModel)
1713 def _read_nodes(
1714 self,
1715 nodes: Iterable[uuid.UUID] | None,
1716 address_index: int,
1717 mb_name: str,
1718 model_type: type[ProvenanceDatasetModel] | type[ProvenanceQuantumModel],
1719 ) -> None:
1720 node: ProvenanceDatasetModel | ProvenanceQuantumModel | None
1721 if nodes is None:
1722 self.address_reader.read_all()
1723 nodes = self.address_reader.rows.keys()
1724 for node in MultiblockReader.read_all_models_in_zip(
1725 self.zf,
1726 mb_name,
1727 model_type,
1728 self.decompressor,
1729 int_size=self.header.int_size,
1730 page_size=self.page_size,
1731 ):
1732 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(node.node_id, {}):
1733 # Use the old node to reduce memory usage (since it might
1734 # also have other outstanding reference holders).
1735 continue
1736 node._add_to_graph(self.graph)
1737 else:
1738 with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1739 for node_id_or_index in nodes:
1740 address_row = self.address_reader.find(node_id_or_index)
1741 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1742 # Use the old node to reduce memory usage (since it
1743 # might also have other outstanding reference holders).
1744 continue
1745 node = mb_reader.read_model(
1746 address_row.addresses[address_index], model_type, self.decompressor
1747 )
1748 if node is not None:
1749 node._add_to_graph(self.graph)
1751 def fetch_logs(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[ButlerLogRecords | None]]:
1752 """Fetch log datasets.
1754 Parameters
1755 ----------
1756 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1757 UUIDs of the log datasets themselves or of the quanta they
1758 correspond to.
1760 Returns
1761 -------
1762 logs : `dict` [ `uuid.UUID`, `list` [\
1763 `lsst.daf.butler.ButlerLogRecords` or `None`] ]
1764 Logs for the given IDs. Each value is a list of
1765 `lsst.daf.butler.ButlerLogRecords` instances representing different
1766 execution attempts, ordered chronologically from first to last.
1767 Attempts where logs were missing will have `None` in this list.
1768 """
1769 result: dict[uuid.UUID, list[ButlerLogRecords | None]] = {}
1770 with MultiblockReader.open_in_zip(self.zf, LOG_MB_NAME, int_size=self.header.int_size) as mb_reader:
1771 for node_id_or_index in nodes:
1772 address_row = self.address_reader.find(node_id_or_index)
1773 logs_by_attempt = mb_reader.read_model(
1774 address_row.addresses[LOG_ADDRESS_INDEX], ProvenanceLogRecordsModel, self.decompressor
1775 )
1776 if logs_by_attempt is not None:
1777 result[node_id_or_index] = [
1778 ButlerLogRecords.from_records(attempt_logs) if attempt_logs is not None else None
1779 for attempt_logs in logs_by_attempt.attempts
1780 ]
1781 return result
1783 def fetch_metadata(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[TaskMetadata | None]]:
1784 """Fetch metadata datasets.
1786 Parameters
1787 ----------
1788 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1789 UUIDs of the metadata datasets themselves or of the quanta they
1790 correspond to.
1792 Returns
1793 -------
1794 metadata : `dict` [ `uuid.UUID`, `list` [`.TaskMetadata`] ]
1795 Metadata for the given IDs. Each value is a list of
1796 `.TaskMetadata` instances representing different execution
1797 attempts, ordered chronologically from first to last. Attempts
1798 where metadata was missing (not written even in the fallback extra
1799 provenance in the logs) will have `None` in this list.
1800 """
1801 result: dict[uuid.UUID, list[TaskMetadata | None]] = {}
1802 with MultiblockReader.open_in_zip(
1803 self.zf, METADATA_MB_NAME, int_size=self.header.int_size
1804 ) as mb_reader:
1805 for node_id_or_index in nodes:
1806 address_row = self.address_reader.find(node_id_or_index)
1807 metadata_by_attempt = mb_reader.read_model(
1808 address_row.addresses[METADATA_ADDRESS_INDEX],
1809 ProvenanceTaskMetadataModel,
1810 self.decompressor,
1811 )
1812 if metadata_by_attempt is not None:
1813 result[node_id_or_index] = metadata_by_attempt.attempts
1814 return result
1816 def fetch_packages(self) -> Packages:
1817 """Fetch package version information."""
1818 data = self._read_single_block_raw("packages")
1819 return Packages.fromBytes(data, format="json")
1822class ProvenanceQuantumGraphWriter:
1823 """A struct of low-level writer objects for the main components of a
1824 provenance quantum graph.
1826 Parameters
1827 ----------
1828 output_path : `str`
1829 Path to write the graph to.
1830 exit_stack : `contextlib.ExitStack`
1831 Object that can be used to manage multiple context managers.
1832 log_on_close : `LogOnClose`
1833 Factory for context managers that log when closed.
1834 predicted : `.PredictedQuantumGraphComponents`
1835 Components of the predicted graph.
1836 zstd_level : `int`, optional
1837 Compression level.
1838 cdict_data : `bytes` or `None`, optional
1839 Bytes representation of the compression dictionary used by the
1840 compressor.
1841 loop_wrapper : `~collections.abc.Callable`, optional
1842 A callable that takes an iterable and returns an equivalent one, to be
1843 used in all potentially-large loops. This can be used to add progress
1844 reporting or check for cancelation signals.
1845 log : `LsstLogAdapter`, optional
1846 Logger to use for debug messages.
1847 """
1849 def __init__(
1850 self,
1851 output_path: str,
1852 *,
1853 exit_stack: ExitStack,
1854 log_on_close: LogOnClose,
1855 predicted: PredictedQuantumGraphComponents | PredictedQuantumGraph,
1856 zstd_level: int = 10,
1857 cdict_data: bytes | None = None,
1858 loop_wrapper: LoopWrapper = pass_through,
1859 log: LsstLogAdapter | None = None,
1860 ) -> None:
1861 header = predicted.header.model_copy()
1862 header.graph_type = "provenance"
1863 if log is None:
1864 log = _LOG
1865 self.log = log
1866 self._base_writer = exit_stack.enter_context(
1867 log_on_close.wrap(
1868 BaseQuantumGraphWriter.open(
1869 output_path,
1870 header,
1871 predicted.pipeline_graph,
1872 address_filename="nodes",
1873 zstd_level=zstd_level,
1874 cdict_data=cdict_data,
1875 ),
1876 "Finishing writing provenance quantum graph.",
1877 )
1878 )
1879 self._base_writer.address_writer.addresses = [{}, {}, {}, {}]
1880 self._log_writer = exit_stack.enter_context(
1881 log_on_close.wrap(
1882 MultiblockWriter.open_in_zip(
1883 self._base_writer.zf, LOG_MB_NAME, header.int_size, use_tempfile=True
1884 ),
1885 "Copying logs into zip archive.",
1886 ),
1887 )
1888 self._base_writer.address_writer.addresses[LOG_ADDRESS_INDEX] = self._log_writer.addresses
1889 self._metadata_writer = exit_stack.enter_context(
1890 log_on_close.wrap(
1891 MultiblockWriter.open_in_zip(
1892 self._base_writer.zf, METADATA_MB_NAME, header.int_size, use_tempfile=True
1893 ),
1894 "Copying metadata into zip archive.",
1895 )
1896 )
1897 self._base_writer.address_writer.addresses[METADATA_ADDRESS_INDEX] = self._metadata_writer.addresses
1898 self._dataset_writer = exit_stack.enter_context(
1899 log_on_close.wrap(
1900 MultiblockWriter.open_in_zip(
1901 self._base_writer.zf, DATASET_MB_NAME, header.int_size, use_tempfile=True
1902 ),
1903 "Copying dataset provenance into zip archive.",
1904 )
1905 )
1906 self._base_writer.address_writer.addresses[DATASET_ADDRESS_INDEX] = self._dataset_writer.addresses
1907 self._quantum_writer = exit_stack.enter_context(
1908 log_on_close.wrap(
1909 MultiblockWriter.open_in_zip(
1910 self._base_writer.zf, QUANTUM_MB_NAME, header.int_size, use_tempfile=True
1911 ),
1912 "Copying quantum provenance into zip archive.",
1913 )
1914 )
1915 self._base_writer.address_writer.addresses[QUANTUM_ADDRESS_INDEX] = self._quantum_writer.addresses
1916 self._init_predicted_quanta(predicted)
1917 self._populate_xgraph_and_inputs(loop_wrapper)
1918 self._existing_init_outputs: set[uuid.UUID] = set()
1920 def _init_predicted_quanta(
1921 self, predicted: PredictedQuantumGraph | PredictedQuantumGraphComponents
1922 ) -> None:
1923 self._predicted_init_quanta: list[PredictedQuantumDatasetsModel] = []
1924 self._predicted_quanta: dict[uuid.UUID, PredictedQuantumDatasetsModel] = {}
1925 if isinstance(predicted, PredictedQuantumGraph):
1926 self._predicted_init_quanta.extend(predicted._init_quanta.values())
1927 self._predicted_quanta.update(predicted._quantum_datasets)
1928 else:
1929 self._predicted_init_quanta.extend(predicted.init_quanta.root)
1930 self._predicted_quanta.update(predicted.quantum_datasets)
1931 self._predicted_quanta.update({q.quantum_id: q for q in self._predicted_init_quanta})
1933 def _populate_xgraph_and_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1934 self._xgraph = networkx.DiGraph()
1935 self._overall_inputs: dict[uuid.UUID, PredictedDatasetModel] = {}
1936 output_dataset_ids: set[uuid.UUID] = set()
1937 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1938 if not predicted_quantum.task_label:
1939 # Skip the 'packages' producer quantum.
1940 continue
1941 output_dataset_ids.update(predicted_quantum.iter_output_dataset_ids())
1942 for predicted_quantum in loop_wrapper(self._predicted_quanta.values()):
1943 if not predicted_quantum.task_label:
1944 # Skip the 'packages' producer quantum.
1945 continue
1946 for predicted_input in itertools.chain.from_iterable(predicted_quantum.inputs.values()):
1947 self._xgraph.add_edge(predicted_input.dataset_id, predicted_quantum.quantum_id)
1948 if predicted_input.dataset_id not in output_dataset_ids:
1949 self._overall_inputs.setdefault(predicted_input.dataset_id, predicted_input)
1950 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
1951 self._xgraph.add_edge(predicted_quantum.quantum_id, predicted_output.dataset_id)
1953 @property
1954 def compressor(self) -> Compressor:
1955 """Object that should be used to compress all JSON blocks."""
1956 return self._base_writer.compressor
1958 def write_packages(self) -> None:
1959 """Write package version information to the provenance graph."""
1960 packages = Packages.fromSystem(include_all=True)
1961 data = packages.toBytes("json")
1962 self._base_writer.write_single_block("packages", data)
1964 def write_overall_inputs(self, loop_wrapper: LoopWrapper = pass_through) -> None:
1965 """Write provenance for overall-input datasets.
1967 Parameters
1968 ----------
1969 loop_wrapper : `~collections.abc.Callable`, optional
1970 A callable that takes an iterable and returns an equivalent one, to
1971 be used in all potentially-large loops. This can be used to add
1972 progress reporting or check for cancelation signals.
1973 """
1974 for predicted_input in loop_wrapper(self._overall_inputs.values()):
1975 if predicted_input.dataset_id not in self._dataset_writer.addresses:
1976 self._dataset_writer.write_model(
1977 predicted_input.dataset_id,
1978 ProvenanceDatasetModel.from_predicted(
1979 predicted_input,
1980 producer=None,
1981 consumers=self._xgraph.successors(predicted_input.dataset_id),
1982 ),
1983 self.compressor,
1984 )
1985 del self._overall_inputs
1987 def write_init_outputs(self, assume_existence: bool = True) -> None:
1988 """Write provenance for init-output datasets and init-quanta.
1990 Parameters
1991 ----------
1992 assume_existence : `bool`, optional
1993 If `True`, just assume all init-outputs exist.
1994 """
1995 init_quanta = ProvenanceInitQuantaModel()
1996 for predicted_init_quantum in self._predicted_init_quanta:
1997 if not predicted_init_quantum.task_label:
1998 # Skip the 'packages' producer quantum.
1999 continue
2000 for predicted_output in itertools.chain.from_iterable(predicted_init_quantum.outputs.values()):
2001 provenance_output = ProvenanceDatasetModel.from_predicted(
2002 predicted_output,
2003 producer=predicted_init_quantum.quantum_id,
2004 consumers=self._xgraph.successors(predicted_output.dataset_id),
2005 )
2006 provenance_output.produced = assume_existence or (
2007 provenance_output.dataset_id in self._existing_init_outputs
2008 )
2009 self._dataset_writer.write_model(
2010 provenance_output.dataset_id, provenance_output, self.compressor
2011 )
2012 init_quanta.root.append(ProvenanceInitQuantumModel.from_predicted(predicted_init_quantum))
2013 self._base_writer.write_single_model("init_quanta", init_quanta)
2015 def write_quantum_provenance(
2016 self, quantum_id: uuid.UUID, metadata: TaskMetadata | None, logs: ButlerLogRecords | None
2017 ) -> None:
2018 """Gather and write provenance for a quantum.
2020 Parameters
2021 ----------
2022 quantum_id : `uuid.UUID`
2023 Unique ID for the quantum.
2024 metadata : `..TaskMetadata` or `None`
2025 Task metadata.
2026 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
2027 Task logs.
2028 """
2029 predicted_quantum = self._predicted_quanta[quantum_id]
2030 provenance_models = ProvenanceQuantumScanModels.from_metadata_and_logs(
2031 predicted_quantum, metadata, logs, incomplete=False
2032 )
2033 scan_data = provenance_models.to_scan_data(predicted_quantum, compressor=self.compressor)
2034 self.write_scan_data(scan_data)
2036 def write_scan_data(self, scan_data: ProvenanceQuantumScanData) -> None:
2037 """Write the output of a quantum provenance scan to disk.
2039 Parameters
2040 ----------
2041 scan_data : `ProvenanceQuantumScanData`
2042 Result of a quantum provenance scan.
2043 """
2044 if scan_data.status is ProvenanceQuantumScanStatus.INIT:
2045 self.log.debug("Handling init-output scan for %s.", scan_data.quantum_id)
2046 self._existing_init_outputs.update(scan_data.existing_outputs)
2047 return
2048 self.log.debug("Handling quantum scan for %s.", scan_data.quantum_id)
2049 # We shouldn't need this predicted quantum after this method runs; pop
2050 # from the dict it in the hopes that'll free up some memory when we're
2051 # done.
2052 predicted_quantum = self._predicted_quanta.pop(scan_data.quantum_id)
2053 outputs: dict[uuid.UUID, bytes] = {}
2054 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
2055 provenance_output = ProvenanceDatasetModel.from_predicted(
2056 predicted_output,
2057 producer=predicted_quantum.quantum_id,
2058 consumers=self._xgraph.successors(predicted_output.dataset_id),
2059 )
2060 provenance_output.produced = provenance_output.dataset_id in scan_data.existing_outputs
2061 outputs[provenance_output.dataset_id] = self.compressor.compress(
2062 provenance_output.model_dump_json().encode()
2063 )
2064 if not scan_data.quantum:
2065 scan_data.quantum = (
2066 ProvenanceQuantumModel.from_predicted(predicted_quantum).model_dump_json().encode()
2067 )
2068 if scan_data.is_compressed:
2069 scan_data.quantum = self.compressor.compress(scan_data.quantum)
2070 if not scan_data.is_compressed:
2071 scan_data.quantum = self.compressor.compress(scan_data.quantum)
2072 if scan_data.metadata:
2073 scan_data.metadata = self.compressor.compress(scan_data.metadata)
2074 if scan_data.logs:
2075 scan_data.logs = self.compressor.compress(scan_data.logs)
2076 self.log.debug("Writing quantum %s.", scan_data.quantum_id)
2077 self._quantum_writer.write_bytes(scan_data.quantum_id, scan_data.quantum)
2078 for dataset_id, dataset_data in outputs.items():
2079 self._dataset_writer.write_bytes(dataset_id, dataset_data)
2080 if scan_data.metadata:
2081 (metadata_output,) = predicted_quantum.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
2082 address = self._metadata_writer.write_bytes(scan_data.quantum_id, scan_data.metadata)
2083 self._metadata_writer.addresses[metadata_output.dataset_id] = address
2084 if scan_data.logs:
2085 (log_output,) = predicted_quantum.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
2086 address = self._log_writer.write_bytes(scan_data.quantum_id, scan_data.logs)
2087 self._log_writer.addresses[log_output.dataset_id] = address
2090class ProvenanceQuantumScanStatus(enum.Enum):
2091 """Status enum for quantum scanning.
2093 Note that this records the status for the *scanning* which is distinct
2094 from the status of the quantum's execution.
2095 """
2097 INCOMPLETE = enum.auto()
2098 """The quantum is not necessarily done running, and cannot be scanned
2099 conclusively yet.
2100 """
2102 ABANDONED = enum.auto()
2103 """The quantum's execution appears to have failed but we cannot rule out
2104 the possibility that it could be recovered, but we've also waited long
2105 enough (according to `ScannerTimeConfigDict.retry_timeout`) that it's time
2106 to stop trying for now.
2108 This state means `ProvenanceQuantumScanModels.from_metadata_and_logs` must
2109 be run again with ``incomplete=False``.
2110 """
2112 SUCCESSFUL = enum.auto()
2113 """The quantum was conclusively scanned and was executed successfully,
2114 unblocking scans for downstream quanta.
2115 """
2117 FAILED = enum.auto()
2118 """The quantum was conclusively scanned and failed execution, blocking
2119 scans for downstream quanta.
2120 """
2122 BLOCKED = enum.auto()
2123 """A quantum upstream of this one failed."""
2125 INIT = enum.auto()
2126 """Init quanta need special handling, because they don't have logs and
2127 metadata.
2128 """
2131@dataclasses.dataclass
2132class ProvenanceQuantumScanModels:
2133 """A struct that represents provenance information for a single quantum."""
2135 quantum_id: uuid.UUID
2136 """Unique ID for the quantum."""
2138 status: ProvenanceQuantumScanStatus = ProvenanceQuantumScanStatus.INCOMPLETE
2139 """Combined status for the scan and the execution of the quantum."""
2141 attempts: list[ProvenanceQuantumAttemptModel] = dataclasses.field(default_factory=list)
2142 """Provenance information about each attempt to run the quantum."""
2144 output_existence: dict[uuid.UUID, bool] = dataclasses.field(default_factory=dict)
2145 """Unique IDs of the output datasets mapped to whether they were actually
2146 produced.
2147 """
2149 metadata: ProvenanceTaskMetadataModel = dataclasses.field(default_factory=ProvenanceTaskMetadataModel)
2150 """Task metadata information for each attempt.
2151 """
2153 logs: ProvenanceLogRecordsModel = dataclasses.field(default_factory=ProvenanceLogRecordsModel)
2154 """Log records for each attempt.
2155 """
2157 @classmethod
2158 def from_metadata_and_logs(
2159 cls,
2160 predicted: PredictedQuantumDatasetsModel,
2161 metadata: TaskMetadata | None,
2162 logs: ButlerLogRecords | None,
2163 *,
2164 incomplete: bool = False,
2165 ) -> ProvenanceQuantumScanModels:
2166 """Construct provenance information from task metadata and logs.
2168 Parameters
2169 ----------
2170 predicted : `PredictedQuantumDatasetsModel`
2171 Information about the predicted quantum.
2172 metadata : `..TaskMetadata` or `None`
2173 Task metadata.
2174 logs : `lsst.daf.butler.logging.ButlerLogRecords` or `None`
2175 Task logs.
2176 incomplete : `bool`, optional
2177 If `True`, treat execution failures as possibly-incomplete quanta
2178 and do not fully process them; instead just set the status to
2179 `ProvenanceQuantumScanStatus.ABANDONED` and return.
2181 Returns
2182 -------
2183 scan_models : `ProvenanceQuantumScanModels`
2184 Struct of models that describe quantum provenance.
2186 Notes
2187 -----
2188 This method does not necessarily fully populate the `output_existence`
2189 field; it does what it can given the information in the metadata and
2190 logs, but the caller is responsible for filling in the existence status
2191 for any predicted outputs that are not present at all in that `dict`.
2192 """
2193 self = ProvenanceQuantumScanModels(predicted.quantum_id)
2194 last_attempt = ProvenanceQuantumAttemptModel()
2195 self._process_logs(predicted, logs, last_attempt, incomplete=incomplete)
2196 self._process_metadata(predicted, metadata, last_attempt, incomplete=incomplete)
2197 if self.status is ProvenanceQuantumScanStatus.ABANDONED:
2198 return self
2199 self._reconcile_attempts(last_attempt)
2200 self._extract_output_existence(predicted)
2201 return self
2203 def _process_logs(
2204 self,
2205 predicted: PredictedQuantumDatasetsModel,
2206 logs: ButlerLogRecords | None,
2207 last_attempt: ProvenanceQuantumAttemptModel,
2208 *,
2209 incomplete: bool,
2210 ) -> None:
2211 (predicted_log_dataset,) = predicted.outputs[acc.LOG_OUTPUT_CONNECTION_NAME]
2212 if logs is None:
2213 self.output_existence[predicted_log_dataset.dataset_id] = False
2214 if incomplete:
2215 self.status = ProvenanceQuantumScanStatus.ABANDONED
2216 else:
2217 self.status = ProvenanceQuantumScanStatus.FAILED
2218 else:
2219 # Set the attempt's run status to FAILED, since the default is
2220 # UNKNOWN (i.e. logs *and* metadata are missing) and we now know
2221 # the logs exist. This will usually get replaced by SUCCESSFUL
2222 # when we look for metadata next.
2223 last_attempt.status = QuantumAttemptStatus.FAILED
2224 self.output_existence[predicted_log_dataset.dataset_id] = True
2225 if logs.extra:
2226 log_extra = _ExecutionLogRecordsExtra.model_validate(logs.extra)
2227 self._extract_from_log_extra(log_extra, last_attempt=last_attempt)
2228 self.logs.attempts.append(list(logs))
2230 def _extract_from_log_extra(
2231 self,
2232 log_extra: _ExecutionLogRecordsExtra,
2233 last_attempt: ProvenanceQuantumAttemptModel | None,
2234 ) -> None:
2235 for previous_attempt_log_extra in log_extra.previous_attempts:
2236 self._extract_from_log_extra(
2237 previous_attempt_log_extra,
2238 last_attempt=None,
2239 )
2240 quantum_attempt: ProvenanceQuantumAttemptModel
2241 if last_attempt is None:
2242 # This is not the last attempt, so it must be a failure.
2243 quantum_attempt = ProvenanceQuantumAttemptModel(
2244 attempt=len(self.attempts), status=QuantumAttemptStatus.FAILED
2245 )
2246 # We also need to get the logs from this extra provenance, since
2247 # they won't be the main section of the log records.
2248 self.logs.attempts.append(log_extra.logs)
2249 # The special last attempt is only appended after we attempt to
2250 # read metadata later, but we have to append this one now.
2251 self.attempts.append(quantum_attempt)
2252 else:
2253 assert not log_extra.logs, "Logs for the last attempt should not be stored in the extra JSON."
2254 quantum_attempt = last_attempt
2255 if log_extra.exception is not None or log_extra.metadata is not None or last_attempt is None:
2256 # We won't be getting a separate metadata dataset, so anything we
2257 # might get from the metadata has to come from this extra
2258 # provenance in the logs.
2259 quantum_attempt.exception = log_extra.exception
2260 if log_extra.metadata is not None:
2261 quantum_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(log_extra.metadata)
2262 self.metadata.attempts.append(log_extra.metadata)
2263 else:
2264 self.metadata.attempts.append(None)
2265 # Regardless of whether this is the last attempt or not, we can only
2266 # get the previous_process_quanta from the log extra.
2267 quantum_attempt.previous_process_quanta.extend(log_extra.previous_process_quanta)
2269 def _process_metadata(
2270 self,
2271 predicted: PredictedQuantumDatasetsModel,
2272 metadata: TaskMetadata | None,
2273 last_attempt: ProvenanceQuantumAttemptModel,
2274 *,
2275 incomplete: bool,
2276 ) -> None:
2277 (predicted_metadata_dataset,) = predicted.outputs[acc.METADATA_OUTPUT_CONNECTION_NAME]
2278 if metadata is None:
2279 self.output_existence[predicted_metadata_dataset.dataset_id] = False
2280 if incomplete:
2281 self.status = ProvenanceQuantumScanStatus.ABANDONED
2282 else:
2283 self.status = ProvenanceQuantumScanStatus.FAILED
2284 else:
2285 self.status = ProvenanceQuantumScanStatus.SUCCESSFUL
2286 self.output_existence[predicted_metadata_dataset.dataset_id] = True
2287 last_attempt.status = QuantumAttemptStatus.SUCCESSFUL
2288 try:
2289 # Int conversion guards against spurious conversion to
2290 # float that can apparently sometimes happen in
2291 # TaskMetadata.
2292 last_attempt.caveats = QuantumSuccessCaveats(int(metadata["quantum"]["caveats"]))
2293 except LookupError:
2294 pass
2295 try:
2296 last_attempt.exception = ExceptionInfo._from_metadata(
2297 metadata[predicted.task_label]["failure"]
2298 )
2299 except LookupError:
2300 pass
2301 last_attempt.resource_usage = QuantumResourceUsage.from_task_metadata(metadata)
2302 self.metadata.attempts.append(metadata)
2304 def _reconcile_attempts(self, last_attempt: ProvenanceQuantumAttemptModel) -> None:
2305 last_attempt.attempt = len(self.attempts)
2306 self.attempts.append(last_attempt)
2307 assert self.status is not ProvenanceQuantumScanStatus.INCOMPLETE
2308 assert self.status is not ProvenanceQuantumScanStatus.ABANDONED
2309 if len(self.logs.attempts) < len(self.attempts):
2310 # Logs were not found for this attempt; must have been a hard error
2311 # that kept the `finally` block from running or otherwise
2312 # interrupted the writing of the logs.
2313 self.logs.attempts.append(None)
2314 if self.status is ProvenanceQuantumScanStatus.SUCCESSFUL:
2315 # But we found the metadata! Either that hard error happened
2316 # at a very unlucky time (in between those two writes), or
2317 # something even weirder happened.
2318 self.attempts[-1].status = QuantumAttemptStatus.ABORTED_SUCCESS
2319 else:
2320 self.attempts[-1].status = QuantumAttemptStatus.FAILED
2321 if len(self.metadata.attempts) < len(self.attempts):
2322 # Metadata missing usually just means a failure. In any case, the
2323 # status will already be correct, either because it was set to a
2324 # failure when we read the logs, or left at UNKNOWN if there were
2325 # no logs. Note that scanners never process BLOCKED quanta at all.
2326 self.metadata.attempts.append(None)
2327 assert len(self.logs.attempts) == len(self.attempts) or len(self.metadata.attempts) == len(
2328 self.attempts
2329 ), (
2330 "The only way we can add more than one quantum attempt is by "
2331 "extracting info stored with the logs, and that always appends "
2332 "a log attempt and a metadata attempt, so this must be a bug in "
2333 "this class."
2334 )
2336 def _extract_output_existence(self, predicted: PredictedQuantumDatasetsModel) -> None:
2337 try:
2338 outputs_put = self.metadata.attempts[-1]["quantum"].getArray("outputs") # type: ignore[index]
2339 except (
2340 IndexError, # metadata.attempts is empty
2341 TypeError, # metadata.attempts[-1] is None
2342 LookupError, # no 'quantum' entry in metadata or 'outputs' in that
2343 ):
2344 pass
2345 else:
2346 for id_str in ensure_iterable(outputs_put):
2347 self.output_existence[uuid.UUID(id_str)] = True
2348 # If the metadata told us what it wrote, anything not in that
2349 # list was not written.
2350 for predicted_output in itertools.chain.from_iterable(predicted.outputs.values()):
2351 self.output_existence.setdefault(predicted_output.dataset_id, False)
2353 def to_scan_data(
2354 self: ProvenanceQuantumScanModels,
2355 predicted_quantum: PredictedQuantumDatasetsModel,
2356 compressor: Compressor | None = None,
2357 ) -> ProvenanceQuantumScanData:
2358 """Convert these models to JSON data.
2360 Parameters
2361 ----------
2362 predicted_quantum : `PredictedQuantumDatasetsModel`
2363 Information about the predicted quantum.
2364 compressor : `Compressor`
2365 Object that can compress bytes.
2367 Returns
2368 -------
2369 scan_data : `ProvenanceQuantumScanData`
2370 Scan information ready for serialization.
2371 """
2372 quantum: ProvenanceInitQuantumModel | ProvenanceQuantumModel
2373 if self.status is ProvenanceQuantumScanStatus.INIT:
2374 quantum = ProvenanceInitQuantumModel.from_predicted(predicted_quantum)
2375 else:
2376 quantum = ProvenanceQuantumModel.from_predicted(predicted_quantum)
2377 quantum.attempts = self.attempts
2378 for predicted_output in itertools.chain.from_iterable(predicted_quantum.outputs.values()):
2379 if predicted_output.dataset_id not in self.output_existence:
2380 raise RuntimeError(
2381 "Logic bug in provenance gathering or execution invariants: "
2382 f"no existence information for output {predicted_output.dataset_id} "
2383 f"({predicted_output.dataset_type_name}@{predicted_output.data_coordinate})."
2384 )
2385 data = ProvenanceQuantumScanData(
2386 self.quantum_id,
2387 self.status,
2388 existing_outputs={
2389 dataset_id for dataset_id, was_produced in self.output_existence.items() if was_produced
2390 },
2391 quantum=quantum.model_dump_json().encode(),
2392 logs=self.logs.model_dump_json().encode() if self.logs.attempts else b"",
2393 metadata=self.metadata.model_dump_json().encode() if self.metadata.attempts else b"",
2394 )
2395 if compressor is not None:
2396 data.compress(compressor)
2397 return data
2400@dataclasses.dataclass
2401class ProvenanceQuantumScanData:
2402 """A struct that represents ready-for-serialization provenance information
2403 for a single quantum.
2404 """
2406 quantum_id: uuid.UUID
2407 """Unique ID for the quantum."""
2409 status: ProvenanceQuantumScanStatus
2410 """Combined status for the scan and the execution of the quantum."""
2412 existing_outputs: set[uuid.UUID] = dataclasses.field(default_factory=set)
2413 """Unique IDs of the output datasets that were actually written."""
2415 quantum: bytes = b""
2416 """Serialized quantum provenance model.
2418 This may be empty for quanta that had no attempts.
2419 """
2421 metadata: bytes = b""
2422 """Serialized task metadata."""
2424 logs: bytes = b""
2425 """Serialized logs."""
2427 is_compressed: bool = False
2428 """Whether the ``quantum``, ``metadata``, and ``log`` attributes are
2429 compressed.
2430 """
2432 def compress(self, compressor: Compressor) -> None:
2433 """Compress the data in this struct if it has not been compressed
2434 already.
2436 Parameters
2437 ----------
2438 compressor : `Compressor`
2439 Object with a ``compress`` method that takes and returns `bytes`.
2440 """
2441 if not self.is_compressed:
2442 self.quantum = compressor.compress(self.quantum)
2443 self.logs = compressor.compress(self.logs) if self.logs else b""
2444 self.metadata = compressor.compress(self.metadata) if self.metadata else b""
2445 self.is_compressed = True