Coverage for python/lsst/pipe/base/graph/_versionDeserializers.py: 26%

257 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-03-22 02:08 -0700

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 program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21from __future__ import annotations 

22 

23__all__ = ("DESERIALIZER_MAP",) 

24 

25import json 

26import lzma 

27import pickle 

28import struct 

29import uuid 

30from abc import ABC, abstractmethod 

31from collections import defaultdict 

32from dataclasses import dataclass 

33from types import SimpleNamespace 

34from typing import ( 

35 TYPE_CHECKING, 

36 Callable, 

37 ClassVar, 

38 DefaultDict, 

39 Dict, 

40 List, 

41 Optional, 

42 Set, 

43 Tuple, 

44 Type, 

45 cast, 

46) 

47 

48import networkx as nx 

49from lsst.daf.butler import ( 

50 DatasetRef, 

51 DatasetType, 

52 DimensionConfig, 

53 DimensionRecord, 

54 DimensionUniverse, 

55 Quantum, 

56 SerializedDimensionRecord, 

57) 

58from lsst.utils import doImportType 

59 

60from ..config import PipelineTaskConfig 

61from ..pipeline import TaskDef 

62from ..pipelineTask import PipelineTask 

63from ._implDetails import DatasetTypeName, _DatasetTracker 

64from .quantumNode import QuantumNode, SerializedQuantumNode 

65 

66if TYPE_CHECKING: 66 ↛ 67line 66 didn't jump to line 67, because the condition on line 66 was never true

67 from .graph import QuantumGraph 

68 

69 

70class StructSizeDescriptor: 

71 """This is basically a class level property. It exists to report the size 

72 (number of bytes) of whatever the formatter string is for a deserializer 

73 """ 

74 

75 def __get__(self, inst: Optional[DeserializerBase], owner: Type[DeserializerBase]) -> int: 

76 return struct.calcsize(owner.FMT_STRING()) 

77 

78 

79@dataclass 

80class DeserializerBase(ABC): 

81 @classmethod 

82 @abstractmethod 

83 def FMT_STRING(cls) -> str: # noqa: N805 # flake8 wants self 

84 raise NotImplementedError("Base class does not implement this method") 

85 

86 structSize: ClassVar[StructSizeDescriptor] 

87 

88 preambleSize: int 

89 sizeBytes: bytes 

90 

91 def __init_subclass__(cls) -> None: 

92 # attach the size decriptor 

93 cls.structSize = StructSizeDescriptor() 

94 super().__init_subclass__() 

95 

96 def unpackHeader(self, rawHeader: bytes) -> Optional[str]: 

97 """Transforms the raw bytes corresponding to the header of a save into 

98 a string of the header information. Returns none if the save format has 

99 no header string implementation (such as save format 1 that is all 

100 pickle) 

101 

102 Parameters 

103 ---------- 

104 rawheader : bytes 

105 The bytes that are to be parsed into the header information. These 

106 are the bytes after the preamble and structsize number of bytes 

107 and before the headerSize bytes 

108 """ 

109 raise NotImplementedError("Base class does not implement this method") 

110 

111 @property 

112 def headerSize(self) -> int: 

113 """Returns the number of bytes from the beginning of the file to the 

114 end of the metadata. 

115 """ 

116 raise NotImplementedError("Base class does not implement this method") 

117 

118 def readHeaderInfo(self, rawHeader: bytes) -> SimpleNamespace: 

119 """Parse the supplied raw bytes into the header information and 

120 byte ranges of specific TaskDefs and QuantumNodes 

121 

122 Parameters 

123 ---------- 

124 rawheader : bytes 

125 The bytes that are to be parsed into the header information. These 

126 are the bytes after the preamble and structsize number of bytes 

127 and before the headerSize bytes 

128 """ 

129 raise NotImplementedError("Base class does not implement this method") 

130 

131 def constructGraph( 

132 self, 

133 nodes: set[uuid.UUID], 

134 _readBytes: Callable[[int, int], bytes], 

135 universe: Optional[DimensionUniverse] = None, 

136 ) -> QuantumGraph: 

