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