Coverage for python / lsst / pipe / base / quantum_graph / _predicted.py: 21%
615 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-20 01:37 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = (
31 "PredictedDatasetInfo",
32 "PredictedDatasetModel",
33 "PredictedInitQuantaModel",
34 "PredictedQuantumDatasetsModel",
35 "PredictedQuantumGraph",
36 "PredictedQuantumGraphComponents",
37 "PredictedQuantumGraphReader",
38 "PredictedQuantumInfo",
39 "PredictedThinGraphModel",
40 "PredictedThinQuantumModel",
41)
43import dataclasses
44import itertools
45import logging
46import sys
47import uuid
48import warnings
49from collections import defaultdict
50from collections.abc import Iterable, Iterator, Mapping, Sequence
51from contextlib import AbstractContextManager, contextmanager
52from typing import TYPE_CHECKING, Any, TypeVar, cast
54import networkx
55import networkx.algorithms.bipartite
56import pydantic
57import zstandard
59from lsst.daf.butler import (
60 Config,
61 DataCoordinate,
62 DataIdValue,
63 DatasetRef,
64 DatasetType,
65 DimensionDataAttacher,
66 DimensionDataExtractor,
67 DimensionGroup,
68 DimensionRecordSetDeserializer,
69 DimensionUniverse,
70 LimitedButler,
71 Quantum,
72 QuantumBackedButler,
73 SerializableDimensionData,
74)
75from lsst.daf.butler._rubin import generate_uuidv7
76from lsst.daf.butler.datastore.record_data import DatastoreRecordData, SerializedDatastoreRecordData
77from lsst.daf.butler.registry import ConflictingDefinitionError
78from lsst.resources import ResourcePath, ResourcePathExpression
79from lsst.utils.packages import Packages
81from .. import automatic_connection_constants as acc
82from ..pipeline import TaskDef
83from ..pipeline_graph import (
84 PipelineGraph,
85 TaskImportMode,
86 TaskInitNode,
87 TaskNode,
88 compare_packages,
89 log_config_mismatch,
90)
91from ._common import (
92 FORMAT_VERSION,
93 BaseQuantumGraph,
94 BaseQuantumGraphReader,
95 BaseQuantumGraphWriter,
96 ConnectionName,
97 DataCoordinateValues,
98 DatasetInfo,
99 DatasetTypeName,
100 DatastoreName,
101 HeaderModel,
102 IncompleteQuantumGraphError,
103 QuantumIndex,
104 QuantumInfo,
105 TaskLabel,
106)
107from ._multiblock import DEFAULT_PAGE_SIZE, AddressRow, MultiblockReader, MultiblockWriter
109if TYPE_CHECKING:
110 from ..config import PipelineTaskConfig
111 from ..graph import QgraphSummary, QuantumGraph
113_LOG = logging.getLogger(__name__)
116_T = TypeVar("_T", bound=pydantic.BaseModel)
119class _PredictedThinQuantumModelV0(pydantic.BaseModel):
120 """Data model for a quantum data ID and internal integer ID in a predicted
121 quantum graph.
122 """
124 quantum_index: QuantumIndex
125 """Internal integer ID for this quantum."""
127 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
128 """Full (required and implied) data coordinate values for this quantum."""
131class PredictedThinQuantumModel(pydantic.BaseModel):
132 """Data model for a quantum data ID and UUID in a predicted
133 quantum graph.
134 """
136 quantum_id: uuid.UUID
137 """Universally unique ID for this quantum."""
139 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
140 """Full (required and implied) data coordinate values for this quantum."""
142 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
143 # when we inherit those docstrings in our public classes.
144 if "sphinx" in sys.modules and not TYPE_CHECKING:
146 def copy(self, *args: Any, **kwargs: Any) -> Any:
147 """See `pydantic.BaseModel.copy`."""
148 return super().copy(*args, **kwargs)
150 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
151 """See `pydantic.BaseModel.model_dump`."""
152 return super().model_dump(*args, **kwargs)
154 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
155 """See `pydantic.BaseModel.model_dump_json`."""
156 return super().model_dump(*args, **kwargs)
158 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
159 """See `pydantic.BaseModel.model_copy`."""
160 return super().model_copy(*args, **kwargs)
162 @classmethod
163 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
164 """See `pydantic.BaseModel.model_construct`."""
165 return super().model_construct(*args, **kwargs)
167 @classmethod
168 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
169 """See `pydantic.BaseModel.model_json_schema`."""
170 return super().model_json_schema(*args, **kwargs)
172 @classmethod
173 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
174 """See `pydantic.BaseModel.model_validate`."""
175 return super().model_validate(*args, **kwargs)
177 @classmethod
178 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
179 """See `pydantic.BaseModel.model_validate_json`."""
180 return super().model_validate_json(*args, **kwargs)
182 @classmethod
183 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
184 """See `pydantic.BaseModel.model_validate_strings`."""
185 return super().model_validate_strings(*args, **kwargs)
188class _PredictedThinGraphModelV0(pydantic.BaseModel):
189 """Data model for the predicted quantum graph component that maps each
190 task label to the data IDs and internal integer IDs of its quanta.
191 """
193 quanta: dict[TaskLabel, list[_PredictedThinQuantumModelV0]] = pydantic.Field(default_factory=dict)
194 """Minimal descriptions of all quanta, grouped by task label."""
196 edges: list[tuple[QuantumIndex, QuantumIndex]] = pydantic.Field(default_factory=list)
197 """Pairs of (predecessor, successor) internal integer quantum IDs."""
199 def _upgraded(self, address_rows: Mapping[uuid.UUID, AddressRow]) -> PredictedThinGraphModel:
200 """Convert to the v1+ model."""
201 uuid_by_index = {v.index: k for k, v in address_rows.items()}
202 return PredictedThinGraphModel.model_construct(
203 quanta={
204 task_label: [
205 PredictedThinQuantumModel.model_construct(
206 quantum_id=uuid_by_index[q.quantum_index], data_coordinate=q.data_coordinate
207 )
208 for q in quanta
209 ]
210 for task_label, quanta in self.quanta.items()
211 },
212 edges=[(uuid_by_index[index1], uuid_by_index[index2]) for index1, index2 in self.edges],
213 )
216class PredictedThinGraphModel(pydantic.BaseModel):
217 """Data model for the predicted quantum graph component that maps each
218 task label to the data IDs and UUIDs of its quanta.
219 """
221 quanta: dict[TaskLabel, list[PredictedThinQuantumModel]] = pydantic.Field(default_factory=dict)
222 """Minimal descriptions of all quanta, grouped by task label."""
224 edges: list[tuple[uuid.UUID, uuid.UUID]] = pydantic.Field(default_factory=list)
225 """Pairs of (predecessor, successor) quantum IDs."""
227 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
228 # when we inherit those docstrings in our public classes.
229 if "sphinx" in sys.modules and not TYPE_CHECKING:
231 def copy(self, *args: Any, **kwargs: Any) -> Any:
232 """See `pydantic.BaseModel.copy`."""
233 return super().copy(*args, **kwargs)
235 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
236 """See `pydantic.BaseModel.model_dump`."""
237 return super().model_dump(*args, **kwargs)
239 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
240 """See `pydantic.BaseModel.model_dump_json`."""
241 return super().model_dump(*args, **kwargs)
243 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
244 """See `pydantic.BaseModel.model_copy`."""
245 return super().model_copy(*args, **kwargs)
247 @classmethod
248 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
249 """See `pydantic.BaseModel.model_construct`."""
250 return super().model_construct(*args, **kwargs)
252 @classmethod
253 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
254 """See `pydantic.BaseModel.model_json_schema`."""
255 return super().model_json_schema(*args, **kwargs)
257 @classmethod
258 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
259 """See `pydantic.BaseModel.model_validate`."""
260 return super().model_validate(*args, **kwargs)
262 @classmethod
263 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
264 """See `pydantic.BaseModel.model_validate_json`."""
265 return super().model_validate_json(*args, **kwargs)
267 @classmethod
268 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
269 """See `pydantic.BaseModel.model_validate_strings`."""
270 return super().model_validate_strings(*args, **kwargs)
273class PredictedDatasetModel(pydantic.BaseModel):
274 """Data model for the datasets in a predicted quantum graph file."""
276 dataset_id: uuid.UUID
277 """Universally unique ID for the dataset."""
279 dataset_type_name: DatasetTypeName
280 """Name of the type of this dataset.
282 This is always a parent dataset type name, not a component.
284 Note that full dataset type definitions are stored in the pipeline graph.
285 """
287 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
288 """The full values (required and implied) of this dataset's data ID."""
290 run: str
291 """This dataset's RUN collection name."""
293 @classmethod
294 def from_dataset_ref(cls, ref: DatasetRef) -> PredictedDatasetModel:
295 """Construct from a butler `~lsst.daf.butler.DatasetRef`.
297 Parameters
298 ----------
299 ref : `lsst.daf.butler.DatasetRef`
300 Dataset reference.
302 Returns
303 -------
304 model : `PredictedDatasetModel`
305 Model for the dataset.
306 """
307 dataset_type_name, _ = DatasetType.splitDatasetTypeName(ref.datasetType.name)
308 return cls.model_construct(
309 dataset_id=ref.id,
310 dataset_type_name=dataset_type_name,
311 data_coordinate=list(ref.dataId.full_values),
312 run=ref.run,
313 )
315 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
316 # when we inherit those docstrings in our public classes.
317 if "sphinx" in sys.modules and not TYPE_CHECKING:
319 def copy(self, *args: Any, **kwargs: Any) -> Any:
320 """See `pydantic.BaseModel.copy`."""
321 return super().copy(*args, **kwargs)
323 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
324 """See `pydantic.BaseModel.model_dump`."""
325 return super().model_dump(*args, **kwargs)
327 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
328 """See `pydantic.BaseModel.model_dump_json`."""
329 return super().model_dump(*args, **kwargs)
331 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
332 """See `pydantic.BaseModel.model_copy`."""
333 return super().model_copy(*args, **kwargs)
335 @classmethod
336 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
337 """See `pydantic.BaseModel.model_construct`."""
338 return super().model_construct(*args, **kwargs)
340 @classmethod
341 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
342 """See `pydantic.BaseModel.model_json_schema`."""
343 return super().model_json_schema(*args, **kwargs)
345 @classmethod
346 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
347 """See `pydantic.BaseModel.model_validate`."""
348 return super().model_validate(*args, **kwargs)
350 @classmethod
351 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
352 """See `pydantic.BaseModel.model_validate_json`."""
353 return super().model_validate_json(*args, **kwargs)
355 @classmethod
356 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
357 """See `pydantic.BaseModel.model_validate_strings`."""
358 return super().model_validate_strings(*args, **kwargs)
361class PredictedQuantumDatasetsModel(pydantic.BaseModel):
362 """Data model for a description of a single predicted quantum that includes
363 its inputs and outputs.
364 """
366 quantum_id: uuid.UUID
367 """Universally unique ID for the quantum."""
369 task_label: TaskLabel
370 """Label of the task.
372 Note that task label definitions are stored in the pipeline graph.
373 """
375 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)
376 """The full values (required and implied) of this quantum's data ID."""
378 inputs: dict[ConnectionName, list[PredictedDatasetModel]] = pydantic.Field(default_factory=dict)
379 """The input datasets to this quantum, grouped by connection name."""
381 outputs: dict[ConnectionName, list[PredictedDatasetModel]] = pydantic.Field(default_factory=dict)
382 """The datasets output by this quantum, grouped by connection name."""
384 datastore_records: dict[DatastoreName, SerializedDatastoreRecordData] = pydantic.Field(
385 default_factory=dict
386 )
387 """Datastore records for inputs to this quantum that are already present in
388 the data repository.
389 """
391 def iter_input_dataset_ids(self) -> Iterator[uuid.UUID]:
392 """Return an iterator over the UUIDs of all datasets consumed by this
393 quantum.
395 Returns
396 -------
397 iter : `~collections.abc.Iterator` [ `uuid.UUID` ]
398 Iterator over dataset IDs.
399 """
400 for datasets in self.inputs.values():
401 for dataset in datasets:
402 yield dataset.dataset_id
404 def iter_output_dataset_ids(self) -> Iterator[uuid.UUID]:
405 """Return an iterator over the UUIDs of all datasets produced by this
406 quantum.
408 Returns
409 -------
410 iter : `~collections.abc.Iterator` [ `uuid.UUID` ]
411 Iterator over dataset IDs.
412 """
413 for datasets in self.outputs.values():
414 for dataset in datasets:
415 yield dataset.dataset_id
417 def iter_dataset_ids(self) -> Iterator[uuid.UUID]:
418 """Return an iterator over the UUIDs of all datasets referenced by this
419 quantum.
421 Returns
422 -------
423 iter : `~collections.abc.Iterator` [ `uuid.UUID` ]
424 Iterator over dataset IDs.
425 """
426 yield from self.iter_input_dataset_ids()
427 yield from self.iter_output_dataset_ids()
429 def deserialize_datastore_records(self) -> dict[DatastoreName, DatastoreRecordData]:
430 """Deserialize the mapping of datastore records."""
431 return {
432 datastore_name: DatastoreRecordData.from_simple(serialized_records)
433 for datastore_name, serialized_records in self.datastore_records.items()
434 }
436 @classmethod
437 def from_execution_quantum(
438 cls, task_node: TaskNode, quantum: Quantum, quantum_id: uuid.UUID
439 ) -> PredictedQuantumDatasetsModel:
440 """Construct from an `lsst.daf.butler.Quantum` instance.
442 Parameters
443 ----------
444 task_node : `.pipeline_graph.TaskNode`
445 Task node from the pipeline graph.
446 quantum : `lsst.daf.butler.quantum`
447 Quantum object.
448 quantum_id : `uuid.UUID`
449 ID for this quantum.
451 Returns
452 -------
453 model : `PredictedFullQuantumModel`
454 Model for this quantum.
455 """
456 result: PredictedQuantumDatasetsModel = cls.model_construct(
457 quantum_id=quantum_id,
458 task_label=task_node.label,
459 data_coordinate=list(cast(DataCoordinate, quantum.dataId).full_values),
460 )
461 for read_edge in task_node.iter_all_inputs():
462 refs = sorted(quantum.inputs[read_edge.dataset_type_name], key=lambda ref: ref.dataId)
463 result.inputs[read_edge.connection_name] = [
464 PredictedDatasetModel.from_dataset_ref(ref) for ref in refs
465 ]
466 for write_edge in task_node.iter_all_outputs():
467 refs = sorted(quantum.outputs[write_edge.dataset_type_name], key=lambda ref: ref.dataId)
468 result.outputs[write_edge.connection_name] = [
469 PredictedDatasetModel.from_dataset_ref(ref) for ref in refs
470 ]
471 result.datastore_records = {
472 store_name: records.to_simple() for store_name, records in quantum.datastore_records.items()
473 }
474 return result
476 @classmethod
477 def from_old_quantum_graph_init(
478 cls, task_init_node: TaskInitNode, old_quantum_graph: QuantumGraph
479 ) -> PredictedQuantumDatasetsModel:
480 """Construct from the init-input and init-output dataset types of a
481 task in an old `QuantumGraph` instance.
483 Parameters
484 ----------
485 task_init_node : `.pipeline_graph.TaskNode`
486 Task init node from the pipeline graph.
487 old_quantum_graph : `QuantumGraph`
488 Quantum graph.
490 Returns
491 -------
492 model : `PredictedFullQuantumModel`
493 Model for this "init" quantum.
494 """
495 task_def = old_quantum_graph.findTaskDefByLabel(task_init_node.label)
496 assert task_def is not None
497 init_input_refs = {
498 ref.datasetType.name: ref for ref in (old_quantum_graph.initInputRefs(task_def) or [])
499 }
500 init_output_refs = {
501 ref.datasetType.name: ref for ref in (old_quantum_graph.initOutputRefs(task_def) or [])
502 }
503 init_input_ids = {ref.id for ref in init_input_refs.values()}
504 result: PredictedQuantumDatasetsModel = cls.model_construct(
505 quantum_id=generate_uuidv7(), task_label=task_init_node.label
506 )
507 for read_edge in task_init_node.iter_all_inputs():
508 ref = init_input_refs[read_edge.dataset_type_name]
509 result.inputs[read_edge.connection_name] = [PredictedDatasetModel.from_dataset_ref(ref)]
510 for write_edge in task_init_node.iter_all_outputs():
511 ref = init_output_refs[write_edge.dataset_type_name]
512 result.outputs[write_edge.connection_name] = [PredictedDatasetModel.from_dataset_ref(ref)]
513 datastore_records: dict[str, DatastoreRecordData] = {}
514 for quantum in old_quantum_graph.get_task_quanta(task_init_node.label).values():
515 for store_name, records in quantum.datastore_records.items():
516 subset = records.subset(init_input_ids)
517 if subset is not None:
518 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset)
519 break # All quanta have same init-inputs, so we only need one.
520 result.datastore_records = {
521 store_name: records.to_simple() for store_name, records in datastore_records.items()
522 }
523 return result
525 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
526 # when we inherit those docstrings in our public classes.
527 if "sphinx" in sys.modules and not TYPE_CHECKING:
529 def copy(self, *args: Any, **kwargs: Any) -> Any:
530 """See `pydantic.BaseModel.copy`."""
531 return super().copy(*args, **kwargs)
533 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
534 """See `pydantic.BaseModel.model_dump`."""
535 return super().model_dump(*args, **kwargs)
537 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
538 """See `pydantic.BaseModel.model_dump_json`."""
539 return super().model_dump(*args, **kwargs)
541 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
542 """See `pydantic.BaseModel.model_copy`."""
543 return super().model_copy(*args, **kwargs)
545 @classmethod
546 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
547 """See `pydantic.BaseModel.model_construct`."""
548 return super().model_construct(*args, **kwargs)
550 @classmethod
551 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
552 """See `pydantic.BaseModel.model_json_schema`."""
553 return super().model_json_schema(*args, **kwargs)
555 @classmethod
556 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
557 """See `pydantic.BaseModel.model_validate`."""
558 return super().model_validate(*args, **kwargs)
560 @classmethod
561 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
562 """See `pydantic.BaseModel.model_validate_json`."""
563 return super().model_validate_json(*args, **kwargs)
565 @classmethod
566 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
567 """See `pydantic.BaseModel.model_validate_strings`."""
568 return super().model_validate_strings(*args, **kwargs)
571class PredictedInitQuantaModel(pydantic.RootModel):
572 """Data model for the init-inputs and init-outputs of a predicted quantum
573 graph.
574 """
576 root: list[PredictedQuantumDatasetsModel] = pydantic.Field(default_factory=list)
577 """List of special "init" quanta: one for each task, and another for global
578 init-outputs.
579 """
581 def update_from_old_quantum_graph(self, old_quantum_graph: QuantumGraph) -> None:
582 """Update this model in-place by extracting from an old `QuantumGraph`
583 instance.
585 Parameters
586 ----------
587 old_quantum_graph : `QuantumGraph`
588 Quantum graph.
589 """
590 global_init_quantum = PredictedQuantumDatasetsModel.model_construct(
591 quantum_id=generate_uuidv7(), task_label=""
592 )
593 for ref in old_quantum_graph.globalInitOutputRefs():
594 global_init_quantum.outputs[ref.datasetType.name] = [PredictedDatasetModel.from_dataset_ref(ref)]
595 self.root.append(global_init_quantum)
596 for task_node in old_quantum_graph.pipeline_graph.tasks.values():
597 self.root.append(
598 PredictedQuantumDatasetsModel.from_old_quantum_graph_init(task_node.init, old_quantum_graph)
599 )
601 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
602 # when we inherit those docstrings in our public classes.
603 if "sphinx" in sys.modules and not TYPE_CHECKING:
605 def copy(self, *args: Any, **kwargs: Any) -> Any:
606 """See `pydantic.BaseModel.copy`."""
607 return super().copy(*args, **kwargs)
609 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
610 """See `pydantic.BaseModel.model_dump`."""
611 return super().model_dump(*args, **kwargs)
613 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
614 """See `pydantic.BaseModel.model_dump_json`."""
615 return super().model_dump(*args, **kwargs)
617 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
618 """See `pydantic.BaseModel.model_copy`."""
619 return super().model_copy(*args, **kwargs)
621 @classmethod
622 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
623 """See `pydantic.BaseModel.model_construct`."""
624 return super().model_construct(*args, **kwargs)
626 @classmethod
627 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
628 """See `pydantic.BaseModel.model_json_schema`."""
629 return super().model_json_schema(*args, **kwargs)
631 @classmethod
632 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
633 """See `pydantic.BaseModel.model_validate`."""
634 return super().model_validate(*args, **kwargs)
636 @classmethod
637 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
638 """See `pydantic.BaseModel.model_validate_json`."""
639 return super().model_validate_json(*args, **kwargs)
641 @classmethod
642 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
643 """See `pydantic.BaseModel.model_validate_strings`."""
644 return super().model_validate_strings(*args, **kwargs)
647class PredictedQuantumInfo(QuantumInfo):
648 """A typed dictionary that annotates the attributes of the NetworkX graph
649 node data for a predicted quantum.
651 Since NetworkX types are not generic over their node mapping type, this has
652 to be used explicitly, e.g.::
654 node_data: PredictedQuantumInfo = xgraph.nodes[quantum_id]
656 where ``xgraph`` can be either `PredictedQuantumGraph.quantum_only_xgraph`
657 or `PredictedQuantumGraph.bipartite_xgraph`.
658 """
660 quantum: Quantum
661 """Quantum object that can be passed directly to an executor.
663 This attribute is only present if
664 `PredictedQuantumGraph.build_execution_quanta` has been run on this node's
665 quantum ID already.
666 """
669class PredictedDatasetInfo(DatasetInfo):
670 """A typed dictionary that annotates the attributes of the NetworkX graph
671 node data for a dataset.
673 Since NetworkX types are not generic over their node mapping type, this has
674 to be used explicitly, e.g.::
676 node_data: PredictedDatasetInfo = xgraph.nodes[dataset_ids]
678 where ``xgraph`` is from the `PredictedQuantumGraph.bipartite_xgraph`
679 property.
680 """
683class PredictedQuantumGraph(BaseQuantumGraph):
684 """A directed acyclic graph that predicts a processing run and supports it
685 during execution.
687 Parameters
688 ----------
689 components : `PredictedQuantumGraphComponents`
690 A struct of components used to construct the graph.
692 Notes
693 -----
694 Iteration over a `PredictedQuantumGraph` yields loaded quantum IDs in
695 deterministic topological order (but the tiebreaker is unspecified). The
696 `len` of a `PredictedQuantumGraph` is the number of loaded non-init quanta,
697 i.e. the same as the number of quanta iterated over.
698 """
700 def __init__(self, components: PredictedQuantumGraphComponents):
701 if not components.header.graph_type == "predicted":
702 raise TypeError(f"Header is for a {components.header.graph_type!r} graph, not 'predicted'.")
703 super().__init__(components.header, components.pipeline_graph)
704 self._quantum_only_xgraph = networkx.DiGraph()
705 self._bipartite_xgraph = networkx.DiGraph()
706 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {
707 task_label: {} for task_label in self.pipeline_graph.tasks.keys()
708 }
709 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {
710 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()
711 }
712 self._datasets_by_type[self.pipeline_graph.packages_dataset_type.name] = {}
713 self._dimension_data = components.dimension_data
714 self._add_init_quanta(components.init_quanta)
715 self._quantum_datasets: dict[uuid.UUID, PredictedQuantumDatasetsModel] = {}
716 self._expanded_data_ids: dict[DataCoordinate, DataCoordinate] = {}
717 self._add_thin_graph(components.thin_graph)
718 for quantum_datasets in components.quantum_datasets.values():
719 self._add_quantum_datasets(quantum_datasets)
720 if not components.thin_graph.edges:
721 # If we loaded the thin_graph, we've already populated this graph.
722 self._quantum_only_xgraph.update(
723 networkx.algorithms.bipartite.projected_graph(
724 networkx.DiGraph(self._bipartite_xgraph),
725 self._quantum_only_xgraph.nodes.keys(),
726 )
727 )
728 if _LOG.isEnabledFor(logging.DEBUG):
729 for quantum_id in self:
730 _LOG.debug(
731 "%s: %s @ %s",
732 quantum_id,
733 self._quantum_only_xgraph.nodes[quantum_id]["task_label"],
734 self._quantum_only_xgraph.nodes[quantum_id]["data_id"].required,
735 )
737 def _add_init_quanta(self, component: PredictedInitQuantaModel) -> None:
738 self._init_quanta = {q.task_label: q for q in component.root}
739 empty_data_id = DataCoordinate.make_empty(self.pipeline_graph.universe)
740 for quantum_datasets in self._init_quanta.values():
741 for init_datasets in itertools.chain(
742 quantum_datasets.inputs.values(), quantum_datasets.outputs.values()
743 ):
744 for init_dataset in init_datasets:
745 self._datasets_by_type[init_dataset.dataset_type_name][empty_data_id] = (
746 init_dataset.dataset_id
747 )
748 _LOG.debug(
749 "%s: %s @ init",
750 quantum_datasets.quantum_id,
751 quantum_datasets.task_label,
752 )
754 def _add_thin_graph(self, component: PredictedThinGraphModel) -> None:
755 self._quantum_only_xgraph.add_edges_from(component.edges)
756 for task_label, thin_quanta_for_task in component.quanta.items():
757 for thin_quantum in thin_quanta_for_task:
758 self._add_quantum(thin_quantum.quantum_id, task_label, thin_quantum.data_coordinate)
760 def _add_quantum_datasets(self, quantum_datasets: PredictedQuantumDatasetsModel) -> None:
761 self._quantum_datasets[quantum_datasets.quantum_id] = quantum_datasets
762 self._add_quantum(
763 quantum_datasets.quantum_id, quantum_datasets.task_label, quantum_datasets.data_coordinate
764 )
765 task_node = self.pipeline_graph.tasks[quantum_datasets.task_label]
766 for connection_name, input_datasets in quantum_datasets.inputs.items():
767 pipeline_edge = task_node.get_input_edge(connection_name)
768 for input_dataset in input_datasets:
769 self._add_dataset(input_dataset)
770 self._bipartite_xgraph.add_edge(
771 input_dataset.dataset_id,
772 quantum_datasets.quantum_id,
773 key=connection_name,
774 is_read=True,
775 )
776 # There might be multiple input connections for the same
777 # dataset type.
778 self._bipartite_xgraph.edges[
779 input_dataset.dataset_id, quantum_datasets.quantum_id
780 ].setdefault("pipeline_edges", []).append(pipeline_edge)
781 for connection_name, output_datasets in quantum_datasets.outputs.items():
782 pipeline_edges = [task_node.get_output_edge(connection_name)]
783 for output_dataset in output_datasets:
784 self._add_dataset(output_dataset)
785 self._bipartite_xgraph.add_edge(
786 quantum_datasets.quantum_id,
787 output_dataset.dataset_id,
788 key=connection_name,
789 is_read=False,
790 pipeline_edges=pipeline_edges,
791 )
793 def _add_quantum(
794 self, quantum_id: uuid.UUID, task_label: str, data_coordinate_values: Sequence[DataIdValue]
795 ) -> None:
796 task_node = self.pipeline_graph.tasks[task_label]
797 self._quantum_only_xgraph.add_node(quantum_id, task_label=task_label, pipeline_node=task_node)
798 self._bipartite_xgraph.add_node(quantum_id, task_label=task_label, pipeline_node=task_node)
799 data_coordinate_values = tuple(data_coordinate_values)
800 dimensions = self.pipeline_graph.tasks[task_label].dimensions
801 data_id = DataCoordinate.from_full_values(dimensions, tuple(data_coordinate_values))
802 self._quantum_only_xgraph.nodes[quantum_id].setdefault("data_id", data_id)
803 self._bipartite_xgraph.nodes[quantum_id].setdefault("data_id", data_id)
804 self._quanta_by_task_label[task_label][data_id] = quantum_id
806 def _add_dataset(self, model: PredictedDatasetModel) -> None:
807 dataset_type_node = self.pipeline_graph.dataset_types[model.dataset_type_name]
808 data_id = DataCoordinate.from_full_values(dataset_type_node.dimensions, tuple(model.data_coordinate))
809 self._bipartite_xgraph.add_node(
810 model.dataset_id,
811 dataset_type_name=dataset_type_node.name,
812 pipeline_node=dataset_type_node,
813 run=model.run,
814 )
815 self._bipartite_xgraph.nodes[model.dataset_id].setdefault("data_id", data_id)
816 self._datasets_by_type[model.dataset_type_name][data_id] = model.dataset_id
818 @classmethod
819 def open(
820 cls,
821 uri: ResourcePathExpression,
822 page_size: int = DEFAULT_PAGE_SIZE,
823 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,
824 ) -> AbstractContextManager[PredictedQuantumGraphReader]:
825 """Open a quantum graph and return a reader to load from it.
827 Parameters
828 ----------
829 uri : convertible to `lsst.resources.ResourcePath`
830 URI to open. Should have a ``.qg`` extension.
831 page_size : `int`, optional
832 Approximate number of bytes to read at once from address files.
833 Note that this does not set a page size for *all* reads, but it
834 does affect the smallest, most numerous reads.
835 import_mode : `.pipeline_graph.TaskImportMode`, optional
836 How to handle importing the task classes referenced in the pipeline
837 graph.
839 Returns
840 -------
841 reader : `contextlib.AbstractContextManager` [ \
842 `PredictedQuantumGraphReader` ]
843 A context manager that returns the reader when entered.
844 """
845 return PredictedQuantumGraphReader.open(uri, page_size=page_size, import_mode=import_mode)
847 @classmethod
848 def read_execution_quanta(
849 cls,
850 uri: ResourcePathExpression,
851 quantum_ids: Iterable[uuid.UUID] | None = None,
852 page_size: int = DEFAULT_PAGE_SIZE,
853 ) -> PredictedQuantumGraph:
854 """Read one or more executable quanta from a quantum graph file.
856 Parameters
857 ----------
858 uri : convertible to `lsst.resources.ResourcePath`
859 URI to open. Should have a ``.qg`` extension for new quantum graph
860 files, or ``.qgraph`` for the old format.
861 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional
862 Iterable of quantum IDs to load. If not provided, all quanta will
863 be loaded. The UUIDs of special init quanta will be ignored.
864 page_size : `int`, optional
865 Approximate number of bytes to read at once from address files.
866 Note that this does not set a page size for *all* reads, but it
867 does affect the smallest, most numerous reads.
869 Returns
870 -------
871 quantum_graph : `PredictedQuantumGraph` ]
872 A quantum graph that can build execution quanta for all of the
873 given IDs.
874 """
875 return PredictedQuantumGraphComponents.read_execution_quanta(
876 uri,
877 quantum_ids,
878 page_size=page_size,
879 ).assemble()
881 @classmethod
882 def make_empty(
883 cls,
884 universe: DimensionUniverse,
885 *,
886 output_run: str,
887 inputs: Iterable[str] = (),
888 output: str | None = None,
889 add_packages: bool = True,
890 ) -> PredictedQuantumGraph:
891 """Make an empty quantum graph with no tasks.
893 Parameters
894 ----------
895 universe : `lsst.daf.butler.DimensionUniverse`
896 Definitions for all butler dimensions.
897 output_run : `str`
898 Output run collection.
899 inputs : `~collections.abc.Iterable` [`str`], optional
900 Iterable of input collection names.
901 output : `str` or `None`, optional
902 Output chained collection.
903 add_packages : `bool`, optional
904 Whether to add the special init quantum that writes the 'packages'
905 dataset. The default (`True`) is consistent with
906 `~..quantum_graph_builder.QuantumGraphBuilder` behavior when there
907 are no regular quanta generated.
909 Returns
910 -------
911 quantum_graph : `PredictedQuantumGraph`
912 An empty quantum graph.
913 """
914 return cls(
915 PredictedQuantumGraphComponents.make_empty(
916 universe,
917 output_run=output_run,
918 inputs=inputs,
919 output=output,
920 add_packages=add_packages,
921 )
922 )
924 @property
925 def quanta_by_task(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:
926 """A nested mapping of all quanta, keyed first by task name and then by
927 data ID.
929 Notes
930 -----
931 This is populated by the ``thin_graph`` component (all quanta are
932 added) and the `quantum_datasets`` component (only loaded quanta are
933 added). All tasks in the pipeline graph are included, even if none of
934 their quanta were loaded (i.e. nested mappings may be empty).
936 The returned object may be an internal dictionary; as the type
937 annotation indicates, it should not be modified in place.
938 """
939 return self._quanta_by_task_label
941 @property
942 def datasets_by_type(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:
943 """A nested mapping of all datasets, keyed first by dataset type name
944 and then by data ID.
946 Notes
947 -----
948 This is populated only by the ``quantum_datasets`` and ``init_quanta``
949 components, and only datasets referenced by loaded quanta are present.
950 All dataset types in the pipeline graph are included, even if none of
951 their datasets were loaded (i.e. nested mappings may be empty).
953 The returned object may be an internal dictionary; as the type
954 annotation indicates, it should not be modified in place.
955 """
956 return self._datasets_by_type
958 @property
959 def quantum_only_xgraph(self) -> networkx.DiGraph:
960 """A directed acyclic graph with quanta as nodes and datasets elided.
962 Notes
963 -----
964 Node keys are quantum UUIDs, and are populated by the ``thin_graph``
965 component (all nodes and edges) and ``quantum_datasets`` component
966 (only those that were loaded).
968 Node state dictionaries are described by the
969 `PredictedQuantumInfo` type.
971 The returned object is a read-only view of an internal one.
972 """
973 return self._quantum_only_xgraph.copy(as_view=True)
975 @property
976 def bipartite_xgraph(self) -> networkx.MultiDiGraph:
977 """A directed acyclic graph with quantum and dataset nodes.
979 This graph never includes init-input and init-output datasets.
981 Notes
982 -----
983 Node keys are quantum or dataset UUIDs. Nodes for quanta are present
984 if the ``thin_graph`` component is loaded (all nodes) or if the
985 ``quantum_datasets`` component is loaded (just loaded quanta). Edges
986 and dataset nodes are only present for quanta whose
987 ``quantum_datasets`` were loaded.
989 Node state dictionaries are described by the
990 `PredictedQuantumInfo` and `PredictedDatasetInfo` types.
992 The returned object is a read-only view of an internal one.
993 """
994 return self._bipartite_xgraph.copy(as_view=True)
996 @property
997 def dimension_data(self) -> DimensionDataAttacher | None:
998 """All dimension records needed to expand the data IDS in the graph.
1000 This may be `None` if the dimension data was not loaded. If all
1001 execution quanta have been built, all records are guaranteed to have
1002 been deserialized and the ``records`` attribute is complete. In other
1003 cases some records may still only be present in the ``deserializers``
1004 attribute.
1005 """
1006 return self._dimension_data
1008 def __iter__(self) -> Iterator[uuid.UUID]:
1009 for quanta_for_task in self.quanta_by_task.values():
1010 for data_id in sorted(quanta_for_task.keys()):
1011 yield quanta_for_task[data_id]
1013 def __len__(self) -> int:
1014 return len(self._quantum_only_xgraph)
1016 def get_init_inputs(self, task_label: str) -> dict[ConnectionName, DatasetRef]:
1017 """Return the init-input datasets for the given task.
1019 Parameters
1020 ----------
1021 task_label : `str`
1022 Label of the task.
1024 Returns
1025 -------
1026 init_inputs : `dict` [ `str`, `lsst.daf.butler.DatasetRef` ]
1027 Dataset references for init-input datasets, keyed by connection
1028 name. Dataset types storage classes match the task connection
1029 declarations, not necessarily the data repository, and may be
1030 components.
1031 """
1032 if self._init_quanta is None:
1033 raise IncompleteQuantumGraphError("The init_quanta component was not loaded.")
1034 task_init_node = self.pipeline_graph.tasks[task_label].init
1035 return {
1036 connection_name: task_init_node.inputs[connection_name].adapt_dataset_ref(
1037 self._make_init_ref(datasets[0])
1038 )
1039 for connection_name, datasets in self._init_quanta[task_label].inputs.items()
1040 }
1042 def get_init_outputs(self, task_label: str) -> dict[ConnectionName, DatasetRef]:
1043 """Return the init-output datasets for the given task.
1045 Parameters
1046 ----------
1047 task_label : `str`
1048 Label of the task. ``""`` may be used to get global init-outputs.
1050 Returns
1051 -------
1052 init_outputs : `dict` [ `str`, `lsst.daf.butler.DatasetRef` ]
1053 Dataset references for init-outputs datasets, keyed by connection
1054 name. Dataset types storage classes match the task connection
1055 declarations, not necessarily the data repository.
1056 """
1057 if self._init_quanta is None:
1058 raise IncompleteQuantumGraphError("The init_quanta component was not loaded.")
1059 if not task_label:
1060 (datasets,) = self._init_quanta[""].outputs.values()
1061 return {
1062 acc.PACKAGES_INIT_OUTPUT_NAME: DatasetRef(
1063 self.pipeline_graph.packages_dataset_type,
1064 DataCoordinate.make_empty(self.pipeline_graph.universe),
1065 run=datasets[0].run,
1066 id=datasets[0].dataset_id,
1067 conform=False,
1068 )
1069 }
1070 task_init_node = self.pipeline_graph.tasks[task_label].init
1071 result: dict[ConnectionName, DatasetRef] = {}
1072 for connection_name, datasets in self._init_quanta[task_label].outputs.items():
1073 if connection_name == acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME:
1074 edge = task_init_node.config_output
1075 else:
1076 edge = task_init_node.outputs[connection_name]
1077 result[connection_name] = edge.adapt_dataset_ref(self._make_init_ref(datasets[0]))
1078 return result
1080 def _make_init_ref(self, dataset: PredictedDatasetModel) -> DatasetRef:
1081 dataset_type = self.pipeline_graph.dataset_types[dataset.dataset_type_name].dataset_type
1082 return DatasetRef(
1083 dataset_type,
1084 DataCoordinate.make_empty(self.pipeline_graph.universe),
1085 run=dataset.run,
1086 id=dataset.dataset_id,
1087 conform=False,
1088 )
1090 def build_execution_quanta(
1091 self,
1092 quantum_ids: Iterable[uuid.UUID] | None = None,
1093 task_label: str | None = None,
1094 ) -> dict[uuid.UUID, Quantum]:
1095 """Build `lsst.daf.butler.Quantum` objects suitable for executing
1096 tasks.
1098 In addition to returning the quantum objects directly, this also causes
1099 the `quantum_only_xgraph` and `bipartite_xgraph` graphs to include a
1100 ``quantum`` attribute for the affected quanta.
1102 Parameters
1103 ----------
1104 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional
1105 IDs of all quanta to return. If not provided, all quanta for the
1106 given task label (if given) or graph are returned.
1107 task_label : `str`, optional
1108 Task label whose quanta should be generated. Ignored if
1109 ``quantum_ids`` is not `None`.
1111 Returns
1112 -------
1113 quanta : `dict` [ `uuid.UUID`, `lsst.daf.butler.Quantum` ]
1114 Mapping of quanta, keyed by UUID. All dataset types are adapted to
1115 the task's storage class declarations and inputs may be components.
1116 All data IDs have dimension records attached.
1117 """
1118 if not self._init_quanta:
1119 raise IncompleteQuantumGraphError(
1120 "Cannot build execution quanta without loading the ``init_quanta`` component."
1121 )
1122 if quantum_ids is None:
1123 if task_label is not None:
1124 quantum_ids = self._quanta_by_task_label[task_label].values()
1125 else:
1126 quantum_ids = self._quantum_only_xgraph.nodes.keys()
1127 else:
1128 # Guard against single-pass iterators.
1129 quantum_ids = list(quantum_ids)
1130 del task_label # make sure we don't accidentally use this.
1131 result: dict[uuid.UUID, Quantum] = {}
1132 self._expand_execution_quantum_data_ids(quantum_ids)
1133 task_init_datastore_records: dict[TaskLabel, dict[DatastoreName, DatastoreRecordData]] = {}
1134 for quantum_id in quantum_ids:
1135 quantum_node_dict: PredictedQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]
1136 if "quantum" in quantum_node_dict:
1137 result[quantum_id] = quantum_node_dict["quantum"]
1138 continue
1139 # We've declare the info dict keys to all be required because that
1140 # saves a lot of casting, but the reality is that they can either
1141 # be fully populated or totally unpopulated. But that makes mypy
1142 # think the check above always succeeds.
1143 try: # type:ignore [unreachable]
1144 quantum_datasets = self._quantum_datasets[quantum_id]
1145 except KeyError:
1146 raise IncompleteQuantumGraphError(
1147 f"Full quantum information for {quantum_id} was not loaded."
1148 ) from None
1149 task_node = self.pipeline_graph.tasks[quantum_datasets.task_label]
1150 quantum_data_id = self._expanded_data_ids[self._bipartite_xgraph.nodes[quantum_id]["data_id"]]
1151 inputs = self._build_execution_quantum_refs(task_node, quantum_datasets.inputs)
1152 outputs = self._build_execution_quantum_refs(task_node, quantum_datasets.outputs)
1153 if task_node.label not in task_init_datastore_records:
1154 task_init_datastore_records[task_node.label] = self._init_quanta[
1155 task_node.label
1156 ].deserialize_datastore_records()
1157 quantum = Quantum(
1158 taskName=task_node.task_class_name,
1159 taskClass=task_node.task_class,
1160 dataId=quantum_data_id,
1161 initInputs={
1162 ref.datasetType: ref for ref in self.get_init_inputs(quantum_datasets.task_label).values()
1163 },
1164 inputs=inputs,
1165 outputs=outputs,
1166 datastore_records=DatastoreRecordData.merge_mappings(
1167 quantum_datasets.deserialize_datastore_records(),
1168 task_init_datastore_records[task_node.label],
1169 ),
1170 )
1171 self._quantum_only_xgraph.nodes[quantum_id]["quantum"] = quantum
1172 self._bipartite_xgraph.nodes[quantum_id]["quantum"] = quantum
1173 result[quantum_id] = quantum
1174 return result
1176 def _expand_execution_quantum_data_ids(self, quantum_ids: Iterable[uuid.UUID]) -> None:
1177 if self._dimension_data is None:
1178 raise IncompleteQuantumGraphError(
1179 "Cannot build execution quanta without loading the ``dimension_data`` component."
1180 )
1181 data_ids_to_expand: dict[DimensionGroup, set[DataCoordinate]] = defaultdict(set)
1182 for quantum_id in quantum_ids:
1183 data_id: DataCoordinate = self._bipartite_xgraph.nodes[quantum_id]["data_id"]
1184 if data_id.hasRecords():
1185 self._expanded_data_ids[data_id] = data_id
1186 else:
1187 data_ids_to_expand[data_id.dimensions].add(data_id)
1188 for dataset_id in itertools.chain(
1189 self._bipartite_xgraph.predecessors(quantum_id),
1190 self._bipartite_xgraph.successors(quantum_id),
1191 ):
1192 data_id = self._bipartite_xgraph.nodes[dataset_id]["data_id"]
1193 if data_id.hasRecords():
1194 self._expanded_data_ids[data_id] = data_id
1195 else:
1196 data_ids_to_expand[data_id.dimensions].add(data_id)
1197 for dimensions, data_ids_for_dimensions in data_ids_to_expand.items():
1198 self._expanded_data_ids.update(
1199 (d, d) for d in self._dimension_data.attach(dimensions, data_ids_for_dimensions)
1200 )
1202 def _build_execution_quantum_refs(
1203 self, task_node: TaskNode, model_mapping: dict[ConnectionName, list[PredictedDatasetModel]]
1204 ) -> dict[DatasetType, list[DatasetRef]]:
1205 results: dict[DatasetType, list[DatasetRef]] = {}
1206 for connection_name, datasets in model_mapping.items():
1207 edge = task_node.get_edge(connection_name)
1208 dataset_type = edge.adapt_dataset_type(
1209 self.pipeline_graph.dataset_types[edge.parent_dataset_type_name].dataset_type
1210 )
1211 results[dataset_type] = [self._make_general_ref(dataset_type, d.dataset_id) for d in datasets]
1212 return results
1214 def _make_general_ref(self, dataset_type: DatasetType, dataset_id: uuid.UUID) -> DatasetRef:
1215 node_state = self._bipartite_xgraph.nodes[dataset_id]
1216 data_id = self._expanded_data_ids[node_state["data_id"]]
1217 return DatasetRef(dataset_type, data_id, run=node_state["run"], id=dataset_id)
1219 def make_init_qbb(
1220 self,
1221 butler_config: Config | ResourcePathExpression,
1222 *,
1223 config_search_paths: Iterable[str] | None = None,
1224 ) -> QuantumBackedButler:
1225 """Construct an quantum-backed butler suitable for reading and writing
1226 init input and init output datasets, respectively.
1228 This only requires the ``init_quanta`` component to have been loaded.
1230 Parameters
1231 ----------
1232 butler_config : `~lsst.daf.butler.Config` or \
1233 `~lsst.resources.ResourcePathExpression`
1234 A butler repository root, configuration filename, or configuration
1235 instance.
1236 config_search_paths : `~collections.abc.Iterable` [ `str` ], optional
1237 Additional search paths for butler configuration.
1239 Returns
1240 -------
1241 qbb : `~lsst.daf.butler.QuantumBackedButler`
1242 A limited butler that can ``get`` init-input datasets and ``put``
1243 init-output datasets.
1244 """
1245 # Collect all init input/output dataset IDs.
1246 predicted_inputs: set[uuid.UUID] = set()
1247 predicted_outputs: set[uuid.UUID] = set()
1248 datastore_record_maps: list[dict[DatastoreName, DatastoreRecordData]] = []
1249 for init_quantum_datasets in self._init_quanta.values():
1250 predicted_inputs.update(
1251 d.dataset_id for d in itertools.chain.from_iterable(init_quantum_datasets.inputs.values())
1252 )
1253 predicted_outputs.update(
1254 d.dataset_id for d in itertools.chain.from_iterable(init_quantum_datasets.outputs.values())
1255 )
1256 datastore_record_maps.append(
1257 {
1258 datastore_name: DatastoreRecordData.from_simple(serialized_records)
1259 for datastore_name, serialized_records in init_quantum_datasets.datastore_records.items()
1260 }
1261 )
1262 # Remove intermediates from inputs.
1263 predicted_inputs -= predicted_outputs
1264 dataset_types = {d.name: d.dataset_type for d in self.pipeline_graph.dataset_types.values()}
1265 # Make butler from everything.
1266 return QuantumBackedButler.from_predicted(
1267 config=butler_config,
1268 predicted_inputs=predicted_inputs,
1269 predicted_outputs=predicted_outputs,
1270 dimensions=self.pipeline_graph.universe,
1271 datastore_records=DatastoreRecordData.merge_mappings(*datastore_record_maps),
1272 search_paths=list(config_search_paths) if config_search_paths is not None else None,
1273 dataset_types=dataset_types,
1274 )
1276 def write_init_outputs(self, butler: LimitedButler, skip_existing: bool = True) -> None:
1277 """Write the init-output datasets for all tasks in the quantum graph.
1279 This only requires the ``init_quanta`` component to have been loaded.
1281 Parameters
1282 ----------
1283 butler : `lsst.daf.butler.LimitedButler`
1284 A limited butler data repository client.
1285 skip_existing : `bool`, optional
1286 If `True` (default) ignore init-outputs that already exist. If
1287 `False`, raise.
1289 Raises
1290 ------
1291 lsst.daf.butler.registry.ConflictingDefinitionError
1292 Raised if an init-output dataset already exists and
1293 ``skip_existing=False``.
1294 """
1295 # Extract init-input and init-output refs from the QG.
1296 input_refs: dict[str, DatasetRef] = {}
1297 output_refs: dict[str, DatasetRef] = {}
1298 for task_node in self.pipeline_graph.tasks.values():
1299 if task_node.label not in self._init_quanta:
1300 continue
1301 input_refs.update(
1302 {ref.datasetType.name: ref for ref in self.get_init_inputs(task_node.label).values()}
1303 )
1304 output_refs.update(
1305 {
1306 ref.datasetType.name: ref
1307 for ref in self.get_init_outputs(task_node.label).values()
1308 if ref.datasetType.name != task_node.init.config_output.dataset_type_name
1309 }
1310 )
1311 for ref, is_stored in butler.stored_many(output_refs.values()).items():
1312 if is_stored:
1313 if not skip_existing:
1314 raise ConflictingDefinitionError(f"Init-output dataset {ref} already exists.")
1315 # We'll `put` whatever's left in output_refs at the end.
1316 del output_refs[ref.datasetType.name]
1317 # Instantiate tasks, reading overall init-inputs and gathering
1318 # init-output in-memory objects.
1319 init_outputs: list[tuple[Any, DatasetType]] = []
1320 self.pipeline_graph.instantiate_tasks(
1321 get_init_input=lambda dataset_type: butler.get(
1322 input_refs[dataset_type.name].overrideStorageClass(dataset_type.storageClass)
1323 ),
1324 init_outputs=init_outputs,
1325 # A task can be in the pipeline graph without having an init
1326 # quantum if it doesn't have any regular quanta either (e.g. they
1327 # were all skipped), and the _init_quanta has a "" entry for global
1328 # init-outputs that we don't want to pass here.
1329 labels=self.pipeline_graph.tasks.keys() & self._init_quanta.keys(),
1330 )
1331 # Write init-outputs that weren't already present.
1332 for obj, dataset_type in init_outputs:
1333 if new_ref := output_refs.get(dataset_type.name):
1334 assert new_ref.datasetType.storageClass_name == dataset_type.storageClass_name, (
1335 "QG init refs should use task connection storage classes."
1336 )
1337 butler.put(obj, new_ref)
1339 def write_configs(self, butler: LimitedButler, compare_existing: bool = True) -> None:
1340 """Write the config datasets for all tasks in the quantum graph.
1342 Parameters
1343 ----------
1344 butler : `lsst.daf.butler.LimitedButler`
1345 A limited butler data repository client.
1346 compare_existing : `bool`, optional
1347 If `True` check configs that already exist for consistency. If
1348 `False`, always raise if configs already exist.
1350 Raises
1351 ------
1352 lsst.daf.butler.registry.ConflictingDefinitionError
1353 Raised if an config dataset already exists and
1354 ``compare_existing=False``, or if the existing config is not
1355 consistent with the config in the quantum graph.
1356 """
1357 to_put: list[tuple[PipelineTaskConfig, DatasetRef]] = []
1358 for task_node in self.pipeline_graph.tasks.values():
1359 if task_node.label not in self._init_quanta:
1360 continue
1361 dataset_type_name = task_node.init.config_output.dataset_type_name
1362 ref = self.get_init_outputs(task_node.label)[acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME]
1363 try:
1364 old_config = butler.get(ref)
1365 except (LookupError, FileNotFoundError):
1366 old_config = None
1367 if old_config is not None:
1368 if not compare_existing:
1369 raise ConflictingDefinitionError(f"Config dataset {ref} already exists.")
1370 if not task_node.config.compare(old_config, shortcut=False, output=log_config_mismatch):
1371 raise ConflictingDefinitionError(
1372 f"Config does not match existing task config {dataset_type_name!r} in "
1373 "butler; tasks configurations must be consistent within the same run collection."
1374 )
1375 else:
1376 to_put.append((task_node.config, ref))
1377 # We do writes at the end to minimize the mess we leave behind when we
1378 # raise an exception.
1379 for config, ref in to_put:
1380 butler.put(config, ref)
1382 def write_packages(self, butler: LimitedButler, compare_existing: bool = True) -> None:
1383 """Write the 'packages' dataset for the currently-active software
1384 versions.
1386 Parameters
1387 ----------
1388 butler : `lsst.daf.butler.LimitedButler`
1389 A limited butler data repository client.
1390 compare_existing : `bool`, optional
1391 If `True` check packages that already exist for consistency. If
1392 `False`, always raise if the packages dataset already exists.
1394 Raises
1395 ------
1396 lsst.daf.butler.registry.ConflictingDefinitionError
1397 Raised if the packages dataset already exists and is not consistent
1398 with the current packages.
1399 """
1400 new_packages = Packages.fromSystem()
1401 (ref,) = self.get_init_outputs("").values()
1402 try:
1403 packages = butler.get(ref)
1404 except (LookupError, FileNotFoundError):
1405 packages = None
1406 if packages is not None:
1407 if not compare_existing:
1408 raise ConflictingDefinitionError(f"Packages dataset {ref} already exists.")
1409 if compare_packages(packages, new_packages):
1410 # have to remove existing dataset first; butler has no
1411 # replace option.
1412 butler.pruneDatasets([ref], unstore=True, purge=True)
1413 butler.put(packages, ref)
1414 else:
1415 butler.put(new_packages, ref)
1417 def init_output_run(self, butler: LimitedButler, existing: bool = True) -> None:
1418 """Initialize a new output RUN collection by writing init-output
1419 datasets (including configs and packages).
1421 Parameters
1422 ----------
1423 butler : `lsst.daf.butler.LimitedButler`
1424 A limited butler data repository client.
1425 existing : `bool`, optional
1426 If `True` check or ignore outputs that already exist. If
1427 `False`, always raise if an output dataset already exists.
1429 Raises
1430 ------
1431 lsst.daf.butler.registry.ConflictingDefinitionError
1432 Raised if there are existing init output datasets, and either
1433 ``existing=False`` or their contents are not compatible with this
1434 graph.
1435 """
1436 self.write_configs(butler, compare_existing=existing)
1437 self.write_packages(butler, compare_existing=existing)
1438 self.write_init_outputs(butler, skip_existing=existing)
1440 @classmethod
1441 def from_old_quantum_graph(cls, old_quantum_graph: QuantumGraph) -> PredictedQuantumGraph:
1442 """Construct from an old `QuantumGraph` instance.
1444 Parameters
1445 ----------
1446 old_quantum_graph : `QuantumGraph`
1447 Quantum graph to transform.
1449 Returns
1450 -------
1451 predicted_quantum_graph : `PredictedQuantumGraph`
1452 A new predicted quantum graph.
1453 """
1454 return PredictedQuantumGraphComponents.from_old_quantum_graph(old_quantum_graph).assemble()
1456 def to_old_quantum_graph(self) -> QuantumGraph:
1457 """Transform into an old `QuantumGraph` instance.
1459 Returns
1460 -------
1461 old_quantum_graph : `QuantumGraph`
1462 Old quantum graph.
1464 Notes
1465 -----
1466 This can only be called on graphs that have loaded all quantum
1467 datasets, init datasets, and dimension records.
1468 """
1469 from ..graph import QuantumGraph
1471 quanta: dict[TaskDef, set[Quantum]] = {}
1472 quantum_to_quantum_id: dict[Quantum, uuid.UUID] = {}
1473 init_inputs: dict[TaskDef, list[DatasetRef]] = {}
1474 init_outputs: dict[TaskDef, list[DatasetRef]] = {}
1475 for task_def in self.pipeline_graph._iter_task_defs():
1476 if not self._quanta_by_task_label.get(task_def.label):
1477 continue
1478 quanta_for_task: set[Quantum] = set()
1479 for quantum_id, quantum in self.build_execution_quanta(task_label=task_def.label).items():
1480 quanta_for_task.add(quantum)
1481 quantum_to_quantum_id[quantum] = quantum_id
1482 quanta[task_def] = quanta_for_task
1483 init_inputs[task_def] = list(self.get_init_inputs(task_def.label).values())
1484 init_outputs[task_def] = list(self.get_init_outputs(task_def.label).values())
1485 global_init_outputs = list(self.get_init_outputs("").values())
1486 registry_dataset_types = [d.dataset_type for d in self.pipeline_graph.dataset_types.values()]
1487 result = object.__new__(QuantumGraph)
1488 result._buildGraphs(
1489 quanta,
1490 _quantumToNodeId=quantum_to_quantum_id,
1491 metadata=self.header.to_old_metadata(),
1492 universe=self.pipeline_graph.universe,
1493 initInputs=init_inputs,
1494 initOutputs=init_outputs,
1495 globalInitOutputs=global_init_outputs,
1496 registryDatasetTypes=registry_dataset_types,
1497 )
1498 return result
1500 def _make_summary(self) -> QgraphSummary:
1501 from ..graph import QgraphSummary, QgraphTaskSummary
1503 summary = QgraphSummary(
1504 cmdLine=self.header.command or None,
1505 creationUTC=str(self.header.timestamp) if self.header.timestamp is not None else None,
1506 inputCollection=self.header.inputs or None,
1507 outputCollection=self.header.output,
1508 outputRun=self.header.output_run,
1509 )
1510 for task_label, quanta_for_task in self.quanta_by_task.items():
1511 task_summary = QgraphTaskSummary(taskLabel=task_label, numQuanta=len(quanta_for_task))
1512 task_node = self.pipeline_graph.tasks[task_label]
1513 for quantum_id in quanta_for_task.values():
1514 quantum_datasets = self._quantum_datasets[quantum_id]
1515 for connection_name, input_datasets in quantum_datasets.inputs.items():
1516 task_summary.numInputs[
1517 task_node.get_input_edge(connection_name).parent_dataset_type_name
1518 ] += len(input_datasets)
1519 for connection_name, output_datasets in quantum_datasets.outputs.items():
1520 task_summary.numOutputs[
1521 task_node.get_output_edge(connection_name).parent_dataset_type_name
1522 ] += len(output_datasets)
1523 summary.qgraphTaskSummaries[task_label] = task_summary
1524 return summary
1527@dataclasses.dataclass(kw_only=True)
1528class PredictedQuantumGraphComponents:
1529 """A helper class for building and writing predicted quantum graphs.
1531 Notes
1532 -----
1533 This class is a simple struct of model classes to allow different tools
1534 that build predicted quantum graphs to assemble them in whatever order they
1535 prefer. It does not enforce any internal invariants (e.g. the quantum and
1536 dataset counts in the header, different representations of quanta, internal
1537 ID sorting, etc.), but it does provide methods that can satisfy them.
1538 """
1540 def __post_init__(self) -> None:
1541 self.header.graph_type = "predicted"
1543 header: HeaderModel = dataclasses.field(default_factory=HeaderModel)
1544 """Basic metadata about the graph."""
1546 pipeline_graph: PipelineGraph
1547 """Description of the pipeline this graph runs, including all task label
1548 and dataset type definitions.
1550 This may include tasks that do not have any quanta (e.g. due to skipping
1551 already-executed tasks).
1553 This also includes the dimension universe used to construct the graph.
1554 """
1556 dimension_data: DimensionDataAttacher | None = None
1557 """Object that can attach dimension records to data IDs.
1558 """
1560 init_quanta: PredictedInitQuantaModel = dataclasses.field(default_factory=PredictedInitQuantaModel)
1561 """A list of special quanta that describe the init-inputs and init-outputs
1562 of the graph.
1564 Tasks that are included in the pipeline graph but do not have any quanta
1565 may or may not have an init quantum, but tasks that do have regular quanta
1566 always have an init quantum as well.
1568 When used to construct a `PredictedQuantumGraph`, this must have either
1569 zero entries or all tasks in the pipeline.
1570 """
1572 thin_graph: PredictedThinGraphModel = dataclasses.field(default_factory=PredictedThinGraphModel)
1573 """A lightweight quantum-quantum DAG with task labels and data IDs only.
1575 This does not include the special "init" quanta.
1576 """
1578 quantum_datasets: dict[uuid.UUID, PredictedQuantumDatasetsModel] = dataclasses.field(default_factory=dict)
1579 """The full descriptions of all quanta, including input and output
1580 dataset, keyed by UUID.
1582 When used to construct a `PredictedQuantumGraph`, this need not have all
1583 entries.
1585 This does not include special "init" quanta.
1586 """
1588 @classmethod
1589 def make_empty(
1590 cls,
1591 universe: DimensionUniverse,
1592 *,
1593 output_run: str,
1594 inputs: Iterable[str] = (),
1595 output: str | None = None,
1596 add_packages: bool = True,
1597 ) -> PredictedQuantumGraphComponents:
1598 """Make components for an empty quantum graph with no tasks.
1600 Parameters
1601 ----------
1602 universe : `lsst.daf.butler.DimensionUniverse`
1603 Definitions for all butler dimensions.
1604 output_run : `str`
1605 Output run collection.
1606 inputs : `~collections.abc.Iterable` [`str`], optional
1607 Iterable of input collection names.
1608 output : `str` or `None`, optional
1609 Output chained collection.
1610 add_packages : `bool`, optional
1611 Whether to add the special init quantum that writes the 'packages'
1612 dataset. The default (`True`) is consistent with
1613 `~..quantum_graph_builder.QuantumGraphBuilder` behavior when there
1614 are no regular quanta generated.
1616 Returns
1617 -------
1618 components : `PredictedQuantumGraphComponents`
1619 Components that can be used to build or write an empty quantum
1620 graph.
1621 """
1622 components = cls(pipeline_graph=PipelineGraph(universe=universe))
1623 components.header.inputs = list(inputs)
1624 components.header.output_run = output_run
1625 components.header.output = output
1626 if add_packages:
1627 components.init_quanta.root = [
1628 PredictedQuantumDatasetsModel.model_construct(
1629 quantum_id=generate_uuidv7(),
1630 task_label="",
1631 outputs={
1632 acc.PACKAGES_INIT_OUTPUT_NAME: [
1633 PredictedDatasetModel(
1634 dataset_id=generate_uuidv7(),
1635 dataset_type_name=acc.PACKAGES_INIT_OUTPUT_NAME,
1636 data_coordinate=[],
1637 run=output_run,
1638 )
1639 ]
1640 },
1641 )
1642 ]
1643 return components
1645 def make_dataset_ref(self, predicted: PredictedDatasetModel) -> DatasetRef:
1646 """Make a `lsst.daf.butler.DatasetRef` from information in the
1647 predicted quantum graph.
1649 Parameters
1650 ----------
1651 predicted : `PredictedDatasetModel`
1652 Model for the dataset in the predicted graph.
1654 Returns
1655 -------
1656 ref : `lsst.daf.butler.DatasetRef`
1657 A dataset reference. Data ID will be expanded if and only if
1658 the dimension data has been loaded.
1659 """
1660 try:
1661 dataset_type = self.pipeline_graph.dataset_types[predicted.dataset_type_name].dataset_type
1662 except KeyError:
1663 if predicted.dataset_type_name == acc.PACKAGES_INIT_OUTPUT_NAME:
1664 dataset_type = self.pipeline_graph.packages_dataset_type
1665 else:
1666 raise
1667 data_id = DataCoordinate.from_full_values(dataset_type.dimensions, tuple(predicted.data_coordinate))
1668 if self.dimension_data is not None:
1669 (data_id,) = self.dimension_data.attach(dataset_type.dimensions, [data_id])
1670 return DatasetRef(
1671 dataset_type,
1672 data_id,
1673 run=predicted.run,
1674 id=predicted.dataset_id,
1675 )
1677 def set_thin_graph(self) -> None:
1678 """Populate the `thin_graph` component from the `pipeline_graph`,
1679 `quantum_datasets` components (which must be complete).
1680 """
1681 bipartite_xgraph = networkx.DiGraph()
1682 self.thin_graph.quanta = {task_label: [] for task_label in self.pipeline_graph.tasks}
1683 graph_quantum_ids: list[uuid.UUID] = []
1684 for quantum_datasets in self.quantum_datasets.values():
1685 self.thin_graph.quanta[quantum_datasets.task_label].append(
1686 PredictedThinQuantumModel.model_construct(
1687 quantum_id=quantum_datasets.quantum_id,
1688 data_coordinate=quantum_datasets.data_coordinate,
1689 )
1690 )
1691 for dataset in itertools.chain.from_iterable(quantum_datasets.inputs.values()):
1692 bipartite_xgraph.add_edge(dataset.dataset_id, quantum_datasets.quantum_id)
1693 for dataset in itertools.chain.from_iterable(quantum_datasets.outputs.values()):
1694 bipartite_xgraph.add_edge(quantum_datasets.quantum_id, dataset.dataset_id)
1695 graph_quantum_ids.append(quantum_datasets.quantum_id)
1696 quantum_only_xgraph: networkx.DiGraph = networkx.bipartite.projected_graph(
1697 bipartite_xgraph, graph_quantum_ids
1698 )
1699 self.thin_graph.edges = list(quantum_only_xgraph.edges)
1701 def set_header_counts(self) -> None:
1702 """Populate the quantum and dataset counts in the header from the
1703 `thin_graph`, `init_quanta`, and `quantum_datasets` components.
1704 """
1705 self.header.n_quanta = len(self.quantum_datasets)
1706 self.header.n_task_quanta = {
1707 task_label: len(thin_quanta) for task_label, thin_quanta in self.thin_graph.quanta.items()
1708 }
1709 all_dataset_ids: set[uuid.UUID] = set()
1710 for quantum_datasets in itertools.chain(self.init_quanta.root, self.quantum_datasets.values()):
1711 all_dataset_ids.update(quantum_datasets.iter_dataset_ids())
1712 self.header.n_datasets = len(all_dataset_ids)
1714 def update_output_run(self, output_run: str) -> None:
1715 """Update the output `~lsst.daf.butler.CollectionType.RUN` collection
1716 name in all datasets and regenerate all output dataset and quantum
1717 UUIDs.
1719 Parameters
1720 ----------
1721 output_run : `str`
1722 New output `~lsst.daf.butler.CollectionType.RUN` collection name.
1723 """
1724 uuid_map: dict[uuid.UUID, uuid.UUID] = {}
1725 # Do all outputs and then all inputs in separate passes so we don't
1726 # need to rely on topological ordering of anything.
1727 for quantum_datasets in itertools.chain(self.init_quanta.root, self.quantum_datasets.values()):
1728 new_quantum_id = generate_uuidv7()
1729 quantum_datasets.quantum_id = new_quantum_id
1730 for output_dataset in itertools.chain.from_iterable(quantum_datasets.outputs.values()):
1731 assert output_dataset.run == self.header.output_run, (
1732 f"Incorrect run {output_dataset.run} for output dataset {output_dataset.dataset_id}."
1733 )
1734 new_dataset_id = generate_uuidv7()
1735 uuid_map[output_dataset.dataset_id] = new_dataset_id
1736 output_dataset.dataset_id = new_dataset_id
1737 output_dataset.run = output_run
1738 for quantum_datasets in itertools.chain(self.init_quanta.root, self.quantum_datasets.values()):
1739 for input_dataset in itertools.chain.from_iterable(quantum_datasets.inputs.values()):
1740 if input_dataset.run == self.header.output_run:
1741 input_dataset.run = output_run
1742 input_dataset.dataset_id = uuid_map.get(
1743 input_dataset.dataset_id,
1744 # This dataset isn't necessary an output of the graph
1745 # just because it's in the output run; the graph could
1746 # have been built with extend_run=True.
1747 input_dataset.dataset_id,
1748 )
1749 # Update the keys of the quantum_datasets dict.
1750 self.quantum_datasets = {qd.quantum_id: qd for qd in self.quantum_datasets.values()}
1751 # Since the UUIDs have changed, the thin graph needs to be rewritten.
1752 self.set_thin_graph()
1753 # Update the header last, since we use it above to get the old run.
1754 self.header.output_run = output_run
1756 def assemble(self) -> PredictedQuantumGraph:
1757 """Construct a `PredictedQuantumGraph` from these components."""
1758 return PredictedQuantumGraph(self)
1760 @classmethod
1761 def read_execution_quanta(
1762 cls,
1763 uri: ResourcePathExpression,
1764 quantum_ids: Iterable[uuid.UUID] | None = None,
1765 page_size: int = DEFAULT_PAGE_SIZE,
1766 ) -> PredictedQuantumGraphComponents:
1767 """Read one or more executable quanta from a quantum graph file.
1769 Parameters
1770 ----------
1771 uri : convertible to `lsst.resources.ResourcePath`
1772 URI to open. Should have a ``.qg`` extension for new quantum graph
1773 files, or ``.qgraph`` for the old format.
1774 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional
1775 Iterable of quantum IDs to load. If not provided, all quanta will
1776 be loaded. The UUIDs of special init quanta will be ignored.
1777 page_size : `int`, optional
1778 Approximate number of bytes to read at once from address files.
1779 Note that this does not set a page size for *all* reads, but it
1780 does affect the smallest, most numerous reads.
1782 Returns
1783 -------
1784 components : `PredictedQuantumGraphComponents` ]
1785 Components for quantum graph that can build execution quanta for
1786 all of the given IDs.
1787 """
1788 uri = ResourcePath(uri)
1789 if uri.getExtension() == ".qgraph":
1790 _LOG.warning(
1791 f"Reading and converting old quantum graph {uri}. "
1792 "Use the '.qg' extension to write in the new format."
1793 )
1794 from ..graph import QuantumGraph
1796 old_qg = QuantumGraph.loadUri(uri, nodes=quantum_ids)
1797 return PredictedQuantumGraphComponents.from_old_quantum_graph(old_qg)
1799 with PredictedQuantumGraph.open(uri, page_size=page_size) as reader:
1800 reader.read_execution_quanta(quantum_ids)
1801 return reader.components
1803 @classmethod
1804 def from_old_quantum_graph(cls, old_quantum_graph: QuantumGraph) -> PredictedQuantumGraphComponents:
1805 """Construct from an old `QuantumGraph` instance.
1807 Parameters
1808 ----------
1809 old_quantum_graph : `QuantumGraph`
1810 Quantum graph to transform.
1812 Returns
1813 -------
1814 components : `PredictedQuantumGraphComponents`
1815 Components for a new predicted quantum graph.
1816 """
1817 header = HeaderModel.from_old_quantum_graph(old_quantum_graph)
1818 result = cls(header=header, pipeline_graph=old_quantum_graph.pipeline_graph)
1819 result.init_quanta.update_from_old_quantum_graph(old_quantum_graph)
1820 dimension_data_extractor = DimensionDataExtractor.from_dimension_group(
1821 old_quantum_graph.pipeline_graph.get_all_dimensions()
1822 )
1823 for task_node in old_quantum_graph.pipeline_graph.tasks.values():
1824 task_quanta = old_quantum_graph.get_task_quanta(task_node.label)
1825 for quantum_id, quantum in task_quanta.items():
1826 result.quantum_datasets[quantum_id] = PredictedQuantumDatasetsModel.from_execution_quantum(
1827 task_node, quantum, quantum_id
1828 )
1829 dimension_data_extractor.update([cast(DataCoordinate, quantum.dataId)])
1830 for refs in itertools.chain(quantum.inputs.values(), quantum.outputs.values()):
1831 dimension_data_extractor.update(ref.dataId for ref in refs)
1832 result.dimension_data = DimensionDataAttacher(
1833 records=dimension_data_extractor.records.values(),
1834 dimensions=result.pipeline_graph.get_all_dimensions(),
1835 )
1836 result.set_thin_graph()
1837 result.set_header_counts()
1838 return result
1840 def write(
1841 self,
1842 uri: ResourcePathExpression,
1843 *,
1844 zstd_level: int = 10,
1845 zstd_dict_size: int = 32768,
1846 zstd_dict_n_inputs: int = 512,
1847 ) -> None:
1848 """Write the graph to a file.
1850 Parameters
1851 ----------
1852 uri : convertible to `lsst.resources.ResourcePath`
1853 Path to write to. Should have a ``.qg`` extension, or ``.qgraph``
1854 to force writing the old format.
1855 zstd_level : `int`, optional
1856 ZStandard compression level to use on JSON blocks.
1857 zstd_dict_size : `int`, optional
1858 Size of a ZStandard dictionary that shares compression information
1859 across components. Set to zero to disable the dictionary.
1860 Dictionary compression is automatically disabled if the number of
1861 quanta is smaller than ``zstd_dict_n_inputs``.
1862 zstd_dict_n_inputs : `int`, optional
1863 Maximum number of `PredictedQuantumDatasetsModel` JSON
1864 representations to feed the ZStandard dictionary training routine.
1866 Notes
1867 -----
1868 Only a complete predicted quantum graph with all components fully
1869 populated should be written.
1870 """
1871 if self.header.n_task_quanta != {
1872 task_label: len(quanta) for task_label, quanta in self.thin_graph.quanta.items()
1873 }:
1874 raise RuntimeError(
1875 "Cannot save graph after partial read of quanta: thin graph is inconsistent with header."
1876 )
1877 # Ensure we record the actual version we're about to write, in case
1878 # we're rewriting an old graph in a new format.
1879 self.header.version = FORMAT_VERSION
1880 uri = ResourcePath(uri)
1881 match uri.getExtension():
1882 case ".qg":
1883 pass
1884 case ".qgraph":
1885 _LOG.warning(
1886 "Converting to an old-format quantum graph.. "
1887 "Use '.qg' instead of '.qgraph' to save in the new format."
1888 )
1889 old_qg = self.assemble().to_old_quantum_graph()
1890 old_qg.saveUri(uri)
1891 return
1892 case ext:
1893 raise ValueError(
1894 f"Unsupported extension {ext!r} for quantum graph; "
1895 "expected '.qg' (or '.qgraph' to force the old format)."
1896 )
1897 cdict_data: bytes | None = None
1898 quantum_datasets_json: dict[uuid.UUID, bytes] = {}
1899 if len(self.quantum_datasets) < zstd_dict_n_inputs:
1900 # ZStandard will fail if we ask to use a compression dict without
1901 # giving it enough data, and it only helps if we have a lot of
1902 # quanta.
1903 zstd_dict_size = 0
1904 if zstd_dict_size:
1905 quantum_datasets_json = {
1906 quantum_model.quantum_id: quantum_model.model_dump_json().encode()
1907 for quantum_model in itertools.islice(self.quantum_datasets.values(), zstd_dict_n_inputs)
1908 }
1909 try:
1910 cdict_data = zstandard.train_dictionary(
1911 zstd_dict_size,
1912 list(quantum_datasets_json.values()),
1913 level=zstd_level,
1914 ).as_bytes()
1915 except zstandard.ZstdError as err:
1916 warnings.warn(f"Not using a compression dictionary: {err}.")
1917 with BaseQuantumGraphWriter.open(
1918 uri,
1919 header=self.header,
1920 pipeline_graph=self.pipeline_graph,
1921 address_filename="quanta",
1922 cdict_data=cdict_data,
1923 zstd_level=zstd_level,
1924 ) as writer:
1925 writer.write_single_model("thin_graph", self.thin_graph)
1926 if self.dimension_data is None:
1927 raise IncompleteQuantumGraphError(
1928 "Cannot save predicted quantum graph with no dimension data."
1929 )
1930 serialized_dimension_data = self.dimension_data.serialized()
1931 writer.write_single_model("dimension_data", serialized_dimension_data)
1932 del serialized_dimension_data
1933 writer.write_single_model("init_quanta", self.init_quanta)
1934 with MultiblockWriter.open_in_zip(
1935 writer.zf, "quantum_datasets", writer.int_size
1936 ) as quantum_datasets_mb:
1937 for quantum_model in self.quantum_datasets.values():
1938 if json_data := quantum_datasets_json.get(quantum_model.quantum_id):
1939 quantum_datasets_mb.write_bytes(
1940 quantum_model.quantum_id, writer.compressor.compress(json_data)
1941 )
1942 else:
1943 quantum_datasets_mb.write_model(
1944 quantum_model.quantum_id, quantum_model, writer.compressor
1945 )
1946 writer.address_writer.addresses.append(quantum_datasets_mb.addresses)
1949@dataclasses.dataclass
1950class PredictedQuantumGraphReader(BaseQuantumGraphReader):
1951 """A helper class for reading predicted quantum graphs."""
1953 components: PredictedQuantumGraphComponents = dataclasses.field(init=False)
1954 """Quantum graph components populated by this reader's methods."""
1956 @classmethod
1957 @contextmanager
1958 def open(
1959 cls,
1960 uri: ResourcePathExpression,
1961 *,
1962 page_size: int = DEFAULT_PAGE_SIZE,
1963 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,
1964 ) -> Iterator[PredictedQuantumGraphReader]:
1965 """Construct a reader from a URI.
1967 Parameters
1968 ----------
1969 uri : convertible to `lsst.resources.ResourcePath`
1970 URI to open. Should have a ``.qg`` extension.
1971 page_size : `int`, optional
1972 Approximate number of bytes to read at once from address files.
1973 Note that this does not set a page size for *all* reads, but it
1974 does affect the smallest, most numerous reads.
1975 import_mode : `.pipeline_graph.TaskImportMode`, optional
1976 How to handle importing the task classes referenced in the pipeline
1977 graph.
1979 Returns
1980 -------
1981 reader : `contextlib.AbstractContextManager` [ \
1982 `PredictedQuantumGraphReader` ]
1983 A context manager that returns the reader when entered.
1984 """
1985 with cls._open(
1986 uri,
1987 graph_type="predicted",
1988 address_filename="quanta",
1989 page_size=page_size,
1990 import_mode=import_mode,
1991 n_addresses=1,
1992 ) as self:
1993 yield self
1995 def __post_init__(self) -> None:
1996 self.components = PredictedQuantumGraphComponents(
1997 header=self.header, pipeline_graph=self.pipeline_graph
1998 )
2000 def finish(self) -> PredictedQuantumGraph:
2001 """Construct a `PredictedQuantumGraph` instance from this reader."""
2002 return self.components.assemble()
2004 def read_all(self) -> None:
2005 """Read all components in full."""
2006 self.read_thin_graph()
2007 self.read_execution_quanta()
2009 def read_thin_graph(self) -> None:
2010 """Read the thin graph.
2012 The thin graph is a quantum-quantum DAG with just task labels and data
2013 IDs as node attributes. It always includes all regular quanta, and
2014 does not include init-input or init-output information.
2015 """
2016 if not self.components.thin_graph.quanta:
2017 if self.header.version > 0:
2018 self.components.thin_graph = self._read_single_block("thin_graph", PredictedThinGraphModel)
2019 else:
2020 self.address_reader.read_all()
2021 thin_graph_v0 = self._read_single_block("thin_graph", _PredictedThinGraphModelV0)
2022 self.components.thin_graph = thin_graph_v0._upgraded(self.address_reader.rows)
2024 def read_init_quanta(self) -> None:
2025 """Read the list of special quanta that represent init-inputs and
2026 init-outputs.
2027 """
2028 if not self.components.init_quanta.root:
2029 self.components.init_quanta = self._read_single_block("init_quanta", PredictedInitQuantaModel)
2031 def read_dimension_data(self) -> None:
2032 """Read all dimension records.
2034 Record data IDs will be immediately deserialized, while other fields
2035 will be left in serialized form until they are needed.
2036 """
2037 if self.components.dimension_data is None:
2038 serializable_dimension_data = self._read_single_block("dimension_data", SerializableDimensionData)
2039 self.components.dimension_data = DimensionDataAttacher(
2040 deserializers=[
2041 DimensionRecordSetDeserializer.from_raw(
2042 self.components.pipeline_graph.universe[element], serialized_records
2043 )
2044 for element, serialized_records in serializable_dimension_data.root.items()
2045 ],
2046 dimensions=DimensionGroup.union(
2047 *self.components.pipeline_graph.group_by_dimensions(prerequisites=True).keys(),
2048 universe=self.components.pipeline_graph.universe,
2049 ),
2050 )
2052 def read_quantum_datasets(self, quantum_ids: Iterable[uuid.UUID] | None = None) -> None:
2053 """Read information about all datasets produced and consumed by the
2054 given quantum IDs.
2056 Parameters
2057 ----------
2058 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional
2059 Iterable of quantum IDs to load. If not provided, all quanta will
2060 be loaded. The UUIDs of special init quanta will be ignored.
2061 """
2062 quantum_datasets: PredictedQuantumDatasetsModel | None
2063 if quantum_ids is None:
2064 if len(self.components.quantum_datasets) != self.header.n_quanta:
2065 for quantum_datasets in MultiblockReader.read_all_models_in_zip(
2066 self.zf,
2067 "quantum_datasets",
2068 PredictedQuantumDatasetsModel,
2069 self.decompressor,
2070 int_size=self.components.header.int_size,
2071 page_size=self.page_size,
2072 ):
2073 self.components.quantum_datasets.setdefault(quantum_datasets.quantum_id, quantum_datasets)
2074 self.address_reader.read_all()
2075 return
2076 with MultiblockReader.open_in_zip(
2077 self.zf, "quantum_datasets", int_size=self.components.header.int_size
2078 ) as mb_reader:
2079 for quantum_id in quantum_ids:
2080 if quantum_id in self.components.quantum_datasets:
2081 continue
2082 address_row = self.address_reader.find(quantum_id)
2083 quantum_datasets = mb_reader.read_model(
2084 address_row.addresses[0], PredictedQuantumDatasetsModel, self.decompressor
2085 )
2086 if quantum_datasets is not None:
2087 self.components.quantum_datasets[address_row.key] = quantum_datasets
2088 return
2090 def read_execution_quanta(self, quantum_ids: Iterable[uuid.UUID] | None = None) -> None:
2091 """Read all information needed to execute the given quanta.
2093 Parameters
2094 ----------
2095 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional
2096 Iterable of quantum IDs to load. If not provided, all quanta will
2097 be loaded. The UUIDs of special init quanta will be ignored.
2098 """
2099 self.read_init_quanta()
2100 self.read_dimension_data()
2101 self.read_quantum_datasets(quantum_ids)