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