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 sys

47import uuid

48import warnings

49from collections import defaultdict

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

51from contextlib import AbstractContextManager, contextmanager

52from typing import TYPE_CHECKING, Any, TypeVar, cast

54import networkx

55import networkx.algorithms.bipartite

56import pydantic

57import zstandard

59from lsst.daf.butler import (

60 Config,

61 DataCoordinate,

62 DataIdValue,

63 DatasetRef,

64 DatasetType,

65 DimensionDataAttacher,

66 DimensionDataExtractor,

67 DimensionGroup,

68 DimensionRecordSetDeserializer,

69 LimitedButler,

70 Quantum,

71 QuantumBackedButler,

72 SerializableDimensionData,

73)

74from lsst.daf.butler._rubin import generate_uuidv7

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

76from lsst.daf.butler.registry import ConflictingDefinitionError

77from lsst.resources import ResourcePath, ResourcePathExpression

78from lsst.utils.packages import Packages

80from .. import automatic_connection_constants as acc

81from ..pipeline import TaskDef

82from ..pipeline_graph import (

83 PipelineGraph,

84 TaskImportMode,

85 TaskInitNode,

86 TaskNode,

87 compare_packages,

88 log_config_mismatch,

89)

90from ._common import (

91 FORMAT_VERSION,

92 BaseQuantumGraph,

93 BaseQuantumGraphReader,

94 BaseQuantumGraphWriter,

95 ConnectionName,

96 DataCoordinateValues,

97 DatasetInfo,

98 DatasetTypeName,

99 DatastoreName,

100 HeaderModel,

101 IncompleteQuantumGraphError,

102 QuantumIndex,

103 QuantumInfo,

104 TaskLabel,

105)

