Coverage for python / lsst / pipe / base / quantum_graph / ingest_graph.py: 23%
147 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-24 01:12 -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/>.
28"""A tool for ingesting provenance quantum graphs (written by the `aggregator`
29module) and [re-]ingesting other datasets (metadata/logs/configs) backed by the
30same file. This "finalizes" the RUN collection, prohibiting (at least
31conceptually) further processing.
33This always proceeds in three steps, so we can resume efficiently:
351. First we ask the butler to "forget" any metadata/log/config datasets that
36 exist in the output RUN collection, removing any record of them from the
37 butler database while preserving their files.
392. Next we ingest the ``run_provenance`` graph dataset itself.
413. Finally, in batches of quanta, we use a
42 `~lsst.daf.butler.QuantumBackedButler` to delete the original
43 metadata/log/config files and ingest new versions of those datasets into the
44 butler.
46Thus, at any point, if the ``run_provenance`` dataset has not been ingested,
47we know any metadata/log/config datasets that have been ingested are backed by
48the original files.
50Moreover, if the ``run_provenance`` dataset has been ingested, any existing
51metadata/log/config datasets must be backed by the graph file, and the original
52files for those datasets will have been deleted.
54We also know that at all times the metadata/log/config *content* is safely
55present in either the original files in the butler storage or in an
56already-ingested ``run_provenance`` dataset.
57"""
59from __future__ import annotations
61__all__ = ("ingest_graph",)
63import dataclasses
64import itertools
65import uuid
66from collections.abc import Iterator
67from contextlib import contextmanager
69from lsst.daf.butler import (
70 Butler,
71 Config,
72 DataCoordinate,
73 DatasetRef,
74 DatasetType,
75 FileDataset,
76 QuantumBackedButler,
77)
78from lsst.daf.butler.registry.sql_registry import SqlRegistry
79from lsst.resources import ResourcePath, ResourcePathExpression
80from lsst.utils.logging import getLogger
82from ..automatic_connection_constants import PROVENANCE_DATASET_TYPE_NAME, PROVENANCE_STORAGE_CLASS
83from ._provenance import (
84 ProvenanceDatasetInfo,
85 ProvenanceInitQuantumInfo,
86 ProvenanceQuantumGraph,
87 ProvenanceQuantumGraphReader,
88 ProvenanceQuantumInfo,
89)
90from .formatter import ProvenanceFormatter
92_LOG = getLogger(__name__)
95def ingest_graph(
96 butler_config: str | Config,
97 uri: ResourcePathExpression | None = None,
98 *,
99 transfer: str | None = "move",
100 batch_size: int = 10000,
101 output_run: str | None = None,
102) -> None:
103 """Ingest a provenance graph into a butler repository.
105 Parameters
106 ----------
107 butler_config : `str`
108 Path or alias for the butler repository, or a butler repository config
109 object.
110 uri : convertible to `lsst.resources.ResourcePath` or `None`, optional
111 Location of the provenance quantum graph to ingest. `None` indicates
112 that the quantum graph has already been ingested, but other ingests
113 and/or deletions failed and need to be resumed.
114 batch_size : `int`, optional
115 Number of datasets to process in each transaction.
116 output_run : `str`, optional
117 Output `~lsst.daf.butler.CollectionType.RUN` collection name. Only
118 needs to be provided if ``uri`` is `None`. If it is provided the
119 output run in the graph is checked against it.
121 Notes
122 -----
123 After this operation, no further processing may be done in the
124 `~lsst.daf.butler.CollectionType.RUN` collection.
126 If this process is interrupted, it can pick up where it left off if run
127 again (at the cost of some duplicate work to figure out how much progress
128 it had made).
129 """
130 with _GraphIngester.open(butler_config, uri, output_run) as helper:
131 helper.fetch_already_ingested_datasets()
132 if not helper.graph_already_ingested:
133 assert uri is not None
134 helper.forget_ingested_datasets(batch_size=batch_size)
135 helper.ingest_graph_dataset(uri, transfer=transfer)
136 helper.clean_and_reingest_datasets(batch_size=batch_size)
139@dataclasses.dataclass
140class _GraphIngester:
141 butler_config: str | Config
142 butler: Butler
143 graph: ProvenanceQuantumGraph
144 graph_already_ingested: bool
145 n_datasets: int
146 datasets_already_ingested: set[uuid.UUID] = dataclasses.field(default_factory=set)
148 @property
149 def output_run(self) -> str:
150 return self.graph.header.output_run
152 @classmethod
153 @contextmanager
154 def open(
155 cls,
156 butler_config: str | Config,
157 uri: ResourcePathExpression | None,
158 output_run: str | None,
159 ) -> Iterator[_GraphIngester]:
160 with Butler.from_config(butler_config, collections=output_run, writeable=True) as butler:
161 butler.registry.registerDatasetType(
162 DatasetType(PROVENANCE_DATASET_TYPE_NAME, butler.dimensions.empty, PROVENANCE_STORAGE_CLASS)
163 )
164 graph, graph_already_ingested = cls.read_graph(butler, uri)
165 if output_run is not None and graph.header.output_run != output_run:
166 raise ValueError(
167 f"Given output run {output_run!r} does not match the graph "
168 f"header {graph.header.output_run!r}."
169 )
170 n_datasets = 2 * len(graph.quantum_only_xgraph) + len(graph.init_quanta)
171 yield cls(
172 butler_config=butler_config,
173 butler=butler,
174 graph=graph,
175 graph_already_ingested=graph_already_ingested,
176 n_datasets=n_datasets,
177 )
179 @staticmethod
180 def read_graph(
181 butler: Butler,
182 uri: ResourcePathExpression | None,
183 ) -> tuple[ProvenanceQuantumGraph, bool]:
184 if uri is not None:
185 _LOG.info("Reading the pre-ingest provenance graph.")
186 with ProvenanceQuantumGraphReader.open(uri) as reader:
187 reader.read_quanta()
188 reader.read_init_quanta()
189 graph = reader.graph
190 already_ingested = (
191 butler.find_dataset(PROVENANCE_DATASET_TYPE_NAME, collections=[graph.header.output_run])
192 is not None
193 )
194 return graph, already_ingested
195 else:
196 _LOG.info("Reading the already-ingested provenance graph.")
197 parameters = {"datasets": [], "read_init_quanta": True}
198 return butler.get(PROVENANCE_DATASET_TYPE_NAME, parameters=parameters), True
200 def fetch_already_ingested_datasets(self) -> None:
201 _LOG.info("Querying for existing datasets in %r.", self.output_run)
202 self.datasets_already_ingested.update(self.butler.registry._fetch_run_dataset_ids(self.output_run))
204 def iter_datasets(self) -> Iterator[tuple[uuid.UUID, ProvenanceDatasetInfo]]:
205 xgraph = self.graph.bipartite_xgraph
206 for task_label, quanta_for_task in self.graph.quanta_by_task.items():
207 _LOG.verbose(
208 "Batching up metadata and log datasets from %d %s quanta.", len(quanta_for_task), task_label
209 )
210 for quantum_id in quanta_for_task.values():
211 quantum_info: ProvenanceQuantumInfo = xgraph.nodes[quantum_id]
212 metadata_id = quantum_info["metadata_id"]
213 yield metadata_id, xgraph.nodes[metadata_id]
214 log_id = quantum_info["log_id"]
215 yield log_id, xgraph.nodes[log_id]
216 _LOG.verbose("Batching up config datasets from %d tasks.", len(self.graph.init_quanta))
217 for task_label, quantum_id in self.graph.init_quanta.items():
218 init_quantum_info: ProvenanceInitQuantumInfo = xgraph.nodes[quantum_id]
219 config_id = init_quantum_info["config_id"]
220 yield config_id, xgraph.nodes[config_id]
222 def forget_ingested_datasets(self, batch_size: int) -> None:
223 _LOG.info(
224 "Dropping database records for metadata/log/config datasets backed by their original files."
225 )
226 to_forget: list[DatasetRef] = []
227 n_forgotten: int = 0
228 n_skipped: int = 0
229 for dataset_id, dataset_info in self.iter_datasets():
230 if dataset_info["produced"] and dataset_id in self.datasets_already_ingested:
231 to_forget.append(self._make_ref_from_info(dataset_id, dataset_info))
232 self.datasets_already_ingested.remove(dataset_id)
233 if len(to_forget) >= batch_size:
234 n_forgotten += self._run_forget(to_forget, n_forgotten + n_skipped)
235 else:
236 n_skipped += 1
237 n_forgotten += self._run_forget(to_forget, n_forgotten + n_skipped)
238 _LOG.info(
239 "Removed database records for %d metadata/log/config datasets, while %d were already absent.",
240 n_forgotten,
241 n_skipped,
242 )
244 def _run_forget(self, to_forget: list[DatasetRef], n_current: int) -> int:
245 if to_forget:
246 _LOG.verbose(
247 "Forgetting a %d-dataset batch; %d/%d forgotten so far or already absent.",
248 len(to_forget),
249 n_current,
250 self.n_datasets,
251 )
252 with self.butler.registry.transaction():
253 self.butler._datastore.forget(to_forget)
254 self.butler.registry.removeDatasets(to_forget)
255 n = len(to_forget)
256 to_forget.clear()
257 return n
259 def ingest_graph_dataset(self, uri: ResourcePathExpression, transfer: str | None) -> None:
260 _LOG.info("Ingesting the provenance quantum graph.")
261 dataset_type = DatasetType(
262 PROVENANCE_DATASET_TYPE_NAME, self.butler.dimensions.empty, PROVENANCE_STORAGE_CLASS
263 )
264 self.butler.registry.registerDatasetType(dataset_type)
265 ref = DatasetRef(dataset_type, DataCoordinate.make_empty(self.butler.dimensions), run=self.output_run)
266 uri = ResourcePath(uri)
267 self.butler.ingest(
268 # We use .abspath() since butler assumes paths are relative to the
269 # repo root, while users expects them to be relative to the CWD in
270 # this context.
271 FileDataset(refs=[ref], path=uri.abspath(), formatter=ProvenanceFormatter),
272 transfer=transfer,
273 )
275 def clean_and_reingest_datasets(self, batch_size: int) -> None:
276 _LOG.info(
277 "Deleting original metadata/log/config files and re-ingesting them with provenance graph backing."
278 )
279 direct_uri = self.butler.getURI(PROVENANCE_DATASET_TYPE_NAME, collections=[self.output_run])
280 qbb = self.make_qbb()
281 to_process: list[DatasetRef] = []
282 n_processed: int = 0
283 n_skipped: int = 0
284 n_not_produced: int = 0
285 for dataset_id, dataset_info in self.iter_datasets():
286 if not dataset_info["produced"]:
287 n_not_produced += 1
288 elif dataset_id not in self.datasets_already_ingested:
289 to_process.append(self._make_ref_from_info(dataset_id, dataset_info))
290 if len(to_process) >= batch_size:
291 n_processed += self._run_clean_and_ingest(
292 qbb, direct_uri, to_process, n_processed + n_skipped
293 )
294 else:
295 n_skipped += 1
296 n_processed += self._run_clean_and_ingest(qbb, direct_uri, to_process, n_processed + n_skipped)
297 _LOG.info(
298 "Deleted and re-ingested %d metadata/log/config datasets "
299 "(%d had already been processed, %d were not produced).",
300 n_processed,
301 n_skipped,
302 n_not_produced,
303 )
305 def _run_clean_and_ingest(
306 self, qbb: QuantumBackedButler, direct_uri: ResourcePath, to_process: list[DatasetRef], n_current: int
307 ) -> int:
308 if not to_process:
309 return 0
310 _LOG.verbose(
311 "Deleting and deleting a %d-dataset batch; %d/%d complete.",
312 len(to_process),
313 n_current,
314 self.n_datasets,
315 )
316 sql_registry: SqlRegistry = self.butler._registry # type: ignore[attr-defined]
317 expanded_refs = sql_registry.expand_refs(to_process)
318 # We need to pass predict=True to keep QBB/FileDatastore from wasting
319 # time doing existence checks, since ResourcePath.mremove will ignore
320 # nonexistent files anyway.
321 original_uris = list(
322 itertools.chain.from_iterable(
323 ref_uris.iter_all() for ref_uris in qbb.get_many_uris(expanded_refs, predict=True).values()
324 )
325 )
326 removal_status = ResourcePath.mremove(original_uris, do_raise=False)
327 for path, status in removal_status.items():
328 if not status.success and not isinstance(status.exception, FileNotFoundError):
329 assert status.exception is not None, "Exception should be set if success=False."
330 status.exception.add_note(f"Attempting to delete original file at {path}.")
331 raise status.exception
332 file_dataset = FileDataset(refs=expanded_refs, path=direct_uri, formatter=ProvenanceFormatter)
333 self.butler.ingest(file_dataset, transfer=None)
334 n = len(to_process)
335 to_process.clear()
336 return n
338 @staticmethod
339 def _make_ref_from_info(dataset_id: uuid.UUID, dataset_info: ProvenanceDatasetInfo) -> DatasetRef:
340 return DatasetRef(
341 dataset_info["pipeline_node"].dataset_type,
342 dataset_info["data_id"],
343 run=dataset_info["run"],
344 id=dataset_id,
345 )
347 def make_qbb(self) -> QuantumBackedButler:
348 dataset_types = {d.name: d.dataset_type for d in self.graph.pipeline_graph.dataset_types.values()}
349 return QuantumBackedButler.from_predicted(
350 config=self.butler_config,
351 predicted_inputs=(),
352 predicted_outputs=(),
353 dimensions=self.butler.dimensions,
354 datastore_records={},
355 dataset_types=dataset_types,
356 )