Coverage for python/lsst/daf/butler/_quantum_backed.py: 92%
192 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-06 01:43 -0800
1# This file is part of daf_butler.
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__ = ("QuantumBackedButler", "QuantumProvenanceData")
32import itertools
33import logging
34import uuid
35from collections import defaultdict
36from collections.abc import Iterable, Mapping
37from typing import TYPE_CHECKING, Any
39import pydantic
41from lsst.resources import ResourcePath, ResourcePathExpression
43from ._butler_config import ButlerConfig
44from ._butler_metrics import ButlerMetrics
45from ._config import Config
46from ._dataset_provenance import DatasetProvenance
47from ._dataset_ref import DatasetId, DatasetRef
48from ._dataset_type import DatasetType
49from ._deferredDatasetHandle import DeferredDatasetHandle
50from ._limited_butler import LimitedButler
51from ._quantum import Quantum
52from ._standalone_datastore import instantiate_standalone_datastore
53from ._storage_class import StorageClass, StorageClassFactory
54from .datastore import Datastore, DuplicateIdError
55from .datastore.record_data import DatastoreRecordData, SerializedDatastoreRecordData
56from .datastores.file_datastore.retrieve_artifacts import retrieve_and_zip
57from .dimensions import DimensionUniverse
58from .registry import ConflictingDefinitionError
59from .registry.interfaces import Database, DatastoreRegistryBridgeManager, OpaqueTableStorageManager
61if TYPE_CHECKING:
62 from ._butler import Butler
64_LOG = logging.getLogger(__name__)
67class QuantumBackedButler(LimitedButler):
68 """An implementation of `LimitedButler` intended to back execution of a
69 single `Quantum`.
71 Parameters
72 ----------
73 predicted_inputs : `~collections.abc.Iterable` [`DatasetId`]
74 Dataset IDs for datasets that can can be read from this butler.
75 predicted_outputs : `~collections.abc.Iterable` [`DatasetId`]
76 Dataset IDs for datasets that can be stored in this butler.
77 dimensions : `DimensionUniverse`
78 Object managing all dimension definitions.
79 datastore : `Datastore`
80 Datastore to use for all dataset I/O and existence checks.
81 storageClasses : `StorageClassFactory`
82 Object managing all storage class definitions.
83 dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`]
84 The registry dataset type definitions, indexed by name.
85 metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
86 Metrics object for tracking butler statistics.
87 database : `Database`, optional
88 Database instance used by datastore. Not required -- only provided
89 to allow database connections to be closed during cleanup.
91 Notes
92 -----
93 Most callers should use the `initialize` `classmethod` to construct new
94 instances instead of calling the constructor directly.
96 `QuantumBackedButler` uses a SQLite database internally, in order to reuse
97 existing `DatastoreRegistryBridge` and `OpaqueTableStorage`
98 implementations that rely SQLAlchemy. If implementations are added in the
99 future that don't rely on SQLAlchemy, it should be possible to swap them
100 in by overriding the type arguments to `initialize` (though at present,
101 `QuantumBackedButler` would still create at least an in-memory SQLite
102 database that would then go unused).`
104 We imagine `QuantumBackedButler` being used during (at least) batch
105 execution to capture `Datastore` records and save them to per-quantum
106 files, which are also a convenient place to store provenance for eventual
107 upload to a SQL-backed `Registry` (once `Registry` has tables to store
108 provenance, that is).
109 These per-quantum files can be written in two ways:
111 - The SQLite file used internally by `QuantumBackedButler` can be used
112 directly but customizing the ``filename`` argument to ``initialize``, and
113 then transferring that file to the object store after execution completes
114 (or fails; a ``try/finally`` pattern probably makes sense here).
116 - A JSON or YAML file can be written by calling `extract_provenance_data`,
117 and using ``pydantic`` methods to write the returned
118 `QuantumProvenanceData` to a file.
120 Note that at present, the SQLite file only contains datastore records, not
121 provenance, but that should be easy to address (if desired) after we
122 actually design a `Registry` schema for provenance. I also suspect that
123 we'll want to explicitly close the SQLite file somehow before trying to
124 transfer it. But I'm guessing we'd prefer to write the per-quantum files
125 as JSON anyway.
126 """
128 def __init__(
129 self,
130 predicted_inputs: Iterable[DatasetId],
131 predicted_outputs: Iterable[DatasetId],
132 dimensions: DimensionUniverse,
133 datastore: Datastore,
134 storageClasses: StorageClassFactory,
135 dataset_types: Mapping[str, DatasetType] | None = None,
136 metrics: ButlerMetrics | None = None,
137 database: Database | None = None,
138 ):
139 self._dimensions = dimensions
140 self._predicted_inputs = set(predicted_inputs)
141 self._predicted_outputs = set(predicted_outputs)
142 self._available_inputs: set[DatasetId] = set()
143 self._unavailable_inputs: set[DatasetId] = set()
144 self._actual_inputs: set[DatasetId] = set()
145 self._actual_output_refs: set[DatasetRef] = set()
146 self._datastore = datastore
147 self.storageClasses = storageClasses
148 self._dataset_types: Mapping[str, DatasetType] = {}
149 self._metrics = metrics if metrics is not None else ButlerMetrics()
150 self._database = database
151 if dataset_types is not None: 151 ↛ 153line 151 didn't jump to line 153 because the condition on line 151 was always true
152 self._dataset_types = dataset_types
153 self._datastore.set_retrieve_dataset_type_method(self._retrieve_dataset_type)
155 @classmethod
156 def initialize(
157 cls,
158 config: Config | ResourcePathExpression,
159 quantum: Quantum,
160 dimensions: DimensionUniverse,
161 filename: str | None = None,
162 OpaqueManagerClass: type[OpaqueTableStorageManager] | None = None,
163 BridgeManagerClass: type[DatastoreRegistryBridgeManager] | None = None,
164 search_paths: list[str] | None = None,
165 dataset_types: Mapping[str, DatasetType] | None = None,
166 metrics: ButlerMetrics | None = None,
167 ) -> QuantumBackedButler:
168 """Construct a new `QuantumBackedButler` from repository configuration
169 and helper types.
171 Parameters
172 ----------
173 config : `Config` or `~lsst.resources.ResourcePathExpression`
174 A butler repository root, configuration filename, or configuration
175 instance.
176 quantum : `Quantum`
177 Object describing the predicted input and output dataset relevant
178 to this butler. This must have resolved `DatasetRef` instances for
179 all inputs and outputs.
180 dimensions : `DimensionUniverse`
181 Object managing all dimension definitions.
182 filename : `str`, optional
183 Name for the SQLite database that will back this butler; defaults
184 to an in-memory database.
185 OpaqueManagerClass : `type`, optional
186 A subclass of `OpaqueTableStorageManager` to use for datastore
187 opaque records. Default is a SQL-backed implementation.
188 BridgeManagerClass : `type`, optional
189 A subclass of `DatastoreRegistryBridgeManager` to use for datastore
190 location records. Default is a SQL-backed implementation.
191 search_paths : `list` of `str`, optional
192 Additional search paths for butler configuration.
193 dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`], \
194 optional
195 Mapping of the dataset type name to its registry definition.
196 metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
197 Metrics object for gathering butler statistics.
198 """
199 predicted_inputs = [ref.id for ref in itertools.chain.from_iterable(quantum.inputs.values())]
200 predicted_inputs += [ref.id for ref in quantum.initInputs.values()]
201 predicted_outputs = [ref.id for ref in itertools.chain.from_iterable(quantum.outputs.values())]
202 return cls._initialize(
203 config=config,
204 predicted_inputs=predicted_inputs,
205 predicted_outputs=predicted_outputs,
206 dimensions=dimensions,
207 filename=filename,
208 datastore_records=quantum.datastore_records,
209 OpaqueManagerClass=OpaqueManagerClass,
210 BridgeManagerClass=BridgeManagerClass,
211 search_paths=search_paths,
212 dataset_types=dataset_types,
213 metrics=metrics,
214 )
216 @classmethod
217 def from_predicted(
218 cls,
219 config: Config | ResourcePathExpression,
220 predicted_inputs: Iterable[DatasetId],
221 predicted_outputs: Iterable[DatasetId],
222 dimensions: DimensionUniverse,
223 datastore_records: Mapping[str, DatastoreRecordData],
224 filename: str | None = None,
225 OpaqueManagerClass: type[OpaqueTableStorageManager] | None = None,
226 BridgeManagerClass: type[DatastoreRegistryBridgeManager] | None = None,
227 search_paths: list[str] | None = None,
228 dataset_types: Mapping[str, DatasetType] | None = None,
229 metrics: ButlerMetrics | None = None,
230 ) -> QuantumBackedButler:
231 """Construct a new `QuantumBackedButler` from sets of input and output
232 dataset IDs.
234 Parameters
235 ----------
236 config : `Config` or `~lsst.resources.ResourcePathExpression`
237 A butler repository root, configuration filename, or configuration
238 instance.
239 predicted_inputs : `~collections.abc.Iterable` [`DatasetId`]
240 Dataset IDs for datasets that can can be read from this butler.
241 predicted_outputs : `~collections.abc.Iterable` [`DatasetId`]
242 Dataset IDs for datasets that can be stored in this butler, must be
243 fully resolved.
244 dimensions : `DimensionUniverse`
245 Object managing all dimension definitions.
246 datastore_records : `dict` [`str`, `DatastoreRecordData`] or `None`
247 Datastore records to import into a datastore.
248 filename : `str`, optional
249 Name for the SQLite database that will back this butler; defaults
250 to an in-memory database.
251 OpaqueManagerClass : `type`, optional
252 A subclass of `OpaqueTableStorageManager` to use for datastore
253 opaque records. Default is a SQL-backed implementation.
254 BridgeManagerClass : `type`, optional
255 A subclass of `DatastoreRegistryBridgeManager` to use for datastore
256 location records. Default is a SQL-backed implementation.
257 search_paths : `list` of `str`, optional
258 Additional search paths for butler configuration.
259 dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`], \
260 optional
261 Mapping of the dataset type name to its registry definition.
262 metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
263 Metrics object for gathering butler statistics.
264 """
265 return cls._initialize(
266 config=config,
267 predicted_inputs=predicted_inputs,
268 predicted_outputs=predicted_outputs,
269 dimensions=dimensions,
270 filename=filename,
271 datastore_records=datastore_records,
272 OpaqueManagerClass=OpaqueManagerClass,
273 BridgeManagerClass=BridgeManagerClass,
274 search_paths=search_paths,
275 dataset_types=dataset_types,
276 metrics=metrics,
277 )
279 @classmethod
280 def _initialize(
281 cls,
282 *,
283 config: Config | ResourcePathExpression,
284 predicted_inputs: Iterable[DatasetId],
285 predicted_outputs: Iterable[DatasetId],
286 dimensions: DimensionUniverse,
287 filename: str | None = None,
288 datastore_records: Mapping[str, DatastoreRecordData] | None = None,
289 OpaqueManagerClass: type[OpaqueTableStorageManager] | None = None,
290 BridgeManagerClass: type[DatastoreRegistryBridgeManager] | None = None,
291 search_paths: list[str] | None = None,
292 dataset_types: Mapping[str, DatasetType] | None = None,
293 metrics: ButlerMetrics | None = None,
294 ) -> QuantumBackedButler:
295 """Initialize quantum-backed butler.
297 Internal method with common implementation used by `initialize` and
298 `for_output`.
300 Parameters
301 ----------
302 config : `Config` or `~lsst.resources.ResourcePathExpression`
303 A butler repository root, configuration filename, or configuration
304 instance.
305 predicted_inputs : `~collections.abc.Iterable` [`DatasetId`]
306 Dataset IDs for datasets that can can be read from this butler.
307 predicted_outputs : `~collections.abc.Iterable` [`DatasetId`]
308 Dataset IDs for datasets that can be stored in this butler.
309 dimensions : `DimensionUniverse`
310 Object managing all dimension definitions.
311 filename : `str`, optional
312 Name for the SQLite database that will back this butler; defaults
313 to an in-memory database.
314 datastore_records : `dict` [`str`, `DatastoreRecordData`] or `None`
315 Datastore records to import into a datastore.
316 OpaqueManagerClass : `type`, optional
317 A subclass of `OpaqueTableStorageManager` to use for datastore
318 opaque records. Default is a SQL-backed implementation.
319 BridgeManagerClass : `type`, optional
320 A subclass of `DatastoreRegistryBridgeManager` to use for datastore
321 location records. Default is a SQL-backed implementation.
322 search_paths : `list` of `str`, optional
323 Additional search paths for butler configuration.
324 dataset_types : `~collections.abc.Mapping` [`str`, `DatasetType`]
325 Mapping of the dataset type name to its registry definition.
326 metrics : `lsst.daf.butler.ButlerMetrics` or `None`, optional
327 Metrics object for gathering butler statistics.
328 """
329 butler_config = ButlerConfig(config, searchPaths=search_paths)
330 datastore, database = instantiate_standalone_datastore(
331 butler_config, dimensions, filename, OpaqueManagerClass, BridgeManagerClass
332 )
334 # TODO: We need to inform `Datastore` here that it needs to support
335 # predictive reads; This only really works for file datastore but
336 # we need to try everything in case there is a chained datastore.
337 datastore._set_trust_mode(True)
339 if datastore_records is not None: 339 ↛ 341line 339 didn't jump to line 341 because the condition on line 339 was always true
340 datastore.import_records(datastore_records)
341 storageClasses = StorageClassFactory()
342 storageClasses.addFromConfig(butler_config)
343 return cls(
344 predicted_inputs,
345 predicted_outputs,
346 dimensions,
347 datastore,
348 storageClasses=storageClasses,
349 dataset_types=dataset_types,
350 metrics=metrics,
351 database=database,
352 )
354 def close(self) -> None:
355 if self._database is not None: 355 ↛ exitline 355 didn't return from function 'close' because the condition on line 355 was always true
356 self._database.dispose()
358 def _retrieve_dataset_type(self, name: str) -> DatasetType | None:
359 """Return DatasetType defined in registry given dataset type name."""
360 return self._dataset_types.get(name)
362 def isWriteable(self) -> bool:
363 # Docstring inherited.
364 return True
366 def get(
367 self,
368 ref: DatasetRef,
369 /,
370 *,
371 parameters: dict[str, Any] | None = None,
372 storageClass: StorageClass | str | None = None,
373 ) -> Any:
374 try:
375 obj = super().get(
376 ref,
377 parameters=parameters,
378 storageClass=storageClass,
379 )
380 except (LookupError, FileNotFoundError, OSError):
381 self._unavailable_inputs.add(ref.id)
382 raise
383 if ref.id in self._predicted_inputs:
384 # do this after delegating to super in case that raises.
385 self._actual_inputs.add(ref.id)
386 self._available_inputs.add(ref.id)
387 return obj
389 def getDeferred(
390 self,
391 ref: DatasetRef,
392 /,
393 *,
394 parameters: dict[str, Any] | None = None,
395 storageClass: str | StorageClass | None = None,
396 ) -> DeferredDatasetHandle:
397 if ref.id in self._predicted_inputs:
398 # Unfortunately, we can't do this after the handle succeeds in
399 # loading, so it's conceivable here that we're marking an input
400 # as "actual" even when it's not even available.
401 self._actual_inputs.add(ref.id)
402 return super().getDeferred(ref, parameters=parameters, storageClass=storageClass)
404 def stored(self, ref: DatasetRef) -> bool:
405 # Docstring inherited.
406 stored = super().stored(ref)
407 if ref.id in self._predicted_inputs:
408 if stored: 408 ↛ 411line 408 didn't jump to line 411 because the condition on line 408 was always true
409 self._available_inputs.add(ref.id)
410 else:
411 self._unavailable_inputs.add(ref.id)
412 return stored
414 def stored_many(
415 self,
416 refs: Iterable[DatasetRef],
417 ) -> dict[DatasetRef, bool]:
418 # Docstring inherited.
419 existence = super().stored_many(refs)
421 for ref, stored in existence.items():
422 if ref.id in self._predicted_inputs:
423 if stored: 423 ↛ 426line 423 didn't jump to line 426 because the condition on line 423 was always true
424 self._available_inputs.add(ref.id)
425 else:
426 self._unavailable_inputs.add(ref.id)
427 return existence
429 def markInputUnused(self, ref: DatasetRef) -> None:
430 # Docstring inherited.
431 self._actual_inputs.discard(ref.id)
433 @property
434 def dimensions(self) -> DimensionUniverse:
435 # Docstring inherited.
436 return self._dimensions
438 def put(self, obj: Any, ref: DatasetRef, /, *, provenance: DatasetProvenance | None = None) -> DatasetRef:
439 # Docstring inherited.
440 if ref.id not in self._predicted_outputs: 440 ↛ 441line 440 didn't jump to line 441 because the condition on line 440 was never true
441 raise RuntimeError("Cannot `put` dataset that was not predicted as an output.")
442 with self._metrics.instrument_put(log=_LOG, msg="Put QBB dataset"):
443 try:
444 self._datastore.put(obj, ref, provenance=provenance)
445 except DuplicateIdError as e:
446 raise ConflictingDefinitionError(f"Datastore already contains dataset: {e}") from e
447 self._actual_output_refs.add(ref)
448 return ref
450 def pruneDatasets(
451 self,
452 refs: Iterable[DatasetRef],
453 *,
454 disassociate: bool = True,
455 unstore: bool = False,
456 tags: Iterable[str] = (),
457 purge: bool = False,
458 ) -> None:
459 # docstring inherited from LimitedButler
461 if purge:
462 if not disassociate:
463 raise TypeError("Cannot pass purge=True without disassociate=True.")
464 if not unstore:
465 raise TypeError("Cannot pass purge=True without unstore=True.")
466 elif disassociate:
467 # No tagged collections for this butler.
468 raise TypeError("Cannot pass disassociate=True without purge=True.")
470 refs = list(refs)
472 # Pruning a component of a DatasetRef makes no sense.
473 for ref in refs:
474 if ref.datasetType.component(): 474 ↛ 475line 474 didn't jump to line 475 because the condition on line 474 was never true
475 raise ValueError(f"Can not prune a component of a dataset (ref={ref})")
477 if unstore: 477 ↛ 479line 477 didn't jump to line 479 because the condition on line 477 was always true
478 self._datastore.trash(refs)
479 if purge:
480 for ref in refs:
481 # We only care about removing them from actual output refs,
482 self._actual_output_refs.discard(ref)
484 if unstore: 484 ↛ exitline 484 didn't return from function 'pruneDatasets' because the condition on line 484 was always true
485 # Point of no return for removing artifacts. Only try to remove
486 # refs associated with this pruning.
487 self._datastore.emptyTrash(refs=refs)
489 def retrieve_artifacts_zip(
490 self,
491 refs: Iterable[DatasetRef],
492 destination: ResourcePathExpression,
493 overwrite: bool = True,
494 ) -> ResourcePath:
495 """Retrieve artifacts from the graph and place in ZIP file.
497 Parameters
498 ----------
499 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
500 The datasets to be included in the zip file.
501 destination : `lsst.resources.ResourcePathExpression`
502 Directory to write the new ZIP file. This directory will
503 also be used as a staging area for the datasets being downloaded
504 from the datastore.
505 overwrite : `bool`, optional
506 If `False` the output Zip will not be written if a file of the
507 same name is already present in ``destination``.
509 Returns
510 -------
511 zip_file : `lsst.resources.ResourcePath`
512 The path to the new ZIP file.
514 Raises
515 ------
516 ValueError
517 Raised if there are no refs to retrieve.
518 """
519 return retrieve_and_zip(refs, destination, self._datastore.retrieveArtifacts, overwrite)
521 def retrieve_artifacts(
522 self,
523 refs: Iterable[DatasetRef],
524 destination: ResourcePathExpression,
525 transfer: str = "auto",
526 preserve_path: bool = True,
527 overwrite: bool = False,
528 ) -> list[ResourcePath]:
529 """Retrieve the artifacts associated with the supplied refs.
531 Parameters
532 ----------
533 refs : `~collections.abc.Iterable` of `DatasetRef`
534 The datasets for which artifacts are to be retrieved.
535 A single ref can result in multiple artifacts. The refs must
536 be resolved.
537 destination : `lsst.resources.ResourcePath` or `str`
538 Location to write the artifacts.
539 transfer : `str`, optional
540 Method to use to transfer the artifacts. Must be one of the options
541 supported by `~lsst.resources.ResourcePath.transfer_from`.
542 "move" is not allowed.
543 preserve_path : `bool`, optional
544 If `True` the full path of the artifact within the datastore
545 is preserved. If `False` the final file component of the path
546 is used.
547 overwrite : `bool`, optional
548 If `True` allow transfers to overwrite existing files at the
549 destination.
551 Returns
552 -------
553 targets : `list` of `lsst.resources.ResourcePath`
554 URIs of file artifacts in destination location. Order is not
555 preserved.
556 """
557 outdir = ResourcePath(destination)
558 artifact_map = self._datastore.retrieveArtifacts(
559 refs,
560 outdir,
561 transfer=transfer,
562 preserve_path=preserve_path,
563 overwrite=overwrite,
564 write_index=True,
565 )
566 return list(artifact_map)
568 def extract_provenance_data(self) -> QuantumProvenanceData:
569 """Extract provenance information and datastore records from this
570 butler.
572 Returns
573 -------
574 provenance : `QuantumProvenanceData`
575 A serializable struct containing input/output dataset IDs and
576 datastore records. This assumes all dataset IDs are UUIDs (just to
577 make it easier for `pydantic` to reason about the struct's types);
578 the rest of this class makes no such assumption, but the approach
579 to processing in which it's useful effectively requires UUIDs
580 anyway.
582 Notes
583 -----
584 `QuantumBackedButler` records this provenance information when its
585 methods are used, which mostly saves `~lsst.pipe.base.PipelineTask`
586 authors from having to worry about while still recording very
587 detailed information. But it has two small weaknesses:
589 - Calling `getDeferred` or `get` is enough to mark a
590 dataset as an "actual input", which may mark some datasets that
591 aren't actually used. We rely on task authors to use
592 `markInputUnused` to address this.
594 - We assume that the execution system will call ``stored``
595 on all predicted inputs prior to execution, in order to populate the
596 "available inputs" set. This is what I envision
597 '`~lsst.ctrl.mpexec.SingleQuantumExecutor` doing after we update it
598 to use this class, but it feels fragile for this class to make such
599 a strong assumption about how it will be used, even if I can't think
600 of any other executor behavior that would make sense.
601 """
602 if not self._actual_inputs.isdisjoint(self._unavailable_inputs): 602 ↛ 603line 602 didn't jump to line 603 because the condition on line 602 was never true
603 _LOG.warning(
604 "Inputs %s were marked as actually used (probably because a DeferredDatasetHandle) "
605 "was obtained, but did not actually exist. This task should be be using markInputUnused "
606 "directly to clarify its provenance.",
607 self._actual_inputs & self._unavailable_inputs,
608 )
609 self._actual_inputs -= self._unavailable_inputs
610 checked_inputs = self._available_inputs | self._unavailable_inputs
611 if self._predicted_inputs != checked_inputs: 611 ↛ 612line 611 didn't jump to line 612 because the condition on line 611 was never true
612 _LOG.warning(
613 "Execution harness did not check predicted inputs %s for existence; available inputs "
614 "recorded in provenance may be incomplete.",
615 self._predicted_inputs - checked_inputs,
616 )
617 datastore_records = self._datastore.export_records(self._actual_output_refs)
618 provenance_records = {
619 datastore_name: records.to_simple() for datastore_name, records in datastore_records.items()
620 }
622 return QuantumProvenanceData(
623 predicted_inputs=self._predicted_inputs,
624 available_inputs=self._available_inputs,
625 actual_inputs=self._actual_inputs,
626 predicted_outputs=self._predicted_outputs,
627 actual_outputs={ref.id for ref in self._actual_output_refs},
628 datastore_records=provenance_records,
629 )
631 def export_predicted_datastore_records(
632 self, refs: Iterable[DatasetRef]
633 ) -> dict[str, DatastoreRecordData]:
634 """Export datastore records for a set of predicted output dataset
635 references.
637 Parameters
638 ----------
639 refs : `~collections.abc.Iterable` [ `DatasetRef` ]
640 Dataset references for which to export datastore records. These
641 refs must be known to this butler.
643 Returns
644 -------
645 records : `dict` [ `str`, `DatastoreRecordData` ]
646 Predicted datastore records indexed by datastore name. No attempt
647 is made to ensure that the associated datasets exist on disk.
648 """
649 unknowns = [
650 ref
651 for ref in refs
652 if (ref.id not in self._predicted_outputs and ref.id not in self._predicted_inputs)
653 ]
654 if unknowns:
655 raise ValueError(f"Cannot export datastore records for unknown outputs: {unknowns}")
657 return self._datastore.export_predicted_records(refs)
660class QuantumProvenanceData(pydantic.BaseModel):
661 """A serializable struct for per-quantum provenance information and
662 datastore records.
664 Notes
665 -----
666 This class slightly duplicates information from the `Quantum` class itself
667 (the ``predicted_inputs`` and ``predicted_outputs`` sets should have the
668 same IDs present in `Quantum.inputs` and `Quantum.outputs`), but overall it
669 assumes the original `Quantum` is also available to reconstruct the
670 complete provenance (e.g. by associating dataset IDs with data IDs,
671 dataset types, and `~CollectionType.RUN` names.
673 Note that ``pydantic`` method ``parse_raw()`` is not going to work
674 correctly for this class, use `direct` method instead.
675 """
677 # This class probably should have information about its execution
678 # environment (anything not controlled and recorded at the
679 # `~CollectionType.RUN` level, such as the compute node ID). but adding it
680 # now is out of scope for this prototype.
682 predicted_inputs: set[uuid.UUID]
683 """Unique IDs of datasets that were predicted as inputs to this quantum
684 when the QuantumGraph was built.
685 """
687 available_inputs: set[uuid.UUID]
688 """Unique IDs of input datasets that were actually present in the datastore
689 when this quantum was executed.
691 This is a subset of ``predicted_inputs``, with the difference generally
692 being datasets were ``predicted_outputs`` but not ``actual_outputs`` of
693 some upstream task.
694 """
696 actual_inputs: set[uuid.UUID]
697 """Unique IDs of datasets that were actually used as inputs by this task.
699 This is a subset of ``available_inputs``.
701 Notes
702 -----
703 The criteria for marking an input as used is that rerunning the quantum
704 with only these ``actual_inputs`` available must yield identical outputs.
705 This means that (for example) even just using an input to help determine
706 an output rejection criteria and then rejecting it as an outlier qualifies
707 that input as actually used.
708 """
710 predicted_outputs: set[uuid.UUID]
711 """Unique IDs of datasets that were predicted as outputs of this quantum
712 when the QuantumGraph was built.
713 """
715 actual_outputs: set[uuid.UUID]
716 """Unique IDs of datasets that were actually written when this quantum
717 was executed.
718 """
720 datastore_records: dict[str, SerializedDatastoreRecordData]
721 """Datastore records indexed by datastore name."""
723 @staticmethod
724 def collect_and_transfer(
725 butler: Butler, quanta: Iterable[Quantum], provenance: Iterable[QuantumProvenanceData]
726 ) -> None:
727 """Transfer output datasets from multiple quanta to a more permanent
728 `Butler` repository.
730 Parameters
731 ----------
732 butler : `Butler`
733 Full butler representing the data repository to transfer datasets
734 to.
735 quanta : `~collections.abc.Iterable` [ `Quantum` ]
736 Iterable of `Quantum` objects that carry information about
737 predicted outputs. May be a single-pass iterator.
738 provenance : `~collections.abc.Iterable` [ `QuantumProvenanceData` ]
739 Provenance and datastore data for each of the given quanta, in the
740 same order. May be a single-pass iterator.
742 Notes
743 -----
744 Input-output provenance data is not actually transferred yet, because
745 `Registry` has no place to store it.
747 This method probably works most efficiently if run on all quanta for a
748 single task label at once, because this will gather all datasets of
749 a particular type together into a single vectorized `Registry` import.
750 It should still behave correctly if run on smaller groups of quanta
751 or even quanta from multiple tasks.
753 Currently this method transfers datastore record data unchanged, with
754 no possibility of actually moving (e.g.) files. Datastores that are
755 present only in execution or only in the more permanent butler are
756 ignored.
757 """
758 grouped_refs = defaultdict(list)
759 summary_records: dict[str, DatastoreRecordData] = {}
760 for quantum, provenance_for_quantum in zip(quanta, provenance, strict=True):
761 quantum_refs_by_id = {
762 ref.id: ref
763 for ref in itertools.chain.from_iterable(quantum.outputs.values())
764 if ref.id in provenance_for_quantum.actual_outputs
765 }
766 for ref in quantum_refs_by_id.values():
767 grouped_refs[ref.datasetType, ref.run].append(ref)
769 # merge datastore records into a summary structure
770 for datastore_name, serialized_records in provenance_for_quantum.datastore_records.items():
771 quantum_records = DatastoreRecordData.from_simple(serialized_records)
772 if (records := summary_records.get(datastore_name)) is not None:
773 records.update(quantum_records)
774 else:
775 summary_records[datastore_name] = quantum_records
777 for refs in grouped_refs.values():
778 butler.registry._importDatasets(refs)
779 butler._datastore.import_records(summary_records)
781 @classmethod
782 def parse_raw(cls, *args: Any, **kwargs: Any) -> QuantumProvenanceData:
783 raise NotImplementedError("parse_raw() is not usable for this class, use direct() instead.")
785 @classmethod
786 def direct(
787 cls,
788 *,
789 predicted_inputs: Iterable[str | uuid.UUID],
790 available_inputs: Iterable[str | uuid.UUID],
791 actual_inputs: Iterable[str | uuid.UUID],
792 predicted_outputs: Iterable[str | uuid.UUID],
793 actual_outputs: Iterable[str | uuid.UUID],
794 datastore_records: Mapping[str, Mapping],
795 ) -> QuantumProvenanceData:
796 """Construct an instance directly without validators.
798 Parameters
799 ----------
800 predicted_inputs : `~collections.abc.Iterable` of `str` or `uuid.UUID`
801 The predicted inputs.
802 available_inputs : `~collections.abc.Iterable` of `str` or `uuid.UUID`
803 The available inputs.
804 actual_inputs : `~collections.abc.Iterable` of `str` or `uuid.UUID`
805 The actual inputs.
806 predicted_outputs : `~collections.abc.Iterable` of `str` or `uuid.UUID`
807 The predicted outputs.
808 actual_outputs : `~collections.abc.Iterable` of `str` or `uuid.UUID`
809 The actual outputs.
810 datastore_records : `~collections.abc.Mapping` [ `str`, \
811 `~collections.abc.Mapping` ]
812 The datastore records.
814 Returns
815 -------
816 provenance : `QuantumProvenanceData`
817 Serializable model of the quantum provenance.
819 Notes
820 -----
821 This differs from the Pydantic "construct" method in that the
822 arguments are explicitly what the model requires, and it will recurse
823 through members, constructing them from their corresponding `direct`
824 methods.
826 This method should only be called when the inputs are trusted.
827 """
829 def _to_uuid_set(uuids: Iterable[str | uuid.UUID]) -> set[uuid.UUID]:
830 """Convert input UUIDs, which could be in string representation to
831 a set of `UUID` instances.
832 """
833 return {uuid.UUID(id) if isinstance(id, str) else id for id in uuids}
835 data = cls.model_construct(
836 predicted_inputs=_to_uuid_set(predicted_inputs),
837 available_inputs=_to_uuid_set(available_inputs),
838 actual_inputs=_to_uuid_set(actual_inputs),
839 predicted_outputs=_to_uuid_set(predicted_outputs),
840 actual_outputs=_to_uuid_set(actual_outputs),
841 datastore_records={
842 key: SerializedDatastoreRecordData.direct(**records)
843 for key, records in datastore_records.items()
844 },
845 )
847 return data