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 "BaseQuantumGraph",

32 "BaseQuantumGraphReader",

33 "BipartiteEdgeInfo",

34 "DatasetInfo",

35 "HeaderModel",

36 "QuantumInfo",

37)

38import dataclasses

39import datetime

40import getpass

41import os

42import sys

43import uuid

44import zipfile

45from abc import ABC, abstractmethod

46from collections.abc import Iterator, Mapping

47from contextlib import contextmanager

48from typing import (

49 TYPE_CHECKING,

50 Any,

51 Self,

52 TypeAlias,

53 TypedDict,

54 TypeVar,

55)

57import networkx

58import networkx.algorithms.bipartite

59import pydantic

60import zstandard

62from lsst.daf.butler import DataCoordinate, DataIdValue

63from lsst.daf.butler._rubin import generate_uuidv7

64from lsst.resources import ResourcePath, ResourcePathExpression

66from ..pipeline_graph import DatasetTypeNode, Edge, PipelineGraph, TaskImportMode, TaskNode

67from ..pipeline_graph.io import SerializedPipelineGraph

68from ._multiblock import (

69 DEFAULT_PAGE_SIZE,

70 AddressReader,

71 AddressWriter,

72 Compressor,

73 Decompressor,

74)

76if TYPE_CHECKING:

77 from ..graph import QuantumGraph

80# These aliases make it a lot easier how the various pydantic models are

81# structured, but they're too verbose to be worth exporting to code outside the

82# quantum_graph subpackage.

83TaskLabel: TypeAlias = str

84DatasetTypeName: TypeAlias = str

85ConnectionName: TypeAlias = str

86DatasetIndex: TypeAlias = int

87QuantumIndex: TypeAlias = int

88DatastoreName: TypeAlias = str

89DimensionElementName: TypeAlias = str

90DataCoordinateValues: TypeAlias = list[DataIdValue]

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

96class IncompleteQuantumGraphError(RuntimeError):

97 pass

100class HeaderModel(pydantic.BaseModel):

101 """Data model for the header of a quantum graph file."""

102

103 version: int = 0

104 """File format / data model version number."""

105

106 graph_type: str = ""

107 """Type of quantum graph stored in this file."""

108

109 inputs: list[str] = pydantic.Field(default_factory=list)

110 """List of input collections used to build the quantum graph."""

111

112 output: str | None = ""

113 """Output CHAINED collection provided when building the quantum graph."""

114

115 output_run: str = ""

116 """Output RUN collection for all output datasets in this graph."""

117

118 user: str = pydantic.Field(default_factory=getpass.getuser)

119 """Username of the process that built this quantum graph."""

120

121 timestamp: datetime.datetime = pydantic.Field(default_factory=datetime.datetime.now)

122 """Timestamp for when this quantum graph was built.

123

124 It is unspecified exactly which point during quantum-graph generation this

125 timestamp is recorded.

126 """

127

128 command: str = pydantic.Field(default_factory=lambda: " ".join(sys.argv))

129 """Command-line invocation that created this graph."""

130

131 metadata: dict[str, Any] = pydantic.Field(default_factory=dict)

132 """Free-form metadata associated with this quantum graph at build time."""

133

134 int_size: int = 8

135 """Number of bytes in the integers used in this file's multi-block and

136 address files.

137 """

138

139 n_quanta: int = 0

140 """Total number of quanta in this graph.

141

142 This does not include special "init" quanta, but it does include quanta

143 that were not loaded in a partial read (except when reading from an old

144 quantum graph file).

145 """

146

147 n_datasets: int = 0

148 """Total number of distinct datasets in the full graph. This includes

149 datasets whose related quanta were not loaded in a partial read (except

150 when reading from an old quantum graph file).

151 """

152

153 n_task_quanta: dict[TaskLabel, int] = pydantic.Field(default_factory=dict)

154 """Number of quanta for each task label.

155

156 This does not include special "init" quanta, but it does include quanta

157 that were not loaded in a partial read (except when reading from an old

158 quantum graph file).

159 """

160

161 provenance_dataset_id: uuid.UUID = pydantic.Field(default_factory=generate_uuidv7)

162 """The dataset ID for provenance quantum graph when it is ingested into

163 a butler repository.

164 """

165

166 @classmethod

167 def from_old_quantum_graph(cls, old_quantum_graph: QuantumGraph) -> HeaderModel:

