Coverage for python / lsst / pipe / base / quantum_graph / _common.py: 67%
194 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-13 01:30 -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 *,
452 address_filename: str,
453 cdict_data: bytes | None = None,
454 zstd_level: int = 10,
455 ) -> Iterator[Self]:
456 uri = ResourcePath(uri)
457 address_writer = AddressWriter()
458 cdict = zstandard.ZstdCompressionDict(cdict_data) if cdict_data is not None else None
459 compressor = zstandard.ZstdCompressor(level=zstd_level, dict_data=cdict)
460 with uri.open(mode="wb") as stream:
461 with zipfile.ZipFile(stream, mode="w", compression=zipfile.ZIP_STORED) as zf:
462 self = cls(zf, compressor, address_writer, header.int_size)
463 self.write_single_model("header", header)
464 if cdict_data is not None:
465 zf.writestr("compression_dict", cdict_data)
466 self.write_single_model("pipeline_graph", SerializedPipelineGraph.serialize(pipeline_graph))
467 yield self
468 address_writer.write_to_zip(zf, address_filename, int_size=self.int_size)
470 def write_single_model(self, name: str, model: pydantic.BaseModel) -> None:
471 """Write a single compressed JSON block as a 'file' in a zip archive.
473 Parameters
474 ----------
475 name : `str`
476 Base name of the file. An extension will be added.
477 model : `pydantic.BaseModel`
478 Pydantic model to convert to JSON.
479 """
480 json_data = model.model_dump_json().encode()
481 self.write_single_block(name, json_data)
483 def write_single_block(self, name: str, json_data: bytes) -> None:
484 """Write a single compressed JSON block as a 'file' in a zip archive.
486 Parameters
487 ----------
488 name : `str`
489 Base name of the file. An extension will be added.
490 json_data : `bytes`
491 Raw JSON to compress and write.
492 """
493 json_data = self.compressor.compress(json_data)
494 self.zf.writestr(f"{name}.json.zst", json_data)
497@dataclasses.dataclass
498class BaseQuantumGraphReader:
499 """A helper class for reading quantum graphs."""
501 header: HeaderModel
502 """Header metadata for the quantum graph."""
504 pipeline_graph: PipelineGraph
505 """Graph of tasks and dataset type names that appear in the quantum
506 graph.
507 """
509 zf: zipfile.ZipFile
510 """The zip archive that represents the quantum graph on disk."""
512 decompressor: Decompressor
513 """A decompressor for all compressed JSON blocks."""
515 address_reader: AddressReader
516 """A helper object for reading addresses into the multi-block files."""
518 page_size: int
519 """Approximate number of bytes to read at a time.
521 Note that this does not set a page size for *all* reads, but it
522 does affect the smallest, most numerous reads.
523 """
525 @classmethod
526 @contextmanager
527 def _open(
528 cls,
529 uri: ResourcePathExpression,
530 *,
531 address_filename: str,
532 graph_type: str,
533 n_addresses: int,
534 page_size: int | None = None,
535 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,
536 ) -> Iterator[Self]:
537 """Construct a reader from a URI.
539 Parameters
540 ----------
541 uri : convertible to `lsst.resources.ResourcePath`
542 URI to open. Should have a ``.qg`` extension.
543 address_filename : `str`
544 Base filename for the address file.
545 graph_type : `str`
546 Value to expect for `HeaderModel.graph_type`.
547 n_addresses : `int`
548 Number of addresses to expect per row in the address file.
549 page_size : `int`, optional
550 Approximate number of bytes to read at once from address files.
551 Note that this does not set a page size for *all* reads, but it
552 does affect the smallest, most numerous reads. When `None`, the
553 ``LSST_QG_PAGE_SIZE`` environment variable is checked before
554 falling back to a default of 5MB.
555 import_mode : `..pipeline_graph.TaskImportMode`, optional
556 How to handle importing the task classes referenced in the pipeline
557 graph.
559 Returns
560 -------
561 reader : `contextlib.AbstractContextManager` [ \
562 `PredictedQuantumGraphReader` ]
563 A context manager that returns the reader when entered.
564 """
565 if page_size is None:
566 page_size = int(os.environ.get("LSST_QG_PAGE_SIZE", DEFAULT_PAGE_SIZE))
567 uri = ResourcePath(uri)
568 cdict: zstandard.ZstdCompressionDict | None = None
569 with uri.open(mode="rb") as zf_stream:
570 with zipfile.ZipFile(zf_stream, "r") as zf:
571 if (cdict_path := zipfile.Path(zf, "compression_dict")).exists():
572 cdict = zstandard.ZstdCompressionDict(cdict_path.read_bytes())
573 decompressor = zstandard.ZstdDecompressor(cdict)
574 header = cls._read_single_block_static("header", HeaderModel, zf, decompressor)
575 if not header.graph_type == graph_type:
576 raise TypeError(f"Header is for a {header.graph_type!r} graph, not {graph_type!r} graph.")
577 serialized_pipeline_graph = cls._read_single_block_static(
578 "pipeline_graph", SerializedPipelineGraph, zf, decompressor
579 )
580 pipeline_graph = serialized_pipeline_graph.deserialize(import_mode)
581 with AddressReader.open_in_zip(
582 zf,
583 address_filename,
584 page_size=page_size,
585 int_size=header.int_size,
586 n_addresses=n_addresses,
587 ) as address_reader:
588 yield cls(
589 header=header,
590 pipeline_graph=pipeline_graph,
591 zf=zf,
592 decompressor=decompressor,
593 address_reader=address_reader,
594 page_size=page_size,
595 )
597 @staticmethod
598 def _read_single_block_static(
599 name: str, model_type: type[_T], zf: zipfile.ZipFile, decompressor: Decompressor
600 ) -> _T:
601 """Read a single compressed JSON block from a 'file' in a zip archive.
603 Parameters
604 ----------
605 zf : `zipfile.ZipFile`
606 Zip archive to read the file from.
607 name : `str`
608 Base name of the file. An extension will be added.
609 model_type : `type` [ `pydantic.BaseModel` ]
610 Pydantic model to validate JSON with.
611 decompressor : `Decompressor`
612 Object with a `decompress` method that takes and returns `bytes`.
614 Returns
615 -------
616 model : `pydantic.BaseModel`
617 Validated model.
618 """
619 compressed_data = zf.read(f"{name}.json.zst")
620 json_data = decompressor.decompress(compressed_data)
621 return model_type.model_validate_json(json_data)
623 def _read_single_block(self, name: str, model_type: type[_T]) -> _T:
624 """Read a single compressed JSON block from a 'file' in a zip archive.
626 Parameters
627 ----------
628 name : `str`
629 Base name of the file. An extension will be added.
630 model_type : `type` [ `pydantic.BaseModel` ]
631 Pydantic model to validate JSON with.
633 Returns
634 -------
635 model : `pydantic.BaseModel`
636 Validated model.
637 """
638 return self._read_single_block_static(name, model_type, self.zf, self.decompressor)
640 def _read_single_block_raw(self, name: str) -> bytes:
641 """Read a single compressed block from a 'file' in a zip archive.
643 Parameters
644 ----------
645 name : `str`
646 Base name of the file. An extension will be added.
648 Returns
649 -------
650 data : `bytes`
651 Decompressed bytes.
652 """
653 compressed_data = self.zf.read(f"{name}.json.zst")
654 return self.decompressor.decompress(compressed_data)