Coverage for python/lsst/pipe/base/quantum_graph/_provenance.py: 40%
318 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-14 01:20 -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 "ProvenanceQuantumInfo",
39 "ProvenanceQuantumModel",
40 "ProvenanceTaskMetadataModel",
41)
44import dataclasses
45import sys
46import uuid
47from collections import Counter
48from collections.abc import Iterable, Iterator, Mapping
49from contextlib import contextmanager
50from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypedDict, TypeVar
52import astropy.table
53import networkx
54import numpy as np
55import pydantic
57from lsst.daf.butler import DataCoordinate
58from lsst.daf.butler.logging import ButlerLogRecord, ButlerLogRecords
59from lsst.resources import ResourcePathExpression
60from lsst.utils.packages import Packages
62from .._status import ExceptionInfo, QuantumAttemptStatus, QuantumSuccessCaveats
63from .._task_metadata import TaskMetadata
64from ..pipeline_graph import PipelineGraph, TaskImportMode, TaskInitNode
65from ..resource_usage import QuantumResourceUsage
66from ._common import (
67 BaseQuantumGraph,
68 BaseQuantumGraphReader,
69 ConnectionName,
70 DataCoordinateValues,
71 DatasetIndex,
72 DatasetInfo,
73 DatasetTypeName,
74 HeaderModel,
75 QuantumIndex,
76 QuantumInfo,
77 TaskLabel,
78)
79from ._multiblock import AddressReader, MultiblockReader
80from ._predicted import PredictedDatasetModel, PredictedQuantumDatasetsModel
82DATASET_ADDRESS_INDEX = 0
83QUANTUM_ADDRESS_INDEX = 1
84LOG_ADDRESS_INDEX = 2
85METADATA_ADDRESS_INDEX = 3
87DATASET_MB_NAME = "datasets"
88QUANTUM_MB_NAME = "quanta"
89LOG_MB_NAME = "logs"
90METADATA_MB_NAME = "metadata"
92_I = TypeVar("_I", bound=uuid.UUID | int)
95class ProvenanceDatasetInfo(DatasetInfo):
96 """A typed dictionary that annotates the attributes of the NetworkX graph
97 node data for a provenance dataset.
99 Since NetworkX types are not generic over their node mapping type, this has
100 to be used explicitly, e.g.::
102 node_data: ProvenanceDatasetInfo = xgraph.nodes[dataset_id]
104 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
105 """
107 dataset_id: uuid.UUID
108 """Unique identifier for the dataset."""
110 produced: bool
111 """Whether this dataset was produced (vs. only predicted).
113 This is always `True` for overall input datasets. It is also `True` for
114 datasets that were produced and then removed before/during transfer back to
115 the central butler repository, so it may not reflect the continued
116 existence of the dataset.
117 """
120class ProvenanceQuantumInfo(QuantumInfo):
121 """A typed dictionary that annotates the attributes of the NetworkX graph
122 node data for a provenance quantum.
124 Since NetworkX types are not generic over their node mapping type, this has
125 to be used explicitly, e.g.::
127 node_data: ProvenanceQuantumInfo = xgraph.nodes[quantum_id]
129 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph` or
130 `ProvenanceQuantumGraph.quantum_only_xgraph`
131 """
133 status: QuantumAttemptStatus
134 """Enumerated status for the quantum.
136 This corresponds to the last attempt to run this quantum, or
137 `QuantumAttemptStatus.BLOCKED` if there were no attempts.
138 """
140 caveats: QuantumSuccessCaveats | None
141 """Flags indicating caveats on successful quanta.
143 This corresponds to the last attempt to run this quantum.
144 """
146 exception: ExceptionInfo | None
147 """Information about an exception raised when the quantum was executing.
149 This corresponds to the last attempt to run this quantum.
150 """
152 resource_usage: QuantumResourceUsage | None
153 """Resource usage information (timing, memory use) for this quantum.
155 This corresponds to the last attempt to run this quantum.
156 """
158 attempts: list[ProvenanceQuantumAttemptModel]
159 """Information about each attempt to run this quantum.
161 An entry is added merely if the quantum *should* have been attempted; an
162 empty `list` is used only for quanta that were blocked by an upstream
163 failure.
164 """
167class ProvenanceInitQuantumInfo(TypedDict):
168 """A typed dictionary that annotates the attributes of the NetworkX graph
169 node data for a provenance init quantum.
171 Since NetworkX types are not generic over their node mapping type, this has
172 to be used explicitly, e.g.::
174 node_data: ProvenanceInitQuantumInfo = xgraph.nodes[quantum_id]
176 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
177 """
179 data_id: DataCoordinate
180 """Data ID of the quantum.
182 This is always an empty ID; this key exists to allow init-quanta and
183 regular quanta to be treated more similarly.
184 """
186 task_label: str
187 """Label of the task for this quantum."""
189 pipeline_node: TaskInitNode
190 """Node in the pipeline graph for this task's init-only step."""
193class ProvenanceDatasetModel(PredictedDatasetModel):
194 """Data model for the datasets in a provenance quantum graph file."""
196 produced: bool
197 """Whether this dataset was produced (vs. only predicted).
199 This is always `True` for overall input datasets. It is also `True` for
200 datasets that were produced and then removed before/during transfer back to
201 the central butler repository, so it may not reflect the continued
202 existence of the dataset.
203 """
205 producer: QuantumIndex | None = None
206 """Internal integer ID of the quantum that produced this dataset.
208 This is `None` for overall inputs to the graph.
209 """
211 consumers: list[QuantumIndex] = pydantic.Field(default_factory=list)
212 """Internal integer IDs of quanta that were predicted to consume this
213 dataset.
214 """
216 @property
217 def node_id(self) -> uuid.UUID:
218 """Alias for the dataset ID."""
219 return self.dataset_id
221 @classmethod
222 def from_predicted(
223 cls,
224 predicted: PredictedDatasetModel,
225 producer: QuantumIndex | None = None,
226 consumers: Iterable[QuantumIndex] = (),
227 ) -> ProvenanceDatasetModel:
228 """Construct from a predicted dataset model.
230 Parameters
231 ----------
232 predicted : `PredictedDatasetModel`
233 Information about the dataset from the predicted graph.
234 producer : `int` or `None`, optional
235 Internal ID of the quantum that was predicted to produce this
236 dataset.
237 consumers : `~collections.abc.Iterable` [`int`], optional
238 Internal IDs of the quanta that were predicted to consume this
239 dataset.
241 Returns
242 -------
243 provenance : `ProvenanceDatasetModel`
244 Provenance dataset model.
246 Notes
247 -----
248 This initializes `produced` to `True` when ``producer is None`` and
249 `False` otherwise, on the assumption that it will be updated later.
250 """
251 return cls.model_construct(
252 dataset_id=predicted.dataset_id,
253 dataset_type_name=predicted.dataset_type_name,
254 data_coordinate=predicted.data_coordinate,
255 run=predicted.run,
256 produced=(producer is None), # if it's not produced by this QG, it's an overall input
257 producer=producer,
258 consumers=list(consumers),
259 )
261 def _add_to_graph(self, graph: ProvenanceQuantumGraph, address_reader: AddressReader) -> None:
262 """Add this dataset and its edges to quanta to a provenance graph.
264 Parameters
265 ----------
266 graph : `ProvenanceQuantumGraph`
267 Graph to update in place.
268 address_reader : `AddressReader`
269 Reader object that can be used to look up UUIDs from integer
270 indexes.
272 Notes
273 -----
274 This method adds:
276 - a ``bipartite_xgraph`` dataset node with full attributes;
277 - ``bipartite_xgraph`` edges to adjacent quanta (which adds quantum
278 nodes with no attributes), without populating edge attributes;
279 - ``quantum_only_xgraph`` edges for each pair of quanta in which one
280 produces this dataset and another consumes it (this also adds quantum
281 nodes with no attributes).
282 """
283 dataset_type_node = graph.pipeline_graph.dataset_types[self.dataset_type_name]
284 data_id = DataCoordinate.from_full_values(dataset_type_node.dimensions, tuple(self.data_coordinate))
285 graph._bipartite_xgraph.add_node(
286 self.dataset_id,
287 data_id=data_id,
288 dataset_type_name=self.dataset_type_name,
289 pipeline_node=dataset_type_node,
290 run=self.run,
291 produced=self.produced,
292 )
293 producer_id: uuid.UUID | None = None
294 if self.producer is not None:
295 producer_id = address_reader.find(self.producer).key
296 graph._bipartite_xgraph.add_edge(producer_id, self.dataset_id)
297 for consumer_index in self.consumers:
298 consumer_id = address_reader.find(consumer_index).key
299 graph._bipartite_xgraph.add_edge(self.dataset_id, consumer_id)
300 if producer_id is not None:
301 graph._quantum_only_xgraph.add_edge(producer_id, consumer_id)
302 graph._datasets_by_type[self.dataset_type_name][data_id] = self.dataset_id
304 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
305 # when we inherit those docstrings in our public classes.
306 if "sphinx" in sys.modules and not TYPE_CHECKING:
308 def copy(self, *args: Any, **kwargs: Any) -> Any:
309 """See `pydantic.BaseModel.copy`."""
310 return super().copy(*args, **kwargs)
312 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
313 """See `pydantic.BaseModel.model_dump`."""
314 return super().model_dump(*args, **kwargs)
316 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
317 """See `pydantic.BaseModel.model_dump_json`."""
318 return super().model_dump(*args, **kwargs)
320 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
321 """See `pydantic.BaseModel.model_copy`."""
322 return super().model_copy(*args, **kwargs)
324 @classmethod
325 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
326 """See `pydantic.BaseModel.model_construct`."""
327 return super().model_construct(*args, **kwargs)
329 @classmethod
330 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
331 """See `pydantic.BaseModel.model_json_schema`."""
332 return super().model_json_schema(*args, **kwargs)
334 @classmethod
335 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
336 """See `pydantic.BaseModel.model_validate`."""
337 return super().model_validate(*args, **kwargs)
339 @classmethod
340 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
341 """See `pydantic.BaseModel.model_validate_json`."""
342 return super().model_validate_json(*args, **kwargs)
344 @classmethod
345 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
346 """See `pydantic.BaseModel.model_validate_strings`."""
347 return super().model_validate_strings(*args, **kwargs)
350class _GenericProvenanceQuantumAttemptModel(pydantic.BaseModel, Generic[_I]):
351 """Data model for a now-superseded attempt to run a quantum in a
352 provenance quantum graph file.
353 """
355 attempt: int = 0
356 """Counter incremented for every attempt to execute this quantum."""
358 status: QuantumAttemptStatus = QuantumAttemptStatus.UNKNOWN
359 """Enumerated status for the quantum."""
361 caveats: QuantumSuccessCaveats | None = None
362 """Flags indicating caveats on successful quanta."""
364 exception: ExceptionInfo | None = None
365 """Information about an exception raised when the quantum was executing."""
367 resource_usage: QuantumResourceUsage | None = None
368 """Resource usage information (timing, memory use) for this quantum."""
370 previous_process_quanta: list[_I] = pydantic.Field(default_factory=list)
371 """The IDs of other quanta previously executed in the same process as this
372 one.
373 """
375 def remap_uuids(
376 self: ProvenanceQuantumAttemptModel, indices: Mapping[uuid.UUID, QuantumIndex]
377 ) -> StorageProvenanceQuantumAttemptModel:
378 return StorageProvenanceQuantumAttemptModel(
379 attempt=self.attempt,
380 status=self.status,
381 caveats=self.caveats,
382 exception=self.exception,
383 resource_usage=self.resource_usage,
384 previous_process_quanta=[indices[q] for q in self.previous_process_quanta],
385 )
387 def remap_indices(
388 self: StorageProvenanceQuantumAttemptModel, address_reader: AddressReader
389 ) -> ProvenanceQuantumAttemptModel:
390 return ProvenanceQuantumAttemptModel(
391 attempt=self.attempt,
392 status=self.status,
393 caveats=self.caveats,
394 exception=self.exception,
395 resource_usage=self.resource_usage,
396 previous_process_quanta=[address_reader.find(q).key for q in self.previous_process_quanta],
397 )
399 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
400 # when we inherit those docstrings in our public classes.
401 if "sphinx" in sys.modules and not TYPE_CHECKING:
403 def copy(self, *args: Any, **kwargs: Any) -> Any:
404 """See `pydantic.BaseModel.copy`."""
405 return super().copy(*args, **kwargs)
407 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
408 """See `pydantic.BaseModel.model_dump`."""
409 return super().model_dump(*args, **kwargs)
411 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
412 """See `pydantic.BaseModel.model_dump_json`."""
413 return super().model_dump(*args, **kwargs)
415 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
416 """See `pydantic.BaseModel.model_copy`."""
417 return super().model_copy(*args, **kwargs)
419 @classmethod
420 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
421 """See `pydantic.BaseModel.model_construct`."""
422 return super().model_construct(*args, **kwargs)
424 @classmethod
425 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
426 """See `pydantic.BaseModel.model_json_schema`."""
427 return super().model_json_schema(*args, **kwargs)
429 @classmethod
430 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
431 """See `pydantic.BaseModel.model_validate`."""
432 return super().model_validate(*args, **kwargs)
434 @classmethod
435 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
436 """See `pydantic.BaseModel.model_validate_json`."""
437 return super().model_validate_json(*args, **kwargs)
439 @classmethod
440 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
441 """See `pydantic.BaseModel.model_validate_strings`."""
442 return super().model_validate_strings(*args, **kwargs)
445StorageProvenanceQuantumAttemptModel: TypeAlias = _GenericProvenanceQuantumAttemptModel[QuantumIndex]
446ProvenanceQuantumAttemptModel: TypeAlias = _GenericProvenanceQuantumAttemptModel[uuid.UUID]
449class ProvenanceLogRecordsModel(pydantic.BaseModel):
450 """Data model for storing execution logs in a provenance quantum graph
451 file.
452 """
454 attempts: list[list[ButlerLogRecord] | None] = pydantic.Field(default_factory=list)
455 """Logs from attempts to run this task, ordered chronologically from first
456 to last.
457 """
459 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
460 # when we inherit those docstrings in our public classes.
461 if "sphinx" in sys.modules and not TYPE_CHECKING:
463 def copy(self, *args: Any, **kwargs: Any) -> Any:
464 """See `pydantic.BaseModel.copy`."""
465 return super().copy(*args, **kwargs)
467 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
468 """See `pydantic.BaseModel.model_dump`."""
469 return super().model_dump(*args, **kwargs)
471 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
472 """See `pydantic.BaseModel.model_dump_json`."""
473 return super().model_dump(*args, **kwargs)
475 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
476 """See `pydantic.BaseModel.model_copy`."""
477 return super().model_copy(*args, **kwargs)
479 @classmethod
480 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
481 """See `pydantic.BaseModel.model_construct`."""
482 return super().model_construct(*args, **kwargs)
484 @classmethod
485 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
486 """See `pydantic.BaseModel.model_json_schema`."""
487 return super().model_json_schema(*args, **kwargs)
489 @classmethod
490 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
491 """See `pydantic.BaseModel.model_validate`."""
492 return super().model_validate(*args, **kwargs)
494 @classmethod
495 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
496 """See `pydantic.BaseModel.model_validate_json`."""
497 return super().model_validate_json(*args, **kwargs)
499 @classmethod
500 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
501 """See `pydantic.BaseModel.model_validate_strings`."""
502 return super().model_validate_strings(*args, **kwargs)
505class ProvenanceTaskMetadataModel(pydantic.BaseModel):
506 """Data model for storing task metadata in a provenance quantum graph
507 file.
508 """
510 attempts: list[TaskMetadata | None] = pydantic.Field(default_factory=list)
511 """Metadata from attempts to run this task, ordered chronologically from
512 first to last.
513 """
515 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
516 # when we inherit those docstrings in our public classes.
517 if "sphinx" in sys.modules and not TYPE_CHECKING:
519 def copy(self, *args: Any, **kwargs: Any) -> Any:
520 """See `pydantic.BaseModel.copy`."""
521 return super().copy(*args, **kwargs)
523 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
524 """See `pydantic.BaseModel.model_dump`."""
525 return super().model_dump(*args, **kwargs)
527 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
528 """See `pydantic.BaseModel.model_dump_json`."""
529 return super().model_dump(*args, **kwargs)
531 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
532 """See `pydantic.BaseModel.model_copy`."""
533 return super().model_copy(*args, **kwargs)
535 @classmethod
536 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
537 """See `pydantic.BaseModel.model_construct`."""
538 return super().model_construct(*args, **kwargs)
540 @classmethod
541 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
542 """See `pydantic.BaseModel.model_json_schema`."""
543 return super().model_json_schema(*args, **kwargs)
545 @classmethod
546 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
547 """See `pydantic.BaseModel.model_validate`."""
548 return super().model_validate(*args, **kwargs)
550 @classmethod
551 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
552 """See `pydantic.BaseModel.model_validate_json`."""
553 return super().model_validate_json(*args, **kwargs)
555 @classmethod
556 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
557 """See `pydantic.BaseModel.model_validate_strings`."""
558 return super().model_validate_strings(*args, **kwargs)
561class ProvenanceQuantumModel(pydantic.BaseModel):
562 """Data model for the quanta in a provenance quantum graph file."""
564 quantum_id: uuid.UUID
565 """Unique identifier for the quantum."""
567 task_label: TaskLabel
568 """Name of the type of this dataset."""
570 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
571 """The full values (required and implied) of this dataset's data ID."""
573 inputs: dict[ConnectionName, list[DatasetIndex]] = pydantic.Field(default_factory=dict)
574 """Internal integer IDs of the datasets predicted to be consumed by this
575 quantum, grouped by connection name.
576 """
578 outputs: dict[ConnectionName, list[DatasetIndex]] = pydantic.Field(default_factory=dict)
579 """Internal integer IDs of the datasets predicted to be produced by this
580 quantum, grouped by connection name.
581 """
583 attempts: list[StorageProvenanceQuantumAttemptModel] = pydantic.Field(default_factory=list)
584 """Provenance for all attempts to execute this quantum, ordered
585 chronologically from first to last.
587 An entry is added merely if the quantum *should* have been attempted; an
588 empty `list` is used only for quanta that were blocked by an upstream
589 failure.
590 """
592 @property
593 def node_id(self) -> uuid.UUID:
594 """Alias for the quantum ID."""
595 return self.quantum_id
597 @classmethod
598 def from_predicted(
599 cls, predicted: PredictedQuantumDatasetsModel, indices: Mapping[uuid.UUID, int]
600 ) -> ProvenanceQuantumModel:
601 """Construct from a predicted quantum model.
603 Parameters
604 ----------
605 predicted : `PredictedQuantumDatasetsModel`
606 Information about the quantum from the predicted graph.
607 indices : `~collections.abc.Mapping [`uuid.UUID`, `int`]
608 Mapping from quantum or dataset UUID to internal integer ID.
610 Returns
611 -------
612 provenance : `ProvenanceQuantumModel`
613 Provenance quantum model.
614 """
615 inputs = {
616 connection_name: [indices[d.dataset_id] for d in predicted_inputs]
617 for connection_name, predicted_inputs in predicted.inputs.items()
618 }
619 outputs = {
620 connection_name: [indices[d.dataset_id] for d in predicted_outputs]
621 for connection_name, predicted_outputs in predicted.outputs.items()
622 }
623 return cls(
624 quantum_id=predicted.quantum_id,
625 task_label=predicted.task_label,
626 data_coordinate=predicted.data_coordinate,
627 inputs=inputs,
628 outputs=outputs,
629 )
631 def _add_to_graph(self, graph: ProvenanceQuantumGraph, address_reader: AddressReader) -> None:
632 """Add this quantum and its edges to datasets to a provenance graph.
634 Parameters
635 ----------
636 graph : `ProvenanceQuantumGraph`
637 Graph to update in place.
638 address_reader : `AddressReader`
639 Reader object that can be used to look up UUIDs from integer
640 indexes.
642 Notes
643 -----
644 This method adds:
646 - a ``bipartite_xgraph`` quantum node with full attributes;
647 - a ``quantum_only_xgraph`` quantum node with full attributes;
648 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
649 nodes with no attributes), while populating those edge attributes;
650 - ``quantum_only_xgraph`` edges to any adjacent quantum that has also
651 already been loaded.
652 """
653 task_node = graph.pipeline_graph.tasks[self.task_label]
654 data_id = DataCoordinate.from_full_values(task_node.dimensions, tuple(self.data_coordinate))
655 last_attempt = (
656 self.attempts[-1]
657 if self.attempts
658 else StorageProvenanceQuantumAttemptModel(status=QuantumAttemptStatus.BLOCKED)
659 )
660 graph._bipartite_xgraph.add_node(
661 self.quantum_id,
662 data_id=data_id,
663 task_label=self.task_label,
664 pipeline_node=task_node,
665 status=last_attempt.status,
666 caveats=last_attempt.caveats,
667 exception=last_attempt.exception,
668 resource_usage=last_attempt.resource_usage,
669 attempts=[a.remap_indices(address_reader) for a in self.attempts],
670 )
671 for connection_name, dataset_indices in self.inputs.items():
672 read_edge = task_node.get_input_edge(connection_name)
673 for dataset_index in dataset_indices:
674 dataset_id = address_reader.find(dataset_index).key
675 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
676 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
677 "pipeline_edges", []
678 ).append(read_edge)
679 for connection_name, dataset_indices in self.outputs.items():
680 write_edge = task_node.get_output_edge(connection_name)
681 for dataset_index in dataset_indices:
682 dataset_id = address_reader.find(dataset_index).key
683 graph._bipartite_xgraph.add_edge(
684 self.quantum_id,
685 dataset_id,
686 is_read=False,
687 # There can only be one pipeline edge for an output.
688 pipeline_edges=[write_edge],
689 )
690 graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
691 graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
692 for dataset_id in graph._bipartite_xgraph.predecessors(self.quantum_id):
693 for upstream_quantum_id in graph._bipartite_xgraph.predecessors(dataset_id):
694 graph._quantum_only_xgraph.add_edge(upstream_quantum_id, self.quantum_id)
695 for dataset_id in graph._bipartite_xgraph.successors(self.quantum_id):
696 for downstream_quantum_id in graph._bipartite_xgraph.successors(dataset_id):
697 graph._quantum_only_xgraph.add_edge(self.quantum_id, downstream_quantum_id)
699 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
700 # when we inherit those docstrings in our public classes.
701 if "sphinx" in sys.modules and not TYPE_CHECKING:
703 def copy(self, *args: Any, **kwargs: Any) -> Any:
704 """See `pydantic.BaseModel.copy`."""
705 return super().copy(*args, **kwargs)
707 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
708 """See `pydantic.BaseModel.model_dump`."""
709 return super().model_dump(*args, **kwargs)
711 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
712 """See `pydantic.BaseModel.model_dump_json`."""
713 return super().model_dump(*args, **kwargs)
715 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
716 """See `pydantic.BaseModel.model_copy`."""
717 return super().model_copy(*args, **kwargs)
719 @classmethod
720 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
721 """See `pydantic.BaseModel.model_construct`."""
722 return super().model_construct(*args, **kwargs)
724 @classmethod
725 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
726 """See `pydantic.BaseModel.model_json_schema`."""
727 return super().model_json_schema(*args, **kwargs)
729 @classmethod
730 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
731 """See `pydantic.BaseModel.model_validate`."""
732 return super().model_validate(*args, **kwargs)
734 @classmethod
735 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
736 """See `pydantic.BaseModel.model_validate_json`."""
737 return super().model_validate_json(*args, **kwargs)
739 @classmethod
740 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
741 """See `pydantic.BaseModel.model_validate_strings`."""
742 return super().model_validate_strings(*args, **kwargs)
745class ProvenanceInitQuantumModel(pydantic.BaseModel):
746 """Data model for the special "init" quanta in a provenance quantum graph
747 file.
748 """
750 quantum_id: uuid.UUID
751 """Unique identifier for the quantum."""
753 task_label: TaskLabel
754 """Name of the type of this dataset.
756 This is always a parent dataset type name, not a component.
758 Note that full dataset type definitions are stored in the pipeline graph.
759 """
761 inputs: dict[ConnectionName, DatasetIndex] = pydantic.Field(default_factory=dict)
762 """Internal integer IDs of the datasets predicted to be consumed by this
763 quantum, grouped by connection name.
764 """
766 outputs: dict[ConnectionName, DatasetIndex] = pydantic.Field(default_factory=dict)
767 """Internal integer IDs of the datasets predicted to be produced by this
768 quantum, grouped by connection name.
769 """
771 @classmethod
772 def from_predicted(
773 cls, predicted: PredictedQuantumDatasetsModel, indices: Mapping[uuid.UUID, int]
774 ) -> ProvenanceInitQuantumModel:
775 """Construct from a predicted quantum model.
777 Parameters
778 ----------
779 predicted : `PredictedQuantumDatasetsModel`
780 Information about the quantum from the predicted graph.
781 indices : `~collections.abc.Mapping [`uuid.UUID`, `int`]
782 Mapping from quantum or dataset UUID to internal integer ID.
784 Returns
785 -------
786 provenance : `ProvenanceInitQuantumModel`
787 Provenance init quantum model.
788 """
789 inputs = {
790 connection_name: indices[predicted_inputs[0].dataset_id]
791 for connection_name, predicted_inputs in predicted.inputs.items()
792 }
793 outputs = {
794 connection_name: indices[predicted_outputs[0].dataset_id]
795 for connection_name, predicted_outputs in predicted.outputs.items()
796 }
797 return cls(
798 quantum_id=predicted.quantum_id,
799 task_label=predicted.task_label,
800 inputs=inputs,
801 outputs=outputs,
802 )
804 def _add_to_graph(
805 self,
806 graph: ProvenanceQuantumGraph,
807 address_reader: AddressReader,
808 empty_data_id: DataCoordinate,
809 ) -> None:
810 """Add this quantum and its edges to datasets to a provenance graph.
812 Parameters
813 ----------
814 graph : `ProvenanceQuantumGraph`
815 Graph to update in place.
816 address_reader : `AddressReader`
817 Reader object that can be used to look up UUIDs from integer
818 indexes.
819 empty_data_id : `lsst.daf.butler.DataCoordinate`
820 The empty data ID for the appropriate dimension universe.
822 Notes
823 -----
824 This method adds:
826 - a ``bipartite_xgraph`` quantum node with full attributes;
827 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
828 nodes with no attributes), while populating those edge attributes;
829 """
830 task_init_node = graph.pipeline_graph.tasks[self.task_label].init
831 graph._bipartite_xgraph.add_node(
832 self.quantum_id, data_id=empty_data_id, task_label=self.task_label, pipeline_node=task_init_node
833 )
834 for connection_name, dataset_index in self.inputs.items():
835 read_edge = task_init_node.get_input_edge(connection_name)
836 dataset_id = address_reader.find(dataset_index).key
837 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
838 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
839 "pipeline_edges", []
840 ).append(read_edge)
841 for connection_name, dataset_index in self.outputs.items():
842 write_edge = task_init_node.get_output_edge(connection_name)
843 dataset_id = address_reader.find(dataset_index).key
844 graph._bipartite_xgraph.add_edge(
845 self.quantum_id,
846 dataset_id,
847 is_read=False,
848 # There can only be one pipeline edge for an output.
849 pipeline_edges=[write_edge],
850 )
851 graph._init_quanta[self.task_label] = self.quantum_id
853 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
854 # when we inherit those docstrings in our public classes.
855 if "sphinx" in sys.modules and not TYPE_CHECKING:
857 def copy(self, *args: Any, **kwargs: Any) -> Any:
858 """See `pydantic.BaseModel.copy`."""
859 return super().copy(*args, **kwargs)
861 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
862 """See `pydantic.BaseModel.model_dump`."""
863 return super().model_dump(*args, **kwargs)
865 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
866 """See `pydantic.BaseModel.model_dump_json`."""
867 return super().model_dump(*args, **kwargs)
869 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
870 """See `pydantic.BaseModel.model_copy`."""
871 return super().model_copy(*args, **kwargs)
873 @classmethod
874 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
875 """See `pydantic.BaseModel.model_construct`."""
876 return super().model_construct(*args, **kwargs)
878 @classmethod
879 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
880 """See `pydantic.BaseModel.model_json_schema`."""
881 return super().model_json_schema(*args, **kwargs)
883 @classmethod
884 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
885 """See `pydantic.BaseModel.model_validate`."""
886 return super().model_validate(*args, **kwargs)
888 @classmethod
889 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
890 """See `pydantic.BaseModel.model_validate_json`."""
891 return super().model_validate_json(*args, **kwargs)
893 @classmethod
894 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
895 """See `pydantic.BaseModel.model_validate_strings`."""
896 return super().model_validate_strings(*args, **kwargs)
899class ProvenanceInitQuantaModel(pydantic.RootModel):
900 """Data model for the init quanta in a provenance graph."""
902 root: list[ProvenanceInitQuantumModel] = pydantic.Field(default_factory=list)
903 """List of special "init" quanta, one for each task."""
905 def _add_to_graph(self, graph: ProvenanceQuantumGraph, address_reader: AddressReader) -> None:
906 """Add this quantum and its edges to datasets to a provenance graph.
908 Parameters
909 ----------
910 graph : `ProvenanceQuantumGraph`
911 Graph to update in place.
912 address_reader : `AddressReader`
913 Reader object that can be used to look up UUIDs from integer
914 indexes.
915 """
916 empty_data_id = DataCoordinate.make_empty(graph.pipeline_graph.universe)
917 for init_quantum in self.root:
918 init_quantum._add_to_graph(graph, address_reader, empty_data_id=empty_data_id)
920 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
921 # when we inherit those docstrings in our public classes.
922 if "sphinx" in sys.modules and not TYPE_CHECKING:
924 def copy(self, *args: Any, **kwargs: Any) -> Any:
925 """See `pydantic.BaseModel.copy`."""
926 return super().copy(*args, **kwargs)
928 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
929 """See `pydantic.BaseModel.model_dump`."""
930 return super().model_dump(*args, **kwargs)
932 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
933 """See `pydantic.BaseModel.model_dump_json`."""
934 return super().model_dump(*args, **kwargs)
936 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
937 """See `pydantic.BaseModel.model_copy`."""
938 return super().model_copy(*args, **kwargs)
940 @classmethod
941 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
942 """See `pydantic.BaseModel.model_construct`."""
943 return super().model_construct(*args, **kwargs)
945 @classmethod
946 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
947 """See `pydantic.BaseModel.model_json_schema`."""
948 return super().model_json_schema(*args, **kwargs)
950 @classmethod
951 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
952 """See `pydantic.BaseModel.model_validate`."""
953 return super().model_validate(*args, **kwargs)
955 @classmethod
956 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
957 """See `pydantic.BaseModel.model_validate_json`."""
958 return super().model_validate_json(*args, **kwargs)
960 @classmethod
961 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
962 """See `pydantic.BaseModel.model_validate_strings`."""
963 return super().model_validate_strings(*args, **kwargs)
966class ProvenanceQuantumGraph(BaseQuantumGraph):
967 """A quantum graph that represents processing that has already been
968 executed.
970 Parameters
971 ----------
972 header : `HeaderModel`
973 General metadata shared with other quantum graph types.
974 pipeline_graph : `.pipeline_graph.PipelineGraph`
975 Graph of tasks and dataset types. May contain a superset of the tasks
976 and dataset types that actually have quanta and datasets in the quantum
977 graph.
979 Notes
980 -----
981 A provenance quantum graph is generally obtained via the
982 `ProvenanceQuantumGraphReader.graph` attribute, which is updated in-place
983 as information is read from disk.
984 """
986 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph) -> None:
987 super().__init__(header, pipeline_graph)
988 self._init_quanta: dict[TaskLabel, uuid.UUID] = {}
989 self._quantum_only_xgraph = networkx.DiGraph()
990 self._bipartite_xgraph = networkx.DiGraph()
991 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {
992 task_label: {} for task_label in self.pipeline_graph.tasks.keys()
993 }
994 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {
995 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
996 }
998 @property
999 def init_quanta(self) -> Mapping[TaskLabel, uuid.UUID]:
1000 """A mapping from task label to the ID of the special init quantum for
1001 that task.
1003 This is populated by the ``init_quanta`` component. Additional
1004 information about each init quantum can be found by using the ID to
1005 look up node attributes in the `bipartite_xgraph`, i.e.::
1007 info: ProvenanceInitQuantumInfo = qg.bipartite_xgraph.nodes[id]
1008 """
1009 return self._init_quanta
1011 @property
1012 def quanta_by_task(self) -> Mapping[TaskLabel, Mapping[DataCoordinate, uuid.UUID]]:
1013 """A nested mapping of all quanta, keyed first by task name and then by
1014 data ID.
1016 Notes
1017 -----
1018 This is populated one quantum at a time as they are read. All tasks in
1019 the pipeline graph are included, even if none of their quanta were
1020 loaded (i.e. nested mappings may be empty).
1022 The returned object may be an internal dictionary; as the type
1023 annotation indicates, it should not be modified in place.
1024 """
1025 return self._quanta_by_task_label
1027 @property
1028 def datasets_by_type(self) -> Mapping[DatasetTypeName, Mapping[DataCoordinate, uuid.UUID]]:
1029 """A nested mapping of all datasets, keyed first by dataset type name
1030 and then by data ID.
1032 Notes
1033 -----
1034 This is populated one dataset at a time as they are read. All dataset
1035 types in the pipeline graph are included, even if none of their
1036 datasets were loaded (i.e. nested mappings may be empty).
1038 The returned object may be an internal dictionary; as the type
1039 annotation indicates, it should not be modified in place.
1040 """
1041 return self._datasets_by_type
1043 @property
1044 def quantum_only_xgraph(self) -> networkx.DiGraph:
1045 """A directed acyclic graph with quanta as nodes (and datasets elided).
1047 Notes
1048 -----
1049 Node keys are quantum UUIDs, and are populated one quantum at a time as
1050 they are loaded. Loading quanta (via
1051 `ProvenanceQuantumGraphReader.read_quanta`) will add the loaded nodes
1052 with full attributes and add edges to adjacent nodes with no
1053 attributes. Loading datasets (via
1054 `ProvenanceQuantumGraphReader.read_datasets`) will also add edges and
1055 nodes with no attributes.
1057 Node attributes are described by the `ProvenanceQuantumInfo` types.
1059 This graph does not include special "init" quanta.
1061 The returned object is a read-only view of an internal one.
1062 """
1063 return self._quantum_only_xgraph.copy(as_view=True)
1065 @property
1066 def bipartite_xgraph(self) -> networkx.DiGraph:
1067 """A directed acyclic graph with quantum and dataset nodes.
1069 Notes
1070 -----
1071 Node keys are quantum or dataset UUIDs, and are populated one quantum
1072 or dataset at a time as they are loaded. Loading quanta (via
1073 `ProvenanceQuantumGraphReader.read_quanta`) or datasets (via
1074 `ProvenanceQuantumGraphReader.read_datasets`) will load those nodes
1075 with full attributes and edges to adjacent nodes with no attributes.
1076 Loading quanta necessary to populate edge attributes.
1078 Node attributes are described by the
1079 `ProvenanceQuantumInfo`, `ProvenanceInitQuantumInfo`, and
1080 `ProvenanceDatasetInfo` types.
1082 This graph includes init-input and init-output datasets, but it does
1083 *not* reflect the dependency between each task's special "init" quantum
1084 and its runtime quanta (as this would require edges between quanta, and
1085 that would break the "bipartite" property).
1087 The returned object is a read-only view of an internal one.
1088 """
1089 return self._bipartite_xgraph.copy(as_view=True)
1091 def make_quantum_table(self) -> astropy.table.Table:
1092 """Construct an `astropy.table.Table` with a tabular summary of the
1093 quanta.
1095 Returns
1096 -------
1097 table : `astropy.table.Table`
1098 A table view of the quantum information. This only includes
1099 counts of status categories and caveats, not any per-data-ID
1100 detail.
1102 Notes
1103 -----
1104 Success caveats in the table are represented by their
1105 `~QuantumSuccessCaveats.concise` form, so when pretty-printing this
1106 table for users, the `~QuantumSuccessCaveats.legend` should generally
1107 be printed as well.
1108 """
1109 rows = []
1110 for task_label, quanta_for_task in self.quanta_by_task.items():
1111 if not self.header.n_task_quanta[task_label]:
1112 continue
1113 status_counts = Counter[QuantumAttemptStatus](
1114 self._quantum_only_xgraph.nodes[q]["status"] for q in quanta_for_task.values()
1115 )
1116 caveat_counts = Counter[QuantumSuccessCaveats | None](
1117 self._quantum_only_xgraph.nodes[q]["caveats"] for q in quanta_for_task.values()
1118 )
1119 caveat_counts.pop(QuantumSuccessCaveats.NO_CAVEATS, None)
1120 caveat_counts.pop(None, None)
1121 if len(caveat_counts) > 1:
1122 caveats = "(multiple)"
1123 elif len(caveat_counts) == 1:
1124 ((code, count),) = caveat_counts.items()
1125 # MyPy can't tell that the pop(None, None) above makes None
1126 # impossible here.
1127 caveats = f"{code.concise()}({count})" # type: ignore[union-attr]
1128 else:
1129 caveats = ""
1130 rows.append(
1131 {
1132 "Task": task_label,
1133 "Unknown": status_counts.get(QuantumAttemptStatus.UNKNOWN, 0),
1134 "Successful": status_counts.get(QuantumAttemptStatus.SUCCESSFUL, 0),
1135 "Caveats": caveats,
1136 "Blocked": status_counts.get(QuantumAttemptStatus.BLOCKED, 0),
1137 "Failed": status_counts.get(QuantumAttemptStatus.FAILED, 0),
1138 "TOTAL": len(quanta_for_task),
1139 "EXPECTED": self.header.n_task_quanta[task_label],
1140 }
1141 )
1142 return astropy.table.Table(rows)
1144 def make_exception_table(self) -> astropy.table.Table:
1145 """Construct an `astropy.table.Table` with counts for each exception
1146 type raised by each task.
1148 At present this only includes information from partial-outputs-error
1149 successes, since exception information for failures is not tracked.
1150 This may change in the future.
1152 Returns
1153 -------
1154 table : `astropy.table.Table`
1155 A table with columns for task label, exception type, and counts.
1156 """
1157 rows = []
1158 for task_label, quanta_for_task in self.quanta_by_task.items():
1159 counts_by_type = Counter(
1160 exc_info.type_name
1161 for q in quanta_for_task.values()
1162 if (exc_info := self._quantum_only_xgraph.nodes[q]["exception"]) is not None
1163 )
1164 for type_name, count in counts_by_type.items():
1165 rows.append({"Task": task_label, "Exception": type_name, "Count": count})
1166 return astropy.table.Table(rows)
1168 def make_task_resource_usage_table(
1169 self, task_label: TaskLabel, include_data_ids: bool = False
1170 ) -> astropy.table.Table:
1171 """Make a table of resource usage for a single task.
1173 Parameters
1174 ----------
1175 task_label : `str`
1176 Label of the task to extract resource usage for.
1177 include_data_ids : `bool`, optional
1178 Whether to also include data ID columns.
1180 Returns
1181 -------
1182 table : `astropy.table.Table`
1183 A table with columns for quantum ID and all fields in
1184 `QuantumResourceUsage`.
1185 """
1186 quanta_for_task = self.quanta_by_task[task_label]
1187 dtype_terms: list[tuple[str, np.dtype]] = [("quantum_id", np.dtype((np.void, 16)))]
1188 if include_data_ids:
1189 dimensions = self.pipeline_graph.tasks[task_label].dimensions
1190 for dimension_name in dimensions.data_coordinate_keys:
1191 dtype = np.dtype(self.pipeline_graph.universe.dimensions[dimension_name].primary_key.pytype)
1192 dtype_terms.append((dimension_name, dtype))
1193 fields = QuantumResourceUsage.get_numpy_fields()
1194 dtype_terms.extend(fields.items())
1195 row_dtype = np.dtype(dtype_terms)
1196 rows: list[object] = []
1197 for data_id, quantum_id in quanta_for_task.items():
1198 info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1199 if (resource_usage := info["resource_usage"]) is not None:
1200 row: tuple[object, ...] = (quantum_id.bytes,)
1201 if include_data_ids:
1202 row += data_id.full_values
1203 row += resource_usage.get_numpy_row()
1204 rows.append(row)
1205 array = np.array(rows, dtype=row_dtype)
1206 return astropy.table.Table(array, units=QuantumResourceUsage.get_units())
1209@dataclasses.dataclass
1210class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1211 """A helper class for reading provenance quantum graphs.
1213 Notes
1214 -----
1215 The `open` context manager should be used to construct new instances.
1216 Instances cannot be used after the context manager exits, except to access
1217 the `graph` attribute`.
1219 The various ``read_*`` methods in this class update the `graph` attribute
1220 in place.
1221 """
1223 graph: ProvenanceQuantumGraph = dataclasses.field(init=False)
1224 """Loaded provenance graph, populated in place as components are read."""
1226 @classmethod
1227 @contextmanager
1228 def open(
1229 cls,
1230 uri: ResourcePathExpression,
1231 *,
1232 page_size: int | None = None,
1233 import_mode: TaskImportMode = TaskImportMode.DO_NOT_IMPORT,
1234 ) -> Iterator[ProvenanceQuantumGraphReader]:
1235 """Construct a reader from a URI.
1237 Parameters
1238 ----------
1239 uri : convertible to `lsst.resources.ResourcePath`
1240 URI to open. Should have a ``.qg`` extension.
1241 page_size : `int`, optional
1242 Approximate number of bytes to read at once from address files and
1243 multi-block files. Note that this does not set a page size for
1244 *all* reads, but it does affect the smallest, most numerous reads.
1245 Can also be set via the ``LSST_QG_PAGE_SIZE`` environment variable.
1246 import_mode : `.pipeline_graph.TaskImportMode`, optional
1247 How to handle importing the task classes referenced in the pipeline
1248 graph.
1250 Returns
1251 -------
1252 reader : `contextlib.AbstractContextManager` [ \
1253 `ProvenanceQuantumGraphReader` ]
1254 A context manager that returns the reader when entered.
1255 """
1256 with cls._open(
1257 uri,
1258 graph_type="provenance",
1259 address_filename="nodes",
1260 page_size=page_size,
1261 import_mode=import_mode,
1262 n_addresses=4,
1263 ) as self:
1264 yield self
1266 def __post_init__(self) -> None:
1267 self.graph = ProvenanceQuantumGraph(self.header, self.pipeline_graph)
1269 def read_init_quanta(self) -> None:
1270 """Read the thin graph, with all edge information and categorization of
1271 quanta by task label.
1272 """
1273 init_quanta = self._read_single_block("init_quanta", ProvenanceInitQuantaModel)
1274 for init_quantum in init_quanta.root:
1275 self.graph._init_quanta[init_quantum.task_label] = init_quantum.quantum_id
1276 init_quanta._add_to_graph(self.graph, self.address_reader)
1278 def read_full_graph(self) -> None:
1279 """Read all bipartite edges and all quantum and dataset node
1280 attributes, fully populating the `graph` attribute.
1282 Notes
1283 -----
1284 This does not read logs, metadata, or packages ; those must always be
1285 fetched explicitly.
1286 """
1287 self.read_init_quanta()
1288 self.read_datasets()
1289 self.read_quanta()
1291 def read_datasets(self, datasets: Iterable[uuid.UUID | DatasetIndex] | None = None) -> None:
1292 """Read information about the given datasets.
1294 Parameters
1295 ----------
1296 datasets : `~collections.abc.Iterable` [`uuid.UUID` or `int`], optional
1297 Iterable of dataset IDs or indices to load. If not provided, all
1298 datasets will be loaded. The UUIDs and indices of quanta will be
1299 ignored.
1300 """
1301 self._read_nodes(datasets, DATASET_ADDRESS_INDEX, DATASET_MB_NAME, ProvenanceDatasetModel)
1303 def read_quanta(self, quanta: Iterable[uuid.UUID | QuantumIndex] | None = None) -> None:
1304 """Read information about the given quanta.
1306 Parameters
1307 ----------
1308 quanta : `~collections.abc.Iterable` [`uuid.UUID` or `int`], optional
1309 Iterable of quantum IDs or indices to load. If not provided, all
1310 quanta will be loaded. The UUIDs and indices of datasets and
1311 special init quanta will be ignored.
1312 """
1313 self._read_nodes(quanta, QUANTUM_ADDRESS_INDEX, QUANTUM_MB_NAME, ProvenanceQuantumModel)
1315 def _read_nodes(
1316 self,
1317 nodes: Iterable[uuid.UUID | int] | None,
1318 address_index: int,
1319 mb_name: str,
1320 model_type: type[ProvenanceDatasetModel] | type[ProvenanceQuantumModel],
1321 ) -> None:
1322 node: ProvenanceDatasetModel | ProvenanceQuantumModel | None
1323 if nodes is None:
1324 self.address_reader.read_all()
1325 nodes = self.address_reader.rows.keys()
1326 for node in MultiblockReader.read_all_models_in_zip(
1327 self.zf,
1328 mb_name,
1329 model_type,
1330 self.decompressor,
1331 int_size=self.header.int_size,
1332 page_size=self.page_size,
1333 ):
1334 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(node.node_id, {}):
1335 # Use the old node to reduce memory usage (since it might
1336 # also have other outstanding reference holders).
1337 continue
1338 node._add_to_graph(self.graph, self.address_reader)
1339 return
1340 with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1341 for node_id_or_index in nodes:
1342 address_row = self.address_reader.find(node_id_or_index)
1343 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1344 # Use the old node to reduce memory usage (since it might
1345 # also have other outstanding reference holders).
1346 continue
1347 node = mb_reader.read_model(
1348 address_row.addresses[address_index], model_type, self.decompressor
1349 )
1350 if node is not None:
1351 node._add_to_graph(self.graph, self.address_reader)
1353 def fetch_logs(
1354 self, nodes: Iterable[uuid.UUID | DatasetIndex | QuantumIndex]
1355 ) -> dict[uuid.UUID | DatasetIndex | QuantumIndex, list[ButlerLogRecords | None]]:
1356 """Fetch log datasets.
1358 Parameters
1359 ----------
1360 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1361 UUIDs or internal integer IDS of the log datasets themselves or of
1362 the quanta they correspond to.
1364 Returns
1365 -------
1366 logs : `dict` [ `uuid.UUID` or `int`, `list` [\
1367 `lsst.daf.butler.ButlerLogRecords` or `None`] ]
1368 Logs for the given IDs. Each value is a list of
1369 `lsst.daf.butler.ButlerLogRecords` instances representing different
1370 execution attempts, ordered chronologically from first to last.
1371 Attempts where logs were missing will have `None` in this list.
1372 """
1373 result: dict[uuid.UUID | DatasetIndex | QuantumIndex, list[ButlerLogRecords | None]] = {}
1374 with MultiblockReader.open_in_zip(self.zf, LOG_MB_NAME, int_size=self.header.int_size) as mb_reader:
1375 for node_id_or_index in nodes:
1376 address_row = self.address_reader.find(node_id_or_index)
1377 logs_by_attempt = mb_reader.read_model(
1378 address_row.addresses[LOG_ADDRESS_INDEX], ProvenanceLogRecordsModel, self.decompressor
1379 )
1380 if logs_by_attempt is not None:
1381 result[node_id_or_index] = [
1382 ButlerLogRecords.from_records(attempt_logs) if attempt_logs is not None else None
1383 for attempt_logs in logs_by_attempt.attempts
1384 ]
1385 return result
1387 def fetch_metadata(
1388 self, nodes: Iterable[uuid.UUID | DatasetIndex | QuantumIndex]
1389 ) -> dict[uuid.UUID | DatasetIndex | QuantumIndex, list[TaskMetadata | None]]:
1390 """Fetch metadata datasets.
1392 Parameters
1393 ----------
1394 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1395 UUIDs or internal integer IDs of the metadata datasets themselves
1396 or of the quanta they correspond to.
1398 Returns
1399 -------
1400 metadata : `dict` [ `uuid.UUID` or `int`, `list` [`.TaskMetadata`] ]
1401 Metadata for the given IDs. Each value is a list of
1402 `.TaskMetadata` instances representing different execution
1403 attempts, ordered chronologically from first to last. Attempts
1404 where metadata was missing (not written even in the fallback extra
1405 provenance in the logs) will have `None` in this list.
1406 """
1407 result: dict[uuid.UUID | DatasetIndex | QuantumIndex, list[TaskMetadata | None]] = {}
1408 with MultiblockReader.open_in_zip(
1409 self.zf, METADATA_MB_NAME, int_size=self.header.int_size
1410 ) as mb_reader:
1411 for node_id_or_index in nodes:
1412 address_row = self.address_reader.find(node_id_or_index)
1413 metadata_by_attempt = mb_reader.read_model(
1414 address_row.addresses[METADATA_ADDRESS_INDEX],
1415 ProvenanceTaskMetadataModel,
1416 self.decompressor,
1417 )
1418 if metadata_by_attempt is not None:
1419 result[node_id_or_index] = metadata_by_attempt.attempts
1420 return result
1422 def fetch_packages(self) -> Packages:
1423 """Fetch package version information."""
1424 data = self._read_single_block_raw("packages")
1425 return Packages.fromBytes(data, format="json")