Coverage for python/lsst/pipe/base/quantum_graph/

1# This file is part of pipe_base.

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.

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 operator

47import sys

48import uuid

49import warnings

50from collections import defaultdict

51from collections.abc import Iterable, Iterator, Mapping, Sequence

52from contextlib import AbstractContextManager, contextmanager

53from typing import TYPE_CHECKING, Any, TypeVar, cast

55import networkx

56import networkx.algorithms.bipartite

57import pydantic

58import zstandard

60from lsst.daf.butler import (

61 Config,

62 DataCoordinate,

63 DataIdValue,

64 DatasetRef,

65 DatasetType,

66 DimensionDataAttacher,

67 DimensionDataExtractor,

68 DimensionGroup,

69 DimensionRecordSetDeserializer,

70 LimitedButler,

71 Quantum,

72 QuantumBackedButler,

73 SerializableDimensionData,

74)

75from lsst.daf.butler._rubin import generate_uuidv7

76from lsst.daf.butler.datastore.record_data import DatastoreRecordData, SerializedDatastoreRecordData

77from lsst.daf.butler.registry import ConflictingDefinitionError

78from lsst.resources import ResourcePath, ResourcePathExpression

79from lsst.utils.packages import Packages

81from .. import automatic_connection_constants as acc

82from ..pipeline import TaskDef

83from ..pipeline_graph import (

84 PipelineGraph,

85 TaskImportMode,

86 TaskInitNode,

87 TaskNode,

88 compare_packages,

89 log_config_mismatch,

90)

91from ._common import (

92 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, MultiblockReader, MultiblockWriter

107

108if TYPE_CHECKING:

109 from ..config import PipelineTaskConfig

110 from ..graph import QgraphSummary, QuantumGraph

111

112_LOG = logging.getLogger(__name__)

113

114

115_T = TypeVar("_T", bound=pydantic.BaseModel)

116

117

118class PredictedThinQuantumModel(pydantic.BaseModel):

119 """Data model for a quantum data ID and internal integer ID in a predicted

120 quantum graph.

121 """

122

123 quantum_index: QuantumIndex

124 """Internal integer ID for this quantum."""

125

126 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)

127 """Full (required and implied) data coordinate values for this quantum."""

128

129 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,

130 # when we inherit those docstrings in our public classes.

131 if "sphinx" in sys.modules and not TYPE_CHECKING:

132

133 def copy(self, *args: Any, **kwargs: Any) -> Any:

134 """See `pydantic.BaseModel.copy`."""

135 return super().copy(*args, **kwargs)

136

137 def model_dump(self, *args: Any, **kwargs: Any) -> Any:

138 """See `pydantic.BaseModel.model_dump`."""

139 return super().model_dump(*args, **kwargs)

140

141 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:

142 """See `pydantic.BaseModel.model_dump_json`."""

143 return super().model_dump(*args, **kwargs)

144

145 def model_copy(self, *args: Any, **kwargs: Any) -> Any:

146 """See `pydantic.BaseModel.model_copy`."""

147 return super().model_copy(*args, **kwargs)

148

149 @classmethod

150 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]

151 """See `pydantic.BaseModel.model_construct`."""

152 return super().model_construct(*args, **kwargs)

153

154 @classmethod

155 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:

156 """See `pydantic.BaseModel.model_json_schema`."""

157 return super().model_json_schema(*args, **kwargs)

158

159 @classmethod

160 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:

161 """See `pydantic.BaseModel.model_validate`."""

162 return super().model_validate(*args, **kwargs)

163

164 @classmethod

165 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:

166 """See `pydantic.BaseModel.model_validate_json`."""

167 return super().model_validate_json(*args, **kwargs)

168

169 @classmethod

170 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:

171 """See `pydantic.BaseModel.model_validate_strings`."""

172 return super().model_validate_strings(*args, **kwargs)

173

174

175class PredictedThinGraphModel(pydantic.BaseModel):

176 """Data model for the predicted quantum graph component that maps each

177 task label to the data IDs and internal integer IDs of its quanta.

178 """

179

180 quanta: dict[TaskLabel, list[PredictedThinQuantumModel]] = pydantic.Field(default_factory=dict)

181 """Minimal descriptions of all quanta, grouped by task label."""

182

183 edges: list[tuple[QuantumIndex, QuantumIndex]] = pydantic.Field(default_factory=list)

184 """Pairs of (predecessor, successor) internal integer quantum IDs."""

185

186 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,

187 # when we inherit those docstrings in our public classes.

188 if "sphinx" in sys.modules and not TYPE_CHECKING:

189

190 def copy(self, *args: Any, **kwargs: Any) -> Any:

191 """See `pydantic.BaseModel.copy`."""

192 return super().copy(*args, **kwargs)

193

194 def model_dump(self, *args: Any, **kwargs: Any) -> Any:

195 """See `pydantic.BaseModel.model_dump`."""

196 return super().model_dump(*args, **kwargs)

197

198 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:

199 """See `pydantic.BaseModel.model_dump_json`."""

200 return super().model_dump(*args, **kwargs)

201

202 def model_copy(self, *args: Any, **kwargs: Any) -> Any:

203 """See `pydantic.BaseModel.model_copy`."""

204 return super().model_copy(*args, **kwargs)

205

206 @classmethod

207 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]

208 """See `pydantic.BaseModel.model_construct`."""

209 return super().model_construct(*args, **kwargs)

210

211 @classmethod

212 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:

213 """See `pydantic.BaseModel.model_json_schema`."""

214 return super().model_json_schema(*args, **kwargs)

215

216 @classmethod

217 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:

218 """See `pydantic.BaseModel.model_validate`."""

219 return super().model_validate(*args, **kwargs)

220

221 @classmethod

222 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:

223 """See `pydantic.BaseModel.model_validate_json`."""

224 return super().model_validate_json(*args, **kwargs)

225

226 @classmethod

227 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:

228 """See `pydantic.BaseModel.model_validate_strings`."""

229 return super().model_validate_strings(*args, **kwargs)

230

231

232class PredictedDatasetModel(pydantic.BaseModel):

233 """Data model for the datasets in a predicted quantum graph file."""

234

235 dataset_id: uuid.UUID

236 """Universally unique ID for the dataset."""

237

238 dataset_type_name: DatasetTypeName

239 """Name of the type of this dataset.

240

241 This is always a parent dataset type name, not a component.

242

243 Note that full dataset type definitions are stored in the pipeline graph.

244 """

245

246 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)

247 """The full values (required and implied) of this dataset's data ID."""

248

249 run: str

250 """This dataset's RUN collection name."""

251

252 @classmethod

253 def from_dataset_ref(cls, ref: DatasetRef) -> PredictedDatasetModel:

254 """Construct from a butler `~lsst.daf.butler.DatasetRef`.

255

256 Parameters

257 ----------

258 ref : `lsst.daf.butler.DatasetRef`

259 Dataset reference.

260

261 Returns

262 -------

263 model : `PredictedDatasetModel`

264 Model for the dataset.

265 """

266 dataset_type_name, _ = DatasetType.splitDatasetTypeName(ref.datasetType.name)

267 return cls.model_construct(

268 dataset_id=ref.id,

269 dataset_type_name=dataset_type_name,

270 data_coordinate=list(ref.dataId.full_values),

271 run=ref.run,

272 )

273

274 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,

275 # when we inherit those docstrings in our public classes.

276 if "sphinx" in sys.modules and not TYPE_CHECKING:

277

278 def copy(self, *args: Any, **kwargs: Any) -> Any:

279 """See `pydantic.BaseModel.copy`."""

280 return super().copy(*args, **kwargs)

281

282 def model_dump(self, *args: Any, **kwargs: Any) -> Any:

283 """See `pydantic.BaseModel.model_dump`."""

284 return super().model_dump(*args, **kwargs)

285

286 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:

287 """See `pydantic.BaseModel.model_dump_json`."""

288 return super().model_dump(*args, **kwargs)

289

290 def model_copy(self, *args: Any, **kwargs: Any) -> Any:

291 """See `pydantic.BaseModel.model_copy`."""

292 return super().model_copy(*args, **kwargs)

293

294 @classmethod

295 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]

296 """See `pydantic.BaseModel.model_construct`."""

297 return super().model_construct(*args, **kwargs)

298

299 @classmethod

300 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:

301 """See `pydantic.BaseModel.model_json_schema`."""

302 return super().model_json_schema(*args, **kwargs)

303

304 @classmethod

305 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:

306 """See `pydantic.BaseModel.model_validate`."""

307 return super().model_validate(*args, **kwargs)

308

309 @classmethod

310 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:

311 """See `pydantic.BaseModel.model_validate_json`."""

312 return super().model_validate_json(*args, **kwargs)

313

314 @classmethod

315 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:

316 """See `pydantic.BaseModel.model_validate_strings`."""

317 return super().model_validate_strings(*args, **kwargs)

318

319

320class PredictedQuantumDatasetsModel(pydantic.BaseModel):

321 """Data model for a description of a single predicted quantum that includes

322 its inputs and outputs.

323 """

324

325 quantum_id: uuid.UUID

326 """Universally unique ID for the quantum."""

327

328 task_label: TaskLabel

329 """Label of the task.

330

331 Note that task label definitions are stored in the pipeline graph.

332 """

333

334 data_coordinate: DataCoordinateValues = pydantic.Field(default_factory=list)

335 """The full values (required and implied) of this quantum's data ID."""

336

337 inputs: dict[ConnectionName, list[PredictedDatasetModel]] = pydantic.Field(default_factory=dict)

338 """The input datasets to this quantum, grouped by connection name."""

339

340 outputs: dict[ConnectionName, list[PredictedDatasetModel]] = pydantic.Field(default_factory=dict)

341 """The datasets output by this quantum, grouped by connection name."""

342

343 datastore_records: dict[DatastoreName, SerializedDatastoreRecordData] = pydantic.Field(

344 default_factory=dict

345 )

346 """Datastore records for inputs to this quantum that are already present in

347 the data repository.

348 """

349

350 def iter_input_dataset_ids(self) -> Iterator[uuid.UUID]:

351 """Return an iterator over the UUIDs of all datasets consumed by this

352 quantum.

353

354 Returns

355 -------

356 iter : `~collections.abc.Iterator` [ `uuid.UUID` ]

357 Iterator over dataset IDs.

358 """