106from ._multiblock import DEFAULT_PAGE_SIZE, AddressRow, MultiblockReader, MultiblockWriter

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 _PredictedThinQuantumModelV0(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

130class PredictedThinQuantumModel(pydantic.BaseModel):

131 """Data model for a quantum data ID and UUID in a predicted

132 quantum graph.

133 """

134

135 quantum_id: uuid.UUID

136 """Universally unique ID for this quantum."""

137

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

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

140

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

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

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

144

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

146 """See `pydantic.BaseModel.copy`."""

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

148

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

150 """See `pydantic.BaseModel.model_dump`."""

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

152

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

154 """See `pydantic.BaseModel.model_dump_json`."""

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

156

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

158 """See `pydantic.BaseModel.model_copy`."""

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

160

161 @classmethod

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

163 """See `pydantic.BaseModel.model_construct`."""

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

165

166 @classmethod

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

168 """See `pydantic.BaseModel.model_json_schema`."""

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

170

171 @classmethod

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

173 """See `pydantic.BaseModel.model_validate`."""

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

175

176 @classmethod

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

178 """See `pydantic.BaseModel.model_validate_json`."""

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

180

181 @classmethod

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

183 """See `pydantic.BaseModel.model_validate_strings`."""

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

185

186

187class _PredictedThinGraphModelV0(pydantic.BaseModel):

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

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

190 """

191

192 quanta: dict[TaskLabel, list[_PredictedThinQuantumModelV0]] = pydantic.Field(default_factory=dict)

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

194

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

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

197

198 def _upgraded(self, address_rows: Mapping[uuid.UUID, AddressRow]) -> PredictedThinGraphModel:

199 """Convert to the v1+ model."""

200 uuid_by_index = {v.index: k for k, v in address_rows.items()}

201 return PredictedThinGraphModel.model_construct(

202 quanta={

203 task_label: [

204 PredictedThinQuantumModel.model_construct(

205 quantum_id=uuid_by_index[q.quantum_index], data_coordinate=q.data_coordinate

206 )

207 for q in quanta

208 ]

209 for task_label, quanta in self.quanta.items()

210 },

211 edges=[(uuid_by_index[index1], uuid_by_index[index2]) for index1, index2 in self.edges],

212 )

213

214

215class PredictedThinGraphModel(pydantic.BaseModel):

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

217 task label to the data IDs and UUIDs of its quanta.

218 """

219

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

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

222

223 edges: list[tuple[uuid.UUID, uuid.UUID]] = pydantic.Field(default_factory=list)

224 """Pairs of (predecessor, successor) quantum IDs."""

225

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

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

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

229

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

231 """See `pydantic.BaseModel.copy`."""

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

233

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

235 """See `pydantic.BaseModel.model_dump`."""

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

237

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

239 """See `pydantic.BaseModel.model_dump_json`."""

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

241

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

243 """See `pydantic.BaseModel.model_copy`."""

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

245

246 @classmethod

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

248 """See `pydantic.BaseModel.model_construct`."""

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

250

251 @classmethod

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

253 """See `pydantic.BaseModel.model_json_schema`."""

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

255

256 @classmethod

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

258 """See `pydantic.BaseModel.model_validate`."""

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

260

261 @classmethod

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

263 """See `pydantic.BaseModel.model_validate_json`."""

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

265

266 @classmethod

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

268 """See `pydantic.BaseModel.model_validate_strings`."""

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

270

271

272class PredictedDatasetModel(pydantic.BaseModel):

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

274

275 dataset_id: uuid.UUID

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

277

278 dataset_type_name: DatasetTypeName

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

280

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

282

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

284 """

285

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

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

288

289 run: str

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

291

292 @classmethod

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

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

295

296 Parameters

297 ----------

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

299 Dataset reference.

300

301 Returns

302 -------

303 model : `PredictedDatasetModel`

304 Model for the dataset.

305 """

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

307 return cls.model_construct(

308 dataset_id=ref.id,

309 dataset_type_name=dataset_type_name,

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

311 run=ref.run,

312 )

313

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

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

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

317

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

319 """See `pydantic.BaseModel.copy`."""

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

321

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

323 """See `pydantic.BaseModel.model_dump`."""

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

325

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

327 """See `pydantic.BaseModel.model_dump_json`."""

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

329

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

331 """See `pydantic.BaseModel.model_copy`."""

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

333

334 @classmethod

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

336 """See `pydantic.BaseModel.model_construct`."""

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

338

339 @classmethod

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

341 """See `pydantic.BaseModel.model_json_schema`."""

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

343

344 @classmethod

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

346 """See `pydantic.BaseModel.model_validate`."""

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

348

349 @classmethod

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

351 """See `pydantic.BaseModel.model_validate_json`."""

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

353

354 @classmethod

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

356 """See `pydantic.BaseModel.model_validate_strings`."""

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

358

359

360class PredictedQuantumDatasetsModel(pydantic.BaseModel):

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

362 its inputs and outputs.

363 """

364

365 quantum_id: uuid.UUID

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

367

368 task_label: TaskLabel

369 """Label of the task.

370

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

372 """

373

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

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

376

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

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

379

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

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

382

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

384 default_factory=dict

385 )

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

387 the data repository.

388 """

389

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

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

392 quantum.

393

394 Returns

395 -------

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

397 Iterator over dataset IDs.

398 """

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

400 for dataset in datasets:

401 yield dataset.dataset_id

402

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

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

405 quantum.

406

407 Returns

408 -------

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

410 Iterator over dataset IDs.

411 """

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

413 for dataset in datasets:

414 yield dataset.dataset_id

415

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

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

418 quantum.

419

420 Returns

421 -------

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

423 Iterator over dataset IDs.

424 """

425 yield from self.iter_input_dataset_ids()

426 yield from self.iter_output_dataset_ids()

427

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

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

430 return {

431 datastore_name: DatastoreRecordData.from_simple(serialized_records)

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

433 }

434

435 @classmethod

436 def from_execution_quantum(

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

438 ) -> PredictedQuantumDatasetsModel:

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

440

441 Parameters

442 ----------

443 task_node : `.pipeline_graph.TaskNode`

444 Task node from the pipeline graph.

445 quantum : `lsst.daf.butler.quantum`

446 Quantum object.

447 quantum_id : `uuid.UUID`

448 ID for this quantum.

449

450 Returns

451 -------

452 model : `PredictedFullQuantumModel`

453 Model for this quantum.

454 """

455 result: PredictedQuantumDatasetsModel = cls.model_construct(

456 quantum_id=quantum_id,

457 task_label=task_node.label,

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

459 )

460 for read_edge in task_node.iter_all_inputs():

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

462 result.inputs[read_edge.connection_name] = [

463 PredictedDatasetModel.from_dataset_ref(ref) for ref in refs

464 ]

465 for write_edge in task_node.iter_all_outputs():

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

467 result.outputs[write_edge.connection_name] = [

468 PredictedDatasetModel.from_dataset_ref(ref) for ref in refs

469 ]

470 result.datastore_records = {

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

472 }

473 return result

474

475 @classmethod

476 def from_old_quantum_graph_init(

477 cls, task_init_node: TaskInitNode, old_quantum_graph: QuantumGraph

478 ) -> PredictedQuantumDatasetsModel:

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

480 task in an old `QuantumGraph` instance.

481

482 Parameters

483 ----------

484 task_init_node : `.pipeline_graph.TaskNode`

485 Task init node from the pipeline graph.

486 old_quantum_graph : `QuantumGraph`

487 Quantum graph.

488

489 Returns

490 -------

491 model : `PredictedFullQuantumModel`

492 Model for this "init" quantum.

493 """

494 task_def = old_quantum_graph.findTaskDefByLabel(task_init_node.label)

495 assert task_def is not None

496 init_input_refs = {

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

498 }

499 init_output_refs = {

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

501 }

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

503 result: PredictedQuantumDatasetsModel = cls.model_construct(

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

505 )

506 for read_edge in task_init_node.iter_all_inputs():

507 ref = init_input_refs[read_edge.dataset_type_name]

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

509 for write_edge in task_init_node.iter_all_outputs():

510 ref = init_output_refs[write_edge.dataset_type_name]

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

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

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

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

515 subset = records.subset(init_input_ids)

516 if subset is not None:

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

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

519 result.datastore_records = {

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

521 }

522 return result

523

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

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

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

527

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

529 """See `pydantic.BaseModel.copy`."""

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

531

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

533 """See `pydantic.BaseModel.model_dump`."""

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

535

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

537 """See `pydantic.BaseModel.model_dump_json`."""

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

539

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

541 """See `pydantic.BaseModel.model_copy`."""

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

543

544 @classmethod

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

546 """See `pydantic.BaseModel.model_construct`."""

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

548

549 @classmethod

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

551 """See `pydantic.BaseModel.model_json_schema`."""

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

553

554 @classmethod

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

556 """See `pydantic.BaseModel.model_validate`."""

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

558

559 @classmethod

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

561 """See `pydantic.BaseModel.model_validate_json`."""

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

563

564 @classmethod

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

566 """See `pydantic.BaseModel.model_validate_strings`."""

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

568

569

570class PredictedInitQuantaModel(pydantic.RootModel):

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

572 graph.

573 """

574

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

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

577 init-outputs.

578 """

579

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

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

582 instance.

583

584 Parameters

585 ----------

586 old_quantum_graph : `QuantumGraph`

587 Quantum graph.

588 """

589 global_init_quantum = PredictedQuantumDatasetsModel.model_construct(

590 quantum_id=generate_uuidv7(), task_label=""

591 )

592 for ref in old_quantum_graph.globalInitOutputRefs():

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

594 self.root.append(global_init_quantum)

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

596 self.root.append(

597 PredictedQuantumDatasetsModel.from_old_quantum_graph_init(task_node.init, old_quantum_graph)

598 )

599

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

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

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

603

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

605 """See `pydantic.BaseModel.copy`."""

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

607

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

609 """See `pydantic.BaseModel.model_dump`."""

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

611

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

613 """See `pydantic.BaseModel.model_dump_json`."""

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

615

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

617 """See `pydantic.BaseModel.model_copy`."""

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

619

620 @classmethod

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

622 """See `pydantic.BaseModel.model_construct`."""

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

624

625 @classmethod

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

627 """See `pydantic.BaseModel.model_json_schema`."""

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

629

630 @classmethod

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

632 """See `pydantic.BaseModel.model_validate`."""

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

634

635 @classmethod

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

637 """See `pydantic.BaseModel.model_validate_json`."""

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

639

640 @classmethod

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

642 """See `pydantic.BaseModel.model_validate_strings`."""

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

644

645

646class PredictedQuantumInfo(QuantumInfo):

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

648 node data for a predicted quantum.

649

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

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

652

653 node_data: PredictedQuantumInfo = xgraph.nodes[quantum_id]

654

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

656 or `PredictedQuantumGraph.bipartite_xgraph`.

657 """

658

659 quantum: Quantum

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

661

662 This attribute is only present if

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

664 quantum ID already.

665 """

666

667

668class PredictedDatasetInfo(DatasetInfo):

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

670 node data for a dataset.

671

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

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

674

675 node_data: PredictedDatasetInfo = xgraph.nodes[dataset_ids]

676

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

678 property.

679 """

680

681

682class PredictedQuantumGraph(BaseQuantumGraph):

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

684 during execution.

685

686 Parameters

687 ----------

688 components : `PredictedQuantumGraphComponents`

689 A struct of components used to construct the graph.

690

691 Notes

692 -----

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

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

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

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

697 """

698

699 def __init__(self, components: PredictedQuantumGraphComponents):

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

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

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

703 self._quantum_only_xgraph = networkx.DiGraph()

704 self._bipartite_xgraph = networkx.DiGraph()

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

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

707 }

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

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

