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