168 """Extract a header from an old `QuantumGraph` instance.

169

170 Parameters

171 ----------

172 old_quantum_graph : `QuantumGraph`

173 Quantum graph to extract a header from.

174

175 Returns

176 -------

177 header : `PredictedHeaderModel`

178 Header for a new predicted quantum graph.

179 """

180 metadata = dict(old_quantum_graph.metadata)

181 metadata.pop("packages", None)

182 if (time_str := metadata.pop("time", None)) is not None:

183 timestamp = datetime.datetime.fromisoformat(time_str)

184 else:

185 timestamp = datetime.datetime.now()

186 return cls(

187 inputs=list(metadata.pop("input", []) or []), # Guard against explicit None and missing key.

188 output=metadata.pop("output", None),

189 output_run=metadata.pop("output_run", ""),

190 user=metadata.pop("user", ""),

191 command=metadata.pop("full_command", ""),

192 timestamp=timestamp,

193 metadata=metadata,

194 )

195

196 def to_old_metadata(self) -> dict[str, Any]:

197 """Return a dictionary using the key conventions used in old quantum

198 graph files.

199 """

200 result = self.metadata.copy()

201 result["input"] = self.inputs

202 result["output"] = self.output

203 result["output_run"] = self.output_run

204 result["full_command"] = self.command

205 result["user"] = self.user

206 result["time"] = str(self.timestamp)

207 return result

208

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

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

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

212

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

214 """See `pydantic.BaseModel.copy`."""

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

216

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

218 """See `pydantic.BaseModel.model_dump`."""

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

220

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

222 """See `pydantic.BaseModel.model_dump_json`."""

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

224

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

226 """See `pydantic.BaseModel.model_copy`."""

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

228

229 @classmethod

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

231 """See `pydantic.BaseModel.model_construct`."""

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

233

234 @classmethod

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

236 """See `pydantic.BaseModel.model_json_schema`."""

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

238

239 @classmethod

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

241 """See `pydantic.BaseModel.model_validate`."""

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

243

244 @classmethod

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

246 """See `pydantic.BaseModel.model_validate_json`."""

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

248

249 @classmethod

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

251 """See `pydantic.BaseModel.model_validate_strings`."""

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

253

254

255class QuantumInfo(TypedDict):

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

257 node data for a quantum.

258

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

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

261

262 node_data: QuantumInfo = xgraph.nodes[quantum_id]

263

264 where ``xgraph`` can be either `BaseQuantumGraph.quantum_only_xgraph`

265 or `BaseQuantumGraph.bipartite_xgraph`.

266 """

267

268 data_id: DataCoordinate

269 """Data ID of the quantum."""

270

271 task_label: str

272 """Label of the task for this quantum."""

273

274 pipeline_node: TaskNode

275 """Node in the pipeline graph for this quantum's task."""

276

277

278class DatasetInfo(TypedDict):

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

280 node data for a dataset.

281

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

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

284

285 node_data: DatasetInfo = xgraph.nodes[dataset_id]

286

287 where ``xgraph`` is from the `BaseQuantumGraph.bipartite_xgraph` property.

288 """

289

290 data_id: DataCoordinate

291 """Data ID of the dataset."""

292

293 dataset_type_name: DatasetTypeName

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

295

296 This is always the general dataset type that matches the data repository

297 storage class, which may differ from any particular task-adapted dataset

298 type whose storage class has been overridden to match the task connections.

299 This means is it never a component.

300 """

301

302 run: str

303 """Name of the `~lsst.daf.butler.CollectionType.RUN` collection that holds

304 or will hold this dataset.

305 """

306

307 pipeline_node: DatasetTypeNode

308 """Node in the pipeline graph for this dataset's type."""

309

310

311class BipartiteEdgeInfo(TypedDict):

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

313 edge data in a bipartite graph.

314 """

315

316 is_read: bool

317 """`True` if this is a dataset -> quantum edge; `False` if it is a

318 quantum -> dataset edge.

319 """

320

321 pipeline_edges: list[Edge]

322 """Corresponding edges in the pipeline graph.

323

324 Note that there may be more than one pipeline edge since a quantum can

325 consume a particular dataset via multiple connections.

326 """

327

328

329class BaseQuantumGraph(ABC):

330 """An abstract base for quantum graphs.

331

332 Parameters

333 ----------

334 header : `HeaderModel`

335 Structured metadata for the graph.

336 pipeline_graph : `.pipeline_graph.PipelineGraph`

337 Graph of tasks and dataset types. May contain a superset of the tasks

338 and dataset types that actually have quanta and datasets in the quantum

339 graph.