710 }

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

712 self._dimension_data = components.dimension_data

713 self._add_init_quanta(components.init_quanta)

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

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

716 self._add_thin_graph(components.thin_graph)

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

718 self._add_quantum_datasets(quantum_datasets)

719 if not components.thin_graph.edges:

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

721 self._quantum_only_xgraph.update(

722 networkx.algorithms.bipartite.projected_graph(

723 networkx.DiGraph(self._bipartite_xgraph),

724 self._quantum_only_xgraph.nodes.keys(),

725 )

726 )

727 if _LOG.isEnabledFor(logging.DEBUG):

728 for quantum_id in self:

729 _LOG.debug(

730 "%s: %s @ %s",

731 quantum_id,

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

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

734 )

735

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

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

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

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

740 for init_datasets in itertools.chain(

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

742 ):

743 for init_dataset in init_datasets:

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

745 init_dataset.dataset_id

746 )

747 _LOG.debug(

748 "%s: %s @ init",

749 quantum_datasets.quantum_id,

750 quantum_datasets.task_label,

751 )

752

753 def _add_thin_graph(self, component: PredictedThinGraphModel) -> None:

754 self._quantum_only_xgraph.add_edges_from(component.edges)

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

756 for thin_quantum in thin_quanta_for_task:

757 self._add_quantum(thin_quantum.quantum_id, task_label, thin_quantum.data_coordinate)

758

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

760 self._quantum_datasets[quantum_datasets.quantum_id] = quantum_datasets

761 self._add_quantum(

762 quantum_datasets.quantum_id, quantum_datasets.task_label, quantum_datasets.data_coordinate

763 )

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

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

766 pipeline_edge = task_node.get_input_edge(connection_name)

767 for input_dataset in input_datasets:

768 self._add_dataset(input_dataset)

769 self._bipartite_xgraph.add_edge(

770 input_dataset.dataset_id,

771 quantum_datasets.quantum_id,

772 key=connection_name,

773 is_read=True,

774 )

775 # There might be multiple input connections for the same

776 # dataset type.