137 """Constructs a graph from the deserialized information. 

138 

139 Parameters 

140 ---------- 

141 nodes : `set` of `uuid.UUID` 

142 The nodes to include in the graph 

143 _readBytes : callable 

144 A callable that can be used to read bytes from the file handle. 

145 The callable will take two ints, start and stop, to use as the 

146 numerical bounds to read and returns a byte stream. 

147 universe : `~lsst.daf.butler.DimensionUniverse` 

148 The singleton of all dimensions known to the middleware registry 

149 """ 

150 raise NotImplementedError("Base class does not implement this method") 

151 

152 def description(self) -> str: 

153 """Return the description of the serialized data format""" 

154 raise NotImplementedError("Base class does not implement this method") 

155 

156 

157Version1Description = """ 

158The save file starts with the first few bytes corresponding to the magic bytes 

159in the QuantumGraph: `qgraph4\xf6\xe8\xa9`. 

160 

161The next few bytes are 2 big endian unsigned 64 bit integers. 

162 

163The first unsigned 64 bit integer corresponds to the number of bytes of a 

164python mapping of TaskDef labels to the byte ranges in the save file where the 

165definition can be loaded. 

166 

167The second unsigned 64 bit integer corrresponds to the number of bytes of a 

168python mapping of QuantumGraph Node number to the byte ranges in the save file 

169where the node can be loaded. The byte range is indexed starting after 

170the `header` bytes of the magic bytes, size bytes, and bytes of the two 

171mappings. 

172 

173Each of the above mappings are pickled and then lzma compressed, so to 

174deserialize the bytes, first lzma decompression must be performed and the 

175results passed to python pickle loader. 

176 

177As stated above, each map contains byte ranges of the corresponding 

178datastructure. Theses bytes are also lzma compressed pickles, and should 

179be deserialized in a similar manner. The byte range is indexed starting after 

180the `header` bytes of the magic bytes, size bytes, and bytes of the two 

181mappings. 

182 

183In addition to the the TaskDef byte locations, the TypeDef map also contains 

184an additional key '__GraphBuildID'. The value associated with this is the 

185unique id assigned to the graph at its creation time. 

186""" 

187 

188 

189@dataclass 

190class DeserializerV1(DeserializerBase): 

191 @classmethod 

192 def FMT_STRING(cls) -> str: 

193 return ">QQ" 

194 

195 def __post_init__(self) -> None: 

196 self.taskDefMapSize, self.nodeMapSize = struct.unpack(self.FMT_STRING(), self.sizeBytes) 

197 

198 @property 

199 def headerSize(self) -> int: 

200 return self.preambleSize + self.structSize + self.taskDefMapSize + self.nodeMapSize 

201 

202 def readHeaderInfo(self, rawHeader: bytes) -> SimpleNamespace: 

203 returnValue = SimpleNamespace() 

204 returnValue.taskDefMap = pickle.loads(rawHeader[: self.taskDefMapSize]) 

205 returnValue._buildId = returnValue.taskDefMap["__GraphBuildID"] 

206 returnValue.map = pickle.loads(rawHeader[self.taskDefMapSize :]) 

207 returnValue.metadata = None 

208 self.returnValue = returnValue 

209 return returnValue 

210 

211 def unpackHeader(self, rawHeader: bytes) -> Optional[str]: 

212 return None 

213 

214 def constructGraph( 

215 self, 

216 nodes: set[uuid.UUID], 

217 _readBytes: Callable[[int, int], bytes], 

218 universe: Optional[DimensionUniverse] = None, 

219 ) -> QuantumGraph: 

220 # need to import here to avoid cyclic imports 

221 from . import QuantumGraph 

222 

223 quanta: DefaultDict[TaskDef, Set[Quantum]] = defaultdict(set) 

224 quantumToNodeId: Dict[Quantum, uuid.UUID] = {} 

225 loadedTaskDef = {} 

226 # loop over the nodes specified above 

227 for node in nodes: 

228 # Get the bytes to read from the map 

229 start, stop = self.returnValue.map[node] 

230 start += self.headerSize 

