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