Coverage for python / lsst / pipe / base / quantum_graph / _provenance.py: 40%
305 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:14 -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, 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 DatasetInfo,
72 DatasetTypeName,
73 HeaderModel,
74 QuantumInfo,
75 TaskLabel,
76)
77from ._multiblock import MultiblockReader
78from ._predicted import PredictedDatasetModel, PredictedQuantumDatasetsModel
80DATASET_ADDRESS_INDEX = 0
81QUANTUM_ADDRESS_INDEX = 1
82LOG_ADDRESS_INDEX = 2
83METADATA_ADDRESS_INDEX = 3
85DATASET_MB_NAME = "datasets"
86QUANTUM_MB_NAME = "quanta"
87LOG_MB_NAME = "logs"
88METADATA_MB_NAME = "metadata"
90_I = TypeVar("_I", bound=uuid.UUID | int)
93class ProvenanceDatasetInfo(DatasetInfo):
94 """A typed dictionary that annotates the attributes of the NetworkX graph
95 node data for a provenance dataset.
97 Since NetworkX types are not generic over their node mapping type, this has
98 to be used explicitly, e.g.::
100 node_data: ProvenanceDatasetInfo = xgraph.nodes[dataset_id]
102 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
103 """
105 dataset_id: uuid.UUID
106 """Unique identifier for the dataset."""
108 produced: bool
109 """Whether this dataset was produced (vs. only predicted).
111 This is always `True` for overall input datasets. It is also `True` for
112 datasets that were produced and then removed before/during transfer back to
113 the central butler repository, so it may not reflect the continued
114 existence of the dataset.
115 """
118class ProvenanceQuantumInfo(QuantumInfo):
119 """A typed dictionary that annotates the attributes of the NetworkX graph
120 node data for a provenance quantum.
122 Since NetworkX types are not generic over their node mapping type, this has
123 to be used explicitly, e.g.::
125 node_data: ProvenanceQuantumInfo = xgraph.nodes[quantum_id]
127 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph` or
128 `ProvenanceQuantumGraph.quantum_only_xgraph`
129 """
131 status: QuantumAttemptStatus
132 """Enumerated status for the quantum.
134 This corresponds to the last attempt to run this quantum, or
135 `QuantumAttemptStatus.BLOCKED` if there were no attempts.
136 """
138 caveats: QuantumSuccessCaveats | None
139 """Flags indicating caveats on successful quanta.
141 This corresponds to the last attempt to run this quantum.
142 """
144 exception: ExceptionInfo | None
145 """Information about an exception raised when the quantum was executing.
147 This corresponds to the last attempt to run this quantum.
148 """
150 resource_usage: QuantumResourceUsage | None
151 """Resource usage information (timing, memory use) for this quantum.
153 This corresponds to the last attempt to run this quantum.
154 """
156 attempts: list[ProvenanceQuantumAttemptModel]
157 """Information about each attempt to run this quantum.
159 An entry is added merely if the quantum *should* have been attempted; an
160 empty `list` is used only for quanta that were blocked by an upstream
161 failure.
162 """
165class ProvenanceInitQuantumInfo(TypedDict):
166 """A typed dictionary that annotates the attributes of the NetworkX graph
167 node data for a provenance init quantum.
169 Since NetworkX types are not generic over their node mapping type, this has
170 to be used explicitly, e.g.::
172 node_data: ProvenanceInitQuantumInfo = xgraph.nodes[quantum_id]
174 where ``xgraph`` is `ProvenanceQuantumGraph.bipartite_xgraph`.
175 """
177 data_id: DataCoordinate
178 """Data ID of the quantum.
180 This is always an empty ID; this key exists to allow init-quanta and
181 regular quanta to be treated more similarly.
182 """
184 task_label: str
185 """Label of the task for this quantum."""
187 pipeline_node: TaskInitNode
188 """Node in the pipeline graph for this task's init-only step."""
191class ProvenanceDatasetModel(PredictedDatasetModel):
192 """Data model for the datasets in a provenance quantum graph file."""
194 produced: bool
195 """Whether this dataset was produced (vs. only predicted).
197 This is always `True` for overall input datasets. It is also `True` for
198 datasets that were produced and then removed before/during transfer back to
199 the central butler repository, so it may not reflect the continued
200 existence of the dataset.
201 """
203 producer: uuid.UUID | None = None
204 """ID of the quantum that produced this dataset.
206 This is `None` for overall inputs to the graph.
207 """
209 consumers: list[uuid.UUID] = pydantic.Field(default_factory=list)
210 """IDs of quanta that were predicted to consume this dataset."""
212 @property
213 def node_id(self) -> uuid.UUID:
214 """Alias for the dataset ID."""
215 return self.dataset_id
217 @classmethod
218 def from_predicted(
219 cls,
220 predicted: PredictedDatasetModel,
221 producer: uuid.UUID | None = None,
222 consumers: Iterable[uuid.UUID] = (),
223 ) -> ProvenanceDatasetModel:
224 """Construct from a predicted dataset model.
226 Parameters
227 ----------
228 predicted : `PredictedDatasetModel`
229 Information about the dataset from the predicted graph.
230 producer : `uuid.UUID` or `None`, optional
231 ID of the quantum that was predicted to produce this dataset.
232 consumers : `~collections.abc.Iterable` [`uuid.UUID`], optional
233 IDs of the quanta that were predicted to consume this dataset.
235 Returns
236 -------
237 provenance : `ProvenanceDatasetModel`
238 Provenance dataset model.
240 Notes
241 -----
242 This initializes `produced` to `True` when ``producer is None`` and
243 `False` otherwise, on the assumption that it will be updated later.
244 """
245 return cls.model_construct(
246 dataset_id=predicted.dataset_id,
247 dataset_type_name=predicted.dataset_type_name,
248 data_coordinate=predicted.data_coordinate,
249 run=predicted.run,
250 produced=(producer is None), # if it's not produced by this QG, it's an overall input
251 producer=producer,
252 consumers=list(consumers),
253 )
255 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
256 """Add this dataset and its edges to quanta to a provenance graph.
258 Parameters
259 ----------
260 graph : `ProvenanceQuantumGraph`
261 Graph to update in place.
263 Notes
264 -----
265 This method adds:
267 - a ``bipartite_xgraph`` dataset node with full attributes;
268 - ``bipartite_xgraph`` edges to adjacent quanta (which adds quantum
269 nodes with no attributes), without populating edge attributes;
270 - ``quantum_only_xgraph`` edges for each pair of quanta in which one
271 produces this dataset and another consumes it (this also adds quantum
272 nodes with no attributes).
273 """
274 dataset_type_node = graph.pipeline_graph.dataset_types[self.dataset_type_name]
275 data_id = DataCoordinate.from_full_values(dataset_type_node.dimensions, tuple(self.data_coordinate))
276 graph._bipartite_xgraph.add_node(
277 self.dataset_id,
278 data_id=data_id,
279 dataset_type_name=self.dataset_type_name,
280 pipeline_node=dataset_type_node,
281 run=self.run,
282 produced=self.produced,
283 )
284 if self.producer is not None:
285 graph._bipartite_xgraph.add_edge(self.producer, self.dataset_id)
286 for consumer_id in self.consumers:
287 graph._bipartite_xgraph.add_edge(self.dataset_id, consumer_id)
288 if self.producer is not None:
289 graph._quantum_only_xgraph.add_edge(self.producer, consumer_id)
290 graph._datasets_by_type[self.dataset_type_name][data_id] = self.dataset_id
292 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
293 # when we inherit those docstrings in our public classes.
294 if "sphinx" in sys.modules and not TYPE_CHECKING:
296 def copy(self, *args: Any, **kwargs: Any) -> Any:
297 """See `pydantic.BaseModel.copy`."""
298 return super().copy(*args, **kwargs)
300 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
301 """See `pydantic.BaseModel.model_dump`."""
302 return super().model_dump(*args, **kwargs)
304 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
305 """See `pydantic.BaseModel.model_dump_json`."""
306 return super().model_dump(*args, **kwargs)
308 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
309 """See `pydantic.BaseModel.model_copy`."""
310 return super().model_copy(*args, **kwargs)
312 @classmethod
313 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
314 """See `pydantic.BaseModel.model_construct`."""
315 return super().model_construct(*args, **kwargs)
317 @classmethod
318 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
319 """See `pydantic.BaseModel.model_json_schema`."""
320 return super().model_json_schema(*args, **kwargs)
322 @classmethod
323 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
324 """See `pydantic.BaseModel.model_validate`."""
325 return super().model_validate(*args, **kwargs)
327 @classmethod
328 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
329 """See `pydantic.BaseModel.model_validate_json`."""
330 return super().model_validate_json(*args, **kwargs)
332 @classmethod
333 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
334 """See `pydantic.BaseModel.model_validate_strings`."""
335 return super().model_validate_strings(*args, **kwargs)
338class ProvenanceQuantumAttemptModel(pydantic.BaseModel):
339 """Data model for a now-superseded attempt to run a quantum in a
340 provenance quantum graph file.
341 """
343 attempt: int = 0
344 """Counter incremented for every attempt to execute this quantum."""
346 status: QuantumAttemptStatus = QuantumAttemptStatus.UNKNOWN
347 """Enumerated status for the quantum."""
349 caveats: QuantumSuccessCaveats | None = None
350 """Flags indicating caveats on successful quanta."""
352 exception: ExceptionInfo | None = None
353 """Information about an exception raised when the quantum was executing."""
355 resource_usage: QuantumResourceUsage | None = None
356 """Resource usage information (timing, memory use) for this quantum."""
358 previous_process_quanta: list[uuid.UUID] = pydantic.Field(default_factory=list)
359 """The IDs of other quanta previously executed in the same process as this
360 one.
361 """
363 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
364 # when we inherit those docstrings in our public classes.
365 if "sphinx" in sys.modules and not TYPE_CHECKING:
367 def copy(self, *args: Any, **kwargs: Any) -> Any:
368 """See `pydantic.BaseModel.copy`."""
369 return super().copy(*args, **kwargs)
371 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
372 """See `pydantic.BaseModel.model_dump`."""
373 return super().model_dump(*args, **kwargs)
375 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
376 """See `pydantic.BaseModel.model_dump_json`."""
377 return super().model_dump(*args, **kwargs)
379 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
380 """See `pydantic.BaseModel.model_copy`."""
381 return super().model_copy(*args, **kwargs)
383 @classmethod
384 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
385 """See `pydantic.BaseModel.model_construct`."""
386 return super().model_construct(*args, **kwargs)
388 @classmethod
389 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
390 """See `pydantic.BaseModel.model_json_schema`."""
391 return super().model_json_schema(*args, **kwargs)
393 @classmethod
394 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
395 """See `pydantic.BaseModel.model_validate`."""
396 return super().model_validate(*args, **kwargs)
398 @classmethod
399 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
400 """See `pydantic.BaseModel.model_validate_json`."""
401 return super().model_validate_json(*args, **kwargs)
403 @classmethod
404 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
405 """See `pydantic.BaseModel.model_validate_strings`."""
406 return super().model_validate_strings(*args, **kwargs)
409class ProvenanceLogRecordsModel(pydantic.BaseModel):
410 """Data model for storing execution logs in a provenance quantum graph
411 file.
412 """
414 attempts: list[list[ButlerLogRecord] | None] = pydantic.Field(default_factory=list)
415 """Logs from attempts to run this task, ordered chronologically from first
416 to last.
417 """
419 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
420 # when we inherit those docstrings in our public classes.
421 if "sphinx" in sys.modules and not TYPE_CHECKING:
423 def copy(self, *args: Any, **kwargs: Any) -> Any:
424 """See `pydantic.BaseModel.copy`."""
425 return super().copy(*args, **kwargs)
427 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
428 """See `pydantic.BaseModel.model_dump`."""
429 return super().model_dump(*args, **kwargs)
431 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
432 """See `pydantic.BaseModel.model_dump_json`."""
433 return super().model_dump(*args, **kwargs)
435 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
436 """See `pydantic.BaseModel.model_copy`."""
437 return super().model_copy(*args, **kwargs)
439 @classmethod
440 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
441 """See `pydantic.BaseModel.model_construct`."""
442 return super().model_construct(*args, **kwargs)
444 @classmethod
445 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
446 """See `pydantic.BaseModel.model_json_schema`."""
447 return super().model_json_schema(*args, **kwargs)
449 @classmethod
450 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
451 """See `pydantic.BaseModel.model_validate`."""
452 return super().model_validate(*args, **kwargs)
454 @classmethod
455 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
456 """See `pydantic.BaseModel.model_validate_json`."""
457 return super().model_validate_json(*args, **kwargs)
459 @classmethod
460 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
461 """See `pydantic.BaseModel.model_validate_strings`."""
462 return super().model_validate_strings(*args, **kwargs)
465class ProvenanceTaskMetadataModel(pydantic.BaseModel):
466 """Data model for storing task metadata in a provenance quantum graph
467 file.
468 """
470 attempts: list[TaskMetadata | None] = pydantic.Field(default_factory=list)
471 """Metadata from attempts to run this task, ordered chronologically from
472 first to last.
473 """
475 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
476 # when we inherit those docstrings in our public classes.
477 if "sphinx" in sys.modules and not TYPE_CHECKING:
479 def copy(self, *args: Any, **kwargs: Any) -> Any:
480 """See `pydantic.BaseModel.copy`."""
481 return super().copy(*args, **kwargs)
483 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
484 """See `pydantic.BaseModel.model_dump`."""
485 return super().model_dump(*args, **kwargs)
487 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
488 """See `pydantic.BaseModel.model_dump_json`."""
489 return super().model_dump(*args, **kwargs)
491 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
492 """See `pydantic.BaseModel.model_copy`."""
493 return super().model_copy(*args, **kwargs)
495 @classmethod
496 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
497 """See `pydantic.BaseModel.model_construct`."""
498 return super().model_construct(*args, **kwargs)
500 @classmethod
501 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
502 """See `pydantic.BaseModel.model_json_schema`."""
503 return super().model_json_schema(*args, **kwargs)
505 @classmethod
506 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
507 """See `pydantic.BaseModel.model_validate`."""
508 return super().model_validate(*args, **kwargs)
510 @classmethod
511 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
512 """See `pydantic.BaseModel.model_validate_json`."""
513 return super().model_validate_json(*args, **kwargs)
515 @classmethod
516 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
517 """See `pydantic.BaseModel.model_validate_strings`."""
518 return super().model_validate_strings(*args, **kwargs)
521class ProvenanceQuantumModel(pydantic.BaseModel):
522 """Data model for the quanta in a provenance quantum graph file."""
524 quantum_id: uuid.UUID
525 """Unique identifier for the quantum."""
527 task_label: TaskLabel
528 """Name of the type of this dataset."""
530 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
531 """The full values (required and implied) of this dataset's data ID."""
533 inputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
534 """IDs of the datasets predicted to be consumed by this quantum, grouped by
535 connection name.
536 """
538 outputs: dict[ConnectionName, list[uuid.UUID]] = pydantic.Field(default_factory=dict)
539 """IDs of the datasets predicted to be produced by this quantum, grouped by
540 connection name.
541 """
543 attempts: list[ProvenanceQuantumAttemptModel] = pydantic.Field(default_factory=list)
544 """Provenance for all attempts to execute this quantum, ordered
545 chronologically from first to last.
547 An entry is added merely if the quantum *should* have been attempted; an
548 empty `list` is used only for quanta that were blocked by an upstream
549 failure.
550 """
552 @property
553 def node_id(self) -> uuid.UUID:
554 """Alias for the quantum ID."""
555 return self.quantum_id
557 @classmethod
558 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceQuantumModel:
559 """Construct from a predicted quantum model.
561 Parameters
562 ----------
563 predicted : `PredictedQuantumDatasetsModel`
564 Information about the quantum from the predicted graph.
566 Returns
567 -------
568 provenance : `ProvenanceQuantumModel`
569 Provenance quantum model.
570 """
571 inputs = {
572 connection_name: [d.dataset_id for d in predicted_inputs]
573 for connection_name, predicted_inputs in predicted.inputs.items()
574 }
575 outputs = {
576 connection_name: [d.dataset_id for d in predicted_outputs]
577 for connection_name, predicted_outputs in predicted.outputs.items()
578 }
579 return cls(
580 quantum_id=predicted.quantum_id,
581 task_label=predicted.task_label,
582 data_coordinate=predicted.data_coordinate,
583 inputs=inputs,
584 outputs=outputs,
585 )
587 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
588 """Add this quantum and its edges to datasets to a provenance graph.
590 Parameters
591 ----------
592 graph : `ProvenanceQuantumGraph`
593 Graph to update in place.
595 Notes
596 -----
597 This method adds:
599 - a ``bipartite_xgraph`` quantum node with full attributes;
600 - a ``quantum_only_xgraph`` quantum node with full attributes;
601 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
602 nodes with no attributes), while populating those edge attributes;
603 - ``quantum_only_xgraph`` edges to any adjacent quantum that has also
604 already been loaded.
605 """
606 task_node = graph.pipeline_graph.tasks[self.task_label]
607 data_id = DataCoordinate.from_full_values(task_node.dimensions, tuple(self.data_coordinate))
608 last_attempt = (
609 self.attempts[-1]
610 if self.attempts
611 else ProvenanceQuantumAttemptModel(status=QuantumAttemptStatus.BLOCKED)
612 )
613 graph._bipartite_xgraph.add_node(
614 self.quantum_id,
615 data_id=data_id,
616 task_label=self.task_label,
617 pipeline_node=task_node,
618 status=last_attempt.status,
619 caveats=last_attempt.caveats,
620 exception=last_attempt.exception,
621 resource_usage=last_attempt.resource_usage,
622 attempts=self.attempts,
623 )
624 for connection_name, dataset_ids in self.inputs.items():
625 read_edge = task_node.get_input_edge(connection_name)
626 for dataset_id in dataset_ids:
627 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
628 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
629 "pipeline_edges", []
630 ).append(read_edge)
631 for connection_name, dataset_ids in self.outputs.items():
632 write_edge = task_node.get_output_edge(connection_name)
633 for dataset_id in dataset_ids:
634 graph._bipartite_xgraph.add_edge(
635 self.quantum_id,
636 dataset_id,
637 is_read=False,
638 # There can only be one pipeline edge for an output.
639 pipeline_edges=[write_edge],
640 )
641 graph._quanta_by_task_label[self.task_label][data_id] = self.quantum_id
642 graph._quantum_only_xgraph.add_node(self.quantum_id, **graph._bipartite_xgraph.nodes[self.quantum_id])
643 for dataset_id in graph._bipartite_xgraph.predecessors(self.quantum_id):
644 for upstream_quantum_id in graph._bipartite_xgraph.predecessors(dataset_id):
645 graph._quantum_only_xgraph.add_edge(upstream_quantum_id, self.quantum_id)
646 for dataset_id in graph._bipartite_xgraph.successors(self.quantum_id):
647 for downstream_quantum_id in graph._bipartite_xgraph.successors(dataset_id):
648 graph._quantum_only_xgraph.add_edge(self.quantum_id, downstream_quantum_id)
650 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
651 # when we inherit those docstrings in our public classes.
652 if "sphinx" in sys.modules and not TYPE_CHECKING:
654 def copy(self, *args: Any, **kwargs: Any) -> Any:
655 """See `pydantic.BaseModel.copy`."""
656 return super().copy(*args, **kwargs)
658 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
659 """See `pydantic.BaseModel.model_dump`."""
660 return super().model_dump(*args, **kwargs)
662 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
663 """See `pydantic.BaseModel.model_dump_json`."""
664 return super().model_dump(*args, **kwargs)
666 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
667 """See `pydantic.BaseModel.model_copy`."""
668 return super().model_copy(*args, **kwargs)
670 @classmethod
671 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
672 """See `pydantic.BaseModel.model_construct`."""
673 return super().model_construct(*args, **kwargs)
675 @classmethod
676 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
677 """See `pydantic.BaseModel.model_json_schema`."""
678 return super().model_json_schema(*args, **kwargs)
680 @classmethod
681 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
682 """See `pydantic.BaseModel.model_validate`."""
683 return super().model_validate(*args, **kwargs)
685 @classmethod
686 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
687 """See `pydantic.BaseModel.model_validate_json`."""
688 return super().model_validate_json(*args, **kwargs)
690 @classmethod
691 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
692 """See `pydantic.BaseModel.model_validate_strings`."""
693 return super().model_validate_strings(*args, **kwargs)
696class ProvenanceInitQuantumModel(pydantic.BaseModel):
697 """Data model for the special "init" quanta in a provenance quantum graph
698 file.
699 """
701 quantum_id: uuid.UUID
702 """Unique identifier for the quantum."""
704 task_label: TaskLabel
705 """Name of the type of this dataset.
707 This is always a parent dataset type name, not a component.
709 Note that full dataset type definitions are stored in the pipeline graph.
710 """
712 inputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
713 """IDs of the datasets predicted to be consumed by this quantum, grouped by
714 connection name.
715 """
717 outputs: dict[ConnectionName, uuid.UUID] = pydantic.Field(default_factory=dict)
718 """IDs of the datasets predicted to be produced by this quantum, grouped by
719 connection name.
720 """
722 @classmethod
723 def from_predicted(cls, predicted: PredictedQuantumDatasetsModel) -> ProvenanceInitQuantumModel:
724 """Construct from a predicted quantum model.
726 Parameters
727 ----------
728 predicted : `PredictedQuantumDatasetsModel`
729 Information about the quantum from the predicted graph.
731 Returns
732 -------
733 provenance : `ProvenanceInitQuantumModel`
734 Provenance init quantum model.
735 """
736 inputs = {
737 connection_name: predicted_inputs[0].dataset_id
738 for connection_name, predicted_inputs in predicted.inputs.items()
739 }
740 outputs = {
741 connection_name: predicted_outputs[0].dataset_id
742 for connection_name, predicted_outputs in predicted.outputs.items()
743 }
744 return cls(
745 quantum_id=predicted.quantum_id,
746 task_label=predicted.task_label,
747 inputs=inputs,
748 outputs=outputs,
749 )
751 def _add_to_graph(self, graph: ProvenanceQuantumGraph, empty_data_id: DataCoordinate) -> None:
752 """Add this quantum and its edges to datasets to a provenance graph.
754 Parameters
755 ----------
756 graph : `ProvenanceQuantumGraph`
757 Graph to update in place.
758 empty_data_id : `lsst.daf.butler.DataCoordinate`
759 The empty data ID for the appropriate dimension universe.
761 Notes
762 -----
763 This method adds:
765 - a ``bipartite_xgraph`` quantum node with full attributes;
766 - ``bipartite_xgraph`` edges to adjacent datasets (which adds datasets
767 nodes with no attributes), while populating those edge attributes;
768 """
769 task_init_node = graph.pipeline_graph.tasks[self.task_label].init
770 graph._bipartite_xgraph.add_node(
771 self.quantum_id, data_id=empty_data_id, task_label=self.task_label, pipeline_node=task_init_node
772 )
773 for connection_name, dataset_id in self.inputs.items():
774 read_edge = task_init_node.get_input_edge(connection_name)
775 graph._bipartite_xgraph.add_edge(dataset_id, self.quantum_id, is_read=True)
776 graph._bipartite_xgraph.edges[dataset_id, self.quantum_id].setdefault(
777 "pipeline_edges", []
778 ).append(read_edge)
779 for connection_name, dataset_id in self.outputs.items():
780 write_edge = task_init_node.get_output_edge(connection_name)
781 graph._bipartite_xgraph.add_edge(
782 self.quantum_id,
783 dataset_id,
784 is_read=False,
785 # There can only be one pipeline edge for an output.
786 pipeline_edges=[write_edge],
787 )
788 graph._init_quanta[self.task_label] = self.quantum_id
790 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
791 # when we inherit those docstrings in our public classes.
792 if "sphinx" in sys.modules and not TYPE_CHECKING:
794 def copy(self, *args: Any, **kwargs: Any) -> Any:
795 """See `pydantic.BaseModel.copy`."""
796 return super().copy(*args, **kwargs)
798 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
799 """See `pydantic.BaseModel.model_dump`."""
800 return super().model_dump(*args, **kwargs)
802 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
803 """See `pydantic.BaseModel.model_dump_json`."""
804 return super().model_dump(*args, **kwargs)
806 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
807 """See `pydantic.BaseModel.model_copy`."""
808 return super().model_copy(*args, **kwargs)
810 @classmethod
811 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
812 """See `pydantic.BaseModel.model_construct`."""
813 return super().model_construct(*args, **kwargs)
815 @classmethod
816 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
817 """See `pydantic.BaseModel.model_json_schema`."""
818 return super().model_json_schema(*args, **kwargs)
820 @classmethod
821 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
822 """See `pydantic.BaseModel.model_validate`."""
823 return super().model_validate(*args, **kwargs)
825 @classmethod
826 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
827 """See `pydantic.BaseModel.model_validate_json`."""
828 return super().model_validate_json(*args, **kwargs)
830 @classmethod
831 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
832 """See `pydantic.BaseModel.model_validate_strings`."""
833 return super().model_validate_strings(*args, **kwargs)
836class ProvenanceInitQuantaModel(pydantic.RootModel):
837 """Data model for the init quanta in a provenance graph."""
839 root: list[ProvenanceInitQuantumModel] = pydantic.Field(default_factory=list)
840 """List of special "init" quanta, one for each task."""
842 def _add_to_graph(self, graph: ProvenanceQuantumGraph) -> None:
843 """Add this quantum and its edges to datasets to a provenance graph.
845 Parameters
846 ----------
847 graph : `ProvenanceQuantumGraph`
848 Graph to update in place.
849 """
850 empty_data_id = DataCoordinate.make_empty(graph.pipeline_graph.universe)
851 for init_quantum in self.root:
852 init_quantum._add_to_graph(graph, empty_data_id=empty_data_id)
854 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
855 # when we inherit those docstrings in our public classes.
856 if "sphinx" in sys.modules and not TYPE_CHECKING:
858 def copy(self, *args: Any, **kwargs: Any) -> Any:
859 """See `pydantic.BaseModel.copy`."""
860 return super().copy(*args, **kwargs)
862 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
863 """See `pydantic.BaseModel.model_dump`."""
864 return super().model_dump(*args, **kwargs)
866 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
867 """See `pydantic.BaseModel.model_dump_json`."""
868 return super().model_dump(*args, **kwargs)
870 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
871 """See `pydantic.BaseModel.model_copy`."""
872 return super().model_copy(*args, **kwargs)
874 @classmethod
875 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
876 """See `pydantic.BaseModel.model_construct`."""
877 return super().model_construct(*args, **kwargs)
879 @classmethod
880 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
881 """See `pydantic.BaseModel.model_json_schema`."""
882 return super().model_json_schema(*args, **kwargs)
884 @classmethod
885 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
886 """See `pydantic.BaseModel.model_validate`."""
887 return super().model_validate(*args, **kwargs)
889 @classmethod
890 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
891 """See `pydantic.BaseModel.model_validate_json`."""
892 return super().model_validate_json(*args, **kwargs)
894 @classmethod
895 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
896 """See `pydantic.BaseModel.model_validate_strings`."""
897 return super().model_validate_strings(*args, **kwargs)
900class ProvenanceQuantumGraph(BaseQuantumGraph):
901 """A quantum graph that represents processing that has already been
902 executed.
904 Parameters
905 ----------
906 header : `HeaderModel`
907 General metadata shared with other quantum graph types.
908 pipeline_graph : `.pipeline_graph.PipelineGraph`
909 Graph of tasks and dataset types. May contain a superset of the tasks
910 and dataset types that actually have quanta and datasets in the quantum
911 graph.
913 Notes
914 -----
915 A provenance quantum graph is generally obtained via the
916 `ProvenanceQuantumGraphReader.graph` attribute, which is updated in-place
917 as information is read from disk.
918 """
920 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph) -> None:
921 super().__init__(header, pipeline_graph)
922 self._init_quanta: dict[TaskLabel, uuid.UUID] = {}
923 self._quantum_only_xgraph = networkx.DiGraph()
924 self._bipartite_xgraph = networkx.DiGraph()
925 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {
926 task_label: {} for task_label in self.pipeline_graph.tasks.keys()
927 }
928 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {
929 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
930 }
932 @property
933 def init_quanta(self) -> Mapping[TaskLabel, uuid.UUID]:
934 """A mapping from task label to the ID of the special init quantum for
935 that task.
937 This is populated by the ``init_quanta`` component. Additional
938 information about each init quantum can be found by using the ID to
939 look up node attributes in the `bipartite_xgraph`, i.e.::
941 info: ProvenanceInitQuantumInfo = qg.bipartite_xgraph.nodes[id]
942 """
943 return self._init_quanta
945 @property
946 def quanta_by_task(self) -> Mapping[TaskLabel, Mapping[DataCoordinate, uuid.UUID]]:
947 """A nested mapping of all quanta, keyed first by task name and then by
948 data ID.
950 Notes
951 -----
952 This is populated one quantum at a time as they are read. All tasks in
953 the pipeline graph are included, even if none of their quanta were
954 loaded (i.e. nested mappings may be empty).
956 The returned object may be an internal dictionary; as the type
957 annotation indicates, it should not be modified in place.
958 """
959 return self._quanta_by_task_label
961 @property
962 def datasets_by_type(self) -> Mapping[DatasetTypeName, Mapping[DataCoordinate, uuid.UUID]]:
963 """A nested mapping of all datasets, keyed first by dataset type name
964 and then by data ID.
966 Notes
967 -----
968 This is populated one dataset at a time as they are read. All dataset
969 types in the pipeline graph are included, even if none of their
970 datasets were loaded (i.e. nested mappings may be empty).
972 The returned object may be an internal dictionary; as the type
973 annotation indicates, it should not be modified in place.
974 """
975 return self._datasets_by_type
977 @property
978 def quantum_only_xgraph(self) -> networkx.DiGraph:
979 """A directed acyclic graph with quanta as nodes (and datasets elided).
981 Notes
982 -----
983 Node keys are quantum UUIDs, and are populated one quantum at a time as
984 they are loaded. Loading quanta (via
985 `ProvenanceQuantumGraphReader.read_quanta`) will add the loaded nodes
986 with full attributes and add edges to adjacent nodes with no
987 attributes. Loading datasets (via
988 `ProvenanceQuantumGraphReader.read_datasets`) will also add edges and
989 nodes with no attributes.
991 Node attributes are described by the `ProvenanceQuantumInfo` types.
993 This graph does not include special "init" quanta.
995 The returned object is a read-only view of an internal one.
996 """
997 return self._quantum_only_xgraph.copy(as_view=True)
999 @property
1000 def bipartite_xgraph(self) -> networkx.DiGraph:
1001 """A directed acyclic graph with quantum and dataset nodes.
1003 Notes
1004 -----
1005 Node keys are quantum or dataset UUIDs, and are populated one quantum
1006 or dataset at a time as they are loaded. Loading quanta (via
1007 `ProvenanceQuantumGraphReader.read_quanta`) or datasets (via
1008 `ProvenanceQuantumGraphReader.read_datasets`) will load those nodes
1009 with full attributes and edges to adjacent nodes with no attributes.
1010 Loading quanta necessary to populate edge attributes.
1012 Node attributes are described by the
1013 `ProvenanceQuantumInfo`, `ProvenanceInitQuantumInfo`, and
1014 `ProvenanceDatasetInfo` types.
1016 This graph includes init-input and init-output datasets, but it does
1017 *not* reflect the dependency between each task's special "init" quantum
1018 and its runtime quanta (as this would require edges between quanta, and
1019 that would break the "bipartite" property).
1021 The returned object is a read-only view of an internal one.
1022 """
1023 return self._bipartite_xgraph.copy(as_view=True)
1025 def make_quantum_table(self) -> astropy.table.Table:
1026 """Construct an `astropy.table.Table` with a tabular summary of the
1027 quanta.
1029 Returns
1030 -------
1031 table : `astropy.table.Table`
1032 A table view of the quantum information. This only includes
1033 counts of status categories and caveats, not any per-data-ID
1034 detail.
1036 Notes
1037 -----
1038 Success caveats in the table are represented by their
1039 `~QuantumSuccessCaveats.concise` form, so when pretty-printing this
1040 table for users, the `~QuantumSuccessCaveats.legend` should generally
1041 be printed as well.
1042 """
1043 rows = []
1044 for task_label, quanta_for_task in self.quanta_by_task.items():
1045 if not self.header.n_task_quanta[task_label]:
1046 continue
1047 status_counts = Counter[QuantumAttemptStatus](
1048 self._quantum_only_xgraph.nodes[q]["status"] for q in quanta_for_task.values()
1049 )
1050 caveat_counts = Counter[QuantumSuccessCaveats | None](
1051 self._quantum_only_xgraph.nodes[q]["caveats"] for q in quanta_for_task.values()
1052 )
1053 caveat_counts.pop(QuantumSuccessCaveats.NO_CAVEATS, None)
1054 caveat_counts.pop(None, None)
1055 if len(caveat_counts) > 1:
1056 caveats = "(multiple)"
1057 elif len(caveat_counts) == 1:
1058 ((code, count),) = caveat_counts.items()
1059 # MyPy can't tell that the pop(None, None) above makes None
1060 # impossible here.
1061 caveats = f"{code.concise()}({count})" # type: ignore[union-attr]
1062 else:
1063 caveats = ""
1064 rows.append(
1065 {
1066 "Task": task_label,
1067 "Unknown": status_counts.get(QuantumAttemptStatus.UNKNOWN, 0),
1068 "Successful": status_counts.get(QuantumAttemptStatus.SUCCESSFUL, 0),
1069 "Caveats": caveats,
1070 "Blocked": status_counts.get(QuantumAttemptStatus.BLOCKED, 0),
1071 "Failed": status_counts.get(QuantumAttemptStatus.FAILED, 0),
1072 "TOTAL": len(quanta_for_task),
1073 "EXPECTED": self.header.n_task_quanta[task_label],
1074 }
1075 )
1076 return astropy.table.Table(rows)
1078 def make_exception_table(self) -> astropy.table.Table:
1079 """Construct an `astropy.table.Table` with counts for each exception
1080 type raised by each task.
1082 At present this only includes information from partial-outputs-error
1083 successes, since exception information for failures is not tracked.
1084 This may change in the future.
1086 Returns
1087 -------
1088 table : `astropy.table.Table`
1089 A table with columns for task label, exception type, and counts.
1090 """
1091 rows = []
1092 for task_label, quanta_for_task in self.quanta_by_task.items():
1093 counts_by_type = Counter(
1094 exc_info.type_name
1095 for q in quanta_for_task.values()
1096 if (exc_info := self._quantum_only_xgraph.nodes[q]["exception"]) is not None
1097 )
1098 for type_name, count in counts_by_type.items():
1099 rows.append({"Task": task_label, "Exception": type_name, "Count": count})
1100 return astropy.table.Table(rows)
1102 def make_task_resource_usage_table(
1103 self, task_label: TaskLabel, include_data_ids: bool = False
1104 ) -> astropy.table.Table:
1105 """Make a table of resource usage for a single task.
1107 Parameters
1108 ----------
1109 task_label : `str`
1110 Label of the task to extract resource usage for.
1111 include_data_ids : `bool`, optional
1112 Whether to also include data ID columns.
1114 Returns
1115 -------
1116 table : `astropy.table.Table`
1117 A table with columns for quantum ID and all fields in
1118 `QuantumResourceUsage`.
1119 """
1120 quanta_for_task = self.quanta_by_task[task_label]
1121 dtype_terms: list[tuple[str, np.dtype]] = [("quantum_id", np.dtype((np.void, 16)))]
1122 if include_data_ids:
1123 dimensions = self.pipeline_graph.tasks[task_label].dimensions
1124 for dimension_name in dimensions.data_coordinate_keys:
1125 dtype = np.dtype(self.pipeline_graph.universe.dimensions[dimension_name].primary_key.pytype)
1126 dtype_terms.append((dimension_name, dtype))
1127 fields = QuantumResourceUsage.get_numpy_fields()
1128 dtype_terms.extend(fields.items())
1129 row_dtype = np.dtype(dtype_terms)
1130 rows: list[object] = []
1131 for data_id, quantum_id in quanta_for_task.items():
1132 info: ProvenanceQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1133 if (resource_usage := info["resource_usage"]) is not None:
1134 row: tuple[object, ...] = (quantum_id.bytes,)
1135 if include_data_ids:
1136 row += data_id.full_values
1137 row += resource_usage.get_numpy_row()
1138 rows.append(row)
1139 array = np.array(rows, dtype=row_dtype)
1140 return astropy.table.Table(array, units=QuantumResourceUsage.get_units())
1143@dataclasses.dataclass
1144class ProvenanceQuantumGraphReader(BaseQuantumGraphReader):
1145 """A helper class for reading provenance quantum graphs.
1147 Notes
1148 -----
1149 The `open` context manager should be used to construct new instances.
1150 Instances cannot be used after the context manager exits, except to access
1151 the `graph` attribute`.
1153 The various ``read_*`` methods in this class update the `graph` attribute
1154 in place.
1155 """
1157 graph: ProvenanceQuantumGraph = dataclasses.field(init=False)
1158 """Loaded provenance graph, populated in place as components are read."""
1160 @classmethod
1161 @contextmanager
1162 def open(
1163 cls,
1164 uri: ResourcePathExpression,
1165 *,
1166 page_size: int | None = None,
1167 import_mode: TaskImportMode = TaskImportMode.DO_NOT_IMPORT,
1168 ) -> Iterator[ProvenanceQuantumGraphReader]:
1169 """Construct a reader from a URI.
1171 Parameters
1172 ----------
1173 uri : convertible to `lsst.resources.ResourcePath`
1174 URI to open. Should have a ``.qg`` extension.
1175 page_size : `int`, optional
1176 Approximate number of bytes to read at once from address files and
1177 multi-block files. Note that this does not set a page size for
1178 *all* reads, but it does affect the smallest, most numerous reads.
1179 Can also be set via the ``LSST_QG_PAGE_SIZE`` environment variable.
1180 import_mode : `.pipeline_graph.TaskImportMode`, optional
1181 How to handle importing the task classes referenced in the pipeline
1182 graph.
1184 Returns
1185 -------
1186 reader : `contextlib.AbstractContextManager` [ \
1187 `ProvenanceQuantumGraphReader` ]
1188 A context manager that returns the reader when entered.
1189 """
1190 with cls._open(
1191 uri,
1192 graph_type="provenance",
1193 address_filename="nodes",
1194 page_size=page_size,
1195 import_mode=import_mode,
1196 n_addresses=4,
1197 ) as self:
1198 yield self
1200 def __post_init__(self) -> None:
1201 self.graph = ProvenanceQuantumGraph(self.header, self.pipeline_graph)
1203 def read_init_quanta(self) -> None:
1204 """Read the thin graph, with all edge information and categorization of
1205 quanta by task label.
1206 """
1207 init_quanta = self._read_single_block("init_quanta", ProvenanceInitQuantaModel)
1208 for init_quantum in init_quanta.root:
1209 self.graph._init_quanta[init_quantum.task_label] = init_quantum.quantum_id
1210 init_quanta._add_to_graph(self.graph)
1212 def read_full_graph(self) -> None:
1213 """Read all bipartite edges and all quantum and dataset node
1214 attributes, fully populating the `graph` attribute.
1216 Notes
1217 -----
1218 This does not read logs, metadata, or packages ; those must always be
1219 fetched explicitly.
1220 """
1221 self.read_init_quanta()
1222 self.read_datasets()
1223 self.read_quanta()
1225 def read_datasets(self, datasets: Iterable[uuid.UUID] | None = None) -> None:
1226 """Read information about the given datasets.
1228 Parameters
1229 ----------
1230 datasets : `~collections.abc.Iterable` [`uuid.UUID`], optional
1231 Iterable of dataset IDs to load. If not provided, all datasets
1232 will be loaded. The UUIDs and indices of quanta will be ignored.
1233 """
1234 self._read_nodes(datasets, DATASET_ADDRESS_INDEX, DATASET_MB_NAME, ProvenanceDatasetModel)
1236 def read_quanta(self, quanta: Iterable[uuid.UUID] | None = None) -> None:
1237 """Read information about the given quanta.
1239 Parameters
1240 ----------
1241 quanta : `~collections.abc.Iterable` [`uuid.UUID`], optional
1242 Iterable of quantum IDs to load. If not provided, all quanta will
1243 be loaded. The UUIDs and indices of datasets and special init
1244 quanta will be ignored.
1245 """
1246 self._read_nodes(quanta, QUANTUM_ADDRESS_INDEX, QUANTUM_MB_NAME, ProvenanceQuantumModel)
1248 def _read_nodes(
1249 self,
1250 nodes: Iterable[uuid.UUID] | None,
1251 address_index: int,
1252 mb_name: str,
1253 model_type: type[ProvenanceDatasetModel] | type[ProvenanceQuantumModel],
1254 ) -> None:
1255 node: ProvenanceDatasetModel | ProvenanceQuantumModel | None
1256 if nodes is None:
1257 self.address_reader.read_all()
1258 nodes = self.address_reader.rows.keys()
1259 for node in MultiblockReader.read_all_models_in_zip(
1260 self.zf,
1261 mb_name,
1262 model_type,
1263 self.decompressor,
1264 int_size=self.header.int_size,
1265 page_size=self.page_size,
1266 ):
1267 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(node.node_id, {}):
1268 # Use the old node to reduce memory usage (since it might
1269 # also have other outstanding reference holders).
1270 continue
1271 node._add_to_graph(self.graph)
1272 return
1273 with MultiblockReader.open_in_zip(self.zf, mb_name, int_size=self.header.int_size) as mb_reader:
1274 for node_id_or_index in nodes:
1275 address_row = self.address_reader.find(node_id_or_index)
1276 if "pipeline_node" in self.graph._bipartite_xgraph.nodes.get(address_row.key, {}):
1277 # Use the old node to reduce memory usage (since it might
1278 # also have other outstanding reference holders).
1279 continue
1280 node = mb_reader.read_model(
1281 address_row.addresses[address_index], model_type, self.decompressor
1282 )
1283 if node is not None:
1284 node._add_to_graph(self.graph)
1286 def fetch_logs(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[ButlerLogRecords | None]]:
1287 """Fetch log datasets.
1289 Parameters
1290 ----------
1291 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1292 UUIDs of the log datasets themselves or of the quanta they
1293 correspond to.
1295 Returns
1296 -------
1297 logs : `dict` [ `uuid.UUID`, `list` [\
1298 `lsst.daf.butler.ButlerLogRecords` or `None`] ]
1299 Logs for the given IDs. Each value is a list of
1300 `lsst.daf.butler.ButlerLogRecords` instances representing different
1301 execution attempts, ordered chronologically from first to last.
1302 Attempts where logs were missing will have `None` in this list.
1303 """
1304 result: dict[uuid.UUID, list[ButlerLogRecords | None]] = {}
1305 with MultiblockReader.open_in_zip(self.zf, LOG_MB_NAME, int_size=self.header.int_size) as mb_reader:
1306 for node_id_or_index in nodes:
1307 address_row = self.address_reader.find(node_id_or_index)
1308 logs_by_attempt = mb_reader.read_model(
1309 address_row.addresses[LOG_ADDRESS_INDEX], ProvenanceLogRecordsModel, self.decompressor
1310 )
1311 if logs_by_attempt is not None:
1312 result[node_id_or_index] = [
1313 ButlerLogRecords.from_records(attempt_logs) if attempt_logs is not None else None
1314 for attempt_logs in logs_by_attempt.attempts
1315 ]
1316 return result
1318 def fetch_metadata(self, nodes: Iterable[uuid.UUID]) -> dict[uuid.UUID, list[TaskMetadata | None]]:
1319 """Fetch metadata datasets.
1321 Parameters
1322 ----------
1323 nodes : `~collections.abc.Iterable` [ `uuid.UUID` ]
1324 UUIDs of the metadata datasets themselves or of the quanta they
1325 correspond to.
1327 Returns
1328 -------
1329 metadata : `dict` [ `uuid.UUID`, `list` [`.TaskMetadata`] ]
1330 Metadata for the given IDs. Each value is a list of
1331 `.TaskMetadata` instances representing different execution
1332 attempts, ordered chronologically from first to last. Attempts
1333 where metadata was missing (not written even in the fallback extra
1334 provenance in the logs) will have `None` in this list.
1335 """
1336 result: dict[uuid.UUID, list[TaskMetadata | None]] = {}
1337 with MultiblockReader.open_in_zip(
1338 self.zf, METADATA_MB_NAME, int_size=self.header.int_size
1339 ) as mb_reader:
1340 for node_id_or_index in nodes:
1341 address_row = self.address_reader.find(node_id_or_index)
1342 metadata_by_attempt = mb_reader.read_model(
1343 address_row.addresses[METADATA_ADDRESS_INDEX],
1344 ProvenanceTaskMetadataModel,
1345 self.decompressor,
1346 )
1347 if metadata_by_attempt is not None:
1348 result[node_id_or_index] = metadata_by_attempt.attempts
1349 return result
1351 def fetch_packages(self) -> Packages:
1352 """Fetch package version information."""
1353 data = self._read_single_block_raw("packages")
1354 return Packages.fromBytes(data, format="json")