231 stop += self.headerSize 

232 

233 # read the specified bytes, will be overloaded by subclasses 

234 # bytes are compressed, so decompress them 

235 dump = lzma.decompress(_readBytes(start, stop)) 

236 

237 # reconstruct node 

238 qNode = pickle.loads(dump) 

239 object.__setattr__(qNode, "nodeId", uuid.uuid4()) 

240 

241 # read the saved node, name. If it has been loaded, attach it, if 

242 # not read in the taskDef first, and then load it 

243 nodeTask = qNode.taskDef 

244 if nodeTask not in loadedTaskDef: 

245 # Get the byte ranges corresponding to this taskDef 

246 start, stop = self.returnValue.taskDefMap[nodeTask] 

247 start += self.headerSize 

248 stop += self.headerSize 

249 

250 # load the taskDef, this method call will be overloaded by 

251 # subclasses. 

252 # bytes are compressed, so decompress them 

253 taskDef = pickle.loads(lzma.decompress(_readBytes(start, stop))) 

254 loadedTaskDef[nodeTask] = taskDef 

255 # Explicitly overload the "frozen-ness" of nodes to attach the 

256 # taskDef back into the un-persisted node 

257 object.__setattr__(qNode, "taskDef", loadedTaskDef[nodeTask]) 

258 quanta[qNode.taskDef].add(qNode.quantum) 

259 

260 # record the node for later processing 

261 quantumToNodeId[qNode.quantum] = qNode.nodeId 

262 

263 # construct an empty new QuantumGraph object, and run the associated 

264 # creation method with the un-persisted data 

265 qGraph = object.__new__(QuantumGraph) 

266 qGraph._buildGraphs( 

267 quanta, 

268 _quantumToNodeId=quantumToNodeId, 

269 _buildId=self.returnValue._buildId, 

270 metadata=self.returnValue.metadata, 

271 universe=universe, 

272 ) 

273 return qGraph 

274 

275 def description(self) -> str: 

276 return Version1Description 

277 

278 

279Version2Description = """ 

280The save file starts with the first few bytes corresponding to the magic bytes 

281in the QuantumGraph: `qgraph4\xf6\xe8\xa9`. 

282 

283The next few bytes are a big endian unsigned long long. 

284 

285The unsigned long long corresponds to the number of bytes of a python mapping 

286of header information. This mapping is encoded into json and then lzma 

287compressed, meaning the operations must be performed in the opposite order to 

288deserialize. 

289 

290The json encoded header mapping contains 4 fields: TaskDefs, GraphBuildId, 

291Nodes, and Metadata. 

292 

293The `TaskDefs` key corresponds to a value which is a mapping of Task label to 

294task data. The task data is a mapping of key to value, where the only key is 

295`bytes` and it corresponds to a tuple of a byte range of the start, stop 

296bytes (indexed after all the header bytes) 

297 

298The `GraphBuildId` corresponds with a string that is the unique id assigned to 

299this graph when it was created. 

300 

301The `Nodes` key is like the `TaskDefs` key except it corresponds to 

302QuantumNodes instead of TaskDefs. Another important difference is that JSON 

303formatting does not allow using numbers as keys, and this mapping is keyed by 

304the node number. Thus it is stored in JSON as two equal length lists, the first 

305being the keys, and the second the values associated with those keys. 

306 

307The `Metadata` key is a mapping of strings to associated values. This metadata 

308may be anything that is important to be transported alongside the graph. 

309 

310As stated above, each map contains byte ranges of the corresponding 

311datastructure. Theses bytes are also lzma compressed pickles, and should 

312be deserialized in a similar manner. 

313""" 

314 

315 

316@dataclass 

317class DeserializerV2(DeserializerBase): 

318 @classmethod 

319 def FMT_STRING(cls) -> str: 

320 return ">Q" 

321 

322 def __post_init__(self) -> None: 

323 (self.mapSize,) = struct.unpack(self.FMT_STRING(), self.sizeBytes) 

324 

325 @property 

326 def headerSize(self) -> int: 

327 return self.preambleSize + self.structSize + self.mapSize 

328 