359 for datasets in self.inputs.values():

360 for dataset in datasets:

361 yield dataset.dataset_id

362

363 def iter_output_dataset_ids(self) -> Iterator[uuid.UUID]:

364 """Return an iterator over the UUIDs of all datasets produced by this

365 quantum.

366

367 Returns

368 -------

369 iter : `~collections.abc.Iterator` [ `uuid.UUID` ]

370 Iterator over dataset IDs.

371 """

372 for datasets in self.outputs.values():

373 for dataset in datasets:

374 yield dataset.dataset_id

375

376 def iter_dataset_ids(self) -> Iterator[uuid.UUID]:

377 """Return an iterator over the UUIDs of all datasets referenced by this

378 quantum.

379

380 Returns

381 -------

382 iter : `~collections.abc.Iterator` [ `uuid.UUID` ]

383 Iterator over dataset IDs.

384 """

385 yield from self.iter_input_dataset_ids()

386 yield from self.iter_output_dataset_ids()

387

388 def deserialize_datastore_records(self) -> dict[DatastoreName, DatastoreRecordData]:

389 """Deserialize the mapping of datastore records."""

390 return {

391 datastore_name: DatastoreRecordData.from_simple(serialized_records)

392 for datastore_name, serialized_records in self.datastore_records.items()

393 }

394

395 @classmethod

396 def from_execution_quantum(

397 cls, task_node: TaskNode, quantum: Quantum, quantum_id: uuid.UUID

398 ) -> PredictedQuantumDatasetsModel:

399 """Construct from an `lsst.daf.butler.Quantum` instance.

400

401 Parameters

402 ----------

403 task_node : `.pipeline_graph.TaskNode`

404 Task node from the pipeline graph.

405 quantum : `lsst.daf.butler.quantum`

406 Quantum object.

407 quantum_id : `uuid.UUID`

408 ID for this quantum.

409

410 Returns

411 -------

412 model : `PredictedFullQuantumModel`

413 Model for this quantum.

414 """

415 result: PredictedQuantumDatasetsModel = cls.model_construct(

416 quantum_id=quantum_id,

417 task_label=task_node.label,

418 data_coordinate=list(cast(DataCoordinate, quantum.dataId).full_values),

419 )

420 for read_edge in task_node.iter_all_inputs():

421 refs = sorted(quantum.inputs[read_edge.dataset_type_name], key=lambda ref: ref.dataId)

422 result.inputs[read_edge.connection_name] = [

423 PredictedDatasetModel.from_dataset_ref(ref) for ref in refs

424 ]

425 for write_edge in task_node.iter_all_outputs():

426 refs = sorted(quantum.outputs[write_edge.dataset_type_name], key=lambda ref: ref.dataId)

427 result.outputs[write_edge.connection_name] = [

428 PredictedDatasetModel.from_dataset_ref(ref) for ref in refs

429 ]

430 result.datastore_records = {

431 store_name: records.to_simple() for store_name, records in quantum.datastore_records.items()

432 }

433 return result

434

435 @classmethod

436 def from_old_quantum_graph_init(

437 cls, task_init_node: TaskInitNode, old_quantum_graph: QuantumGraph

438 ) -> PredictedQuantumDatasetsModel:

439 """Construct from the init-input and init-output dataset types of a

440 task in an old `QuantumGraph` instance.

441

442 Parameters

443 ----------

444 task_init_node : `.pipeline_graph.TaskNode`

445 Task init node from the pipeline graph.

446 old_quantum_graph : `QuantumGraph`

447 Quantum graph.

448

449 Returns

450 -------

451 model : `PredictedFullQuantumModel`

452 Model for this "init" quantum.

453 """

454 task_def = old_quantum_graph.findTaskDefByLabel(task_init_node.label)

455 assert task_def is not None

456 init_input_refs = {

457 ref.datasetType.name: ref for ref in (old_quantum_graph.initInputRefs(task_def) or [])

458 }

459 init_output_refs = {

460 ref.datasetType.name: ref for ref in (old_quantum_graph.initOutputRefs(task_def) or [])

461 }

462 init_input_ids = {ref.id for ref in init_input_refs.values()}

463 result: PredictedQuantumDatasetsModel = cls.model_construct(

464 quantum_id=generate_uuidv7(), task_label=task_init_node.label

465 )

466 for read_edge in task_init_node.iter_all_inputs():

467 ref = init_input_refs[read_edge.dataset_type_name]

468 result.inputs[read_edge.connection_name] = [PredictedDatasetModel.from_dataset_ref(ref)]

469 for write_edge in task_init_node.iter_all_outputs():

470 ref = init_output_refs[write_edge.dataset_type_name]

471 result.outputs[write_edge.connection_name] = [PredictedDatasetModel.from_dataset_ref(ref)]

472 datastore_records: dict[str, DatastoreRecordData] = {}

473 for quantum in old_quantum_graph.get_task_quanta(task_init_node.label).values():

474 for store_name, records in quantum.datastore_records.items():

475 subset = records.subset(init_input_ids)

476 if subset is not None:

477 datastore_records.setdefault(store_name, DatastoreRecordData()).update(subset)

478 break # All quanta have same init-inputs, so we only need one.

479 result.datastore_records = {

480 store_name: records.to_simple() for store_name, records in datastore_records.items()

481 }

482 return result

483

484 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,

485 # when we inherit those docstrings in our public classes.

486 if "sphinx" in sys.modules and not TYPE_CHECKING:

487

488 def copy(self, *args: Any, **kwargs: Any) -> Any:

489 """See `pydantic.BaseModel.copy`."""

490 return super().copy(*args, **kwargs)

491

492 def model_dump(self, *args: Any, **kwargs: Any) -> Any:

493 """See `pydantic.BaseModel.model_dump`."""

494 return super().model_dump(*args, **kwargs)

495

496 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:

497 """See `pydantic.BaseModel.model_dump_json`."""

498 return super().model_dump(*args, **kwargs)

499

500 def model_copy(self, *args: Any, **kwargs: Any) -> Any:

501 """See `pydantic.BaseModel.model_copy`."""

502 return super().model_copy(*args, **kwargs)

503

504 @classmethod

505 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]

506 """See `pydantic.BaseModel.model_construct`."""

507 return super().model_construct(*args, **kwargs)

508

509 @classmethod

510 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:

511 """See `pydantic.BaseModel.model_json_schema`."""

512 return super().model_json_schema(*args, **kwargs)

513

514 @classmethod

515 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:

516 """See `pydantic.BaseModel.model_validate`."""

517 return super().model_validate(*args, **kwargs)

518

519 @classmethod

520 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:

521 """See `pydantic.BaseModel.model_validate_json`."""

522 return super().model_validate_json(*args, **kwargs)

523

524 @classmethod

525 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:

526 """See `pydantic.BaseModel.model_validate_strings`."""

527 return super().model_validate_strings(*args, **kwargs)

528

529

530class PredictedInitQuantaModel(pydantic.RootModel):

531 """Data model for the init-inputs and init-outputs of a predicted quantum

532 graph.

533 """

534

535 root: list[PredictedQuantumDatasetsModel] = pydantic.Field(default_factory=list)

536 """List of special "init" quanta: one for each task, and another for global

537 init-outputs.

538 """

539

540 def update_from_old_quantum_graph(self, old_quantum_graph: QuantumGraph) -> None:

541 """Update this model in-place by extracting from an old `QuantumGraph`

542 instance.

543

544 Parameters

545 ----------

546 old_quantum_graph : `QuantumGraph`

547 Quantum graph.

548 """

549 global_init_quantum = PredictedQuantumDatasetsModel.model_construct(

550 quantum_id=generate_uuidv7(), task_label=""

551 )

552 for ref in old_quantum_graph.globalInitOutputRefs():

553 global_init_quantum.outputs[ref.datasetType.name] = [PredictedDatasetModel.from_dataset_ref(ref)]

554 self.root.append(global_init_quantum)

555 for task_node in old_quantum_graph.pipeline_graph.tasks.values():

556 self.root.append(

557 PredictedQuantumDatasetsModel.from_old_quantum_graph_init(task_node.init, old_quantum_graph)

558 )

559

560 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,

561 # when we inherit those docstrings in our public classes.

562 if "sphinx" in sys.modules and not TYPE_CHECKING:

563

564 def copy(self, *args: Any, **kwargs: Any) -> Any:

565 """See `pydantic.BaseModel.copy`."""

566 return super().copy(*args, **kwargs)

567

568 def model_dump(self, *args: Any, **kwargs: Any) -> Any:

569 """See `pydantic.BaseModel.model_dump`."""

570 return super().model_dump(*args, **kwargs)

571

572 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:

573 """See `pydantic.BaseModel.model_dump_json`."""

574 return super().model_dump(*args, **kwargs)

575

576 def model_copy(self, *args: Any, **kwargs: Any) -> Any:

577 """See `pydantic.BaseModel.model_copy`."""

578 return super().model_copy(*args, **kwargs)

579

580 @classmethod

581 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]

582 """See `pydantic.BaseModel.model_construct`."""

583 return super().model_construct(*args, **kwargs)

584

585 @classmethod

586 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:

587 """See `pydantic.BaseModel.model_json_schema`."""

588 return super().model_json_schema(*args, **kwargs)

589

590 @classmethod

591 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:

592 """See `pydantic.BaseModel.model_validate`."""

593 return super().model_validate(*args, **kwargs)

594

595 @classmethod

596 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:

597 """See `pydantic.BaseModel.model_validate_json`."""

598 return super().model_validate_json(*args, **kwargs)

599

600 @classmethod

601 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:

602 """See `pydantic.BaseModel.model_validate_strings`."""

603 return super().model_validate_strings(*args, **kwargs)

604

605

606class PredictedQuantumInfo(QuantumInfo):

607 """A typed dictionary that annotates the attributes of the NetworkX graph

608 node data for a predicted quantum.

609

610 Since NetworkX types are not generic over their node mapping type, this has

611 to be used explicitly, e.g.::

612

613 node_data: PredictedQuantumInfo = xgraph.nodes[quantum_id]

614

615 where ``xgraph`` can be either `PredictedQuantumGraph.quantum_only_xgraph`

616 or `PredictedQuantumGraph.bipartite_xgraph`.

617 """

