Coverage for python / lsst / pipe / base / quantum_graph / _common.py: 68%
192 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:13 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = (
31 "FORMAT_VERSION",
32 "BaseQuantumGraph",
33 "BaseQuantumGraphReader",
34 "BipartiteEdgeInfo",
35 "DatasetInfo",
36 "HeaderModel",
37 "QuantumInfo",
38)
39import dataclasses
40import datetime
41import getpass
42import os
43import sys
44import uuid
45import zipfile
46from abc import ABC, abstractmethod
47from collections.abc import Iterator, Mapping
48from contextlib import contextmanager
49from typing import (
50 TYPE_CHECKING,
51 Any,
52 Self,
53 TypeAlias,
54 TypedDict,
55 TypeVar,
56)
58import networkx
59import networkx.algorithms.bipartite
60import pydantic
61import zstandard
63from lsst.daf.butler import DataCoordinate, DataIdValue
64from lsst.daf.butler._rubin import generate_uuidv7
65from lsst.resources import ResourcePath, ResourcePathExpression
67from ..pipeline_graph import DatasetTypeNode, Edge, PipelineGraph, TaskImportMode, TaskNode
68from ..pipeline_graph.io import SerializedPipelineGraph
69from ._multiblock import (
70 DEFAULT_PAGE_SIZE,
71 AddressReader,
72 AddressWriter,
73 Compressor,
74 Decompressor,
75)
77if TYPE_CHECKING:
78 from ..graph import QuantumGraph
81# These aliases make it a lot easier how the various pydantic models are
82# structured, but they're too verbose to be worth exporting to code outside the
83# quantum_graph subpackage.
84TaskLabel: TypeAlias = str
85DatasetTypeName: TypeAlias = str
86ConnectionName: TypeAlias = str
87DatasetIndex: TypeAlias = int
88QuantumIndex: TypeAlias = int
89DatastoreName: TypeAlias = str
90DimensionElementName: TypeAlias = str
91DataCoordinateValues: TypeAlias = list[DataIdValue]
94_T = TypeVar("_T", bound=pydantic.BaseModel)
96FORMAT_VERSION: int = 1
97"""
98File format version number for new files.
100This applies to both predicted and provenance QGs, since they usually change
101in concert.
103CHANGELOG:
105- 0: Initial version.
106- 1: Switched from internal integer IDs to UUIDs in all models.
107"""
110class IncompleteQuantumGraphError(RuntimeError):
111 pass
114class HeaderModel(pydantic.BaseModel):
115 """Data model for the header of a quantum graph file."""
117 version: int = FORMAT_VERSION
118 """File format / data model version number."""
120 graph_type: str = ""
121 """Type of quantum graph stored in this file."""
123 inputs: list[str] = pydantic.Field(default_factory=list)
124 """List of input collections used to build the quantum graph."""
126 output: str | None = ""
127 """Output CHAINED collection provided when building the quantum graph."""
129 output_run: str = ""
130 """Output RUN collection for all output datasets in this graph."""
132 user: str = pydantic.Field(default_factory=getpass.getuser)
133 """Username of the process that built this quantum graph."""
135 timestamp: datetime.datetime = pydantic.Field(default_factory=datetime.datetime.now)
136 """Timestamp for when this quantum graph was built.
138 It is unspecified exactly which point during quantum-graph generation this
139 timestamp is recorded.
140 """
142 command: str = pydantic.Field(default_factory=lambda: " ".join(sys.argv))
143 """Command-line invocation that created this graph."""
145 metadata: dict[str, Any] = pydantic.Field(default_factory=dict)
146 """Free-form metadata associated with this quantum graph at build time."""
148 int_size: int = 8
149 """Number of bytes in the integers used in this file's multi-block and
150 address files.
151 """
153 n_quanta: int = 0
154 """Total number of quanta in this graph.
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 """
161 n_datasets: int = 0
162 """Total number of distinct datasets in the full graph. This includes
163 datasets whose related quanta were not loaded in a partial read (except
164 when reading from an old quantum graph file).
165 """
167 n_task_quanta: dict[TaskLabel, int] = pydantic.Field(default_factory=dict)
168 """Number of quanta for each task label.
170 This does not include special "init" quanta, but it does include quanta
171 that were not loaded in a partial read (except when reading from an old
172 quantum graph file).
173 """
175 provenance_dataset_id: uuid.UUID = pydantic.Field(default_factory=generate_uuidv7)
176 """The dataset ID for provenance quantum graph when it is ingested into
177 a butler repository.
178 """
180 @classmethod
181 def from_old_quantum_graph(cls, old_quantum_graph: QuantumGraph) -> HeaderModel:
182 """Extract a header from an old `QuantumGraph` instance.
184 Parameters
185 ----------
186 old_quantum_graph : `QuantumGraph`
187 Quantum graph to extract a header from.
189 Returns
190 -------
191 header : `PredictedHeaderModel`
192 Header for a new predicted quantum graph.
193 """
194 metadata = dict(old_quantum_graph.metadata)
195 metadata.pop("packages", None)
196 if (time_str := metadata.pop("time", None)) is not None:
197 timestamp = datetime.datetime.fromisoformat(time_str)
198 else:
199 timestamp = datetime.datetime.now()
200 return cls(
201 inputs=list(metadata.pop("input", []) or []), # Guard against explicit None and missing key.
202 output=metadata.pop("output", None),
203 output_run=metadata.pop("output_run", ""),
204 user=metadata.pop("user", ""),
205 command=metadata.pop("full_command", ""),
206 timestamp=timestamp,
207 metadata=metadata,
208 )
210 def to_old_metadata(self) -> dict[str, Any]:
211 """Return a dictionary using the key conventions used in old quantum
212 graph files.
213 """
214 result = self.metadata.copy()
215 result["input"] = self.inputs
216 result["output"] = self.output
217 result["output_run"] = self.output_run
218 result["full_command"] = self.command
219 result["user"] = self.user
220 result["time"] = str(self.timestamp)
221 return result
223 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
224 # when we inherit those docstrings in our public classes.
225 if "sphinx" in sys.modules and not TYPE_CHECKING:
227 def copy(self, *args: Any, **kwargs: Any) -> Any:
228 """See `pydantic.BaseModel.copy`."""
229 return super().copy(*args, **kwargs)
231 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
232 """See `pydantic.BaseModel.model_dump`."""
233 return super().model_dump(*args, **kwargs)
235 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
236 """See `pydantic.BaseModel.model_dump_json`."""
237 return super().model_dump(*args, **kwargs)
239 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
240 """See `pydantic.BaseModel.model_copy`."""
241 return super().model_copy(*args, **kwargs)
243 @classmethod
244 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
245 """See `pydantic.BaseModel.model_construct`."""
246 return super().model_construct(*args, **kwargs)
248 @classmethod
249 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
250 """See `pydantic.BaseModel.model_json_schema`."""
251 return super().model_json_schema(*args, **kwargs)
253 @classmethod
254 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
255 """See `pydantic.BaseModel.model_validate`."""
256 return super().model_validate(*args, **kwargs)
258 @classmethod
259 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
260 """See `pydantic.BaseModel.model_validate_json`."""
261 return super().model_validate_json(*args, **kwargs)
263 @classmethod
264 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
265 """See `pydantic.BaseModel.model_validate_strings`."""
266 return super().model_validate_strings(*args, **kwargs)
269class QuantumInfo(TypedDict):
270 """A typed dictionary that annotates the attributes of the NetworkX graph
271 node data for a quantum.
273 Since NetworkX types are not generic over their node mapping type, this has
274 to be used explicitly, e.g.::
276 node_data: QuantumInfo = xgraph.nodes[quantum_id]
278 where ``xgraph`` can be either `BaseQuantumGraph.quantum_only_xgraph`
279 or `BaseQuantumGraph.bipartite_xgraph`.
280 """
282 data_id: DataCoordinate
283 """Data ID of the quantum."""
285 task_label: str
286 """Label of the task for this quantum."""
288 pipeline_node: TaskNode
289 """Node in the pipeline graph for this quantum's task."""
292class DatasetInfo(TypedDict):
293 """A typed dictionary that annotates the attributes of the NetworkX graph
294 node data for a dataset.
296 Since NetworkX types are not generic over their node mapping type, this has
297 to be used explicitly, e.g.::
299 node_data: DatasetInfo = xgraph.nodes[dataset_id]
301 where ``xgraph`` is from the `BaseQuantumGraph.bipartite_xgraph` property.
302 """
304 data_id: DataCoordinate
305 """Data ID of the dataset."""
307 dataset_type_name: DatasetTypeName
308 """Name of the type of this dataset.
310 This is always the general dataset type that matches the data repository
311 storage class, which may differ from any particular task-adapted dataset
312 type whose storage class has been overridden to match the task connections.
313 This means is it never a component.
314 """
316 run: str
317 """Name of the `~lsst.daf.butler.CollectionType.RUN` collection that holds
318 or will hold this dataset.
319 """
321 pipeline_node: DatasetTypeNode
322 """Node in the pipeline graph for this dataset's type."""
325class BipartiteEdgeInfo(TypedDict):
326 """A typed dictionary that annotates the attributes of the NetworkX graph
327 edge data in a bipartite graph.
328 """
330 is_read: bool
331 """`True` if this is a dataset -> quantum edge; `False` if it is a
332 quantum -> dataset edge.
333 """
335 pipeline_edges: list[Edge]
336 """Corresponding edges in the pipeline graph.
338 Note that there may be more than one pipeline edge since a quantum can
339 consume a particular dataset via multiple connections.
340 """
343class BaseQuantumGraph(ABC):
344 """An abstract base for quantum graphs.
346 Parameters
347 ----------
348 header : `HeaderModel`
349 Structured metadata for the graph.
350 pipeline_graph : `.pipeline_graph.PipelineGraph`
351 Graph of tasks and dataset types. May contain a superset of the tasks
352 and dataset types that actually have quanta and datasets in the quantum
353 graph.
354 """
356 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph):
357 self.header = header
358 self.pipeline_graph = pipeline_graph
360 @property
361 @abstractmethod
362 def quanta_by_task(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:
363 """A nested mapping of all quanta, keyed first by task name and then by
364 data ID.
366 Notes
367 -----
368 Partial loads may not fully populate this mapping, but it can always
369 be accessed.
371 The returned object may be an internal dictionary; as the type
372 annotation indicates, it should not be modified in place.
373 """
374 raise NotImplementedError()
376 @property
377 @abstractmethod
378 def datasets_by_type(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:
379 """A nested mapping of all datasets, keyed first by dataset type name
380 and then by data ID.
382 Notes
383 -----
384 Partial loads may not fully populate this mapping, but it can always
385 be accessed.
387 The returned object may be an internal dictionary; as the type
388 annotation indicates, it should not be modified in place.
389 """
390 raise NotImplementedError()
392 @property
393 @abstractmethod
394 def quantum_only_xgraph(self) -> networkx.DiGraph:
395 """A directed acyclic graph with quanta as nodes and datasets elided.
397 Notes
398 -----
399 Partial loads may not fully populate this graph, but it can always be
400 accessed.
402 Node state dictionaries are described by the `QuantumInfo` type
403 (or a subtype thereof).
405 The returned object is a read-only view of an internal one.
406 """
407 raise NotImplementedError()
409 @property
410 @abstractmethod
411 def bipartite_xgraph(self) -> networkx.DiGraph:
412 """A directed acyclic graph with quantum and dataset nodes.
414 Notes
415 -----
416 Partial loads may not fully populate this graph, but it can always be
417 accessed.
419 Node state dictionaries are described by the `QuantumInfo` and
420 `DatasetInfo` types (or a subtypes thereof). Edges have state
421 dictionaries described by `BipartiteEdgeInfo`.
423 The returned object is a read-only view of an internal one.
424 """
425 raise NotImplementedError()
428@dataclasses.dataclass
429class BaseQuantumGraphWriter:
430 """A helper class for writing quantum graphs."""
432 zf: zipfile.ZipFile
433 """The zip archive that represents the quantum graph on disk."""
435 compressor: Compressor
436 """A compressor for all compressed JSON blocks."""
438 address_writer: AddressWriter
439 """A helper object for reading addresses into the multi-block files."""
441 int_size: int
442 """Size (in bytes) used to write integers to binary files."""
444 @classmethod
445 @contextmanager
446 def open(
447 cls,
448 uri: ResourcePathExpression,
449 header: HeaderModel,
450 pipeline_graph: PipelineGraph,
451 indices: dict[uuid.UUID, int],
452 *,
453 address_filename: str,
454 compressor: Compressor,
455 cdict_data: bytes | None = None,
456 ) -> Iterator[Self]:
457 uri = ResourcePath(uri)
458 address_writer = AddressWriter(indices)
459 with uri.open(mode="wb") as stream:
460 with zipfile.ZipFile(stream, mode="w", compression=zipfile.ZIP_STORED) as zf:
461 self = cls(zf, compressor, address_writer, header.int_size)
462 self.write_single_model("header", header)
463 if cdict_data is not None:
464 zf.writestr("compression_dict", cdict_data)
465 self.write_single_model("pipeline_graph", SerializedPipelineGraph.serialize(pipeline_graph))
466 yield self
467 address_writer.write_to_zip(zf, address_filename, int_size=self.int_size)
469 def write_single_model(self, name: str, model: pydantic.BaseModel) -> None:
470 """Write a single compressed JSON block as a 'file' in a zip archive.
472 Parameters
473 ----------
474 name : `str`
475 Base name of the file. An extension will be added.
476 model : `pydantic.BaseModel`
477 Pydantic model to convert to JSON.
478 """
479 json_data = model.model_dump_json().encode()
480 self.write_single_block(name, json_data)
482 def write_single_block(self, name: str, json_data: bytes) -> None:
483 """Write a single compressed JSON block as a 'file' in a zip archive.
485 Parameters
486 ----------
487 name : `str`
488 Base name of the file. An extension will be added.
489 json_data : `bytes`
490 Raw JSON to compress and write.
491 """
492 json_data = self.compressor.compress(json_data)
493 self.zf.writestr(f"{name}.json.zst", json_data)
496@dataclasses.dataclass
497class BaseQuantumGraphReader:
498 """A helper class for reading quantum graphs."""
500 header: HeaderModel
501 """Header metadata for the quantum graph."""
503 pipeline_graph: PipelineGraph
504 """Graph of tasks and dataset type names that appear in the quantum
505 graph.
506 """
508 zf: zipfile.ZipFile
509 """The zip archive that represents the quantum graph on disk."""
511 decompressor: Decompressor
512 """A decompressor for all compressed JSON blocks."""
514 address_reader: AddressReader
515 """A helper object for reading addresses into the multi-block files."""
517 page_size: int
518 """Approximate number of bytes to read at a time.
520 Note that this does not set a page size for *all* reads, but it
521 does affect the smallest, most numerous reads.
522 """
524 @classmethod
525 @contextmanager
526 def _open(
527 cls,
528 uri: ResourcePathExpression,
529 *,
530 address_filename: str,
531 graph_type: str,
532 n_addresses: int,
533 page_size: int | None = None,
534 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,
535 ) -> Iterator[Self]:
536 """Construct a reader from a URI.
538 Parameters
539 ----------
540 uri : convertible to `lsst.resources.ResourcePath`
541 URI to open. Should have a ``.qg`` extension.
542 address_filename : `str`
543 Base filename for the address file.
544 graph_type : `str`
545 Value to expect for `HeaderModel.graph_type`.
546 n_addresses : `int`
547 Number of addresses to expect per row in the address file.
548 page_size : `int`, optional
549 Approximate number of bytes to read at once from address files.
550 Note that this does not set a page size for *all* reads, but it
551 does affect the smallest, most numerous reads. When `None`, the
552 ``LSST_QG_PAGE_SIZE`` environment variable is checked before
553 falling back to a default of 5MB.
554 import_mode : `..pipeline_graph.TaskImportMode`, optional
555 How to handle importing the task classes referenced in the pipeline
556 graph.
558 Returns
559 -------
560 reader : `contextlib.AbstractContextManager` [ \
561 `PredictedQuantumGraphReader` ]
562 A context manager that returns the reader when entered.
563 """
564 if page_size is None:
565 page_size = int(os.environ.get("LSST_QG_PAGE_SIZE", DEFAULT_PAGE_SIZE))
566 uri = ResourcePath(uri)
567 cdict: zstandard.ZstdCompressionDict | None = None
568 with uri.open(mode="rb") as zf_stream:
569 with zipfile.ZipFile(zf_stream, "r") as zf:
570 if (cdict_path := zipfile.Path(zf, "compression_dict")).exists():
571 cdict = zstandard.ZstdCompressionDict(cdict_path.read_bytes())
572 decompressor = zstandard.ZstdDecompressor(cdict)
573 header = cls._read_single_block_static("header", HeaderModel, zf, decompressor)
574 if not header.graph_type == graph_type:
575 raise TypeError(f"Header is for a {header.graph_type!r} graph, not {graph_type!r} graph.")
576 serialized_pipeline_graph = cls._read_single_block_static(
577 "pipeline_graph", SerializedPipelineGraph, zf, decompressor
578 )
579 pipeline_graph = serialized_pipeline_graph.deserialize(import_mode)
580 with AddressReader.open_in_zip(
581 zf,
582 address_filename,
583 page_size=page_size,
584 int_size=header.int_size,
585 n_addresses=n_addresses,
586 ) as address_reader:
587 yield cls(
588 header=header,
589 pipeline_graph=pipeline_graph,
590 zf=zf,
591 decompressor=decompressor,
592 address_reader=address_reader,
593 page_size=page_size,
594 )
596 @staticmethod
597 def _read_single_block_static(
598 name: str, model_type: type[_T], zf: zipfile.ZipFile, decompressor: Decompressor
599 ) -> _T:
600 """Read a single compressed JSON block from a 'file' in a zip archive.
602 Parameters
603 ----------
604 zf : `zipfile.ZipFile`
605 Zip archive to read the file from.
606 name : `str`
607 Base name of the file. An extension will be added.
608 model_type : `type` [ `pydantic.BaseModel` ]
609 Pydantic model to validate JSON with.
610 decompressor : `Decompressor`
611 Object with a `decompress` method that takes and returns `bytes`.
613 Returns
614 -------
615 model : `pydantic.BaseModel`
616 Validated model.
617 """
618 compressed_data = zf.read(f"{name}.json.zst")
619 json_data = decompressor.decompress(compressed_data)
620 return model_type.model_validate_json(json_data)
622 def _read_single_block(self, name: str, model_type: type[_T]) -> _T:
623 """Read a single compressed JSON block from a 'file' in a zip archive.
625 Parameters
626 ----------
627 name : `str`
628 Base name of the file. An extension will be added.
629 model_type : `type` [ `pydantic.BaseModel` ]
630 Pydantic model to validate JSON with.
632 Returns
633 -------
634 model : `pydantic.BaseModel`
635 Validated model.
636 """
637 return self._read_single_block_static(name, model_type, self.zf, self.decompressor)
639 def _read_single_block_raw(self, name: str) -> bytes:
640 """Read a single compressed block from a 'file' in a zip archive.
642 Parameters
643 ----------
644 name : `str`
645 Base name of the file. An extension will be added.
647 Returns
648 -------
649 data : `bytes`
650 Decompressed bytes.
651 """
652 compressed_data = self.zf.read(f"{name}.json.zst")
653 return self.decompressor.decompress(compressed_data)