777 self._bipartite_xgraph.edges[

778 input_dataset.dataset_id, quantum_datasets.quantum_id

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

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

781 pipeline_edges = [task_node.get_output_edge(connection_name)]

782 for output_dataset in output_datasets:

783 self._add_dataset(output_dataset)

784 self._bipartite_xgraph.add_edge(

785 quantum_datasets.quantum_id,

786 output_dataset.dataset_id,

787 key=connection_name,

788 is_read=False,

789 pipeline_edges=pipeline_edges,

790 )

791

792 def _add_quantum(

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

794 ) -> None:

795 task_node = self.pipeline_graph.tasks[task_label]

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

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

798 data_coordinate_values = tuple(data_coordinate_values)

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

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

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

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

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

804

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

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

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

808 self._bipartite_xgraph.add_node(

809 model.dataset_id,

810 dataset_type_name=dataset_type_node.name,

811 pipeline_node=dataset_type_node,

812 run=model.run,

813 )

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

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

816

817 @classmethod

818 def open(

819 cls,

820 uri: ResourcePathExpression,

821 page_size: int = DEFAULT_PAGE_SIZE,

822 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,

823 ) -> AbstractContextManager[PredictedQuantumGraphReader]:

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

825

826 Parameters

827 ----------

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

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

830 page_size : `int`, optional

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

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

833 does affect the smallest, most numerous reads.

834 import_mode : `.pipeline_graph.TaskImportMode`, optional

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

836 graph.

837

838 Returns

839 -------

840 reader : `contextlib.AbstractContextManager` [ \

841 `PredictedQuantumGraphReader` ]

842 A context manager that returns the reader when entered.

843 """

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

845

846 @classmethod

847 def read_execution_quanta(

848 cls,

849 uri: ResourcePathExpression,

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

851 page_size: int = DEFAULT_PAGE_SIZE,

852 ) -> PredictedQuantumGraph:

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

854

855 Parameters

856 ----------

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

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

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

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

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

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

863 page_size : `int`, optional

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

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

866 does affect the smallest, most numerous reads.

867

868 Returns

869 -------

870 quantum_graph : `PredictedQuantumGraph` ]

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

872 given IDs.

873 """

874 return PredictedQuantumGraphComponents.read_execution_quanta(

875 uri,

876 quantum_ids,

877 page_size=page_size,

878 ).assemble()

879

880 @property

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

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

883 data ID.

884

885 Notes

886 -----

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

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

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

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

891

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

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

894 """

895 return self._quanta_by_task_label

896

897 @property

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

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

900 and then by data ID.

901

902 Notes

903 -----

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

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

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

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

908

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

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

911 """

912 return self._datasets_by_type

913

914 @property

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

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

917

918 Notes

919 -----

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

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

922 (only those that were loaded).

923

924 Node state dictionaries are described by the

925 `PredictedQuantumInfo` type.

926

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

928 """

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

930

931 @property

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

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

934

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

936

937 Notes

938 -----

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

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

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

942 and dataset nodes are only present for quanta whose

943 ``quantum_datasets`` were loaded.

944

945 Node state dictionaries are described by the

946 `PredictedQuantumInfo` and `PredictedDatasetInfo` types.

947

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

949 """

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

951

952 @property

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

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

955

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

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

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

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

960 attribute.

961 """

962 return self._dimension_data

963

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

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

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

967 yield quanta_for_task[data_id]

968

969 def __len__(self) -> int:

970 return len(self._quantum_only_xgraph)

971

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

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

974

975 Parameters

976 ----------

977 task_label : `str`

978 Label of the task.

979

980 Returns

981 -------

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

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

984 name. Dataset types storage classes match the task connection

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

986 components.

987 """

988 if self._init_quanta is None:

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

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

991 return {

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

993 self._make_init_ref(datasets[0])

994 )

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

996 }

997

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

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

1000

1001 Parameters

1002 ----------

1003 task_label : `str`

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

1005

1006 Returns

1007 -------

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

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

1010 name. Dataset types storage classes match the task connection

1011 declarations, not necessarily the data repository.

1012 """

1013 if self._init_quanta is None:

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

1015 if not task_label:

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

1017 return {

1018 acc.PACKAGES_INIT_OUTPUT_NAME: DatasetRef(

1019 self.pipeline_graph.packages_dataset_type,

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

1021 run=datasets[0].run,

1022 id=datasets[0].dataset_id,

1023 conform=False,

1024 )

1025 }

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

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

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

1029 if connection_name == acc.CONFIG_INIT_OUTPUT_CONNECTION_NAME:

1030 edge = task_init_node.config_output

1031 else:

1032 edge = task_init_node.outputs[connection_name]

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

1034 return result

1035

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

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

1038 return DatasetRef(

1039 dataset_type,

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

1041 run=dataset.run,

1042 id=dataset.dataset_id,

1043 conform=False,

1044 )

1045

1046 def build_execution_quanta(

1047 self,

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

1049 task_label: str | None = None,

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

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

1052 tasks.

1053

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

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

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

1057

1058 Parameters

1059 ----------

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

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

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

1063 task_label : `str`, optional

1064 Task label whose quanta should be generated. Ignored if

1065 ``quantum_ids`` is not `None`.

1066

1067 Returns

1068 -------

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

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

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

1072 All data IDs have dimension records attached.

1073 """

1074 if not self._init_quanta:

1075 raise IncompleteQuantumGraphError(

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

1077 )

1078 if quantum_ids is None:

1079 if task_label is not None:

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

1081 else:

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

1083 else:

1084 # Guard against single-pass iterators.

1085 quantum_ids = list(quantum_ids)

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

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

1088 self._expand_execution_quantum_data_ids(quantum_ids)

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

1090 for quantum_id in quantum_ids:

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

1092 if "quantum" in quantum_node_dict:

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

1094 continue

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

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

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

1098 # think the check above always succeeds.

1099 try: # type:ignore [unreachable]

1100 quantum_datasets = self._quantum_datasets[quantum_id]

1101 except KeyError:

1102 raise IncompleteQuantumGraphError(

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

1104 ) from None

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

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

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

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

1109 if task_node.label not in task_init_datastore_records:

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

1111 task_node.label

1112 ].deserialize_datastore_records()

1113 quantum = Quantum(

1114 taskName=task_node.task_class_name,

1115 taskClass=task_node.task_class,

1116 dataId=quantum_data_id,

1117 initInputs={

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

1119 },

1120 inputs=inputs,

1121 outputs=outputs,

1122 datastore_records=DatastoreRecordData.merge_mappings(

1123 quantum_datasets.deserialize_datastore_records(),

1124 task_init_datastore_records[task_node.label],

1125 ),

1126 )

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

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