618

619 quantum: Quantum

620 """Quantum object that can be passed directly to an executor.

621

622 This attribute is only present if

623 `PredictedQuantumGraph.build_execution_quanta` has been run on this node's

624 quantum ID already.

625 """

626

627

628class PredictedDatasetInfo(DatasetInfo):

629 """A typed dictionary that annotates the attributes of the NetworkX graph

630 node data for a dataset.

631

632 Since NetworkX types are not generic over their node mapping type, this has

633 to be used explicitly, e.g.::

634

635 node_data: PredictedDatasetInfo = xgraph.nodes[dataset_ids]

636

637 where ``xgraph`` is from the `PredictedQuantumGraph.bipartite_xgraph`

638 property.

639 """

640

641

642class PredictedQuantumGraph(BaseQuantumGraph):

643 """A directed acyclic graph that predicts a processing run and supports it

644 during execution.

645

646 Parameters

647 ----------

648 components : `PredictedQuantumGraphComponents`

649 A struct of components used to construct the graph.

650

651 Notes

652 -----

653 Iteration over a `PredictedQuantumGraph` yields loaded quantum IDs in

654 deterministic topological order (but the tiebreaker is unspecified). The

655 `len` of a `PredictedQuantumGraph` is the number of loaded non-init quanta,

656 i.e. the same as the number of quanta iterated over.

657 """

658

659 def __init__(self, components: PredictedQuantumGraphComponents):

660 if not components.header.graph_type == "predicted":

661 raise TypeError(f"Header is for a {components.header.graph_type!r} graph, not 'predicted'.")

662 super().__init__(components.header, components.pipeline_graph)

663 self._quantum_only_xgraph = networkx.DiGraph()

664 self._bipartite_xgraph = networkx.DiGraph()

665 self._quanta_by_task_label: dict[str, dict[DataCoordinate, uuid.UUID]] = {

666 task_label: {} for task_label in self.pipeline_graph.tasks.keys()

667 }

668 self._datasets_by_type: dict[str, dict[DataCoordinate, uuid.UUID]] = {

669 dataset_type_name: {} for dataset_type_name in self.pipeline_graph.dataset_types.keys()

670 }

671 self._datasets_by_type[self.pipeline_graph.packages_dataset_type.name] = {}

672 self._dimension_data = components.dimension_data

673 self._add_init_quanta(components.init_quanta)

674 self._quantum_datasets: dict[uuid.UUID, PredictedQuantumDatasetsModel] = {}

675 self._expanded_data_ids: dict[DataCoordinate, DataCoordinate] = {}

676 self._add_thin_graph(components.thin_graph, components.quantum_indices)

677 for quantum_datasets in components.quantum_datasets.values():

678 self._add_quantum_datasets(quantum_datasets)

679 if not components.thin_graph.edges:

680 # If we loaded the thin_graph, we've already populated this graph.

681 self._quantum_only_xgraph.update(

682 networkx.algorithms.bipartite.projected_graph(

683 networkx.DiGraph(self._bipartite_xgraph),

684 self._quantum_only_xgraph.nodes.keys(),

685 )

686 )

687 if _LOG.isEnabledFor(logging.DEBUG):

688 for quantum_id in self:

689 _LOG.debug(

690 "%s: %s @ %s",

691 quantum_id,

692 self._quantum_only_xgraph.nodes[quantum_id]["task_label"],

693 self._quantum_only_xgraph.nodes[quantum_id]["data_id"].required,

694 )

695

696 def _add_init_quanta(self, component: PredictedInitQuantaModel) -> None:

697 self._init_quanta = {q.task_label: q for q in component.root}

698 empty_data_id = DataCoordinate.make_empty(self.pipeline_graph.universe)

699 for quantum_datasets in self._init_quanta.values():

700 for init_datasets in itertools.chain(

701 quantum_datasets.inputs.values(), quantum_datasets.outputs.values()

702 ):

703 for init_dataset in init_datasets:

704 self._datasets_by_type[init_dataset.dataset_type_name][empty_data_id] = (

705 init_dataset.dataset_id

706 )

707 _LOG.debug(

708 "%s: %s @ init",

709 quantum_datasets.quantum_id,

710 quantum_datasets.task_label,

711 )

712

713 def _add_thin_graph(

714 self, component: PredictedThinGraphModel, indices: Mapping[uuid.UUID, QuantumIndex]

715 ) -> None:

716 uuid_by_index = {v: k for k, v in indices.items()}

717 for index1, index2 in component.edges:

718 self._quantum_only_xgraph.add_edge(uuid_by_index[index1], uuid_by_index[index2])

719 for task_label, thin_quanta_for_task in component.quanta.items():

720 for thin_quantum in thin_quanta_for_task:

721 self._add_quantum(

722 uuid_by_index[thin_quantum.quantum_index],

723 task_label,

724 thin_quantum.data_coordinate,

725 )

726

727 def _add_quantum_datasets(self, quantum_datasets: PredictedQuantumDatasetsModel) -> None:

728 self._quantum_datasets[quantum_datasets.quantum_id] = quantum_datasets

729 self._add_quantum(

730 quantum_datasets.quantum_id, quantum_datasets.task_label, quantum_datasets.data_coordinate

731 )

732 task_node = self.pipeline_graph.tasks[quantum_datasets.task_label]

733 for connection_name, input_datasets in quantum_datasets.inputs.items():

734 pipeline_edge = task_node.get_input_edge(connection_name)

735 for input_dataset in input_datasets:

736 self._add_dataset(input_dataset)

737 self._bipartite_xgraph.add_edge(

738 input_dataset.dataset_id,

739 quantum_datasets.quantum_id,

740 key=connection_name,

741 is_read=True,

742 )

743 # There might be multiple input connections for the same

744 # dataset type.

745 self._bipartite_xgraph.edges[

746 input_dataset.dataset_id, quantum_datasets.quantum_id

747 ].setdefault("pipeline_edges", []).append(pipeline_edge)

748 for connection_name, output_datasets in quantum_datasets.outputs.items():

749 pipeline_edges = [task_node.get_output_edge(connection_name)]

750 for output_dataset in output_datasets:

751 self._add_dataset(output_dataset)

752 self._bipartite_xgraph.add_edge(

753 quantum_datasets.quantum_id,

754 output_dataset.dataset_id,

755 key=connection_name,

756 is_read=False,

757 pipeline_edges=pipeline_edges,

758 )

759

760 def _add_quantum(

761 self, quantum_id: uuid.UUID, task_label: str, data_coordinate_values: Sequence[DataIdValue]

762 ) -> None:

763 task_node = self.pipeline_graph.tasks[task_label]

764 self._quantum_only_xgraph.add_node(quantum_id, task_label=task_label, pipeline_node=task_node)

765 self._bipartite_xgraph.add_node(quantum_id, task_label=task_label, pipeline_node=task_node)

766 data_coordinate_values = tuple(data_coordinate_values)

767 dimensions = self.pipeline_graph.tasks[task_label].dimensions

768 data_id = DataCoordinate.from_full_values(dimensions, tuple(data_coordinate_values))

769 self._quantum_only_xgraph.nodes[quantum_id].setdefault("data_id", data_id)

770 self._bipartite_xgraph.nodes[quantum_id].setdefault("data_id", data_id)

771 self._quanta_by_task_label[task_label][data_id] = quantum_id

772

773 def _add_dataset(self, model: PredictedDatasetModel) -> None:

774 dataset_type_node = self.pipeline_graph.dataset_types[model.dataset_type_name]

775 data_id = DataCoordinate.from_full_values(dataset_type_node.dimensions, tuple(model.data_coordinate))

776 self._bipartite_xgraph.add_node(

777 model.dataset_id,

778 dataset_type_name=dataset_type_node.name,

779 pipeline_node=dataset_type_node,

780 run=model.run,

781 )

782 self._bipartite_xgraph.nodes[model.dataset_id].setdefault("data_id", data_id)

783 self._datasets_by_type[model.dataset_type_name][data_id] = model.dataset_id

784

785 @classmethod

786 def open(

787 cls,

788 uri: ResourcePathExpression,

789 page_size: int = DEFAULT_PAGE_SIZE,

790 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,

791 ) -> AbstractContextManager[PredictedQuantumGraphReader]:

792 """Open a quantum graph and return a reader to load from it.

793

794 Parameters

795 ----------

796 uri : convertible to `lsst.resources.ResourcePath`

797 URI to open. Should have a ``.qg`` extension.

798 page_size : `int`, optional

799 Approximate number of bytes to read at once from address files.

800 Note that this does not set a page size for *all* reads, but it

801 does affect the smallest, most numerous reads.

802 import_mode : `.pipeline_graph.TaskImportMode`, optional

803 How to handle importing the task classes referenced in the pipeline

804 graph.

805

806 Returns

807 -------

808 reader : `contextlib.AbstractContextManager` [ \

809 `PredictedQuantumGraphReader` ]

810 A context manager that returns the reader when entered.

811 """

812 return PredictedQuantumGraphReader.open(uri, page_size=page_size, import_mode=import_mode)

813

814 @classmethod

815 def read_execution_quanta(

816 cls,

817 uri: ResourcePathExpression,

818 quantum_ids: Iterable[uuid.UUID] | None = None,

819 page_size: int = DEFAULT_PAGE_SIZE,

820 ) -> PredictedQuantumGraph:

821 """Read one or more executable quanta from a quantum graph file.

822

823 Parameters

824 ----------

825 uri : convertible to `lsst.resources.ResourcePath`

826 URI to open. Should have a ``.qg`` extension for new quantum graph

827 files, or ``.qgraph`` for the old format.

828 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional

829 Iterable of quantum IDs to load. If not provided, all quanta will

830 be loaded. The UUIDs of special init quanta will be ignored.

831 page_size : `int`, optional

832 Approximate number of bytes to read at once from address files.

833 Note that this does not set a page size for *all* reads, but it

834 does affect the smallest, most numerous reads.

835

836 Returns

837 -------

838 quantum_graph : `PredictedQuantumGraph` ]

839 A quantum graph that can build execution quanta for all of the

840 given IDs.

841 """