329 def readHeaderInfo(self, rawHeader: bytes) -> SimpleNamespace: 

330 uncompressedHeaderMap = self.unpackHeader(rawHeader) 

331 if uncompressedHeaderMap is None: 

332 raise ValueError( 

333 "This error is not possible because self.unpackHeader cannot return None," 

334 " but is done to satisfy type checkers" 

335 ) 

336 header = json.loads(uncompressedHeaderMap) 

337 returnValue = SimpleNamespace() 

338 returnValue.taskDefMap = header["TaskDefs"] 

339 returnValue._buildId = header["GraphBuildID"] 

340 returnValue.map = dict(header["Nodes"]) 

341 returnValue.metadata = header["Metadata"] 

342 self.returnValue = returnValue 

343 return returnValue 

344 

345 def unpackHeader(self, rawHeader: bytes) -> Optional[str]: 

346 return lzma.decompress(rawHeader).decode() 

347 

348 def constructGraph( 

349 self, 

350 nodes: set[uuid.UUID], 

351 _readBytes: Callable[[int, int], bytes], 

352 universe: Optional[DimensionUniverse] = None, 

353 ) -> QuantumGraph: 

354 # need to import here to avoid cyclic imports 

355 from . import QuantumGraph 

356 

357 quanta: DefaultDict[TaskDef, Set[Quantum]] = defaultdict(set) 

358 quantumToNodeId: Dict[Quantum, uuid.UUID] = {} 

359 loadedTaskDef = {} 

360 # loop over the nodes specified above 

361 for node in nodes: 

362 # Get the bytes to read from the map 

363 start, stop = self.returnValue.map[node]["bytes"] 

364 start += self.headerSize 

365 stop += self.headerSize 

366 

367 # read the specified bytes, will be overloaded by subclasses 

368 # bytes are compressed, so decompress them 

369 dump = lzma.decompress(_readBytes(start, stop)) 

370 

371 # reconstruct node 

372 qNode = pickle.loads(dump) 

373 object.__setattr__(qNode, "nodeId", uuid.uuid4()) 

374 

375 # read the saved node, name. If it has been loaded, attach it, if 

376 # not read in the taskDef first, and then load it 

377 nodeTask = qNode.taskDef 

378 if nodeTask not in loadedTaskDef: 

379 # Get the byte ranges corresponding to this taskDef 

380 start, stop = self.returnValue.taskDefMap[nodeTask]["bytes"] 

381 start += self.headerSize 

382 stop += self.headerSize 

383 

384 # load the taskDef, this method call will be overloaded by 

385 # subclasses. 

386 # bytes are compressed, so decompress them 

387 taskDef = pickle.loads(lzma.decompress(_readBytes(start, stop))) 

388 loadedTaskDef[nodeTask] = taskDef 

389 # Explicitly overload the "frozen-ness" of nodes to attach the 

390 # taskDef back into the un-persisted node 

391 object.__setattr__(qNode, "taskDef", loadedTaskDef[nodeTask]) 

392 quanta[qNode.taskDef].add(qNode.quantum) 

393 

394 # record the node for later processing 

395 quantumToNodeId[qNode.quantum] = qNode.nodeId 

396 

397 # construct an empty new QuantumGraph object, and run the associated 

398 # creation method with the un-persisted data 

399 qGraph = object.__new__(QuantumGraph) 

400 qGraph._buildGraphs( 

401 quanta, 

402 _quantumToNodeId=quantumToNodeId, 

403 _buildId=self.returnValue._buildId, 

404 metadata=self.returnValue.metadata, 

405 universe=universe, 

406 ) 

407 return qGraph 

408 

409 def description(self) -> str: 

410 return Version2Description 

411 

412 