1129 result[quantum_id] = quantum

1130 return result

1131

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

1133 if self._dimension_data is None:

1134 raise IncompleteQuantumGraphError(

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

1136 )

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

1138 for quantum_id in quantum_ids:

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

1140 if data_id.hasRecords():

1141 self._expanded_data_ids[data_id] = data_id

1142 else:

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

1144 for dataset_id in itertools.chain(

1145 self._bipartite_xgraph.predecessors(quantum_id),

1146 self._bipartite_xgraph.successors(quantum_id),

1147 ):

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

1149 if data_id.hasRecords():

1150 self._expanded_data_ids[data_id] = data_id

1151 else:

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

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

1154 self._expanded_data_ids.update(

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

1156 )

1157

1158 def _build_execution_quantum_refs(

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

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

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

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

1163 edge = task_node.get_edge(connection_name)

1164 dataset_type = edge.adapt_dataset_type(

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

1166 )

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

1168 return results

1169

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

1171 node_state = self._bipartite_xgraph.nodes[dataset_id]

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

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

1174

1175 def make_init_qbb(

1176 self,

1177 butler_config: Config | ResourcePathExpression,

1178 *,

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

1180 ) -> QuantumBackedButler:

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

1182 init input and init output datasets, respectively.

1183

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

1185

1186 Parameters

1187 ----------

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

1189 `~lsst.resources.ResourcePathExpression`

1190 A butler repository root, configuration filename, or configuration

1191 instance.

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

1193 Additional search paths for butler configuration.

1194

1195 Returns

1196 -------

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

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

1199 init-output datasets.

1200 """

1201 # Collect all init input/output dataset IDs.

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

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

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

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

1206 predicted_inputs.update(

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

1208 )

1209 predicted_outputs.update(

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

1211 )

1212 datastore_record_maps.append(

1213 {

1214 datastore_name: DatastoreRecordData.from_simple(serialized_records)

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

1216 }

1217 )

1218 # Remove intermediates from inputs.

1219 predicted_inputs -= predicted_outputs

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

1221 # Make butler from everything.

1222 return QuantumBackedButler.from_predicted(

1223 config=butler_config,

1224 predicted_inputs=predicted_inputs,

1225 predicted_outputs=predicted_outputs,

1226 dimensions=self.pipeline_graph.universe,

1227 datastore_records=DatastoreRecordData.merge_mappings(*datastore_record_maps),

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

1229 dataset_types=dataset_types,

1230 )

1231

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

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

1234

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

1236

1237 Parameters

1238 ----------

1239 butler : `lsst.daf.butler.LimitedButler`

1240 A limited butler data repository client.

1241 skip_existing : `bool`, optional

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

1243 `False`, raise.

1244

1245 Raises

1246 ------

1247 lsst.daf.butler.registry.ConflictingDefinitionError

1248 Raised if an init-output dataset already exists and

1249 ``skip_existing=False``.

1250 """

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

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

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

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

1255 if task_node.label not in self._init_quanta:

1256 continue

1257 input_refs.update(

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

1259 )

1260 output_refs.update(

1261 {

1262 ref.datasetType.name: ref

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

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

1265 }

1266 )

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

1268 if is_stored:

1269 if not skip_existing:

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

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

1272 del output_refs[ref.datasetType.name]

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

1274 # init-output in-memory objects.

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

1276 self.pipeline_graph.instantiate_tasks(

1277 get_init_input=lambda dataset_type: butler.get(

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

1279 ),

1280 init_outputs=init_outputs,

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

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

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

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

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

1286 )

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

1288 for obj, dataset_type in init_outputs:

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

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

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

1292 )

1293 butler.put(obj, new_ref)

1294

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

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

1297

1298 Parameters

1299 ----------

1300 butler : `lsst.daf.butler.LimitedButler`

1301 A limited butler data repository client.

1302 compare_existing : `bool`, optional

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

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

1305

1306 Raises

1307 ------

1308 lsst.daf.butler.registry.ConflictingDefinitionError

1309 Raised if an config dataset already exists and

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

1311 consistent with the config in the quantum graph.

1312 """

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

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

1315 if task_node.label not in self._init_quanta:

1316 continue

1317 dataset_type_name = task_node.init.config_output.dataset_type_name

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

1319 try:

1320 old_config = butler.get(ref)

1321 except (LookupError, FileNotFoundError):

1322 old_config = None

1323 if old_config is not None:

1324 if not compare_existing:

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

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

1327 raise ConflictingDefinitionError(

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

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

1330 )

1331 else:

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

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

1334 # raise an exception.

1335 for config, ref in to_put:

1336 butler.put(config, ref)

1337

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

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

1340 versions.

1341

1342 Parameters

1343 ----------

1344 butler : `lsst.daf.butler.LimitedButler`

1345 A limited butler data repository client.

1346 compare_existing : `bool`, optional

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

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

1349

1350 Raises

1351 ------

1352 lsst.daf.butler.registry.ConflictingDefinitionError

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

1354 with the current packages.

1355 """

1356 new_packages = Packages.fromSystem()

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

1358 try:

1359 packages = butler.get(ref)

1360 except (LookupError, FileNotFoundError):

1361 packages = None

1362 if packages is not None:

1363 if not compare_existing:

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

1365 if compare_packages(packages, new_packages):

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

1367 # replace option.

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

1369 butler.put(packages, ref)

1370 else:

1371 butler.put(new_packages, ref)

1372

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

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

1375 datasets (including configs and packages).

1376

1377 Parameters

1378 ----------