340 """

341

342 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph):

343 self.header = header

344 self.pipeline_graph = pipeline_graph

345

346 @property

347 @abstractmethod

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

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

350 data ID.

351

352 Notes

353 -----

354 Partial loads may not fully populate this mapping, but it can always

355 be accessed.

356

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

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

359 """

360 raise NotImplementedError()

361

362 @property

363 @abstractmethod

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

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

366 and then by data ID.

367

368 Notes

369 -----

370 Partial loads may not fully populate this mapping, but it can always

371 be accessed.

372

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

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

375 """

376 raise NotImplementedError()

377

378 @property

379 @abstractmethod

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

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

382

383 Notes

384 -----

385 Partial loads may not fully populate this graph, but it can always be

386 accessed.

387

388 Node state dictionaries are described by the `QuantumInfo` type

389 (or a subtype thereof).

390

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

392 """

393 raise NotImplementedError()

394

395 @property

396 @abstractmethod

397 def bipartite_xgraph(self) -> networkx.DiGraph:

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

399

400 Notes

401 -----

402 Partial loads may not fully populate this graph, but it can always be

403 accessed.

404

405 Node state dictionaries are described by the `QuantumInfo` and

406 `DatasetInfo` types (or a subtypes thereof). Edges have state

407 dictionaries described by `BipartiteEdgeInfo`.

408

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

410 """

411 raise NotImplementedError()

412

413

414@dataclasses.dataclass

415class BaseQuantumGraphWriter:

416 """A helper class for writing quantum graphs."""

417

418 zf: zipfile.ZipFile

419 """The zip archive that represents the quantum graph on disk."""

420

421 compressor: Compressor

422 """A compressor for all compressed JSON blocks."""

423

424 address_writer: AddressWriter

425 """A helper object for reading addresses into the multi-block files."""

426

427 int_size: int

428 """Size (in bytes) used to write integers to binary files."""

429

430 @classmethod

431 @contextmanager

432 def open(

433 cls,

434 uri: ResourcePathExpression,

435 header: HeaderModel,

436 pipeline_graph: PipelineGraph,

437 indices: dict[uuid.UUID, int],

438 *,

439 address_filename: str,

440 compressor: Compressor,

441 cdict_data: bytes | None = None,

442 ) -> Iterator[Self]:

443 uri = ResourcePath(uri)

444 address_writer = AddressWriter(indices)

445 with uri.open(mode="wb") as stream:

446 with zipfile.ZipFile(stream, mode="w", compression=zipfile.ZIP_STORED) as zf:

447 self = cls(zf, compressor, address_writer, header.int_size)

448 self.write_single_model("header", header)

449 if cdict_data is not None:

450 zf.writestr("compression_dict", cdict_data)

451 self.write_single_model("pipeline_graph", SerializedPipelineGraph.serialize(pipeline_graph))

452 yield self

453 address_writer.write_to_zip(zf, address_filename, int_size=self.int_size)

454

455 def write_single_model(self, name: str, model: pydantic.BaseModel) -> None:

456 """Write a single compressed JSON block as a 'file' in a zip archive.

457

458 Parameters

459 ----------

460 name : `str`

461 Base name of the file. An extension will be added.

462 model : `pydantic.BaseModel`

463 Pydantic model to convert to JSON.

464 """

465 json_data = model.model_dump_json().encode()

466 self.write_single_block(name, json_data)

467

468 def write_single_block(self, name: str, json_data: bytes) -> None:

469 """Write a single compressed JSON block as a 'file' in a zip archive.

470

471 Parameters

472 ----------

473 name : `str`

474 Base name of the file. An extension will be added.

475 json_data : `bytes`

476 Raw JSON to compress and write.

477 """

478 json_data = self.compressor.compress(json_data)

479 self.zf.writestr(f"{name}.json.zst", json_data)

480

481

482@dataclasses.dataclass

483class BaseQuantumGraphReader:

484 """A helper class for reading quantum graphs."""

485

486 header: HeaderModel

487 """Header metadata for the quantum graph."""

488

489 pipeline_graph: PipelineGraph

490 """Graph of tasks and dataset type names that appear in the quantum

491 graph.

492 """

493

494 zf: zipfile.ZipFile

495 """The zip archive that represents the quantum graph on disk."""

496

497 decompressor: Decompressor

498 """A decompressor for all compressed JSON blocks."""

499

500 address_reader: AddressReader

501 """A helper object for reading addresses into the multi-block files."""

502

503 page_size: int

504 """Approximate number of bytes to read at a time.

505

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

507 does affect the smallest, most numerous reads.