413Version3Description = """ 

414The save file starts with the first few bytes corresponding to the magic bytes 

415in the QuantumGraph: `qgraph4\xf6\xe8\xa9`. 

416 

417The next few bytes are a big endian unsigned long long. 

418 

419The unsigned long long corresponds to the number of bytes of a mapping 

420of header information. This mapping is encoded into json and then lzma 

421compressed, meaning the operations must be performed in the opposite order to 

422deserialize. 

423 

424The json encoded header mapping contains 5 fields: GraphBuildId, TaskDefs, 

425Nodes, Metadata, and DimensionRecords. 

426 

427The `GraphBuildId` key corresponds with a string that is the unique id assigned 

428to this graph when it was created. 

429 

430The `TaskDefs` key corresponds to a value which is a mapping of Task label to 

431task data. The task data is a mapping of key to value. The keys of this mapping 

432are `bytes`, `inputs`, and `outputs`. 

433 

434The `TaskDefs` `bytes` key corresponds to a tuple of a byte range of the 

435start, stop bytes (indexed after all the header bytes). This byte rage 

436corresponds to a lzma compressed json mapping. This mapping has keys of 

437`taskName`, corresponding to a fully qualified python class, `config` a 

438pex_config string that is used to configure the class, and `label` which 

439corresponds to a string that uniquely identifies the task within a given 

440execution pipeline. 

441 

442The `TaskDefs` `inputs` key is associated with a list of tuples where each 

443tuple is a label of a task that is considered coming before a given task, and 

444the name of the dataset that is shared between the tasks (think node and edge 

445in a graph sense). 

446 

447The `TaskDefs` `outputs` key is like inputs except the values in a list 

448correspond to all the output connections of a task. 

449 

450The `Nodes` key is also a json mapping with keys corresponding to the UUIDs of 

451QuantumNodes. The values associated with these keys is another mapping with 

452the keys `bytes`, `inputs`, and `outputs`. 

453 

454`Nodes` key `bytes` corresponds to a tuple of a byte range of the start, stop 

455bytes (indexed after all the header bytes). These bytes are a lzma compressed 

456json mapping which contains many sub elements, this mapping will be referred to 

457as the SerializedQuantumNode (related to the python class it corresponds to). 

458 

459SerializedQUantumNodes have 3 keys, `quantum` corresponding to a json mapping 

460(described below) referred to as a SerializedQuantum, `taskLabel` a string 

461which corresponds to a label in the `TaskDefs` mapping, and `nodeId. 

462 

463A SerializedQuantum has many keys; taskName, dataId, datasetTypeMapping, 

464initInputs, inputs, outputs, dimensionRecords. 

465 

466like the `TaskDefs` key except it corresponds to 

467QuantumNodes instead of TaskDefs, and the keys of the mappings are string 

468representations of the UUIDs of the QuantumNodes. 

469 

470The `Metadata` key is a mapping of strings to associated values. This metadata 

471may be anything that is important to be transported alongside the graph. 

472 

473As stated above, each map contains byte ranges of the corresponding 

474datastructure. Theses bytes are also lzma compressed pickles, and should 

475be deserialized in a similar manner. 

476""" 

477 

478 

479@dataclass 

480class DeserializerV3(DeserializerBase): 

481 @classmethod 

482 def FMT_STRING(cls) -> str: 

483 return ">Q" 

484 

485 def __post_init__(self) -> None: 

486 self.infoSize: int 

487 (self.infoSize,) = struct.unpack(self.FMT_STRING(), self.sizeBytes) 

488 

489 @property 

490 def headerSize(self) -> int: 

491 return self.preambleSize + self.structSize + self.infoSize 

492 

493 def readHeaderInfo(self, rawHeader: bytes) -> SimpleNamespace: 

494 uncompressedinfoMap = self.unpackHeader(rawHeader) 

495 assert uncompressedinfoMap is not None # for python typing, this variant can't be None 

496 infoMap = json.loads(uncompressedinfoMap) 

497 infoMappings = SimpleNamespace() 

498 infoMappings.taskDefMap = infoMap["TaskDefs"] 

499 infoMappings._buildId = infoMap["GraphBuildID"] 

500 infoMappings.map = {uuid.UUID(k): v for k, v in infoMap["Nodes"]} 

501 infoMappings.metadata = infoMap["Metadata"] 

502 infoMappings.dimensionRecords = {} 

503 for k, v in infoMap["DimensionRecords"].items(): 