842 return PredictedQuantumGraphComponents.read_execution_quanta(

843 uri,

844 quantum_ids,

845 page_size=page_size,

846 ).assemble()

847

848 @property

849 def quanta_by_task(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:

850 """A nested mapping of all quanta, keyed first by task name and then by

851 data ID.

852

853 Notes

854 -----

855 This is populated by the ``thin_graph`` component (all quanta are

856 added) and the `quantum_datasets`` component (only loaded quanta are

857 added). All tasks in the pipeline graph are included, even if none of

858 their quanta were loaded (i.e. nested mappings may be empty).

859

860 The returned object may be an internal dictionary; as the type

861 annotation indicates, it should not be modified in place.

862 """

863 return self._quanta_by_task_label

864

865 @property

866 def datasets_by_type(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:

867 """A nested mapping of all datasets, keyed first by dataset type name

868 and then by data ID.

869

870 Notes

871 -----

872 This is populated only by the ``quantum_datasets`` and ``init_quanta``

873 components, and only datasets referenced by loaded quanta are present.

874 All dataset types in the pipeline graph are included, even if none of

875 their datasets were loaded (i.e. nested mappings may be empty).

876

877 The returned object may be an internal dictionary; as the type

878 annotation indicates, it should not be modified in place.

879 """

880 return self._datasets_by_type

881

882 @property

883 def quantum_only_xgraph(self) -> networkx.DiGraph:

884 """A directed acyclic graph with quanta as nodes and datasets elided.

885

886 Notes

887 -----

888 Node keys are quantum UUIDs, and are populated by the ``thin_graph``

889 component (all nodes and edges) and ``quantum_datasets`` component

890 (only those that were loaded).

891

892 Node state dictionaries are described by the

893 `PredictedQuantumInfo` type.

894

895 The returned object is a read-only view of an internal one.

896 """

897 return self._quantum_only_xgraph.copy(as_view=True)

898

899 @property

900 def bipartite_xgraph(self) -> networkx.MultiDiGraph:

901 """A directed acyclic graph with quantum and dataset nodes.

902

903 This graph never includes init-input and init-output datasets.

904

905 Notes

906 -----

907 Node keys are quantum or dataset UUIDs. Nodes for quanta are present

908 if the ``thin_graph`` component is loaded (all nodes) or if the

909 ``quantum_datasets`` component is loaded (just loaded quanta). Edges

910 and dataset nodes are only present for quanta whose

911 ``quantum_datasets`` were loaded.

912

913 Node state dictionaries are described by the

914 `PredictedQuantumInfo` and `PredictedDatasetInfo` types.

915

916 The returned object is a read-only view of an internal one.

917 """

918 return self._bipartite_xgraph.copy(as_view=True)

919

920 @property

921 def dimension_data(self) -> DimensionDataAttacher | None:

922 """All dimension records needed to expand the data IDS in the graph.

923

924 This may be `None` if the dimension data was not loaded. If all

925 execution quanta have been built, all records are guaranteed to have

926 been deserialized and the ``records`` attribute is complete. In other

927 cases some records may still only be present in the ``deserializers``

928 attribute.

929 """

930 return self._dimension_data

931

932 def __iter__(self) -> Iterator[uuid.UUID]:

933 for quanta_for_task in self.quanta_by_task.values():

934 for data_id in sorted(quanta_for_task.keys()):

935 yield quanta_for_task[data_id]

936

937 def __len__(self) -> int:

938 return len(self._quantum_only_xgraph)

939

940 def get_init_inputs(self, task_label: str) -> dict[ConnectionName, DatasetRef]:

941 """Return the init-input datasets for the given task.

942

943 Parameters

944 ----------

945 task_label : `str`

946 Label of the task.

947

948 Returns

949 -------

950 init_inputs : `dict` [ `str`, `lsst.daf.butler.DatasetRef` ]

951 Dataset references for init-input datasets, keyed by connection

952 name. Dataset types storage classes match the task connection

953 declarations, not necessarily the data repository, and may be

954 components.

955 """

956 if self._init_quanta is None:

957 raise IncompleteQuantumGraphError("The init_quanta component was not loaded.")

958 task_init_node = self.pipeline_graph.tasks[task_label].init

959 return {

960 connection_name: task_init_node.inputs[connection_name].adapt_dataset_ref(

961 self._make_init_ref(datasets[0])

962 )

963 for connection_name, datasets in self._init_quanta[task_label].inputs.items()

964 }

965

966 def get_init_outputs(self, task_label: str) -> dict[ConnectionName, DatasetRef]:

967 """Return the init-output datasets for the given task.

968

969 Parameters

970 ----------

971 task_label : `str`

972 Label of the task. ``""`` may be used to get global init-outputs.

973

974 Returns

975 -------

976 init_outputs : `dict` [ `str`, `lsst.daf.butler.DatasetRef` ]

977 Dataset references for init-outputs datasets, keyed by connection

978 name. Dataset types storage classes match the task connection

979 declarations, not necessarily the data repository.

980 """

981 if self._init_quanta is None:

982 raise IncompleteQuantumGraphError("The init_quanta component was not loaded.")

983 if not task_label:

984 (datasets,) = self._init_quanta[""].outputs.values()

985 return {

986 acc.PACKAGES_INIT_OUTPUT_NAME: DatasetRef(

987 self.pipeline_graph.packages_dataset_type,

988 DataCoordinate.make_empty(self.pipeline_graph.universe),

989 run=datasets[0].run,

990 id=datasets[0].dataset_id,

991 conform=False,

992 )

993 }

994 task_init_node = self.pipeline_graph.tasks[task_label].init

995 result: dict[ConnectionName, DatasetRef] = {}

996 for connection_name, datasets in self._init_quanta[task_label].outputs.items():

997 if connection_name == acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME:

998 edge = task_init_node.config_output

999 else:

1000 edge = task_init_node.outputs[connection_name]

1001 result[connection_name] = edge.adapt_dataset_ref(self._make_init_ref(datasets[0]))

1002 return result

1003

1004 def _make_init_ref(self, dataset: PredictedDatasetModel) -> DatasetRef:

1005 dataset_type = self.pipeline_graph.dataset_types[dataset.dataset_type_name].dataset_type

1006 return DatasetRef(

1007 dataset_type,

1008 DataCoordinate.make_empty(self.pipeline_graph.universe),

1009 run=dataset.run,

1010 id=dataset.dataset_id,

1011 conform=False,

1012 )

1013

1014 def build_execution_quanta(

1015 self,

1016 quantum_ids: Iterable[uuid.UUID] | None = None,

1017 task_label: str | None = None,

1018 ) -> dict[uuid.UUID, Quantum]:

1019 """Build `lsst.daf.butler.Quantum` objects suitable for executing

1020 tasks.

1021

1022 In addition to returning the quantum objects directly, this also causes

1023 the `quantum_only_xgraph` and `bipartite_xgraph` graphs to include a

1024 ``quantum`` attribute for the affected quanta.

1025

1026 Parameters

1027 ----------

1028 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional

1029 IDs of all quanta to return. If not provided, all quanta for the

1030 given task label (if given) or graph are returned.

1031 task_label : `str`, optional

1032 Task label whose quanta should be generated. Ignored if

1033 ``quantum_ids`` is not `None`.

1034

1035 Returns

1036 -------

1037 quanta : `dict` [ `uuid.UUID`, `lsst.daf.butler.Quantum` ]

1038 Mapping of quanta, keyed by UUID. All dataset types are adapted to

1039 the task's storage class declarations and inputs may be components.

1040 All data IDs have dimension records attached.

1041 """

1042 if not self._init_quanta:

1043 raise IncompleteQuantumGraphError(

1044 "Cannot build execution quanta without loading the ``init_quanta`` component."

1045 )

1046 if quantum_ids is None:

1047 if task_label is not None:

1048 quantum_ids = self._quanta_by_task_label[task_label].values()

1049 else:

1050 quantum_ids = self._quantum_only_xgraph.nodes.keys()

1051 else:

1052 # Guard against single-pass iterators.

1053 quantum_ids = list(quantum_ids)

1054 del task_label # make sure we don't accidentally use this.

1055 result: dict[uuid.UUID, Quantum] = {}

1056 self._expand_execution_quantum_data_ids(quantum_ids)

1057 task_init_datastore_records: dict[TaskLabel, dict[DatastoreName, DatastoreRecordData]] = {}

1058 for quantum_id in quantum_ids:

1059 quantum_node_dict: PredictedQuantumInfo = self._quantum_only_xgraph.nodes[quantum_id]

1060 if "quantum" in quantum_node_dict:

1061 result[quantum_id] = quantum_node_dict["quantum"]

1062 continue

1063 # We've declare the info dict keys to all be required because that

1064 # saves a lot of casting, but the reality is that they can either

1065 # be fully populated or totally unpopulated. But that makes mypy

1066 # think the check above always succeeds.

1067 try: # type:ignore [unreachable]

1068 quantum_datasets = self._quantum_datasets[quantum_id]

1069 except KeyError:

1070 raise IncompleteQuantumGraphError(

1071 f"Full quantum information for {quantum_id} was not loaded."

1072 ) from None

1073 task_node = self.pipeline_graph.tasks[quantum_datasets.task_label]

1074 quantum_data_id = self._expanded_data_ids[self._bipartite_xgraph.nodes[quantum_id]["data_id"]]

1075 inputs = self._build_execution_quantum_refs(task_node, quantum_datasets.inputs)

1076 outputs = self._build_execution_quantum_refs(task_node, quantum_datasets.outputs)

1077 if task_node.label not in task_init_datastore_records:

1078 task_init_datastore_records[task_node.label] = self._init_quanta[

1079 task_node.label

1080 ].deserialize_datastore_records()

1081 quantum = Quantum(

1082 taskName=task_node.task_class_name,

1083 taskClass=task_node.task_class,

1084 dataId=quantum_data_id,

1085 initInputs={

1086 ref.datasetType: ref for ref in self.get_init_inputs(quantum_datasets.task_label).values()

1087 },

1088 inputs=inputs,

1089 outputs=outputs,

1090 datastore_records=DatastoreRecordData.merge_mappings(

1091 quantum_datasets.deserialize_datastore_records(),

1092 task_init_datastore_records[task_node.label],

1093 ),

1094 )

1095 self._quantum_only_xgraph.nodes[quantum_id]["quantum"] = quantum

1096 self._bipartite_xgraph.nodes[quantum_id]["quantum"] = quantum

1097 result[quantum_id] = quantum

1098 return result

1099

1100 def _expand_execution_quantum_data_ids(self, quantum_ids: Iterable[uuid.UUID]) -> None:

1101 if self._dimension_data is None:

1102 raise IncompleteQuantumGraphError(

1103 "Cannot build execution quanta without loading the ``dimension_data`` component."

1104 )

1105 data_ids_to_expand: dict[DimensionGroup, set[DataCoordinate]] = defaultdict(set)

1106 for quantum_id in quantum_ids:

1107 data_id: DataCoordinate = self._bipartite_xgraph.nodes[quantum_id]["data_id"]

1108 if data_id.hasRecords():

1109 self._expanded_data_ids[data_id] = data_id

1110 else:

1111 data_ids_to_expand[data_id.dimensions].add(data_id)

1112 for dataset_id in itertools.chain(

1113 self._bipartite_xgraph.predecessors(quantum_id),

1114 self._bipartite_xgraph.successors(quantum_id),

1115 ):

1116 data_id = self._bipartite_xgraph.nodes[dataset_id]["data_id"]

1117 if data_id.hasRecords():

1118 self._expanded_data_ids[data_id] = data_id

1119 else:

1120 data_ids_to_expand[data_id.dimensions].add(data_id)

1121 for dimensions, data_ids_for_dimensions in data_ids_to_expand.items():

1122 self._expanded_data_ids.update(

1123 (d, d) for d in self._dimension_data.attach(dimensions, data_ids_for_dimensions)

1124 )

1125

1126 def _build_execution_quantum_refs(

1127 self, task_node: TaskNode, model_mapping: dict[ConnectionName, list[PredictedDatasetModel]]

1128 ) -> dict[DatasetType, list[DatasetRef]]:

1129 results: dict[DatasetType, list[DatasetRef]] = {}

1130 for connection_name, datasets in model_mapping.items():

1131 edge = task_node.get_edge(connection_name)

1132 dataset_type = edge.adapt_dataset_type(

1133 self.pipeline_graph.dataset_types[edge.parent_dataset_type_name].dataset_type

1134 )

1135 results[dataset_type] = [self._make_general_ref(dataset_type, d.dataset_id) for d in datasets]

1136 return results

1137

1138 def _make_general_ref(self, dataset_type: DatasetType, dataset_id: uuid.UUID) -> DatasetRef:

1139 node_state = self._bipartite_xgraph.nodes[dataset_id]

1140 data_id = self._expanded_data_ids[node_state["data_id"]]

1141 return DatasetRef(dataset_type, data_id, run=node_state["run"], id=dataset_id)

1142

1143 def make_init_qbb(

1144 self,

1145 butler_config: Config | ResourcePathExpression,

1146 *,

1147 config_search_paths: Iterable[str] | None = None,

1148 ) -> QuantumBackedButler:

1149 """Construct an quantum-backed butler suitable for reading and writing

1150 init input and init output datasets, respectively.

1151

1152 This only requires the ``init_quanta`` component to have been loaded.

1153

1154 Parameters

1155 ----------

1156 butler_config : `~lsst.daf.butler.Config` or \

1157 `~lsst.resources.ResourcePathExpression`

1158 A butler repository root, configuration filename, or configuration

1159 instance.

1160 config_search_paths : `~collections.abc.Iterable` [ `str` ], optional

1161 Additional search paths for butler configuration.

1162

1163 Returns

1164 -------

1165 qbb : `~lsst.daf.butler.QuantumBackedButler`

1166 A limited butler that can ``get`` init-input datasets and ``put``

1167 init-output datasets.

1168 """

1169 # Collect all init input/output dataset IDs.

1170 predicted_inputs: set[uuid.UUID] = set()

1171 predicted_outputs: set[uuid.UUID] = set()

1172 datastore_record_maps: list[dict[DatastoreName, DatastoreRecordData]] = []

1173 for init_quantum_datasets in self._init_quanta.values():

1174 predicted_inputs.update(

1175 d.dataset_id for d in itertools.chain.from_iterable(init_quantum_datasets.inputs.values())

1176 )

1177 predicted_outputs.update(

1178 d.dataset_id for d in itertools.chain.from_iterable(init_quantum_datasets.outputs.values())

1179 )

1180 datastore_record_maps.append(

1181 {

1182 datastore_name: DatastoreRecordData.from_simple(serialized_records)

1183 for datastore_name, serialized_records in init_quantum_datasets.datastore_records.items()

1184 }

1185 )

1186 # Remove intermediates from inputs.

1187 predicted_inputs -= predicted_outputs

1188 dataset_types = {d.name: d.dataset_type for d in self.pipeline_graph.dataset_types.values()}

1189 # Make butler from everything.

1190 return QuantumBackedButler.from_predicted(

1191 config=butler_config,

1192 predicted_inputs=predicted_inputs,

1193 predicted_outputs=predicted_outputs,

1194 dimensions=self.pipeline_graph.universe,

1195 datastore_records=DatastoreRecordData.merge_mappings(*datastore_record_maps),

1196 search_paths=list(config_search_paths) if config_search_paths is not None else None,

1197 dataset_types=dataset_types,

1198 )

1199

1200 def write_init_outputs(self, butler: LimitedButler, skip_existing: bool = True) -> None:

1201 """Write the init-output datasets for all tasks in the quantum graph.

1202

1203 This only requires the ``init_quanta`` component to have been loaded.

1204

1205 Parameters

1206 ----------

1207 butler : `lsst.daf.butler.LimitedButler`

1208 A limited butler data repository client.

1209 skip_existing : `bool`, optional

1210 If `True` (default) ignore init-outputs that already exist. If

1211 `False`, raise.

1212

1213 Raises

1214 ------

1215 lsst.daf.butler.registry.ConflictingDefinitionError

1216 Raised if an init-output dataset already exists and

1217 ``skip_existing=False``.

1218 """

1219 # Extract init-input and init-output refs from the QG.

1220 input_refs: dict[str, DatasetRef] = {}

1221 output_refs: dict[str, DatasetRef] = {}

1222 for task_node in self.pipeline_graph.tasks.values():

1223 if task_node.label not in self._init_quanta:

1224 continue

1225 input_refs.update(

1226 {ref.datasetType.name: ref for ref in self.get_init_inputs(task_node.label).values()}

1227 )

1228 output_refs.update(

1229 {

1230 ref.datasetType.name: ref

1231 for ref in self.get_init_outputs(task_node.label).values()

1232 if ref.datasetType.name != task_node.init.config_output.dataset_type_name

1233 }

1234 )

1235 for ref, is_stored in butler.stored_many(output_refs.values()).items():

1236 if is_stored:

1237 if not skip_existing:

1238 raise ConflictingDefinitionError(f"Init-output dataset {ref} already exists.")

1239 # We'll `put` whatever's left in output_refs at the end.

1240 del output_refs[ref.datasetType.name]

1241 # Instantiate tasks, reading overall init-inputs and gathering

1242 # init-output in-memory objects.

1243 init_outputs: list[tuple[Any, DatasetType]] = []

1244 self.pipeline_graph.instantiate_tasks(

1245 get_init_input=lambda dataset_type: butler.get(

1246 input_refs[dataset_type.name].overrideStorageClass(dataset_type.storageClass)

1247 ),

1248 init_outputs=init_outputs,

1249 # A task can be in the pipeline graph without having an init

1250 # quantum if it doesn't have any regular quanta either (e.g. they

1251 # were all skipped), and the _init_quanta has a "" entry for global

1252 # init-outputs that we don't want to pass here.

1253 labels=self.pipeline_graph.tasks.keys() & self._init_quanta.keys(),

1254 )

1255 # Write init-outputs that weren't already present.

1256 for obj, dataset_type in init_outputs:

1257 if new_ref := output_refs.get(dataset_type.name):

1258 assert new_ref.datasetType.storageClass_name == dataset_type.storageClass_name, (

1259 "QG init refs should use task connection storage classes."

1260 )

1261 butler.put(obj, new_ref)

1262

1263 def write_configs(self, butler: LimitedButler, compare_existing: bool = True) -> None:

1264 """Write the config datasets for all tasks in the quantum graph.

1265

1266 Parameters

1267 ----------

1268 butler : `lsst.daf.butler.LimitedButler`

1269 A limited butler data repository client.

1270 compare_existing : `bool`, optional

1271 If `True` check configs that already exist for consistency. If

1272 `False`, always raise if configs already exist.

1273

1274 Raises

1275 ------

1276 lsst.daf.butler.registry.ConflictingDefinitionError

1277 Raised if an config dataset already exists and

1278 ``compare_existing=False``, or if the existing config is not

1279 consistent with the config in the quantum graph.

1280 """

1281 to_put: list[tuple[PipelineTaskConfig, DatasetRef]] = []

1282 for task_node in self.pipeline_graph.tasks.values():

1283 if task_node.label not in self._init_quanta:

1284 continue

1285 dataset_type_name = task_node.init.config_output.dataset_type_name

1286 ref = self.get_init_outputs(task_node.label)[acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME]

1287 try:

1288 old_config = butler.get(ref)

1289 except (LookupError, FileNotFoundError):

1290 old_config = None

1291 if old_config is not None:

1292 if not compare_existing:

1293 raise ConflictingDefinitionError(f"Config dataset {ref} already exists.")

1294 if not task_node.config.compare(old_config, shortcut=False, output=log_config_mismatch):

1295 raise ConflictingDefinitionError(

1296 f"Config does not match existing task config {dataset_type_name!r} in "

1297 "butler; tasks configurations must be consistent within the same run collection."

1298 )

1299 else:

1300 to_put.append((task_node.config, ref))

1301 # We do writes at the end to minimize the mess we leave behind when we

1302 # raise an exception.

1303 for config, ref in to_put:

1304 butler.put(config, ref)

1305

1306 def write_packages(self, butler: LimitedButler, compare_existing: bool = True) -> None:

1307 """Write the 'packages' dataset for the currently-active software

1308 versions.

1309

1310 Parameters

1311 ----------

1312 butler : `lsst.daf.butler.LimitedButler`

1313 A limited butler data repository client.

1314 compare_existing : `bool`, optional

1315 If `True` check packages that already exist for consistency. If

1316 `False`, always raise if the packages dataset already exists.

1317

1318 Raises

1319 ------

1320 lsst.daf.butler.registry.ConflictingDefinitionError

1321 Raised if the packages dataset already exists and is not consistent

1322 with the current packages.

1323 """

1324 new_packages = Packages.fromSystem()

1325 (ref,) = self.get_init_outputs("").values()

1326 try:

1327 packages = butler.get(ref)

1328 except (LookupError, FileNotFoundError):

1329 packages = None

1330 if packages is not None:

1331 if not compare_existing:

1332 raise ConflictingDefinitionError(f"Packages dataset {ref} already exists.")

1333 if compare_packages(packages, new_packages):

1334 # have to remove existing dataset first; butler has no

1335 # replace option.

1336 butler.pruneDatasets([ref], unstore=True, purge=True)

1337 butler.put(packages, ref)

1338 else:

1339 butler.put(new_packages, ref)

1340

1341 def init_output_run(self, butler: LimitedButler, existing: bool = True) -> None:

1342 """Initialize a new output RUN collection by writing init-output

1343 datasets (including configs and packages).

1344

1345 Parameters

1346 ----------

1347 butler : `lsst.daf.butler.LimitedButler`

1348 A limited butler data repository client.

1349 existing : `bool`, optional

1350 If `True` check or ignore outputs that already exist. If

1351 `False`, always raise if an output dataset already exists.

1352

1353 Raises

1354 ------

1355 lsst.daf.butler.registry.ConflictingDefinitionError

1356 Raised if there are existing init output datasets, and either

1357 ``existing=False`` or their contents are not compatible with this

1358 graph.

1359 """

1360 self.write_configs(butler, compare_existing=existing)

1361 self.write_packages(butler, compare_existing=existing)

1362 self.write_init_outputs(butler, skip_existing=existing)

1363

1364 @classmethod

1365 def from_old_quantum_graph(cls, old_quantum_graph: QuantumGraph) -> PredictedQuantumGraph:

1366 """Construct from an old `QuantumGraph` instance.

1367

1368 Parameters

1369 ----------

1370 old_quantum_graph : `QuantumGraph`

1371 Quantum graph to transform.

1372

1373 Returns

1374 -------

1375 predicted_quantum_graph : `PredictedQuantumGraph`

1376 A new predicted quantum graph.

1377 """

1378 return PredictedQuantumGraphComponents.from_old_quantum_graph(old_quantum_graph).assemble()

1379

1380 def to_old_quantum_graph(self) -> QuantumGraph:

1381 """Transform into an old `QuantumGraph` instance.

1382

1383 Returns

1384 -------

1385 old_quantum_graph : `QuantumGraph`

1386 Old quantum graph.

1387

1388 Notes

1389 -----

1390 This can only be called on graphs that have loaded all quantum

1391 datasets, init datasets, and dimension records.

1392 """

1393 from ..graph import QuantumGraph

1394

1395 quanta: dict[TaskDef, set[Quantum]] = {}

1396 quantum_to_quantum_id: dict[Quantum, uuid.UUID] = {}

1397 init_inputs: dict[TaskDef, list[DatasetRef]] = {}

1398 init_outputs: dict[TaskDef, list[DatasetRef]] = {}

1399 for task_def in self.pipeline_graph._iter_task_defs():

1400 if not self._quanta_by_task_label.get(task_def.label):

1401 continue

1402 quanta_for_task: set[Quantum] = set()

1403 for quantum_id, quantum in self.build_execution_quanta(task_label=task_def.label).items():

1404 quanta_for_task.add(quantum)

1405 quantum_to_quantum_id[quantum] = quantum_id

1406 quanta[task_def] = quanta_for_task

1407 init_inputs[task_def] = list(self.get_init_inputs(task_def.label).values())

1408 init_outputs[task_def] = list(self.get_init_outputs(task_def.label).values())

1409 global_init_outputs = list(self.get_init_outputs("").values())

1410 registry_dataset_types = [d.dataset_type for d in self.pipeline_graph.dataset_types.values()]

1411 result = object.__new__(QuantumGraph)

1412 result._buildGraphs(

1413 quanta,

1414 _quantumToNodeId=quantum_to_quantum_id,

1415 metadata=self.header.to_old_metadata(),

1416 universe=self.pipeline_graph.universe,

1417 initInputs=init_inputs,

1418 initOutputs=init_outputs,

1419 globalInitOutputs=global_init_outputs,

1420 registryDatasetTypes=registry_dataset_types,

1421 )

1422 return result

1423

1424 def _make_summary(self) -> QgraphSummary:

1425 from ..graph import QgraphSummary, QgraphTaskSummary

1426

1427 summary = QgraphSummary(

1428 cmdLine=self.header.command or None,

1429 creationUTC=str(self.header.timestamp) if self.header.timestamp is not None else None,

1430 inputCollection=self.header.inputs or None,

1431 outputCollection=self.header.output,

1432 outputRun=self.header.output_run,

1433 )

1434 for task_label, quanta_for_task in self.quanta_by_task.items():

1435 task_summary = QgraphTaskSummary(taskLabel=task_label, numQuanta=len(quanta_for_task))

1436 task_node = self.pipeline_graph.tasks[task_label]

1437 for quantum_id in quanta_for_task.values():

1438 quantum_datasets = self._quantum_datasets[quantum_id]

1439 for connection_name, input_datasets in quantum_datasets.inputs.items():

1440 task_summary.numInputs[

1441 task_node.get_input_edge(connection_name).parent_dataset_type_name

1442 ] += len(input_datasets)

1443 for connection_name, output_datasets in quantum_datasets.outputs.items():

1444 task_summary.numOutputs[

1445 task_node.get_output_edge(connection_name).parent_dataset_type_name

1446 ] += len(output_datasets)

1447 summary.qgraphTaskSummaries[task_label] = task_summary

1448 return summary

1449

1450

1451@dataclasses.dataclass(kw_only=True)

1452class PredictedQuantumGraphComponents:

1453 """A helper class for building and writing predicted quantum graphs.

1454

1455 Notes

1456 -----

1457 This class is a simple struct of model classes to allow different tools

1458 that build predicted quantum graphs to assemble them in whatever order they

1459 prefer. It does not enforce any internal invariants (e.g. the quantum and

1460 dataset counts in the header, different representations of quanta, internal

1461 ID sorting, etc.), but it does provide methods that can satisfy them.

1462 """

1463

1464 def __post_init__(self) -> None:

1465 self.header.graph_type = "predicted"

1466

1467 header: HeaderModel = dataclasses.field(default_factory=HeaderModel)

1468 """Basic metadata about the graph."""

1469

1470 pipeline_graph: PipelineGraph

1471 """Description of the pipeline this graph runs, including all task label

1472 and dataset type definitions.

1473

1474 This may include tasks that do not have any quanta (e.g. due to skipping

1475 already-executed tasks).

1476

1477 This also includes the dimension universe used to construct the graph.

1478 """

1479

1480 dimension_data: DimensionDataAttacher | None = None

1481 """Object that can attach dimension records to data IDs.

1482 """

1483

1484 init_quanta: PredictedInitQuantaModel = dataclasses.field(default_factory=PredictedInitQuantaModel)

1485 """A list of special quanta that describe the init-inputs and init-outputs

1486 of the graph.

1487

1488 Tasks that are included in the pipeline graph but do not have any quanta

1489 may or may not have an init quantum, but tasks that do have regular quanta

1490 always have an init quantum as well.

1491

1492 When used to construct a `PredictedQuantumGraph`, this must have either

1493 zero entries or all tasks in the pipeline.

1494 """

1495

1496 thin_graph: PredictedThinGraphModel = dataclasses.field(default_factory=PredictedThinGraphModel)

1497 """A lightweight quantum-quantum DAG with task labels and data IDs only.

1498

1499 This uses internal integer IDs ("indexes") for node IDs.

1500

1501 This does not include the special "init" quanta.

1502 """

1503

1504 quantum_datasets: dict[uuid.UUID, PredictedQuantumDatasetsModel] = dataclasses.field(default_factory=dict)

1505 """The full descriptions of all quanta, including input and output

1506 dataset, keyed by UUID.

1507

1508 When used to construct a `PredictedQuantumGraph`, this need not have all

1509 entries.

1510

1511 This does not include special "init" quanta.

1512 """

1513

1514 quantum_indices: dict[uuid.UUID, QuantumIndex] = dataclasses.field(default_factory=dict)

1515 """A mapping from external universal quantum ID to internal integer ID.

1516

1517 While this `dict` does not need to be sorted, the internal integer IDs do

1518 need to correspond exactly to ``enumerate(sorted(uuids))``.

1519

1520 When used to construct a `PredictedQuantumGraph`, this must be fully

1521 populated if `thin_graph` is. It can be empty otherwise.

1522

1523 This does include special "init" quanta.

1524 """

1525

1526 def make_dataset_ref(self, predicted: PredictedDatasetModel) -> DatasetRef:

1527 """Make a `lsst.daf.butler.DatasetRef` from information in the

1528 predicted quantum graph.

1529

1530 Parameters

1531 ----------

1532 predicted : `PredictedDatasetModel`

1533 Model for the dataset in the predicted graph.

1534

1535 Returns

1536 -------

1537 ref : `lsst.daf.butler.DatasetRef`

1538 A dataset reference. Data ID will be expanded if and only if

1539 the dimension data has been loaded.

1540 """

1541 try:

1542 dataset_type = self.pipeline_graph.dataset_types[predicted.dataset_type_name].dataset_type

1543 except KeyError:

1544 if predicted.dataset_type_name == acc.PACKAGES_INIT_OUTPUT_NAME:

1545 dataset_type = self.pipeline_graph.packages_dataset_type

1546 else:

1547 raise

1548 data_id = DataCoordinate.from_full_values(dataset_type.dimensions, tuple(predicted.data_coordinate))

1549 if self.dimension_data is not None:

1550 (data_id,) = self.dimension_data.attach(dataset_type.dimensions, [data_id])

1551 return DatasetRef(

1552 dataset_type,

1553 data_id,

1554 run=predicted.run,

1555 id=predicted.dataset_id,

1556 )

1557

1558 def set_quantum_indices(self) -> None:

1559 """Populate the `quantum_indices` component by sorting the UUIDs in the

1560 `init_quanta` and `quantum_datasets` components (which must both be

1561 complete).

1562 """

1563 all_quantum_ids = [q.quantum_id for q in self.init_quanta.root]

1564 all_quantum_ids.extend(self.quantum_datasets.keys())

1565 all_quantum_ids.sort(key=operator.attrgetter("int"))

1566 self.quantum_indices = {quantum_id: index for index, quantum_id in enumerate(all_quantum_ids)}

1567

1568 def set_thin_graph(self) -> None:

1569 """Populate the `thin_graph` component from the `pipeline_graph`,

1570 `quantum_datasets` and `quantum_indices` components (which must all be

1571 complete).

1572 """

1573 bipartite_xgraph = networkx.DiGraph()

1574 self.thin_graph.quanta = {task_label: [] for task_label in self.pipeline_graph.tasks}

1575 graph_quantum_indices = []

1576 for quantum_datasets in self.quantum_datasets.values():

1577 quantum_index = self.quantum_indices[quantum_datasets.quantum_id]

1578 self.thin_graph.quanta[quantum_datasets.task_label].append(

1579 PredictedThinQuantumModel.model_construct(

1580 quantum_index=quantum_index,

1581 data_coordinate=quantum_datasets.data_coordinate,

1582 )

1583 )

1584 for dataset in itertools.chain.from_iterable(quantum_datasets.inputs.values()):

1585 bipartite_xgraph.add_edge(dataset.dataset_id, quantum_index)

1586 for dataset in itertools.chain.from_iterable(quantum_datasets.outputs.values()):

1587 bipartite_xgraph.add_edge(quantum_index, dataset.dataset_id)

1588 graph_quantum_indices.append(quantum_index)

1589 quantum_only_xgraph: networkx.DiGraph = networkx.bipartite.projected_graph(

1590 bipartite_xgraph, graph_quantum_indices

1591 )

1592 self.thin_graph.edges = list(quantum_only_xgraph.edges)

1593

1594 def set_header_counts(self) -> None:

1595 """Populate the quantum and dataset counts in the header from the

1596 `quantum_indices`, `thin_graph`, `init_quanta`, and `quantum_datasets`

1597 components.

1598 """

1599 self.header.n_quanta = len(self.quantum_indices) - len(self.init_quanta.root)

1600 self.header.n_task_quanta = {

1601 task_label: len(thin_quanta) for task_label, thin_quanta in self.thin_graph.quanta.items()

1602 }

1603 all_dataset_ids: set[uuid.UUID] = set()

1604 for quantum_datasets in itertools.chain(self.init_quanta.root, self.quantum_datasets.values()):

1605 all_dataset_ids.update(quantum_datasets.iter_dataset_ids())

1606 self.header.n_datasets = len(all_dataset_ids)

1607

1608 def update_output_run(self, output_run: str) -> None:

1609 """Update the output `~lsst.daf.butler.CollectionType.RUN` collection

1610 name in all datasets and regenerate all output dataset and quantum

1611 UUIDs.

1612

1613 Parameters

1614 ----------

1615 output_run : `str`

1616 New output `~lsst.daf.butler.CollectionType.RUN` collection name.

1617 """

1618 uuid_map: dict[uuid.UUID, uuid.UUID] = {}

1619 # Do all outputs and then all inputs in separate passes so we don't

1620 # need to rely on topological ordering of anything.

1621 for quantum_datasets in itertools.chain(self.init_quanta.root, self.quantum_datasets.values()):

1622 new_quantum_id = generate_uuidv7()

1623 quantum_datasets.quantum_id = new_quantum_id

1624 for output_dataset in itertools.chain.from_iterable(quantum_datasets.outputs.values()):

1625 assert output_dataset.run == self.header.output_run, (

1626 f"Incorrect run {output_dataset.run} for output dataset {output_dataset.dataset_id}."

1627 )

1628 new_dataset_id = generate_uuidv7()

1629 uuid_map[output_dataset.dataset_id] = new_dataset_id

1630 output_dataset.dataset_id = new_dataset_id

1631 output_dataset.run = output_run

1632 for quantum_datasets in itertools.chain(self.init_quanta.root, self.quantum_datasets.values()):

1633 for input_dataset in itertools.chain.from_iterable(quantum_datasets.inputs.values()):

1634 if input_dataset.run == self.header.output_run:

1635 input_dataset.run = output_run

1636 input_dataset.dataset_id = uuid_map.get(

1637 input_dataset.dataset_id,

1638 # This dataset isn't necessary an output of the graph

1639 # just because it's in the output run; the graph could

1640 # have been built with extend_run=True.

1641 input_dataset.dataset_id,

1642 )

1643 # Update the keys of the quantum_datasets dict.

1644 self.quantum_datasets = {qd.quantum_id: qd for qd in self.quantum_datasets.values()}

1645 # Since the UUIDs have changed, the indices need to change, too.

1646 self.set_quantum_indices()

1647 self.set_thin_graph()

1648 # Update the header last, since we use it above to get the old run.

1649 self.header.output_run = output_run

1650

1651 def assemble(self) -> PredictedQuantumGraph:

1652 """Construct a `PredictedQuantumGraph` from these components."""

1653 return PredictedQuantumGraph(self)

1654

1655 @classmethod

1656 def read_execution_quanta(

1657 cls,

1658 uri: ResourcePathExpression,

1659 quantum_ids: Iterable[uuid.UUID] | None = None,

1660 page_size: int = DEFAULT_PAGE_SIZE,

1661 ) -> PredictedQuantumGraphComponents:

1662 """Read one or more executable quanta from a quantum graph file.

1663

1664 Parameters

1665 ----------

1666 uri : convertible to `lsst.resources.ResourcePath`

1667 URI to open. Should have a ``.qg`` extension for new quantum graph

1668 files, or ``.qgraph`` for the old format.

1669 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional

1670 Iterable of quantum IDs to load. If not provided, all quanta will

1671 be loaded. The UUIDs of special init quanta will be ignored.

1672 page_size : `int`, optional

1673 Approximate number of bytes to read at once from address files.

1674 Note that this does not set a page size for *all* reads, but it

1675 does affect the smallest, most numerous reads.

1676

1677 Returns

1678 -------

1679 components : `PredictedQuantumGraphComponents` ]

1680 Components for quantum graph that can build execution quanta for

1681 all of the given IDs.

1682 """

1683 uri = ResourcePath(uri)

1684 if uri.getExtension() == ".qgraph":

1685 _LOG.warning(

1686 f"Reading and converting old quantum graph {uri}. "

1687 "Use the '.qg' extension to write in the new format."

1688 )

1689 from ..graph import QuantumGraph

1690

1691 old_qg = QuantumGraph.loadUri(uri, nodes=quantum_ids)

1692 return PredictedQuantumGraphComponents.from_old_quantum_graph(old_qg)

1693

1694 with PredictedQuantumGraph.open(uri, page_size=page_size) as reader:

1695 reader.read_execution_quanta(quantum_ids)

1696 return reader.components

1697

1698 @classmethod

1699 def from_old_quantum_graph(cls, old_quantum_graph: QuantumGraph) -> PredictedQuantumGraphComponents:

1700 """Construct from an old `QuantumGraph` instance.

1701

1702 Parameters

1703 ----------

1704 old_quantum_graph : `QuantumGraph`

1705 Quantum graph to transform.

1706

1707 Returns

1708 -------

1709 components : `PredictedQuantumGraphComponents`

1710 Components for a new predicted quantum graph.

1711 """

1712 header = HeaderModel.from_old_quantum_graph(old_quantum_graph)

1713 result = cls(header=header, pipeline_graph=old_quantum_graph.pipeline_graph)

1714 result.init_quanta.update_from_old_quantum_graph(old_quantum_graph)

1715 dimension_data_extractor = DimensionDataExtractor.from_dimension_group(

1716 old_quantum_graph.pipeline_graph.get_all_dimensions()

1717 )

1718 for task_node in old_quantum_graph.pipeline_graph.tasks.values():

1719 task_quanta = old_quantum_graph.get_task_quanta(task_node.label)

1720 for quantum_id, quantum in task_quanta.items():

1721 result.quantum_datasets[quantum_id] = PredictedQuantumDatasetsModel.from_execution_quantum(

1722 task_node, quantum, quantum_id

1723 )

1724 dimension_data_extractor.update([cast(DataCoordinate, quantum.dataId)])

1725 for refs in itertools.chain(quantum.inputs.values(), quantum.outputs.values()):

1726 dimension_data_extractor.update(ref.dataId for ref in refs)

1727 result.dimension_data = DimensionDataAttacher(

1728 records=dimension_data_extractor.records.values(),

1729 dimensions=result.pipeline_graph.get_all_dimensions(),

1730 )

1731 result.set_quantum_indices()

1732 result.set_thin_graph()

1733 result.set_header_counts()

1734 return result

1735

1736 def write(

1737 self,

1738 uri: ResourcePathExpression,

1739 *,

1740 zstd_level: int = 10,

1741 zstd_dict_size: int = 32768,

1742 zstd_dict_n_inputs: int = 512,

1743 ) -> None:

1744 """Write the graph to a file.

1745

1746 Parameters

1747 ----------

1748 uri : convertible to `lsst.resources.ResourcePath`

1749 Path to write to. Should have a ``.qg`` extension, or ``.qgraph``

1750 to force writing the old format.

1751 zstd_level : `int`, optional

1752 ZStandard compression level to use on JSON blocks.

1753 zstd_dict_size : `int`, optional

1754 Size of a ZStandard dictionary that shares compression information

1755 across components. Set to zero to disable the dictionary.

1756 Dictionary compression is automatically disabled if the number of

1757 quanta is smaller than ``zstd_dict_n_inputs``.

1758 zstd_dict_n_inputs : `int`, optional

1759 Maximum number of `PredictedQuantumDatasetsModel` JSON

1760 representations to feed the ZStandard dictionary training routine.

1761

1762 Notes

1763 -----

1764 Only a complete predicted quantum graph with all components fully

1765 populated should be written.

1766 """

1767 if self.header.n_quanta + len(self.init_quanta.root) != len(self.quantum_indices):

1768 raise RuntimeError(

1769 f"Cannot save graph after partial read of quanta: expected {self.header.n_quanta}, "

1770 f"got {len(self.quantum_indices)}."

1771 )

1772 uri = ResourcePath(uri)

1773 match uri.getExtension():

1774 case ".qg":

1775 pass

1776 case ".qgraph":

1777 _LOG.warning(

1778 "Converting to an old-format quantum graph.. "

1779 "Use '.qg' instead of '.qgraph' to save in the new format."

1780 )

1781 old_qg = self.assemble().to_old_quantum_graph()

1782 old_qg.saveUri(uri)

1783 return

1784 case ext:

1785 raise ValueError(

1786 f"Unsupported extension {ext!r} for quantum graph; "

1787 "expected '.qg' (or '.qgraph' to force the old format)."

1788 )

1789 cdict: zstandard.ZstdCompressionDict | None = None

1790 cdict_data: bytes | None = None

1791 quantum_datasets_json: dict[uuid.UUID, bytes] = {}

1792 if len(self.quantum_datasets) < zstd_dict_n_inputs:

1793 # ZStandard will fail if we ask to use a compression dict without

1794 # giving it enough data, and it only helps if we have a lot of

1795 # quanta.

1796 zstd_dict_size = 0

1797 if zstd_dict_size:

1798 quantum_datasets_json = {

1799 quantum_model.quantum_id: quantum_model.model_dump_json().encode()

1800 for quantum_model in itertools.islice(self.quantum_datasets.values(), zstd_dict_n_inputs)

1801 }

1802 try:

1803 cdict = zstandard.train_dictionary(

1804 zstd_dict_size,

1805 list(quantum_datasets_json.values()),

1806 level=zstd_level,

1807 )

1808 except zstandard.ZstdError as err:

1809 warnings.warn(f"Not using a compression dictionary: {err}.")

1810 cdict = None

1811 else:

1812 cdict_data = cdict.as_bytes()

1813 compressor = zstandard.ZstdCompressor(level=zstd_level, dict_data=cdict)

1814 with BaseQuantumGraphWriter.open(

1815 uri,

1816 header=self.header,

1817 pipeline_graph=self.pipeline_graph,

1818 indices=self.quantum_indices,

1819 address_filename="quanta",

1820 compressor=compressor,

1821 cdict_data=cdict_data,

1822 ) as writer:

1823 writer.write_single_model("thin_graph", self.thin_graph)

1824 if self.dimension_data is None:

1825 raise IncompleteQuantumGraphError(

1826 "Cannot save predicted quantum graph with no dimension data."

1827 )

1828 serialized_dimension_data = self.dimension_data.serialized()

1829 writer.write_single_model("dimension_data", serialized_dimension_data)

1830 del serialized_dimension_data

1831 writer.write_single_model("init_quanta", self.init_quanta)

1832 with MultiblockWriter.open_in_zip(

1833 writer.zf, "quantum_datasets", writer.int_size

1834 ) as quantum_datasets_mb:

1835 for quantum_model in self.quantum_datasets.values():

1836 if json_data := quantum_datasets_json.get(quantum_model.quantum_id):

1837 quantum_datasets_mb.write_bytes(

1838 quantum_model.quantum_id, writer.compressor.compress(json_data)

1839 )

1840 else:

1841 quantum_datasets_mb.write_model(

1842 quantum_model.quantum_id, quantum_model, writer.compressor

1843 )

1844 writer.address_writer.addresses.append(quantum_datasets_mb.addresses)

1845

1846

1847@dataclasses.dataclass

1848class PredictedQuantumGraphReader(BaseQuantumGraphReader):

1849 """A helper class for reading predicted quantum graphs."""

1850

1851 components: PredictedQuantumGraphComponents = dataclasses.field(init=False)

1852 """Quantum graph components populated by this reader's methods."""

1853

1854 @classmethod

1855 @contextmanager

1856 def open(

1857 cls,

1858 uri: ResourcePathExpression,

1859 *,

1860 page_size: int = DEFAULT_PAGE_SIZE,

1861 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,

1862 ) -> Iterator[PredictedQuantumGraphReader]:

1863 """Construct a reader from a URI.

1864

1865 Parameters

1866 ----------

1867 uri : convertible to `lsst.resources.ResourcePath`

1868 URI to open. Should have a ``.qg`` extension.

1869 page_size : `int`, optional

1870 Approximate number of bytes to read at once from address files.

1871 Note that this does not set a page size for *all* reads, but it

1872 does affect the smallest, most numerous reads.

1873 import_mode : `.pipeline_graph.TaskImportMode`, optional

1874 How to handle importing the task classes referenced in the pipeline

1875 graph.

1876

1877 Returns

1878 -------

1879 reader : `contextlib.AbstractContextManager` [ \

1880 `PredictedQuantumGraphReader` ]

1881 A context manager that returns the reader when entered.

1882 """

1883 with cls._open(

1884 uri,

1885 graph_type="predicted",

1886 address_filename="quanta",

1887 page_size=page_size,

1888 import_mode=import_mode,

1889 n_addresses=1,

1890 ) as self:

1891 yield self

1892

1893 def __post_init__(self) -> None:

1894 self.components = PredictedQuantumGraphComponents(

1895 header=self.header, pipeline_graph=self.pipeline_graph

1896 )

1897

1898 def finish(self) -> PredictedQuantumGraph:

1899 """Construct a `PredictedQuantumGraph` instance from this reader."""

1900 return self.components.assemble()

1901

1902 def read_all(self) -> None:

1903 """Read all components in full."""

1904 self.read_thin_graph()

1905 self.read_execution_quanta()

1906

1907 def read_thin_graph(self) -> None:

1908 """Read the thin graph.

1909

1910 The thin graph is a quantum-quantum DAG with internal integer IDs for

1911 nodes and just task labels and data IDs as node attributes. It always

1912 includes all regular quanta, and does not include init-input or

1913 init-output information.

1914 """

1915 if not self.components.thin_graph.quanta:

1916 self.components.thin_graph = self._read_single_block("thin_graph", PredictedThinGraphModel)

1917 if len(self.components.quantum_indices) != self.components.header.n_quanta:

1918 self.address_reader.read_all()

1919 self.components.quantum_indices.update(

1920 {row.key: row.index for row in self.address_reader.rows.values()}

1921 )

1922

1923 def read_init_quanta(self) -> None:

1924 """Read the list of special quanta that represent init-inputs and

1925 init-outputs.

1926 """

1927 if not self.components.init_quanta.root:

1928 self.components.init_quanta = self._read_single_block("init_quanta", PredictedInitQuantaModel)

1929

1930 def read_dimension_data(self) -> None:

1931 """Read all dimension records.

1932

1933 Record data IDs will be immediately deserialized, while other fields

1934 will be left in serialized form until they are needed.

1935 """

1936 if self.components.dimension_data is None:

1937 serializable_dimension_data = self._read_single_block("dimension_data", SerializableDimensionData)

1938 self.components.dimension_data = DimensionDataAttacher(

1939 deserializers=[

1940 DimensionRecordSetDeserializer.from_raw(

1941 self.components.pipeline_graph.universe[element], serialized_records

1942 )

1943 for element, serialized_records in serializable_dimension_data.root.items()

1944 ],

1945 dimensions=DimensionGroup.union(

1946 *self.components.pipeline_graph.group_by_dimensions(prerequisites=True).keys(),

1947 universe=self.components.pipeline_graph.universe,

1948 ),

1949 )

1950

1951 def read_quantum_datasets(self, quantum_ids: Iterable[uuid.UUID] | None = None) -> None:

1952 """Read information about all datasets produced and consumed by the

1953 given quantum IDs.

1954

1955 Parameters

1956 ----------

1957 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional

1958 Iterable of quantum IDs to load. If not provided, all quanta will

1959 be loaded. The UUIDs of special init quanta will be ignored.

1960 """

1961 quantum_datasets: PredictedQuantumDatasetsModel | None

1962 if quantum_ids is None:

1963 if len(self.components.quantum_datasets) != self.header.n_quanta:

1964 for quantum_datasets in MultiblockReader.read_all_models_in_zip(

1965 self.zf,

1966 "quantum_datasets",

1967 PredictedQuantumDatasetsModel,

1968 self.decompressor,

1969 int_size=self.components.header.int_size,

1970 page_size=self.page_size,

1971 ):

1972 self.components.quantum_datasets.setdefault(quantum_datasets.quantum_id, quantum_datasets)

1973 self.address_reader.read_all()

1974 for address_row in self.address_reader.rows.values():

1975 self.components.quantum_indices[address_row.key] = address_row.index

1976 return

1977 with MultiblockReader.open_in_zip(

1978 self.zf, "quantum_datasets", int_size=self.components.header.int_size

1979 ) as mb_reader:

1980 for quantum_id in quantum_ids:

1981 if quantum_id in self.components.quantum_datasets:

1982 continue

1983 address_row = self.address_reader.find(quantum_id)

1984 self.components.quantum_indices[address_row.key] = address_row.index

1985 quantum_datasets = mb_reader.read_model(

1986 address_row.addresses[0], PredictedQuantumDatasetsModel, self.decompressor

1987 )

1988 if quantum_datasets is not None:

1989 self.components.quantum_datasets[address_row.key] = quantum_datasets

1990 return

1991

1992 def read_execution_quanta(self, quantum_ids: Iterable[uuid.UUID] | None = None) -> None:

1993 """Read all information needed to execute the given quanta.

1994

1995 Parameters

1996 ----------

1997 quantum_ids : `~collections.abc.Iterable` [ `uuid.UUID` ], optional

1998 Iterable of quantum IDs to load. If not provided, all quanta will

1999 be loaded. The UUIDs of special init quanta will be ignored.

2000 """

2001 self.read_init_quanta()

2002 self.read_dimension_data()

2003 self.read_quantum_datasets(quantum_ids)

Coverage for python/lsst/pipe/base/quantum_graph/_predicted.py: 20%

608 statements