Coverage for python/lsst/pipe/base/quantum_graph/_common.py: 67%
187 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-06 01:36 -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 "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.resources import ResourcePath, ResourcePathExpression
65from ..pipeline_graph import DatasetTypeNode, Edge, PipelineGraph, TaskImportMode, TaskNode
66from ..pipeline_graph.io import SerializedPipelineGraph
67from ._multiblock import (
68 DEFAULT_PAGE_SIZE,
69 AddressReader,
70 AddressWriter,
71 Compressor,
72 Decompressor,
73)
75if TYPE_CHECKING:
76 from ..graph import QuantumGraph
79# These aliases make it a lot easier how the various pydantic models are
80# structured, but they're too verbose to be worth exporting to code outside the
81# quantum_graph subpackage.
82TaskLabel: TypeAlias = str
83DatasetTypeName: TypeAlias = str
84ConnectionName: TypeAlias = str
85DatasetIndex: TypeAlias = int
86QuantumIndex: TypeAlias = int
87DatastoreName: TypeAlias = str
88DimensionElementName: TypeAlias = str
89DataCoordinateValues: TypeAlias = list[DataIdValue]
92_T = TypeVar("_T", bound=pydantic.BaseModel)
95class IncompleteQuantumGraphError(RuntimeError):
96 pass
99class HeaderModel(pydantic.BaseModel):
100 """Data model for the header of a quantum graph file."""
102 version: int = 0
103 """File format / data model version number."""
105 graph_type: str = ""
106 """Type of quantum graph stored in this file."""
108 inputs: list[str] = pydantic.Field(default_factory=list)
109 """List of input collections used to build the quantum graph."""
111 output: str | None = ""
112 """Output CHAINED collection provided when building the quantum graph."""
114 output_run: str = ""
115 """Output RUN collection for all output datasets in this graph."""
117 user: str = pydantic.Field(default_factory=getpass.getuser)
118 """Username of the process that built this quantum graph."""
120 timestamp: datetime.datetime = pydantic.Field(default_factory=datetime.datetime.now)
121 """Timestamp for when this quantum graph was built.
123 It is unspecified exactly which point during quantum-graph generation this
124 timestamp is recorded.
125 """
127 command: str = pydantic.Field(default_factory=lambda: " ".join(sys.argv))
128 """Command-line invocation that created this graph."""
130 metadata: dict[str, Any] = pydantic.Field(default_factory=dict)
131 """Free-form metadata associated with this quantum graph at build time."""
133 int_size: int = 8
134 """Number of bytes in the integers used in this file's multi-block and
135 address files.
136 """
138 n_quanta: int = 0
139 """Total number of quanta in this graph.
141 This does not include special "init" quanta, but it does include quanta
142 that were not loaded in a partial read (except when reading from an old
143 quantum graph file).
144 """
146 n_datasets: int = 0
147 """Total number of distinct datasets in the full graph. This includes
148 datasets whose related quanta were not loaded in a partial read (except
149 when reading from an old quantum graph file).
150 """
152 n_task_quanta: dict[TaskLabel, int] = pydantic.Field(default_factory=dict)
153 """Number of quanta for each task label.
155 This does not include special "init" quanta, but it does include quanta
156 that were not loaded in a partial read (except when reading from an old
157 quantum graph file).
158 """
160 @classmethod
161 def from_old_quantum_graph(cls, old_quantum_graph: QuantumGraph) -> HeaderModel:
162 """Extract a header from an old `QuantumGraph` instance.
164 Parameters
165 ----------
166 old_quantum_graph : `QuantumGraph`
167 Quantum graph to extract a header from.
169 Returns
170 -------
171 header : `PredictedHeaderModel`
172 Header for a new predicted quantum graph.
173 """
174 metadata = dict(old_quantum_graph.metadata)
175 metadata.pop("packages", None)
176 if (time_str := metadata.pop("time", None)) is not None:
177 timestamp = datetime.datetime.fromisoformat(time_str)
178 else:
179 timestamp = datetime.datetime.now()
180 return cls(
181 inputs=list(metadata.pop("input", []) or []), # Guard against explicit None and missing key.
182 output=metadata.pop("output", None),
183 output_run=metadata.pop("output_run", ""),
184 user=metadata.pop("user", ""),
185 command=metadata.pop("full_command", ""),
186 timestamp=timestamp,
187 metadata=metadata,
188 )
190 def to_old_metadata(self) -> dict[str, Any]:
191 """Return a dictionary using the key conventions used in old quantum
192 graph files.
193 """
194 result = self.metadata.copy()
195 result["input"] = self.inputs
196 result["output"] = self.output
197 result["output_run"] = self.output_run
198 result["full_command"] = self.command
199 result["user"] = self.user
200 result["time"] = str(self.timestamp)
201 return result
203 # Work around the fact that Sphinx chokes on Pydantic docstring formatting,
204 # when we inherit those docstrings in our public classes.
205 if "sphinx" in sys.modules and not TYPE_CHECKING:
207 def copy(self, *args: Any, **kwargs: Any) -> Any:
208 """See `pydantic.BaseModel.copy`."""
209 return super().copy(*args, **kwargs)
211 def model_dump(self, *args: Any, **kwargs: Any) -> Any:
212 """See `pydantic.BaseModel.model_dump`."""
213 return super().model_dump(*args, **kwargs)
215 def model_dump_json(self, *args: Any, **kwargs: Any) -> Any:
216 """See `pydantic.BaseModel.model_dump_json`."""
217 return super().model_dump(*args, **kwargs)
219 def model_copy(self, *args: Any, **kwargs: Any) -> Any:
220 """See `pydantic.BaseModel.model_copy`."""
221 return super().model_copy(*args, **kwargs)
223 @classmethod
224 def model_construct(cls, *args: Any, **kwargs: Any) -> Any: # type: ignore[misc, override]
225 """See `pydantic.BaseModel.model_construct`."""
226 return super().model_construct(*args, **kwargs)
228 @classmethod
229 def model_json_schema(cls, *args: Any, **kwargs: Any) -> Any:
230 """See `pydantic.BaseModel.model_json_schema`."""
231 return super().model_json_schema(*args, **kwargs)
233 @classmethod
234 def model_validate(cls, *args: Any, **kwargs: Any) -> Any:
235 """See `pydantic.BaseModel.model_validate`."""
236 return super().model_validate(*args, **kwargs)
238 @classmethod
239 def model_validate_json(cls, *args: Any, **kwargs: Any) -> Any:
240 """See `pydantic.BaseModel.model_validate_json`."""
241 return super().model_validate_json(*args, **kwargs)
243 @classmethod
244 def model_validate_strings(cls, *args: Any, **kwargs: Any) -> Any:
245 """See `pydantic.BaseModel.model_validate_strings`."""
246 return super().model_validate_strings(*args, **kwargs)
249class QuantumInfo(TypedDict):
250 """A typed dictionary that annotates the attributes of the NetworkX graph
251 node data for a quantum.
253 Since NetworkX types are not generic over their node mapping type, this has
254 to be used explicitly, e.g.::
256 node_data: QuantumInfo = xgraph.nodes[quantum_id]
258 where ``xgraph`` can be either `BaseQuantumGraph.quantum_only_xgraph`
259 or `BaseQuantumGraph.bipartite_xgraph`.
260 """
262 data_id: DataCoordinate
263 """Data ID of the quantum."""
265 task_label: str
266 """Label of the task for this quantum."""
268 pipeline_node: TaskNode
269 """Node in the pipeline graph for this quantum's task."""
272class DatasetInfo(TypedDict):
273 """A typed dictionary that annotates the attributes of the NetworkX graph
274 node data for a dataset.
276 Since NetworkX types are not generic over their node mapping type, this has
277 to be used explicitly, e.g.::
279 node_data: DatasetInfo = xgraph.nodes[dataset_id]
281 where ``xgraph`` is from the `BaseQuantumGraph.bipartite_xgraph` property.
282 """
284 data_id: DataCoordinate
285 """Data ID of the dataset."""
287 dataset_type_name: DatasetTypeName
288 """Name of the type of this dataset.
290 This is always the general dataset type that matches the data repository
291 storage class, which may differ from any particular task-adapted dataset
292 type whose storage class has been overridden to match the task connections.
293 This means is it never a component.
294 """
296 run: str
297 """Name of the `~lsst.daf.butler.CollectionType.RUN` collection that holds
298 or will hold this dataset.
299 """
301 pipeline_node: DatasetTypeNode
302 """Node in the pipeline graph for this dataset's type."""
305class BipartiteEdgeInfo(TypedDict):
306 """A typed dictionary that annotates the attributes of the NetworkX graph
307 edge data in a bipartite graph.
308 """
310 is_read: bool
311 """`True` if this is a dataset -> quantum edge; `False` if it is a
312 quantum -> dataset edge.
313 """
315 pipeline_edges: list[Edge]
316 """Corresponding edges in the pipeline graph.
318 Note that there may be more than one pipeline edge since a quantum can
319 consume a particular dataset via multiple connections.
320 """
323class BaseQuantumGraph(ABC):
324 """An abstract base for quantum graphs.
326 Parameters
327 ----------
328 header : `HeaderModel`
329 Structured metadata for the graph.
330 pipeline_graph : `.pipeline_graph.PipelineGraph`
331 Graph of tasks and dataset types. May contain a superset of the tasks
332 and dataset types that actually have quanta and datasets in the quantum
333 graph.
334 """
336 def __init__(self, header: HeaderModel, pipeline_graph: PipelineGraph):
337 self.header = header
338 self.pipeline_graph = pipeline_graph
340 @property
341 @abstractmethod
342 def quanta_by_task(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:
343 """A nested mapping of all quanta, keyed first by task name and then by
344 data ID.
346 Notes
347 -----
348 Partial loads may not fully populate this mapping, but it can always
349 be accessed.
351 The returned object may be an internal dictionary; as the type
352 annotation indicates, it should not be modified in place.
353 """
354 raise NotImplementedError()
356 @property
357 @abstractmethod
358 def datasets_by_type(self) -> Mapping[str, Mapping[DataCoordinate, uuid.UUID]]:
359 """A nested mapping of all datasets, keyed first by dataset type name
360 and then by data ID.
362 Notes
363 -----
364 Partial loads may not fully populate this mapping, but it can always
365 be accessed.
367 The returned object may be an internal dictionary; as the type
368 annotation indicates, it should not be modified in place.
369 """
370 raise NotImplementedError()
372 @property
373 @abstractmethod
374 def quantum_only_xgraph(self) -> networkx.DiGraph:
375 """A directed acyclic graph with quanta as nodes and datasets elided.
377 Notes
378 -----
379 Partial loads may not fully populate this graph, but it can always be
380 accessed.
382 Node state dictionaries are described by the `QuantumInfo` type
383 (or a subtype thereof).
385 The returned object is a read-only view of an internal one.
386 """
387 raise NotImplementedError()
389 @property
390 @abstractmethod
391 def bipartite_xgraph(self) -> networkx.DiGraph:
392 """A directed acyclic graph with quantum and dataset nodes.
394 Notes
395 -----
396 Partial loads may not fully populate this graph, but it can always be
397 accessed.
399 Node state dictionaries are described by the `QuantumInfo` and
400 `DatasetInfo` types (or a subtypes thereof). Edges have state
401 dictionaries described by `BipartiteEdgeInfo`.
403 The returned object is a read-only view of an internal one.
404 """
405 raise NotImplementedError()
408@dataclasses.dataclass
409class BaseQuantumGraphWriter:
410 """A helper class for writing quantum graphs."""
412 zf: zipfile.ZipFile
413 """The zip archive that represents the quantum graph on disk."""
415 compressor: Compressor
416 """A compressor for all compressed JSON blocks."""
418 address_writer: AddressWriter
419 """A helper object for reading addresses into the multi-block files."""
421 int_size: int
422 """Size (in bytes) used to write integers to binary files."""
424 @classmethod
425 @contextmanager
426 def open(
427 cls,
428 uri: ResourcePathExpression,
429 header: HeaderModel,
430 pipeline_graph: PipelineGraph,
431 indices: dict[uuid.UUID, int],
432 *,
433 address_filename: str,
434 compressor: Compressor,
435 cdict_data: bytes | None = None,
436 ) -> Iterator[Self]:
437 uri = ResourcePath(uri)
438 address_writer = AddressWriter(indices)
439 with uri.open(mode="wb") as stream:
440 with zipfile.ZipFile(stream, mode="w", compression=zipfile.ZIP_STORED) as zf:
441 self = cls(zf, compressor, address_writer, header.int_size)
442 self.write_single_model("header", header)
443 if cdict_data is not None:
444 zf.writestr("compression_dict", cdict_data)
445 self.write_single_model("pipeline_graph", SerializedPipelineGraph.serialize(pipeline_graph))
446 yield self
447 address_writer.write_to_zip(zf, address_filename, int_size=self.int_size)
449 def write_single_model(self, name: str, model: pydantic.BaseModel) -> None:
450 """Write a single compressed JSON block as a 'file' in a zip archive.
452 Parameters
453 ----------
454 name : `str`
455 Base name of the file. An extension will be added.
456 model : `pydantic.BaseModel`
457 Pydantic model to convert to JSON.
458 """
459 json_data = model.model_dump_json().encode()
460 self.write_single_block(name, json_data)
462 def write_single_block(self, name: str, json_data: bytes) -> None:
463 """Write a single compressed JSON block as a 'file' in a zip archive.
465 Parameters
466 ----------
467 name : `str`
468 Base name of the file. An extension will be added.
469 json_data : `bytes`
470 Raw JSON to compress and write.
471 """
472 json_data = self.compressor.compress(json_data)
473 self.zf.writestr(f"{name}.json.zst", json_data)
476@dataclasses.dataclass
477class BaseQuantumGraphReader:
478 """A helper class for reading quantum graphs."""
480 header: HeaderModel
481 """Header metadata for the quantum graph."""
483 pipeline_graph: PipelineGraph
484 """Graph of tasks and dataset type names that appear in the quantum
485 graph.
486 """
488 zf: zipfile.ZipFile
489 """The zip archive that represents the quantum graph on disk."""
491 decompressor: Decompressor
492 """A decompressor for all compressed JSON blocks."""
494 address_reader: AddressReader
495 """A helper object for reading addresses into the multi-block files."""
497 page_size: int
498 """Approximate number of bytes to read at a time.
500 Note that this does not set a page size for *all* reads, but it
501 does affect the smallest, most numerous reads.
502 """
504 @classmethod
505 @contextmanager
506 def _open(
507 cls,
508 uri: ResourcePathExpression,
509 *,
510 address_filename: str,
511 graph_type: str,
512 n_addresses: int,
513 page_size: int | None = None,
514 import_mode: TaskImportMode = TaskImportMode.ASSUME_CONSISTENT_EDGES,
515 ) -> Iterator[Self]:
516 """Construct a reader from a URI.
518 Parameters
519 ----------
520 uri : convertible to `lsst.resources.ResourcePath`
521 URI to open. Should have a ``.qg`` extension.
522 address_filename : `str`
523 Base filename for the address file.
524 graph_type : `str`
525 Value to expect for `HeaderModel.graph_type`.
526 n_addresses : `int`
527 Number of addresses to expect per row in the address file.
528 page_size : `int`, optional
529 Approximate number of bytes to read at once from address files.
530 Note that this does not set a page size for *all* reads, but it
531 does affect the smallest, most numerous reads. When `None`, the
532 ``LSST_QG_PAGE_SIZE`` environment variable is checked before
533 falling back to a default of 5MB.
534 import_mode : `..pipeline_graph.TaskImportMode`, optional
535 How to handle importing the task classes referenced in the pipeline
536 graph.
538 Returns
539 -------
540 reader : `contextlib.AbstractContextManager` [ \
541 `PredictedQuantumGraphReader` ]
542 A context manager that returns the reader when entered.
543 """
544 if page_size is None:
545 page_size = int(os.environ.get("LSST_QG_PAGE_SIZE", DEFAULT_PAGE_SIZE))
546 uri = ResourcePath(uri)
547 cdict: zstandard.ZstdCompressionDict | None = None
548 with uri.open(mode="rb") as zf_stream:
549 with zipfile.ZipFile(zf_stream, "r") as zf:
550 if (cdict_path := zipfile.Path(zf, "compression_dict")).exists():
551 cdict = zstandard.ZstdCompressionDict(cdict_path.read_bytes())
552 decompressor = zstandard.ZstdDecompressor(cdict)
553 header = cls._read_single_block_static("header", HeaderModel, zf, decompressor)
554 if not header.graph_type == graph_type:
555 raise TypeError(f"Header is for a {header.graph_type!r} graph, not {graph_type!r} graph.")
556 serialized_pipeline_graph = cls._read_single_block_static(
557 "pipeline_graph", SerializedPipelineGraph, zf, decompressor
558 )
559 pipeline_graph = serialized_pipeline_graph.deserialize(import_mode)
560 with AddressReader.open_in_zip(
561 zf,
562 address_filename,
563 page_size=page_size,
564 int_size=header.int_size,
565 n_addresses=n_addresses,
566 ) as address_reader:
567 yield cls(
568 header=header,
569 pipeline_graph=pipeline_graph,
570 zf=zf,
571 decompressor=decompressor,
572 address_reader=address_reader,
573 page_size=page_size,
574 )
576 @staticmethod
577 def _read_single_block_static(
578 name: str, model_type: type[_T], zf: zipfile.ZipFile, decompressor: Decompressor
579 ) -> _T:
580 """Read a single compressed JSON block from a 'file' in a zip archive.
582 Parameters
583 ----------
584 zf : `zipfile.ZipFile`
585 Zip archive to read the file from.
586 name : `str`
587 Base name of the file. An extension will be added.
588 model_type : `type` [ `pydantic.BaseModel` ]
589 Pydantic model to validate JSON with.
590 decompressor : `Decompressor`
591 Object with a `decompress` method that takes and returns `bytes`.
593 Returns
594 -------
595 model : `pydantic.BaseModel`
596 Validated model.
597 """
598 compressed_data = zf.read(f"{name}.json.zst")
599 json_data = decompressor.decompress(compressed_data)
600 return model_type.model_validate_json(json_data)
602 def _read_single_block(self, name: str, model_type: type[_T]) -> _T:
603 """Read a single compressed JSON block from a 'file' in a zip archive.
605 Parameters
606 ----------
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.
612 Returns
613 -------
614 model : `pydantic.BaseModel`
615 Validated model.
616 """
617 return self._read_single_block_static(name, model_type, self.zf, self.decompressor)
619 def _read_single_block_raw(self, name: str) -> bytes:
620 """Read a single compressed block from a 'file' in a zip archive.
622 Parameters
623 ----------
624 name : `str`
625 Base name of the file. An extension will be added.
627 Returns
628 -------
629 data : `bytes`
630 Decompressed bytes.
631 """
632 compressed_data = self.zf.read(f"{name}.json.zst")
633 return self.decompressor.decompress(compressed_data)