504 infoMappings.dimensionRecords[int(k)] = SerializedDimensionRecord(**v) 

505 # This is important to be a get call here, so that it supports versions 

506 # of saved quantum graph that might not have a saved universe without 

507 # changing save format 

508 if (universeConfig := infoMap.get("universe")) is not None: 

509 universe = DimensionUniverse(config=DimensionConfig(universeConfig)) 

510 else: 

511 universe = DimensionUniverse() 

512 infoMappings.universe = universe 

513 infoMappings.globalInitOutputRefs = [] 

514 if (json_refs := infoMap.get("GlobalInitOutputRefs")) is not None: 

515 infoMappings.globalInitOutputRefs = [ 

516 DatasetRef.from_json(json_ref, universe=universe) for json_ref in json_refs 

517 ] 

518 infoMappings.registryDatasetTypes = [] 

519 if (json_refs := infoMap.get("RegistryDatasetTypes")) is not None: 

520 infoMappings.registryDatasetTypes = [ 

521 DatasetType.from_json(json_ref, universe=universe) for json_ref in json_refs 

522 ] 

523 self.infoMappings = infoMappings 

524 return infoMappings 

525 

526 def unpackHeader(self, rawHeader: bytes) -> Optional[str]: 

527 return lzma.decompress(rawHeader).decode() 

528 

529 def constructGraph( 

530 self, 

531 nodes: set[uuid.UUID], 

532 _readBytes: Callable[[int, int], bytes], 

533 universe: Optional[DimensionUniverse] = None, 

534 ) -> QuantumGraph: 

535 # need to import here to avoid cyclic imports 

536 from . import QuantumGraph 

537 

538 graph = nx.DiGraph() 

539 loadedTaskDef: Dict[str, TaskDef] = {} 

540 container = {} 

541 datasetDict = _DatasetTracker[DatasetTypeName, TaskDef](createInverse=True) 

542 taskToQuantumNode: DefaultDict[TaskDef, Set[QuantumNode]] = defaultdict(set) 

543 recontitutedDimensions: Dict[int, Tuple[str, DimensionRecord]] = {} 

544 initInputRefs: Dict[TaskDef, List[DatasetRef]] = {} 

545 initOutputRefs: Dict[TaskDef, List[DatasetRef]] = {} 

546 

547 if universe is not None: 

548 if not universe.isCompatibleWith(self.infoMappings.universe): 

549 saved = self.infoMappings.universe 

550 raise RuntimeError( 

551 f"The saved dimension universe ({saved.namespace}@v{saved.version}) is not " 

552 f"compatible with the supplied universe ({universe.namespace}@v{universe.version})." 

553 ) 

554 else: 

555 universe = self.infoMappings.universe 

556 

557 for node in nodes: 

558 start, stop = self.infoMappings.map[node]["bytes"] 

559 start, stop = start + self.headerSize, stop + self.headerSize 

560 # Read in the bytes corresponding to the node to load and 

561 # decompress it 

562 dump = json.loads(lzma.decompress(_readBytes(start, stop))) 

563 

564 # Turn the json back into the pydandtic model 

565 nodeDeserialized = SerializedQuantumNode.direct(**dump) 

566 # attach the dictionary of dimension records to the pydandtic model 

567 # these are stored seperately because the are stored over and over 

568 # and this saves a lot of space and time. 

569 nodeDeserialized.quantum.dimensionRecords = self.infoMappings.dimensionRecords 

570 # get the label for the current task 

571 nodeTaskLabel = nodeDeserialized.taskLabel 

572 

573 if nodeTaskLabel not in loadedTaskDef: 

574 # Get the byte ranges corresponding to this taskDef 

575 start, stop = self.infoMappings.taskDefMap[nodeTaskLabel]["bytes"] 

576 start, stop = start + self.headerSize, stop + self.headerSize 

577 

578 # bytes are compressed, so decompress them 

579 taskDefDump = json.loads(lzma.decompress(_readBytes(start, stop))) 

580 taskClass: Type[PipelineTask] = doImportType(taskDefDump["taskName"]) 