1379 butler : `lsst.daf.butler.LimitedButler`

1380 A limited butler data repository client.

1381 existing : `bool`, optional

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

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

1384

1385 Raises

1386 ------

1387 lsst.daf.butler.registry.ConflictingDefinitionError

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

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

1390 graph.

1391 """

1392 self.write_configs(butler, compare_existing=existing)

1393 self.write_packages(butler, compare_existing=existing)

1394 self.write_init_outputs(butler, skip_existing=existing)

1395

1396 @classmethod

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

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

1399

1400 Parameters

1401 ----------

1402 old_quantum_graph : `QuantumGraph`

1403 Quantum graph to transform.

1404

1405 Returns

1406 -------

1407 predicted_quantum_graph : `PredictedQuantumGraph`

1408 A new predicted quantum graph.

1409 """

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

1411

1412 def to_old_quantum_graph(self) -> QuantumGraph:

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

1414

1415 Returns

1416 -------

1417 old_quantum_graph : `QuantumGraph`

1418 Old quantum graph.

1419

1420 Notes

1421 -----

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

1423 datasets, init datasets, and dimension records.

1424 """

1425 from ..graph import QuantumGraph

1426

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

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

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

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

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

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

1433 continue

1434 quanta_for_task: set[Quantum] = set()

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

1436 quanta_for_task.add(quantum)

1437 quantum_to_quantum_id[quantum] = quantum_id

1438 quanta[task_def] = quanta_for_task

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

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

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

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

1443 result = object.__new__(QuantumGraph)

1444 result._buildGraphs(

1445 quanta,

1446 _quantumToNodeId=quantum_to_quantum_id,

1447 metadata=self.header.to_old_metadata(),

1448 universe=self.pipeline_graph.universe,

1449 initInputs=init_inputs,

1450 initOutputs=init_outputs,

1451 globalInitOutputs=global_init_outputs,

1452 registryDatasetTypes=registry_dataset_types,

1453 )

1454 return result

1455

1456 def _make_summary(self) -> QgraphSummary:

1457 from ..graph import QgraphSummary, QgraphTaskSummary

1458

1459 summary = QgraphSummary(

1460 cmdLine=self.header.command or None,

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

1462 inputCollection=self.header.inputs or None,

1463 outputCollection=self.header.output,

1464 outputRun=self.header.output_run,

1465 )

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

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

1468 task_node = self.pipeline_graph.tasks[task_label]

1469 for quantum_id in quanta_for_task.values():

1470 quantum_datasets = self._quantum_datasets[quantum_id]

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

1472 task_summary.numInputs[

1473 task_node.get_input_edge(connection_name).parent_dataset_type_name

1474 ] += len(input_datasets)

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

1476 task_summary.numOutputs[

1477 task_node.get_output_edge(connection_name).parent_dataset_type_name

1478 ] += len(output_datasets)

1479 summary.qgraphTaskSummaries[task_label] = task_summary

1480 return summary

1481

1482

1483@dataclasses.dataclass(kw_only=True)

1484class PredictedQuantumGraphComponents:

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

1486

1487 Notes

1488 -----

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

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

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

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

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

1494 """

1495

1496 def __post_init__(self) -> None:

1497 self.header.graph_type = "predicted"

1498

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

1500 """Basic metadata about the graph."""

1501

1502 pipeline_graph: PipelineGraph

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

1504 and dataset type definitions.

1505

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

1507 already-executed tasks).

1508

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

1510 """

1511

1512 dimension_data: DimensionDataAttacher | None = None

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

1514 """

1515

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

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

1518 of the graph.

1519

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

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

1522 always have an init quantum as well.

1523

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

1525 zero entries or all tasks in the pipeline.

1526 """

1527

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

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

1530

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

1532 """

1533

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

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

1536 dataset, keyed by UUID.

1537

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

1539 entries.

1540

1541 This does not include special "init" quanta.

1542 """

1543

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

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

1546 predicted quantum graph.

1547

1548 Parameters

1549 ----------

1550 predicted : `PredictedDatasetModel`

1551 Model for the dataset in the predicted graph.

1552

1553 Returns

1554 -------

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

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

1557 the dimension data has been loaded.

1558 """

1559 try:

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

1561 except KeyError:

1562 if predicted.dataset_type_name == acc.PACKAGES_INIT_OUTPUT_NAME:

1563 dataset_type = self.pipeline_graph.packages_dataset_type

1564 else:

1565 raise

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

1567 if self.dimension_data is not None:

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

1569 return DatasetRef(

1570 dataset_type,

1571 data_id,

1572 run=predicted.run,

1573 id=predicted.dataset_id,

1574 )

1575

1576 def set_thin_graph(self) -> None:

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

1578 `quantum_datasets` components (which must be complete).

1579 """

1580 bipartite_xgraph = networkx.DiGraph()

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

1582 graph_quantum_ids: list[uuid.UUID] = []

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

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

1585 PredictedThinQuantumModel.model_construct(

1586 quantum_id=quantum_datasets.quantum_id,

1587 data_coordinate=quantum_datasets.data_coordinate,

1588 )

1589 )

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

1591 bipartite_xgraph.add_edge(dataset.dataset_id, quantum_datasets.quantum_id)

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

1593 bipartite_xgraph.add_edge(quantum_datasets.quantum_id, dataset.dataset_id)

1594 graph_quantum_ids.append(quantum_datasets.quantum_id)

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

1596 bipartite_xgraph, graph_quantum_ids

1597 )

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

1599

1600 def set_header_counts(self) -> None:

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

1602 `thin_graph`, `init_quanta`, and `quantum_datasets` components.

1603 """

