Add new QuantumGraph.get_refs method

python/lsst/pipe/base/graph/graph.py

@@ -32,6 +32,7 @@
 import getpass
 import io
 import json
+import logging
 import lzma
 import os
 import struct
@@ -75,6 +76,7 @@
 from .quantumNode import BuildId, QuantumNode
 
 _T = TypeVar("_T", bound="QuantumGraph")
+_LOG = logging.getLogger(__name__)
 
 # modify this constant any time the on disk representation of the save file
 # changes, and update the load helpers to behave properly for each version.
@@ -1656,3 +1658,88 @@ def init_output_run(self, butler: LimitedButler, existing: bool = True) -> None:
         self.write_configs(butler, compare_existing=existing)
         self.write_packages(butler, compare_existing=existing)
         self.write_init_outputs(butler, skip_existing=existing)
+
+    def get_refs(
+        self,
+        *,
+        include_inputs: bool = False,
+        include_init_inputs: bool = False,
+        include_init_outputs: bool = False,
+        include_outputs: bool = False,
+        conform_outputs: bool = True,
+    ) -> tuple[set[DatasetRef], dict[str, DatastoreRecordData]]:
+        """Get the requested dataset refs from the graph.
+
+        Parameters
+        ----------
+        include_inputs : `bool`, optional
+            Include inputs.
+        include_init_inputs : `bool`, optional
+            Include init inputs.
+        include_init_outputs : `bool`, optional
+            Include init outpus.
+        include_outputs : `bool`, optional
+            Include outputs.
+        conform_outputs : `bool`, optional
+            Whether any outputs found should have their dataset types conformed
+            with the registry dataset types.
+
+        Returns
+        -------
+        refs : `set` [ `lsst.daf.butler.DatasetRef` ]
+            The requested dataset refs found in the graph.
+        datastore_records : `dict` [ `str`, \
+                `lsst.daf.butler.datastore.record_data.DatastoreRecordData` ]
+            Any datastore records found.
+        """
+        datastore_records: dict[str, DatastoreRecordData] = {}
+        init_input_refs: set[DatasetRef] = set()
+        init_output_refs: set[DatasetRef] = set()
+
+        if include_init_inputs or include_init_outputs:
+            for task_def in self.iterTaskGraph():
+                if include_init_inputs:
+                    if in_refs := self.initInputRefs(task_def):
+                        init_input_refs.update(in_refs)
+                if include_init_outputs:
+                    if out_refs := self.initOutputRefs(task_def):
+                        init_output_refs.update(out_refs)
+
+        input_refs: set[DatasetRef] = set()
+        output_refs: set[DatasetRef] = set()
+
+        for qnode in self:
+            if include_inputs:
+                for other_refs in qnode.quantum.inputs.values():
+                    input_refs.update(other_refs)
+                # Inputs can come with datastore records.
+                for store_name, records in qnode.quantum.datastore_records.items():
+                    datastore_records.setdefault(store_name, DatastoreRecordData()).update(records)
+            if include_outputs:
+                for other_refs in qnode.quantum.outputs.values():
+                    output_refs.update(other_refs)
+
+        if conform_outputs:
+            # Get data repository definitions from the QuantumGraph; these can
+            # have different storage classes than those in the quanta.
+            dataset_types = {dstype.name: dstype for dstype in self.registryDatasetTypes()}
+
+            def _update_ref(ref: DatasetRef) -> DatasetRef:
+                internal_dataset_type = dataset_types.get(ref.datasetType.name, ref.datasetType)
+                if internal_dataset_type.storageClass_name != ref.datasetType.storageClass_name:
+                    ref = ref.overrideStorageClass(internal_dataset_type.storageClass_name)
+                return ref
+
+            # Convert output_refs to the data repository storage classes, too.
+            output_refs = {_update_ref(ref) for ref in output_refs}
+            init_output_refs = {_update_ref(ref) for ref in init_output_refs}
+
+        _LOG.info(
+            "Found the following datasets. InitInputs: %d; Inputs: %d; InitOutputs: %s; Outputs: %d",
+            len(init_input_refs),
+            len(input_refs),
+            len(init_output_refs),
+            len(output_refs),
+        )
+        refs = input_refs | init_input_refs | init_output_refs | output_refs
+        return refs, datastore_records

python/lsst/pipe/base/script/retrieve_artifacts_for_quanta.py

@@ -29,8 +29,7 @@
 
 import logging
 
-from lsst.daf.butler import DatasetRef, QuantumBackedButler
-from lsst.daf.butler.datastore.record_data import DatastoreRecordData
+from lsst.daf.butler import QuantumBackedButler
 from lsst.pipe.base import QuantumGraph
 from lsst.resources import ResourcePath
 
@@ -81,48 +80,18 @@ def retrieve_artifacts_for_quanta(
     nodes = qgraph_node_id or None
     qgraph = QuantumGraph.loadUri(graph, nodes=nodes)
 
+    refs, datastore_records = qgraph.get_refs(
+        include_inputs=include_inputs,
+        include_init_inputs=include_inputs,
+        include_outputs=include_outputs,
+        include_init_outputs=include_outputs,
+        conform_outputs=True,  # Need to look for predicted outputs with correct storage class.
+    )
+
     # Get data repository definitions from the QuantumGraph; these can have
     # different storage classes than those in the quanta.
     dataset_types = {dstype.name: dstype for dstype in qgraph.registryDatasetTypes()}
 
-    datastore_records: dict[str, DatastoreRecordData] = {}
-    refs: set[DatasetRef] = set()
-    if include_inputs:
-        # Collect input refs used by this graph.
-        for task_def in qgraph.iterTaskGraph():
-            if in_refs := qgraph.initInputRefs(task_def):
-                refs.update(in_refs)
-        for qnode in qgraph:
-            for otherRefs in qnode.quantum.inputs.values():
-                refs.update(otherRefs)
-            for store_name, records in qnode.quantum.datastore_records.items():
-                datastore_records.setdefault(store_name, DatastoreRecordData()).update(records)
-    n_inputs = len(refs)
-    if n_inputs:
-        _LOG.info("Found %d input dataset%s.", n_inputs, "" if n_inputs == 1 else "s")
-
-    if include_outputs:
-        # Collect output refs that could be created by this graph.
-        original_output_refs: set[DatasetRef] = set(qgraph.globalInitOutputRefs())
-        for task_def in qgraph.iterTaskGraph():
-            if out_refs := qgraph.initOutputRefs(task_def):
-                original_output_refs.update(out_refs)
-        for qnode in qgraph:
-            for otherRefs in qnode.quantum.outputs.values():
-                original_output_refs.update(otherRefs)
-
-        # Convert output_refs to the data repository storage classes, too.
-        for ref in original_output_refs:
-            internal_dataset_type = dataset_types.get(ref.datasetType.name, ref.datasetType)
-            if internal_dataset_type.storageClass_name != ref.datasetType.storageClass_name:
-                refs.add(ref.overrideStorageClass(internal_dataset_type.storageClass_name))
-            else:
-                refs.add(ref)
-
-    n_outputs = len(refs) - n_inputs
-    if n_outputs:
-        _LOG.info("Found %d output dataset%s.", n_outputs, "" if n_outputs == 1 else "s")
-
     # Make QBB, its config is the same as output Butler.
     qbb = QuantumBackedButler.from_predicted(
         config=repo,