581 config: PipelineTaskConfig = taskClass.ConfigClass() 

582 config.loadFromStream(taskDefDump["config"]) 

583 # Rebuild TaskDef 

584 recreatedTaskDef = TaskDef( 

585 taskName=taskDefDump["taskName"], 

586 taskClass=taskClass, 

587 config=config, 

588 label=taskDefDump["label"], 

589 ) 

590 loadedTaskDef[nodeTaskLabel] = recreatedTaskDef 

591 

592 # initInputRefs and initOutputRefs are optional 

593 if (refs := taskDefDump.get("initInputRefs")) is not None: 

594 initInputRefs[recreatedTaskDef] = [ 

595 cast(DatasetRef, DatasetRef.from_json(ref, universe=universe)) for ref in refs 

596 ] 

597 if (refs := taskDefDump.get("initOutputRefs")) is not None: 

598 initOutputRefs[recreatedTaskDef] = [ 

599 cast(DatasetRef, DatasetRef.from_json(ref, universe=universe)) for ref in refs 

600 ] 

601 

602 # rebuild the mappings that associate dataset type names with 

603 # TaskDefs 

604 for _, input in self.infoMappings.taskDefMap[nodeTaskLabel]["inputs"]: 

605 datasetDict.addConsumer(DatasetTypeName(input), recreatedTaskDef) 

606 

607 added = set() 

608 for outputConnection in self.infoMappings.taskDefMap[nodeTaskLabel]["outputs"]: 

609 typeName = outputConnection[1] 

610 if typeName not in added: 

611 added.add(typeName) 

612 datasetDict.addProducer(DatasetTypeName(typeName), recreatedTaskDef) 

613 

614 # reconstitute the node, passing in the dictionaries for the 

615 # loaded TaskDefs and dimension records. These are used to ensure 

616 # that each unique record is only loaded once 

617 qnode = QuantumNode.from_simple(nodeDeserialized, loadedTaskDef, universe, recontitutedDimensions) 

618 container[qnode.nodeId] = qnode 

619 taskToQuantumNode[loadedTaskDef[nodeTaskLabel]].add(qnode) 

620 

621 # recreate the relations between each node from stored info 

622 graph.add_node(qnode) 

623 for id in self.infoMappings.map[qnode.nodeId]["inputs"]: 

624 # uuid is stored as a string, turn it back into a uuid 

625 id = uuid.UUID(id) 

626 # if the id is not yet in the container, dont make a connection 

627 # this is not an issue, because once it is, that id will add 

628 # the reverse connection 

629 if id in container: 

630 graph.add_edge(container[id], qnode) 

631 for id in self.infoMappings.map[qnode.nodeId]["outputs"]: 

632 # uuid is stored as a string, turn it back into a uuid 

633 id = uuid.UUID(id) 

634 # if the id is not yet in the container, dont make a connection 

635 # this is not an issue, because once it is, that id will add 

636 # the reverse connection 

637 if id in container: 

638 graph.add_edge(qnode, container[id]) 

639 

640 newGraph = object.__new__(QuantumGraph) 

641 newGraph._metadata = self.infoMappings.metadata 

642 newGraph._buildId = self.infoMappings._buildId 

643 newGraph._datasetDict = datasetDict 

644 newGraph._nodeIdMap = container 

645 newGraph._count = len(nodes) 

646 newGraph._taskToQuantumNode = dict(taskToQuantumNode.items()) 

647 newGraph._taskGraph = datasetDict.makeNetworkXGraph() 

648 newGraph._connectedQuanta = graph 

649 newGraph._initInputRefs = initInputRefs 

650 newGraph._initOutputRefs = initOutputRefs 

651 newGraph._globalInitOutputRefs = self.infoMappings.globalInitOutputRefs 

652 newGraph._registryDatasetTypes = self.infoMappings.registryDatasetTypes 

653 newGraph._universe = universe 

654 return newGraph 

655 

656 

657DESERIALIZER_MAP: dict[int, Type[DeserializerBase]] = { 

658 1: DeserializerV1, 

659 2: DeserializerV2, 

660 3: DeserializerV3, 

661}