1604 self.header.n_quanta = len(self.quantum_datasets)

1605 self.header.n_task_quanta = {

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

1607 }

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

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

1610 all_dataset_ids.update(quantum_datasets.iter_dataset_ids())

1611 self.header.n_datasets = len(all_dataset_ids)

1612

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

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

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

1616 UUIDs.

1617

1618 Parameters

1619 ----------

1620 output_run : `str`

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

1622 """

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

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

1625 # need to rely on topological ordering of anything.

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

1627 new_quantum_id = generate_uuidv7()

1628 quantum_datasets.quantum_id = new_quantum_id

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

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

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

1632 )

1633 new_dataset_id = generate_uuidv7()

1634 uuid_map[output_dataset.dataset_id] = new_dataset_id

1635 output_dataset.dataset_id = new_dataset_id

1636 output_dataset.run = output_run

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

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

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

1640 input_dataset.run = output_run

1641 input_dataset.dataset_id = uuid_map.get(

1642 input_dataset.dataset_id,

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

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

1645 # have been built with extend_run=True.

1646 input_dataset.dataset_id,

1647 )

1648 # Update the keys of the quantum_datasets dict.

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

1650 # Since the UUIDs have changed, the thin graph needs to be rewritten.

1651 self.set_thin_graph()

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

1653 self.header.output_run = output_run

1654

1655 def assemble(self) -> PredictedQuantumGraph:

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

1657 return PredictedQuantumGraph(self)

1658

1659 @classmethod

1660 def read_execution_quanta(

1661 cls,

1662 uri: ResourcePathExpression,

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

1664 page_size: int = DEFAULT_PAGE_SIZE,

1665 ) -> PredictedQuantumGraphComponents:

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

1667

1668 Parameters

1669 ----------

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

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

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

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

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

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

1676 page_size : `int`, optional

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

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

1679 does affect the smallest, most numerous reads.

1680

1681 Returns

1682 -------

1683 components : `PredictedQuantumGraphComponents` ]

1684 Components for quantum graph that can build execution quanta for

1685 all of the given IDs.

1686 """

1687 uri = ResourcePath(uri)

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

1689 _LOG.warning(

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

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

1692 )

1693 from ..graph import QuantumGraph

1694

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

1696 return PredictedQuantumGraphComponents.from_old_quantum_graph(old_qg)

1697

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

1699 reader.read_execution_quanta(quantum_ids)

1700 return reader.components

1701

1702 @classmethod

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

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

1705

1706 Parameters

1707 ----------

1708 old_quantum_graph : `QuantumGraph`

1709 Quantum graph to transform.

1710

1711 Returns

1712 -------

1713 components : `PredictedQuantumGraphComponents`

1714 Components for a new predicted quantum graph.

1715 """

1716 header = HeaderModel.from_old_quantum_graph(old_quantum_graph)

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

1718 result.init_quanta.update_from_old_quantum_graph(old_quantum_graph)

1719 dimension_data_extractor = DimensionDataExtractor.from_dimension_group(

1720 old_quantum_graph.pipeline_graph.get_all_dimensions()

1721 )

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

1723 task_quanta = old_quantum_graph.get_task_quanta(task_node.label)

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

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

1726 task_node, quantum, quantum_id

1727 )

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

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

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

1731 result.dimension_data = DimensionDataAttacher(

1732 records=dimension_data_extractor.records.values(),

1733 dimensions=result.pipeline_graph.get_all_dimensions(),

1734 )

1735 result.set_thin_graph()

1736 result.set_header_counts()

1737 return result

1738

1739 def write(

1740 self,

1741 uri: ResourcePathExpression,

1742 *,

1743 zstd_level: int = 10,

1744 zstd_dict_size: int = 32768,

1745 zstd_dict_n_inputs: int = 512,

1746 ) -> None:

1747 """Write the graph to a file.

1748

1749 Parameters

1750 ----------

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

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

1753 to force writing the old format.

1754 zstd_level : `int`, optional

1755 ZStandard compression level to use on JSON blocks.

1756 zstd_dict_size : `int`, optional

1757 Size of a ZStandard dictionary that shares compression information

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

1759 Dictionary compression is automatically disabled if the number of

1760 quanta is smaller than ``zstd_dict_n_inputs``.

1761 zstd_dict_n_inputs : `int`, optional

1762 Maximum number of `PredictedQuantumDatasetsModel` JSON

1763 representations to feed the ZStandard dictionary training routine.

1764

1765 Notes

1766 -----

1767 Only a complete predicted quantum graph with all components fully

1768 populated should be written.