508 """

509

510 @classmethod

511 @contextmanager

512 def _open(

513 cls,

514 uri: ResourcePathExpression,

515 *,

516 address_filename: str,

517 graph_type: str,

518 n_addresses: int,

519 page_size: int | None = None,

520 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,

521 ) -> Iterator[Self]:

522 """Construct a reader from a URI.

523

524 Parameters

525 ----------

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

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

528 address_filename : `str`

529 Base filename for the address file.

530 graph_type : `str`

531 Value to expect for `HeaderModel.graph_type`.

532 n_addresses : `int`

533 Number of addresses to expect per row in the address file.

534 page_size : `int`, optional

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

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

537 does affect the smallest, most numerous reads. When `None`, the

538 ``LSST_QG_PAGE_SIZE`` environment variable is checked before

539 falling back to a default of 5MB.

540 import_mode : `..pipeline_graph.TaskImportMode`, optional

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

542 graph.

543

544 Returns

545 -------

546 reader : `contextlib.AbstractContextManager` [ \

547 `PredictedQuantumGraphReader` ]

548 A context manager that returns the reader when entered.

549 """

550 if page_size is None:

551 page_size = int(os.environ.get("LSST_QG_PAGE_SIZE", DEFAULT_PAGE_SIZE))

552 uri = ResourcePath(uri)

553 cdict: zstandard.ZstdCompressionDict | None = None

554 with uri.open(mode="rb") as zf_stream:

555 with zipfile.ZipFile(zf_stream, "r") as zf:

556 if (cdict_path := zipfile.Path(zf, "compression_dict")).exists():

557 cdict = zstandard.ZstdCompressionDict(cdict_path.read_bytes())

558 decompressor = zstandard.ZstdDecompressor(cdict)

559 header = cls._read_single_block_static("header", HeaderModel, zf, decompressor)

560 if not header.graph_type == graph_type:

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

562 serialized_pipeline_graph = cls._read_single_block_static(

563 "pipeline_graph", SerializedPipelineGraph, zf, decompressor

564 )

565 pipeline_graph = serialized_pipeline_graph.deserialize(import_mode)

566 with AddressReader.open_in_zip(

567 zf,

568 address_filename,

569 page_size=page_size,

570 int_size=header.int_size,

571 n_addresses=n_addresses,

572 ) as address_reader:

573 yield cls(

574 header=header,

575 pipeline_graph=pipeline_graph,

576 zf=zf,

577 decompressor=decompressor,

578 address_reader=address_reader,

579 page_size=page_size,

580 )

581

582 @staticmethod

583 def _read_single_block_static(

584 name: str, model_type: type[_T], zf: zipfile.ZipFile, decompressor: Decompressor

585 ) -> _T:

586 """Read a single compressed JSON block from a 'file' in a zip archive.

587

588 Parameters

589 ----------

590 zf : `zipfile.ZipFile`

591 Zip archive to read the file from.

592 name : `str`

593 Base name of the file. An extension will be added.

594 model_type : `type` [ `pydantic.BaseModel` ]

595 Pydantic model to validate JSON with.

596 decompressor : `Decompressor`

597 Object with a `decompress` method that takes and returns `bytes`.

598

599 Returns

600 -------

601 model : `pydantic.BaseModel`

602 Validated model.

603 """

604 compressed_data = zf.read(f"{name}.json.zst")

605 json_data = decompressor.decompress(compressed_data)

606 return model_type.model_validate_json(json_data)

607

608 def _read_single_block(self, name: str, model_type: type[_T]) -> _T:

609 """Read a single compressed JSON block from a 'file' in a zip archive.

610

611 Parameters

612 ----------

613 name : `str`

614 Base name of the file. An extension will be added.

615 model_type : `type` [ `pydantic.BaseModel` ]

616 Pydantic model to validate JSON with.

617

618 Returns

619 -------

620 model : `pydantic.BaseModel`

621 Validated model.

622 """

623 return self._read_single_block_static(name, model_type, self.zf, self.decompressor)

624

625 def _read_single_block_raw(self, name: str) -> bytes:

626 """Read a single compressed block from a 'file' in a zip archive.

627

628 Parameters

629 ----------

630 name : `str`

631 Base name of the file. An extension will be added.

632

633 Returns

634 -------

635 data : `bytes`

636 Decompressed bytes.

637 """

638 compressed_data = self.zf.read(f"{name}.json.zst")

639 return self.decompressor.decompress(compressed_data)

Coverage for python/lsst/pipe/base/quantum_graph/_common.py: 68%

190 statements