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