1769 """

1770 if self.header.n_task_quanta != {

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

1772 }:

1773 raise RuntimeError(

1774 "Cannot save graph after partial read of quanta: thin graph is inconsistent with header."

1775 )

1776 # Ensure we record the actual version we're about to write, in case

1777 # we're rewriting an old graph in a new format.

1778 self.header.version = FORMAT_VERSION

1779 uri = ResourcePath(uri)

1780 match uri.getExtension():

1781 case ".qg":

1782 pass

1783 case ".qgraph":

1784 _LOG.warning(

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

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

1787 )

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

1789 old_qg.saveUri(uri)

1790 return

1791 case ext:

1792 raise ValueError(

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

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

1795 )

1796 cdict: zstandard.ZstdCompressionDict | None = None

1797 cdict_data: bytes | None = None

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

1799 if len(self.quantum_datasets) < zstd_dict_n_inputs:

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

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

1802 # quanta.

1803 zstd_dict_size = 0

1804 if zstd_dict_size:

1805 quantum_datasets_json = {

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

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

1808 }

1809 try:

1810 cdict = zstandard.train_dictionary(

1811 zstd_dict_size,

1812 list(quantum_datasets_json.values()),

1813 level=zstd_level,

1814 )

1815 except zstandard.ZstdError as err:

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

1817 cdict = None

1818 else:

1819 cdict_data = cdict.as_bytes()

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

1821 indices = {quantum_id: n for n, quantum_id in enumerate(sorted(self.quantum_datasets.keys()))}

1822 with BaseQuantumGraphWriter.open(

1823 uri,

1824 header=self.header,

1825 pipeline_graph=self.pipeline_graph,

1826 indices=indices,

1827 address_filename="quanta",

1828 compressor=compressor,

1829 cdict_data=cdict_data,

1830 ) as writer:

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

1832 if self.dimension_data is None:

1833 raise IncompleteQuantumGraphError(

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

1835 )

1836 serialized_dimension_data = self.dimension_data.serialized()

1837 writer.write_single_model("dimension_data", serialized_dimension_data)

1838 del serialized_dimension_data

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

1840 with MultiblockWriter.open_in_zip(

1841 writer.zf, "quantum_datasets", writer.int_size

1842 ) as quantum_datasets_mb:

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

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

1845 quantum_datasets_mb.write_bytes(

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

1847 )

1848 else:

1849 quantum_datasets_mb.write_model(

1850 quantum_model.quantum_id, quantum_model, writer.compressor

1851 )

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

1853

1854

1855@dataclasses.dataclass

1856class PredictedQuantumGraphReader(BaseQuantumGraphReader):

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

1858

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

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

1861

1862 @classmethod

1863 @contextmanager

1864 def open(

1865 cls,

1866 uri: ResourcePathExpression,

1867 *,

1868 page_size: int = DEFAULT_PAGE_SIZE,

1869 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,

1870 ) -> Iterator[PredictedQuantumGraphReader]:

1871 """Construct a reader from a URI.

1872

1873 Parameters

1874 ----------

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

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

1877 page_size : `int`, optional

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

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

1880 does affect the smallest, most numerous reads.

1881 import_mode : `.pipeline_graph.TaskImportMode`, optional

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

1883 graph.

1884

1885 Returns

1886 -------

1887 reader : `contextlib.AbstractContextManager` [ \

1888 `PredictedQuantumGraphReader` ]

1889 A context manager that returns the reader when entered.

1890 """

1891 with cls._open(

1892 uri,

1893 graph_type="predicted",

1894 address_filename="quanta",

1895 page_size=page_size,

1896 import_mode=import_mode,

1897 n_addresses=1,

1898 ) as self:

1899 yield self

1900

1901 def __post_init__(self) -> None:

1902 self.components = PredictedQuantumGraphComponents(

1903 header=self.header, pipeline_graph=self.pipeline_graph

1904 )

1905

1906 def finish(self) -> PredictedQuantumGraph:

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

1908 return self.components.assemble()

1909

1910 def read_all(self) -> None:

1911 """Read all components in full."""

1912 self.read_thin_graph()

1913 self.read_execution_quanta()

1914

1915 def read_thin_graph(self) -> None:

1916 """Read the thin graph.

1917

1918 The thin graph is a quantum-quantum DAG with just task labels and data

1919 IDs as node attributes. It always includes all regular quanta, and

1920 does not include init-input or init-output information.

1921 """

1922 if not self.components.thin_graph.quanta:

1923 if self.header.version > 0:

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

1925 else:

1926 self.address_reader.read_all()

1927 thin_graph_v0 = self._read_single_block("thin_graph", _PredictedThinGraphModelV0)

1928 self.components.thin_graph = thin_graph_v0._upgraded(self.address_reader.rows)

1929

1930 def read_init_quanta(self) -> None:

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

1932 init-outputs.

1933 """

1934 if not self.components.init_quanta.root:

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

1936

1937 def read_dimension_data(self) -> None:

1938 """Read all dimension records.

1939

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

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

1942 """

1943 if self.components.dimension_data is None:

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

1945 self.components.dimension_data = DimensionDataAttacher(

1946 deserializers=[

1947 DimensionRecordSetDeserializer.from_raw(

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

1949 )

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

1951 ],

1952 dimensions=DimensionGroup.union(

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

1954 universe=self.components.pipeline_graph.universe,

1955 ),

1956 )

1957

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

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

1960 given quantum IDs.

1961

1962 Parameters

1963 ----------

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

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

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

1967 """

1968 quantum_datasets: PredictedQuantumDatasetsModel | None

1969 if quantum_ids is None:

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

1971 for quantum_datasets in MultiblockReader.read_all_models_in_zip(

1972 self.zf,

1973 "quantum_datasets",

1974 PredictedQuantumDatasetsModel,

1975 self.decompressor,

1976 int_size=self.components.header.int_size,

1977 page_size=self.page_size,

1978 ):

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

1980 self.address_reader.read_all()

1981 return

1982 with MultiblockReader.open_in_zip(

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

1984 ) as mb_reader:

1985 for quantum_id in quantum_ids:

1986 if quantum_id in self.components.quantum_datasets:

1987 continue

1988 address_row = self.address_reader.find(quantum_id)

1989 quantum_datasets = mb_reader.read_model(

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

1991 )

1992 if quantum_datasets is not None:

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

1994 return

1995

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

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

1998

1999 Parameters

2000 ----------

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

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

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

2004 """

2005 self.read_init_quanta()

2006 self.read_dimension_data()

2007 self.read_quantum_datasets(quantum_ids)

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

608 statements