Coverage for python / lsst / pipe / base / pipeline_graph / io.py: 45%
213 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-20 01:37 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29__all__ = (
30 "SerializedDatasetTypeNode",
31 "SerializedEdge",
32 "SerializedPipelineGraph",
33 "SerializedTaskInitNode",
34 "SerializedTaskNode",
35 "SerializedTaskSubset",
36 "expect_not_none",
37)
39from collections.abc import Mapping
40from typing import Any, TypeVar
42import networkx
43import pydantic
45from lsst.daf.butler import DatasetType, DimensionConfig, DimensionGroup, DimensionUniverse
47from .. import automatic_connection_constants as acc
48from ._dataset_types import DatasetTypeNode
49from ._edges import Edge, ReadEdge, WriteEdge
50from ._exceptions import PipelineGraphReadError
51from ._nodes import NodeKey, NodeType
52from ._pipeline_graph import PipelineGraph
53from ._task_subsets import StepDefinitions, TaskSubset
54from ._tasks import TaskImportMode, TaskInitNode, TaskNode
56_U = TypeVar("_U")
58_IO_VERSION_INFO = (0, 0, 1)
59"""Version tuple embedded in saved PipelineGraphs.
60"""
63def expect_not_none(value: _U | None, msg: str) -> _U:
64 """Check that a value is not `None` and return it.
66 Parameters
67 ----------
68 value : `~typing.Any`
69 Value to check.
70 msg : `str`
71 Error message for the case where ``value is None``.
73 Returns
74 -------
75 value : `typing.Any`
76 Value, guaranteed not to be `None`.
78 Raises
79 ------
80 PipelineGraphReadError
81 Raised with ``msg`` if ``value is None``.
82 """
83 if value is None:
84 raise PipelineGraphReadError(msg)
85 return value
88class SerializedEdge(pydantic.BaseModel):
89 """Struct used to represent a serialized `Edge` in a `PipelineGraph`.
91 All `ReadEdge` and `WriteEdge` state not included here is instead
92 effectively serialized by the context in which a `SerializedEdge` appears
93 (e.g. the keys of the nested dictionaries in which it serves as the value
94 type).
95 """
97 dataset_type_name: str
98 """Full dataset type name (including component)."""
100 storage_class: str
101 """Name of the storage class."""
103 raw_dimensions: list[str]
104 """Raw dimensions of the dataset type from the task connections."""
106 is_calibration: bool = False
107 """Whether this dataset type can be included in
108 `~lsst.daf.butler.CollectionType.CALIBRATION` collections."""
110 defer_query_constraint: bool = False
111 """If `True`, by default do not include this dataset type's existence as a
112 constraint on the initial data ID query in QuantumGraph generation."""
114 @classmethod
115 def serialize(cls, target: Edge) -> SerializedEdge:
116 """Transform an `Edge` to a `SerializedEdge`.
118 Parameters
119 ----------
120 target : `Edge`
121 The object to serialize.
123 Returns
124 -------
125 `SerializedEdge`
126 Model transformed into something that can be serialized.
127 """
128 return SerializedEdge.model_construct(
129 storage_class=target.storage_class_name,
130 dataset_type_name=target.dataset_type_name,
131 raw_dimensions=sorted(target.raw_dimensions),
132 is_calibration=target.is_calibration,
133 defer_query_constraint=getattr(target, "defer_query_constraint", False),
134 )
136 def deserialize_read_edge(
137 self,
138 task_key: NodeKey,
139 connection_name: str,
140 dataset_type_keys: Mapping[str, NodeKey],
141 is_prerequisite: bool = False,
142 ) -> ReadEdge:
143 """Transform a `SerializedEdge` to a `ReadEdge`.
145 Parameters
146 ----------
147 task_key : `NodeKey`
148 Key for the task node this edge is connected to.
149 connection_name : `str`
150 Internal name for the connection as seen by the task.
151 dataset_type_keys : `~collections.abc.Mapping` [`str`, `NodeKey`]
152 Mapping of dataset type name to node key.
153 is_prerequisite : `bool`, optional
154 Whether this dataset must be present in the data repository prior
155 to `QuantumGraph` generation.
157 Returns
158 -------
159 `ReadEdge`
160 Deserialized object.
161 """
162 parent_dataset_type_name, component = DatasetType.splitDatasetTypeName(self.dataset_type_name)
163 return ReadEdge(
164 dataset_type_key=dataset_type_keys[parent_dataset_type_name],
165 task_key=task_key,
166 storage_class_name=self.storage_class,
167 is_prerequisite=is_prerequisite,
168 component=component,
169 connection_name=connection_name,
170 is_calibration=self.is_calibration,
171 defer_query_constraint=self.defer_query_constraint,
172 raw_dimensions=frozenset(self.raw_dimensions),
173 )
175 def deserialize_write_edge(
176 self,
177 task_key: NodeKey,
178 connection_name: str,
179 dataset_type_keys: Mapping[str, NodeKey],
180 ) -> WriteEdge:
181 """Transform a `SerializedEdge` to a `WriteEdge`.
183 Parameters
184 ----------
185 task_key : `NodeKey`
186 Key for the task node this edge is connected to.
187 connection_name : `str`
188 Internal name for the connection as seen by the task.
189 dataset_type_keys : `~collections.abc.Mapping` [`str`, `NodeKey`]
190 Mapping of dataset type name to node key.
192 Returns
193 -------
194 `WriteEdge`
195 Deserialized object.
196 """
197 return WriteEdge(
198 task_key=task_key,
199 dataset_type_key=dataset_type_keys[self.dataset_type_name],
200 storage_class_name=self.storage_class,
201 connection_name=connection_name,
202 is_calibration=self.is_calibration,
203 raw_dimensions=frozenset(self.raw_dimensions),
204 )
207class SerializedTaskInitNode(pydantic.BaseModel):
208 """Struct used to represent a serialized `TaskInitNode` in a
209 `PipelineGraph`.
211 The task label is serialized by the context in which a
212 `SerializedTaskInitNode` appears (e.g. the keys of the nested dictionary
213 in which it serves as the value type), and the task class name and config
214 string are save with the corresponding `SerializedTaskNode`.
215 """
217 inputs: dict[str, SerializedEdge]
218 """Mapping of serialized init-input edges, keyed by connection name."""
220 outputs: dict[str, SerializedEdge]
221 """Mapping of serialized init-output edges, keyed by connection name."""
223 config_output: SerializedEdge
224 """The serialized config init-output edge."""
226 index: int | None = None
227 """The index of this node in the sorted sequence of `PipelineGraph`.
229 This is `None` if the `PipelineGraph` was not sorted when it was
230 serialized.
231 """
233 @classmethod
234 def serialize(cls, target: TaskInitNode) -> SerializedTaskInitNode:
235 """Transform a `TaskInitNode` to a `SerializedTaskInitNode`.
237 Parameters
238 ----------
239 target : `TaskInitNode`
240 Object to be serialized.
242 Returns
243 -------
244 `SerializedTaskInitNode`
245 Model that can be serialized.
246 """
247 return cls.model_construct(
248 inputs={
249 connection_name: SerializedEdge.serialize(edge)
250 for connection_name, edge in sorted(target.inputs.items())
251 },
252 outputs={
253 connection_name: SerializedEdge.serialize(edge)
254 for connection_name, edge in sorted(target.outputs.items())
255 },
256 config_output=SerializedEdge.serialize(target.config_output),
257 )
259 def deserialize(
260 self,
261 key: NodeKey,
262 task_class_name: str,
263 config_str: str,
264 dataset_type_keys: Mapping[str, NodeKey],
265 ) -> TaskInitNode:
266 """Transform a `SerializedTaskInitNode` to a `TaskInitNode`.
268 Parameters
269 ----------
270 key : `NodeKey`
271 Key that identifies this node in internal and exported networkx
272 graphs.
273 task_class_name : `str`, optional
274 Fully-qualified name of the task class. Must be provided if
275 ``imported_data`` is not.
276 config_str : `str`, optional
277 Configuration for the task as a string of override statements.
278 dataset_type_keys : `~collections.abc.Mapping` [`str`, `NodeKey`]
279 Mapping of dataset type name to node key.
281 Returns
282 -------
283 `TaskInitNode`
284 Deserialized object.
285 """
286 return TaskInitNode(
287 key,
288 inputs={
289 connection_name: serialized_edge.deserialize_read_edge(
290 key, connection_name, dataset_type_keys
291 )
292 for connection_name, serialized_edge in self.inputs.items()
293 },
294 outputs={
295 connection_name: serialized_edge.deserialize_write_edge(
296 key, connection_name, dataset_type_keys
297 )
298 for connection_name, serialized_edge in self.outputs.items()
299 },
300 config_output=self.config_output.deserialize_write_edge(
301 key, acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME, dataset_type_keys
302 ),
303 task_class_name=task_class_name,
304 config_str=config_str,
305 )
308class SerializedTaskNode(pydantic.BaseModel):
309 """Struct used to represent a serialized `TaskNode` in a `PipelineGraph`.
311 The task label is serialized by the context in which a
312 `SerializedTaskNode` appears (e.g. the keys of the nested dictionary in
313 which it serves as the value type).
314 """
316 task_class: str
317 """Fully-qualified name of the task class."""
319 init: SerializedTaskInitNode
320 """Serialized task initialization node."""
322 config_str: str
323 """Configuration for the task as a string of override statements."""
325 prerequisite_inputs: dict[str, SerializedEdge]
326 """Mapping of serialized prerequisiste input edges, keyed by connection
327 name.
328 """
330 inputs: dict[str, SerializedEdge]
331 """Mapping of serialized input edges, keyed by connection name."""
333 outputs: dict[str, SerializedEdge]
334 """Mapping of serialized output edges, keyed by connection name."""
336 metadata_output: SerializedEdge
337 """The serialized metadata output edge."""
339 dimensions: list[str]
340 """The task's dimensions, if they were resolved."""
342 log_output: SerializedEdge | None = None
343 """The serialized log output edge."""
345 index: int | None = None
346 """The index of this node in the sorted sequence of `PipelineGraph`.
348 This is `None` if the `PipelineGraph` was not sorted when it was
349 serialized.
350 """
352 @classmethod
353 def serialize(cls, target: TaskNode) -> SerializedTaskNode:
354 """Transform a `TaskNode` to a `SerializedTaskNode`.
356 Parameters
357 ----------
358 target : `TaskNode`
359 Object to be serialized.
361 Returns
362 -------
363 `SerializedTaskNode`
364 Object that can be serialized.
365 """
366 dimensions = list(target.raw_dimensions)
367 dimensions.sort()
368 return cls.model_construct(
369 task_class=target.task_class_name,
370 init=SerializedTaskInitNode.serialize(target.init),
371 config_str=target.get_config_str(),
372 dimensions=dimensions,
373 prerequisite_inputs={
374 connection_name: SerializedEdge.serialize(edge)
375 for connection_name, edge in sorted(target.prerequisite_inputs.items())
376 },
377 inputs={
378 connection_name: SerializedEdge.serialize(edge)
379 for connection_name, edge in sorted(target.inputs.items())
380 },
381 outputs={
382 connection_name: SerializedEdge.serialize(edge)
383 for connection_name, edge in sorted(target.outputs.items())
384 },
385 metadata_output=SerializedEdge.serialize(target.metadata_output),
386 log_output=(
387 SerializedEdge.serialize(target.log_output) if target.log_output is not None else None
388 ),
389 )
391 def deserialize(
392 self,
393 key: NodeKey,
394 init_key: NodeKey,
395 dataset_type_keys: Mapping[str, NodeKey],
396 universe: DimensionUniverse | None,
397 ) -> TaskNode:
398 """Transform a `SerializedTaskNode` to a `TaskNode`.
400 Parameters
401 ----------
402 key : `NodeKey`
403 Key that identifies this node in internal and exported networkx
404 graphs.
405 init_key : `NodeKey`
406 Key that identifies the init node in internal and exported networkx
407 graphs.
408 dataset_type_keys : `~collections.abc.Mapping` [`str`, `NodeKey`]
409 Mapping of dataset type name to node key.
410 universe : `~lsst.daf.butler.DimensionUniverse` or `None`
411 The dimension universe.
413 Returns
414 -------
415 `TaskNode`
416 Deserialized object.
417 """
418 init = self.init.deserialize(
419 init_key,
420 task_class_name=self.task_class,
421 config_str=expect_not_none(
422 self.config_str, f"No serialized config file for task with label {key.name!r}."
423 ),
424 dataset_type_keys=dataset_type_keys,
425 )
426 inputs = {
427 connection_name: serialized_edge.deserialize_read_edge(key, connection_name, dataset_type_keys)
428 for connection_name, serialized_edge in self.inputs.items()
429 }
430 prerequisite_inputs = {
431 connection_name: serialized_edge.deserialize_read_edge(
432 key, connection_name, dataset_type_keys, is_prerequisite=True
433 )
434 for connection_name, serialized_edge in self.prerequisite_inputs.items()
435 }
436 outputs = {
437 connection_name: serialized_edge.deserialize_write_edge(key, connection_name, dataset_type_keys)
438 for connection_name, serialized_edge in self.outputs.items()
439 }
440 if (serialized_log_output := self.log_output) is not None:
441 log_output = serialized_log_output.deserialize_write_edge(
442 key, acc.LOG_OUTPUT_CONNECTION_NAME, dataset_type_keys
443 )
444 else:
445 log_output = None
446 metadata_output = self.metadata_output.deserialize_write_edge(
447 key, acc.METADATA_OUTPUT_CONNECTION_NAME, dataset_type_keys
448 )
449 dimensions: frozenset[str] | DimensionGroup
450 if universe is not None:
451 dimensions = universe.conform(self.dimensions)
452 else:
453 dimensions = frozenset(self.dimensions)
454 return TaskNode(
455 key=key,
456 init=init,
457 inputs=inputs,
458 prerequisite_inputs=prerequisite_inputs,
459 outputs=outputs,
460 log_output=log_output,
461 metadata_output=metadata_output,
462 dimensions=dimensions,
463 )
466class SerializedDatasetTypeNode(pydantic.BaseModel):
467 """Struct used to represent a serialized `DatasetTypeNode` in a
468 `PipelineGraph`.
470 Unresolved dataset types are serialized as instances with at most the
471 `index` attribute set, and are typically converted to JSON with pydantic's
472 ``exclude_defaults=True`` option to keep this compact.
474 The dataset typename is serialized by the context in which a
475 `SerializedDatasetTypeNode` appears (e.g. the keys of the nested dictionary
476 in which it serves as the value type).
477 """
479 dimensions: list[str] | None = None
480 """Dimensions of the dataset type."""
482 storage_class: str | None = None
483 """Name of the storage class."""
485 is_calibration: bool = False
486 """Whether this dataset type is a calibration."""
488 is_initial_query_constraint: bool = False
489 """Whether this dataset type should be a query constraint during
490 `QuantumGraph` generation."""
492 is_prerequisite: bool = False
493 """Whether datasets of this dataset type must exist in the input collection
494 before `QuantumGraph` generation."""
496 index: int | None = None
497 """The index of this node in the sorted sequence of `PipelineGraph`.
499 This is `None` if the `PipelineGraph` was not sorted when it was
500 serialized.
501 """
503 @classmethod
504 def serialize(cls, target: DatasetTypeNode | None) -> SerializedDatasetTypeNode:
505 """Transform a `DatasetTypeNode` to a `SerializedDatasetTypeNode`.
507 Parameters
508 ----------
509 target : `DatasetTypeNode` or `None`
510 Object to serialize.
512 Returns
513 -------
514 `SerializedDatasetTypeNode`
515 Object in serializable form.
516 """
517 if target is None:
518 return cls.model_construct()
519 return cls.model_construct(
520 dimensions=list(target.dataset_type.dimensions.names),
521 storage_class=target.dataset_type.storageClass_name,
522 is_calibration=target.dataset_type.isCalibration(),
523 is_initial_query_constraint=target.is_initial_query_constraint,
524 is_prerequisite=target.is_prerequisite,
525 )
527 def deserialize(
528 self, key: NodeKey, xgraph: networkx.MultiDiGraph, universe: DimensionUniverse | None
529 ) -> DatasetTypeNode | None:
530 """Transform a `SerializedDatasetTypeNode` to a `DatasetTypeNode`.
532 Parameters
533 ----------
534 key : `NodeKey`
535 Key that identifies this node in internal and exported networkx
536 graphs.
537 xgraph : `networkx.MultiDiGraph`
538 <unknown>.
539 universe : `~lsst.daf.butler.DimensionUniverse` or `None`
540 The dimension universe.
542 Returns
543 -------
544 `DatasetTypeNode`
545 Deserialized object.
546 """
547 if self.dimensions is not None:
548 dataset_type = DatasetType(
549 key.name,
550 expect_not_none(
551 self.dimensions,
552 f"Serialized dataset type {key.name!r} has no dimensions.",
553 ),
554 storageClass=expect_not_none(
555 self.storage_class,
556 f"Serialized dataset type {key.name!r} has no storage class.",
557 ),
558 isCalibration=self.is_calibration,
559 universe=expect_not_none(
560 universe,
561 f"Serialized dataset type {key.name!r} has dimensions, "
562 "but no dimension universe was stored.",
563 ),
564 )
565 producer: str | None = None
566 producing_edge: WriteEdge | None = None
567 for _, _, producing_edge in xgraph.in_edges(key, data="instance"):
568 assert producing_edge is not None, "Should only be None if we never loop."
569 if producer is not None:
570 raise PipelineGraphReadError(
571 f"Serialized dataset type {key.name!r} is produced by both "
572 f"{producing_edge.task_label!r} and {producer!r} in resolved graph."
573 )
574 producer = producing_edge.task_label
575 consuming_edges = tuple(
576 consuming_edge for _, _, consuming_edge in xgraph.in_edges(key, data="instance")
577 )
578 return DatasetTypeNode(
579 dataset_type=dataset_type,
580 is_prerequisite=self.is_prerequisite,
581 is_initial_query_constraint=self.is_initial_query_constraint,
582 producing_edge=producing_edge,
583 consuming_edges=consuming_edges,
584 )
585 return None
588class SerializedTaskSubset(pydantic.BaseModel):
589 """Struct used to represent a serialized `TaskSubset` in a `PipelineGraph`.
591 The subset label is serialized by the context in which a
592 `SerializedDatasetTypeNode` appears (e.g. the keys of the nested dictionary
593 in which it serves as the value type).
594 """
596 description: str
597 """Description of the subset."""
599 tasks: list[str]
600 """Labels of tasks in the subset, sorted lexicographically for
601 determinism.
602 """
604 dimensions: list[str]
605 """Dimensions that can be used to divide this step's quanta into
606 independent groups.
607 """
609 @classmethod
610 def serialize(cls, target: TaskSubset) -> SerializedTaskSubset:
611 """Transform a `TaskSubset` into a `SerializedTaskSubset`.
613 Parameters
614 ----------
615 target : `TaskSubset`
616 Object to serialize.
618 Returns
619 -------
620 `SerializedTaskSubset`
621 Object in serializable form.
622 """
623 dimensions = sorted(target._step_definitions._dimensions_by_label.get(target.label, ()))
624 return cls.model_construct(
625 description=target._description,
626 tasks=sorted(target),
627 dimensions=dimensions,
628 )
630 def deserialize_task_subset(
631 self, label: str, xgraph: networkx.MultiDiGraph, steps: StepDefinitions
632 ) -> TaskSubset:
633 """Transform a `SerializedTaskSubset` into a `TaskSubset`.
635 Parameters
636 ----------
637 label : `str`
638 Subset label.
639 xgraph : `networkx.MultiDiGraph`
640 The under-construction networkx graph that backs the pipeline
641 graph.
642 steps : `StepDefinitions`
643 Step definitions for the pipeline graph. Modified in-place to
644 set sharding dimension if this is a step.
646 Returns
647 -------
648 `TaskSubset`
649 Deserialized object.
650 """
651 members = set(self.tasks)
652 if label in steps:
653 steps._dimensions_by_label[label] = frozenset(self.dimensions)
654 return TaskSubset(xgraph, label, members, self.description, steps)
657class SerializedPipelineGraph(pydantic.BaseModel):
658 """Struct used to represent a serialized `PipelineGraph`."""
660 version: str = ".".join(str(v) for v in _IO_VERSION_INFO)
661 """Serialization version."""
663 description: str
664 """Human-readable description of the pipeline."""
666 tasks: dict[str, SerializedTaskNode] = pydantic.Field(default_factory=dict)
667 """Mapping of serialized tasks, keyed by label."""
669 dataset_types: dict[str, SerializedDatasetTypeNode] = pydantic.Field(default_factory=dict)
670 """Mapping of serialized dataset types, keyed by parent dataset type name.
671 """
673 task_subsets: dict[str, SerializedTaskSubset] = pydantic.Field(default_factory=dict)
674 """Mapping of task subsets, keyed by subset label."""
676 dimensions: dict[str, Any] | None = None
677 """Dimension universe configuration."""
679 data_id: dict[str, Any] = pydantic.Field(default_factory=dict)
680 """Data ID that constrains all quanta generated from this pipeline."""
682 step_labels: list[str] = pydantic.Field(default_factory=list)
683 """List of task subset labels that are steps."""
685 steps_verified: bool
686 """Whether the step definitions in this pipeline were checked for
687 consistency with the task and subset definitions.
688 """
690 @classmethod
691 def serialize(cls, target: PipelineGraph) -> SerializedPipelineGraph:
692 """Transform a `PipelineGraph` into a `SerializedPipelineGraph`.
694 Parameters
695 ----------
696 target : `PipelineGraph`
697 Object to serialize.
699 Returns
700 -------
701 `SerializedPipelineGraph`
702 Object in serializable form.
703 """
704 result = SerializedPipelineGraph.model_construct(
705 description=target.description,
706 tasks={label: SerializedTaskNode.serialize(node) for label, node in target.tasks.items()},
707 dataset_types={
708 name: SerializedDatasetTypeNode().serialize(target.dataset_types.get_if_resolved(name))
709 for name in target.dataset_types
710 },
711 task_subsets={
712 label: SerializedTaskSubset.serialize(subset) for label, subset in target.task_subsets.items()
713 },
714 step_labels=list(target.steps),
715 steps_verified=target.steps.verified,
716 dimensions=target._universe.dimensionConfig.toDict() if target._universe is not None else None,
717 data_id=target._raw_data_id,
718 )
719 if target._sorted_keys:
720 for index, node_key in enumerate(target._sorted_keys):
721 match node_key.node_type:
722 case NodeType.TASK:
723 result.tasks[node_key.name].index = index
724 case NodeType.DATASET_TYPE:
725 result.dataset_types[node_key.name].index = index
726 case NodeType.TASK_INIT:
727 result.tasks[node_key.name].init.index = index
728 return result
730 def deserialize(
731 self,
732 import_mode: TaskImportMode,
733 ) -> PipelineGraph:
734 """Transform a `SerializedPipelineGraph` into a `PipelineGraph`.
736 Parameters
737 ----------
738 import_mode : `TaskImportMode`
739 Import mode.
741 Returns
742 -------
743 `PipelineGraph`
744 Deserialized object.
745 """
746 universe: DimensionUniverse | None = None
747 if self.dimensions is not None:
748 universe = DimensionUniverse(
749 config=DimensionConfig(
750 expect_not_none(
751 self.dimensions,
752 "Serialized pipeline graph has not been resolved; "
753 "load it is a MutablePipelineGraph instead.",
754 )
755 )
756 )
757 xgraph = networkx.MultiDiGraph()
758 sort_index_map: dict[int, NodeKey] = {}
759 # Save the dataset type keys after the first time we make them - these
760 # may be tiny objects, but it's still to have only one copy of each
761 # value floating around the graph.
762 dataset_type_keys: dict[str, NodeKey] = {}
763 for dataset_type_name, serialized_dataset_type in self.dataset_types.items():
764 dataset_type_key = NodeKey(NodeType.DATASET_TYPE, dataset_type_name)
765 # We intentionally don't attach a DatasetTypeNode instance here
766 # yet, since we need edges to do that and those are saved with
767 # the tasks.
768 xgraph.add_node(dataset_type_key, bipartite=NodeType.DATASET_TYPE.value)
769 if serialized_dataset_type.index is not None:
770 sort_index_map[serialized_dataset_type.index] = dataset_type_key
771 dataset_type_keys[dataset_type_name] = dataset_type_key
772 for task_label, serialized_task in self.tasks.items():
773 task_key = NodeKey(NodeType.TASK, task_label)
774 task_init_key = NodeKey(NodeType.TASK_INIT, task_label)
775 task_node = serialized_task.deserialize(task_key, task_init_key, dataset_type_keys, universe)
776 if serialized_task.index is not None:
777 sort_index_map[serialized_task.index] = task_key
778 if serialized_task.init.index is not None:
779 sort_index_map[serialized_task.init.index] = task_init_key
780 xgraph.add_node(task_key, instance=task_node, bipartite=NodeType.TASK.bipartite)
781 xgraph.add_node(task_init_key, instance=task_node.init, bipartite=NodeType.TASK_INIT.bipartite)
782 xgraph.add_edge(task_init_key, task_key, Edge.INIT_TO_TASK_NAME, instance=None)
783 for read_edge in task_node.init.iter_all_inputs():
784 xgraph.add_edge(
785 read_edge.dataset_type_key,
786 read_edge.task_key,
787 read_edge.connection_name,
788 instance=read_edge,
789 )
790 for write_edge in task_node.init.iter_all_outputs():
791 xgraph.add_edge(
792 write_edge.task_key,
793 write_edge.dataset_type_key,
794 write_edge.connection_name,
795 instance=write_edge,
796 )
797 for read_edge in task_node.iter_all_inputs():
798 xgraph.add_edge(
799 read_edge.dataset_type_key,
800 read_edge.task_key,
801 read_edge.connection_name,
802 instance=read_edge,
803 )
804 for write_edge in task_node.iter_all_outputs():
805 xgraph.add_edge(
806 write_edge.task_key,
807 write_edge.dataset_type_key,
808 write_edge.connection_name,
809 instance=write_edge,
810 )
811 # Iterate over dataset types again to add instances.
812 for dataset_type_name, serialized_dataset_type in self.dataset_types.items():
813 dataset_type_key = dataset_type_keys[dataset_type_name]
814 xgraph.nodes[dataset_type_key]["instance"] = serialized_dataset_type.deserialize(
815 dataset_type_key, xgraph, universe
816 )
817 steps = StepDefinitions(
818 universe,
819 dimensions_by_label=dict.fromkeys(self.step_labels, frozenset()),
820 verified=self.steps_verified,
821 )
822 result = PipelineGraph.__new__(PipelineGraph)
823 result._init_from_args(
824 xgraph,
825 sorted_keys=[sort_index_map[i] for i in range(len(xgraph))] if sort_index_map else None,
826 task_subsets={
827 subset_label: serialized_subset.deserialize_task_subset(subset_label, xgraph, steps)
828 for subset_label, serialized_subset in self.task_subsets.items()
829 },
830 description=self.description,
831 step_definitions=steps,
832 universe=universe,
833 data_id=self.data_id,
834 )
835 result._import_and_configure(import_mode)
836 return result