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-27 01:20 -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 : `lsst.resources.ResourcePathExpression` 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 transfer : `str` or `None`, optional
115 Transfer mode to use when ingesting graph. Matches those supported
116 by `lsst.resources.ResourcePath.transfer_from`.
117 batch_size : `int`, optional
118 Number of datasets to process in each transaction.
119 output_run : `str`, optional
120 Output `~lsst.daf.butler.CollectionType.RUN` collection name. Only
121 needs to be provided if ``uri`` is `None`. If it is provided the
122 output run in the graph is checked against it.
124 Notes
125 -----
126 After this operation, no further processing may be done in the
127 `~lsst.daf.butler.CollectionType.RUN` collection.
129 If this process is interrupted, it can pick up where it left off if run
130 again (at the cost of some duplicate work to figure out how much progress
131 it had made).
132 """
133 with _GraphIngester.open(butler_config, uri, output_run) as helper:
134 helper.fetch_already_ingested_datasets()
135 if not helper.graph_already_ingested:
136 assert uri is not None
137 helper.forget_ingested_datasets(batch_size=batch_size)
138 helper.ingest_graph_dataset(uri, transfer=transfer)
139 helper.clean_and_reingest_datasets(batch_size=batch_size)
142@dataclasses.dataclass
143class _GraphIngester:
144 butler_config: str | Config
145 butler: Butler
146 graph: ProvenanceQuantumGraph
147 graph_already_ingested: bool
148 n_datasets: int
149 datasets_already_ingested: set[uuid.UUID] = dataclasses.field(default_factory=set)
151 @property
152 def output_run(self) -> str:
153 return self.graph.header.output_run
155 @classmethod
156 @contextmanager
157 def open(
158 cls,
159 butler_config: str | Config,
160 uri: ResourcePathExpression | None,
161 output_run: str | None,
162 ) -> Iterator[_GraphIngester]:
163 with Butler.from_config(butler_config, collections=output_run, writeable=True) as butler:
164 butler.registry.registerDatasetType(
165 DatasetType(PROVENANCE_DATASET_TYPE_NAME, butler.dimensions.empty, PROVENANCE_STORAGE_CLASS)
166 )
167 graph, graph_already_ingested = cls.read_graph(butler, uri)
168 if output_run is not None and graph.header.output_run != output_run:
169 raise ValueError(
170 f"Given output run {output_run!r} does not match the graph "
171 f"header {graph.header.output_run!r}."
172 )
173 n_datasets = 2 * len(graph.quantum_only_xgraph) + len(graph.init_quanta)
174 yield cls(
175 butler_config=butler_config,
176 butler=butler,
177 graph=graph,
178 graph_already_ingested=graph_already_ingested,
179 n_datasets=n_datasets,
180 )
182 @staticmethod
183 def read_graph(
184 butler: Butler,
185 uri: ResourcePathExpression | None,
186 ) -> tuple[ProvenanceQuantumGraph, bool]:
187 if uri is not None:
188 _LOG.info("Reading the pre-ingest provenance graph.")
189 with ProvenanceQuantumGraphReader.open(uri) as reader:
190 reader.read_quanta()
191 reader.read_init_quanta()
192 graph = reader.graph
193 already_ingested = (
194 butler.find_dataset(PROVENANCE_DATASET_TYPE_NAME, collections=[graph.header.output_run])
195 is not None
196 )
197 return graph, already_ingested
198 else:
199 _LOG.info("Reading the already-ingested provenance graph.")
200 parameters = {"datasets": [], "read_init_quanta": True}
201 return butler.get(PROVENANCE_DATASET_TYPE_NAME, parameters=parameters), True
203 def fetch_already_ingested_datasets(self) -> None:
204 _LOG.info("Querying for existing datasets in %r.", self.output_run)
205 self.datasets_already_ingested.update(self.butler.registry._fetch_run_dataset_ids(self.output_run))
207 def iter_datasets(self) -> Iterator[tuple[uuid.UUID, ProvenanceDatasetInfo]]:
208 xgraph = self.graph.bipartite_xgraph
209 for task_label, quanta_for_task in self.graph.quanta_by_task.items():
210 _LOG.verbose(
211 "Batching up metadata and log datasets from %d %s quanta.", len(quanta_for_task), task_label
212 )
213 for quantum_id in quanta_for_task.values():
214 quantum_info: ProvenanceQuantumInfo = xgraph.nodes[quantum_id]
215 metadata_id = quantum_info["metadata_id"]
216 yield metadata_id, xgraph.nodes[metadata_id]
217 log_id = quantum_info["log_id"]
218 yield log_id, xgraph.nodes[log_id]
219 _LOG.verbose("Batching up config datasets from %d tasks.", len(self.graph.init_quanta))
220 for task_label, quantum_id in self.graph.init_quanta.items():
221 init_quantum_info: ProvenanceInitQuantumInfo = xgraph.nodes[quantum_id]
222 config_id = init_quantum_info["config_id"]
223 yield config_id, xgraph.nodes[config_id]
225 def forget_ingested_datasets(self, batch_size: int) -> None:
226 _LOG.info(
227 "Dropping database records for metadata/log/config datasets backed by their original files."
228 )
229 to_forget: list[DatasetRef] = []
230 n_forgotten: int = 0
231 n_skipped: int = 0
232 for dataset_id, dataset_info in self.iter_datasets():
233 if dataset_info["produced"] and dataset_id in self.datasets_already_ingested:
234 to_forget.append(self._make_ref_from_info(dataset_id, dataset_info))
235 self.datasets_already_ingested.remove(dataset_id)
236 if len(to_forget) >= batch_size:
237 n_forgotten += self._run_forget(to_forget, n_forgotten + n_skipped)
238 else:
239 n_skipped += 1
240 n_forgotten += self._run_forget(to_forget, n_forgotten + n_skipped)
241 _LOG.info(
242 "Removed database records for %d metadata/log/config datasets, while %d were already absent.",
243 n_forgotten,
244 n_skipped,
245 )
247 def _run_forget(self, to_forget: list[DatasetRef], n_current: int) -> int:
248 if to_forget:
249 _LOG.verbose(
250 "Forgetting a %d-dataset batch; %d/%d forgotten so far or already absent.",
251 len(to_forget),
252 n_current,
253 self.n_datasets,
254 )
255 with self.butler.registry.transaction():
256 self.butler._datastore.forget(to_forget)
257 self.butler.registry.removeDatasets(to_forget)
258 n = len(to_forget)
259 to_forget.clear()
260 return n
262 def ingest_graph_dataset(self, uri: ResourcePathExpression, transfer: str | None) -> None:
263 _LOG.info("Ingesting the provenance quantum graph.")
264 dataset_type = DatasetType(
265 PROVENANCE_DATASET_TYPE_NAME, self.butler.dimensions.empty, PROVENANCE_STORAGE_CLASS
266 )
267 self.butler.registry.registerDatasetType(dataset_type)
268 ref = DatasetRef(dataset_type, DataCoordinate.make_empty(self.butler.dimensions), run=self.output_run)
269 uri = ResourcePath(uri)
270 self.butler.ingest(
271 # We use .abspath() since butler assumes paths are relative to the
272 # repo root, while users expects them to be relative to the CWD in
273 # this context.
274 FileDataset(refs=[ref], path=uri.abspath(), formatter=ProvenanceFormatter),
275 transfer=transfer,
276 )
278 def clean_and_reingest_datasets(self, batch_size: int) -> None:
279 _LOG.info(
280 "Deleting original metadata/log/config files and re-ingesting them with provenance graph backing."
281 )
282 direct_uri = self.butler.getURI(PROVENANCE_DATASET_TYPE_NAME, collections=[self.output_run])
283 qbb = self.make_qbb()
284 to_process: list[DatasetRef] = []
285 n_processed: int = 0
286 n_skipped: int = 0
287 n_not_produced: int = 0
288 for dataset_id, dataset_info in self.iter_datasets():
289 if not dataset_info["produced"]:
290 n_not_produced += 1
291 elif dataset_id not in self.datasets_already_ingested:
292 to_process.append(self._make_ref_from_info(dataset_id, dataset_info))
293 if len(to_process) >= batch_size:
294 n_processed += self._run_clean_and_ingest(
295 qbb, direct_uri, to_process, n_processed + n_skipped
296 )
297 else:
298 n_skipped += 1
299 n_processed += self._run_clean_and_ingest(qbb, direct_uri, to_process, n_processed + n_skipped)
300 _LOG.info(
301 "Deleted and re-ingested %d metadata/log/config datasets "
302 "(%d had already been processed, %d were not produced).",
303 n_processed,
304 n_skipped,
305 n_not_produced,
306 )
308 def _run_clean_and_ingest(
309 self, qbb: QuantumBackedButler, direct_uri: ResourcePath, to_process: list[DatasetRef], n_current: int
310 ) -> int:
311 if not to_process:
312 return 0
313 _LOG.verbose(
314 "Deleting and deleting a %d-dataset batch; %d/%d complete.",
315 len(to_process),
316 n_current,
317 self.n_datasets,
318 )
319 sql_registry: SqlRegistry = self.butler._registry # type: ignore[attr-defined]
320 expanded_refs = sql_registry.expand_refs(to_process)
321 # We need to pass predict=True to keep QBB/FileDatastore from wasting
322 # time doing existence checks, since ResourcePath.mremove will ignore
323 # nonexistent files anyway.
324 original_uris = list(
325 itertools.chain.from_iterable(
326 ref_uris.iter_all() for ref_uris in qbb.get_many_uris(expanded_refs, predict=True).values()
327 )
328 )
329 removal_status = ResourcePath.mremove(original_uris, do_raise=False)
330 for path, status in removal_status.items():
331 if not status.success and not isinstance(status.exception, FileNotFoundError):
332 assert status.exception is not None, "Exception should be set if success=False."
333 status.exception.add_note(f"Attempting to delete original file at {path}.")
334 raise status.exception
335 file_dataset = FileDataset(refs=expanded_refs, path=direct_uri, formatter=ProvenanceFormatter)
336 self.butler.ingest(file_dataset, transfer=None)
337 n = len(to_process)
338 to_process.clear()
339 return n
341 @staticmethod
342 def _make_ref_from_info(dataset_id: uuid.UUID, dataset_info: ProvenanceDatasetInfo) -> DatasetRef:
343 return DatasetRef(
344 dataset_info["pipeline_node"].dataset_type,
345 dataset_info["data_id"],
346 run=dataset_info["run"],
347 id=dataset_id,
348 )
350 def make_qbb(self) -> QuantumBackedButler:
351 dataset_types = {d.name: d.dataset_type for d in self.graph.pipeline_graph.dataset_types.values()}
352 return QuantumBackedButler.from_predicted(
353 config=self.butler_config,
354 predicted_inputs=(),
355 predicted_outputs=(),
356 dimensions=self.butler.dimensions,
357 datastore_records={},
358 dataset_types=dataset_types,
359 )