Coverage for python/lsst/pipe/base/connections.py: 39%
382 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-10 01:31 -0700
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"""Module defining connection classes for PipelineTask."""
30from __future__ import annotations
32__all__ = [
33 "AdjustQuantumHelper",
34 "DeferredDatasetRef",
35 "InputQuantizedConnection",
36 "OutputQuantizedConnection",
37 "PipelineTaskConnections",
38 "QuantaAdjuster",
39 "QuantizedConnection",
40 "ScalarError",
41 "ScalarError",
42 "iterConnections",
43]
45import dataclasses
46import itertools
47import string
48import uuid
49import warnings
50from collections import UserDict, defaultdict
51from collections.abc import Collection, Generator, Iterable, Iterator, Mapping, Sequence, Set
52from dataclasses import dataclass
53from types import MappingProxyType, SimpleNamespace
54from typing import TYPE_CHECKING, Any
56from lsst.daf.butler import (
57 Butler,
58 DataCoordinate,
59 DatasetRef,
60 DatasetType,
61 NamedKeyDict,
62 NamedKeyMapping,
63 Quantum,
64)
66from ._status import NoWorkFound
67from .connectionTypes import BaseConnection, BaseInput, Output, PrerequisiteInput
69if TYPE_CHECKING:
70 from .config import PipelineTaskConfig
71 from .pipeline_graph import PipelineGraph, TaskNode
72 from .quantum_graph_skeleton import QuantumGraphSkeleton
75class ScalarError(TypeError):
76 """Exception raised when dataset type is configured as scalar
77 but there are multiple data IDs in a Quantum for that dataset.
78 """
81class PipelineTaskConnectionDict(UserDict):
82 """A special dict class used by `PipelineTaskConnectionMetaclass`.
84 This dict is used in `PipelineTaskConnection` class creation, as the
85 dictionary that is initially used as ``__dict__``. It exists to
86 intercept connection fields declared in a `PipelineTaskConnection`, and
87 what name is used to identify them. The names are then added to class
88 level list according to the connection type of the class attribute. The
89 names are also used as keys in a class level dictionary associated with
90 the corresponding class attribute. This information is a duplicate of
91 what exists in ``__dict__``, but provides a simple place to lookup and
92 iterate on only these variables.
94 Parameters
95 ----------
96 *args : `~typing.Any`
97 Passed to `dict` constructor.
98 **kwargs : `~typing.Any`
99 Passed to `dict` constructor.
100 """
102 def __init__(self, *args: Any, **kwargs: Any):
103 super().__init__(*args, **kwargs)
104 # Initialize class level variables used to track any declared
105 # class level variables that are instances of
106 # connectionTypes.BaseConnection
107 self.data["inputs"] = set()
108 self.data["prerequisiteInputs"] = set()
109 self.data["outputs"] = set()
110 self.data["initInputs"] = set()
111 self.data["initOutputs"] = set()
112 self.data["allConnections"] = {}
114 def __setitem__(self, name: str, value: Any) -> None:
115 if isinstance(value, BaseConnection):
116 if name in { 116 ↛ 126line 116 didn't jump to line 126 because the condition on line 116 was never true
117 "dimensions",
118 "inputs",
119 "prerequisiteInputs",
120 "outputs",
121 "initInputs",
122 "initOutputs",
123 "allConnections",
124 }:
125 # Guard against connections whose names are reserved.
126 raise AttributeError(f"Connection name {name!r} is reserved for internal use.")
127 if (previous := self.data.get(name)) is not None: 127 ↛ 130line 127 didn't jump to line 130 because the condition on line 127 was never true
128 # Guard against changing the type of an in inherited connection
129 # by first removing it from the set it's current in.
130 self.data[previous._connection_type_set].discard(name)
131 object.__setattr__(value, "varName", name)
132 self.data["allConnections"][name] = value
133 self.data[value._connection_type_set].add(name)
134 # defer to the default behavior
135 super().__setitem__(name, value)
138class PipelineTaskConnectionsMetaclass(type):
139 """Metaclass used in the declaration of PipelineTaskConnections classes.
141 Parameters
142 ----------
143 name : `str`
144 Name of connection.
145 bases : `~collections.abc.Collection`
146 Base classes.
147 dct : `~collections.abc.Mapping`
148 Connections dict.
149 **kwargs : `~typing.Any`
150 Additional parameters.
151 """
153 # We can annotate these attributes as `collections.abc.Set` to discourage
154 # undesirable modifications in type-checked code, since the internal code
155 # modifying them is in `PipelineTaskConnectionDict` and that doesn't see
156 # these annotations anyway.
158 dimensions: Set[str]
159 """Set of dimension names that define the unit of work for this task.
161 Required and implied dependencies will automatically be expanded later and
162 need not be provided.
164 This is shadowed by an instance-level attribute on
165 `PipelineTaskConnections` instances.
166 """
168 inputs: Set[str]
169 """Set with the names of all `~connectionTypes.Input` connection
170 attributes.
172 This is updated automatically as class attributes are added. Note that
173 this attribute is shadowed by an instance-level attribute on
174 `PipelineTaskConnections` instances.
175 """
177 prerequisiteInputs: Set[str]
178 """Set with the names of all `~connectionTypes.PrerequisiteInput`
179 connection attributes.
181 See `inputs` for additional information.
182 """
184 outputs: Set[str]
185 """Set with the names of all `~connectionTypes.Output` connection
186 attributes.
188 See `inputs` for additional information.
189 """
191 initInputs: Set[str]
192 """Set with the names of all `~connectionTypes.InitInput` connection
193 attributes.
195 See `inputs` for additional information.
196 """
198 initOutputs: Set[str]
199 """Set with the names of all `~connectionTypes.InitOutput` connection
200 attributes.
202 See `inputs` for additional information.
203 """
205 allConnections: Mapping[str, BaseConnection]
206 """Mapping containing all connection attributes.
208 See `inputs` for additional information.
209 """
211 def __prepare__(name, bases, **kwargs): # noqa: N804
212 # Create an instance of our special dict to catch and track all
213 # variables that are instances of connectionTypes.BaseConnection
214 # Copy any existing connections from a parent class
215 dct = PipelineTaskConnectionDict()
216 for base in bases:
217 if isinstance(base, PipelineTaskConnectionsMetaclass): 217 ↛ 216line 217 didn't jump to line 216 because the condition on line 217 was always true
218 for name, value in base.allConnections.items(): 218 ↛ 219line 218 didn't jump to line 219 because the loop on line 218 never started
219 dct[name] = value
220 return dct
222 def __new__(cls, name, bases, dct, **kwargs):
223 dimensionsValueError = TypeError(
224 "PipelineTaskConnections class must be created with a dimensions "
225 "attribute which is an iterable of dimension names"
226 )
228 if name != "PipelineTaskConnections":
229 # Verify that dimensions are passed as a keyword in class
230 # declaration
231 if "dimensions" not in kwargs: 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true
232 for base in bases:
233 if hasattr(base, "dimensions"):
234 kwargs["dimensions"] = base.dimensions
235 break
236 if "dimensions" not in kwargs:
237 raise dimensionsValueError
238 try:
239 if isinstance(kwargs["dimensions"], str): 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true
240 raise TypeError(
241 "Dimensions must be iterable of dimensions, got str,possibly omitted trailing comma"
242 )
243 if not isinstance(kwargs["dimensions"], Iterable): 243 ↛ 244line 243 didn't jump to line 244 because the condition on line 243 was never true
244 raise TypeError("Dimensions must be iterable of dimensions")
245 dct["dimensions"] = set(kwargs["dimensions"])
246 except TypeError as exc:
247 raise dimensionsValueError from exc
248 # Lookup any python string templates that may have been used in the
249 # declaration of the name field of a class connection attribute
250 allTemplates = set()
251 stringFormatter = string.Formatter()
252 # Loop over all connections
253 for obj in dct["allConnections"].values():
254 nameValue = obj.name
255 # add all the parameters to the set of templates
256 for param in stringFormatter.parse(nameValue):
257 if param[1] is not None:
258 allTemplates.add(param[1])
260 # look up any template from base classes and merge them all
261 # together
262 mergeDict = {}
263 mergeDeprecationsDict = {}
264 for base in bases[::-1]:
265 if hasattr(base, "defaultTemplates"):
266 mergeDict.update(base.defaultTemplates)
267 if hasattr(base, "deprecatedTemplates"):
268 mergeDeprecationsDict.update(base.deprecatedTemplates)
269 if "defaultTemplates" in kwargs:
270 mergeDict.update(kwargs["defaultTemplates"])
271 if "deprecatedTemplates" in kwargs: 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true
272 mergeDeprecationsDict.update(kwargs["deprecatedTemplates"])
273 if len(mergeDict) > 0:
274 kwargs["defaultTemplates"] = mergeDict
275 if len(mergeDeprecationsDict) > 0: 275 ↛ 276line 275 didn't jump to line 276 because the condition on line 275 was never true
276 kwargs["deprecatedTemplates"] = mergeDeprecationsDict
278 # Verify that if templated strings were used, defaults were
279 # supplied as an argument in the declaration of the connection
280 # class
281 if len(allTemplates) > 0 and "defaultTemplates" not in kwargs: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 raise TypeError(
283 "PipelineTaskConnection class contains templated attribute names, but no "
284 "defaut templates were provided, add a dictionary attribute named "
285 "defaultTemplates which contains the mapping between template key and value"
286 )
287 if len(allTemplates) > 0:
288 # Verify all templates have a default, and throw if they do not
289 defaultTemplateKeys = set(kwargs["defaultTemplates"].keys())
290 templateDifference = allTemplates.difference(defaultTemplateKeys)
291 if templateDifference: 291 ↛ 292line 291 didn't jump to line 292 because the condition on line 291 was never true
292 raise TypeError(f"Default template keys were not provided for {templateDifference}")
293 # Verify that templates do not share names with variable names
294 # used for a connection, this is needed because of how
295 # templates are specified in an associated config class.
296 nameTemplateIntersection = allTemplates.intersection(set(dct["allConnections"].keys()))
297 if len(nameTemplateIntersection) > 0: 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true
298 raise TypeError(
299 "Template parameters cannot share names with Class attributes"
300 f" (conflicts are {nameTemplateIntersection})."
301 )
302 dct["defaultTemplates"] = kwargs.get("defaultTemplates", {})
303 dct["deprecatedTemplates"] = kwargs.get("deprecatedTemplates", {})
305 # Convert all the connection containers into frozensets so they cannot
306 # be modified at the class scope
307 for connectionName in ("inputs", "prerequisiteInputs", "outputs", "initInputs", "initOutputs"):
308 dct[connectionName] = frozenset(dct[connectionName])
309 # our custom dict type must be turned into an actual dict to be used in
310 # type.__new__
311 return super().__new__(cls, name, bases, dict(dct))
313 def __init__(cls, name, bases, dct, **kwargs):
314 # This overrides the default init to drop the kwargs argument. Python
315 # metaclasses will have this argument set if any kwargs are passes at
316 # class construction time, but should be consumed before calling
317 # __init__ on the type metaclass. This is in accordance with python
318 # documentation on metaclasses
319 super().__init__(name, bases, dct)
321 def __call__(cls, *, config: PipelineTaskConfig | None = None) -> PipelineTaskConnections:
322 # MyPy appears not to really understand metaclass.__call__ at all, so
323 # we need to tell it to ignore __new__ and __init__ calls here.
324 instance: PipelineTaskConnections = cls.__new__(cls) # type: ignore
326 # Make mutable copies of all set-like class attributes so derived
327 # __init__ implementations can modify them in place.
328 instance.dimensions = set(cls.dimensions)
329 instance.inputs = set(cls.inputs)
330 instance.prerequisiteInputs = set(cls.prerequisiteInputs)
331 instance.outputs = set(cls.outputs)
332 instance.initInputs = set(cls.initInputs)
333 instance.initOutputs = set(cls.initOutputs)
335 # Set self.config. It's a bit strange that we claim to accept None but
336 # really just raise here, but it's not worth changing now.
337 from .config import PipelineTaskConfig # local import to avoid cycle
339 if config is None or not isinstance(config, PipelineTaskConfig):
340 raise ValueError(
341 "PipelineTaskConnections must be instantiated with a PipelineTaskConfig instance"
342 )
343 instance.config = config
345 # Extract the template names that were defined in the config instance
346 # by looping over the keys of the defaultTemplates dict specified at
347 # class declaration time.
348 templateValues = {
349 name: getattr(config.connections, name)
350 for name in cls.defaultTemplates # type: ignore
351 }
353 # We now assemble a mapping of all connection instances keyed by
354 # internal name, applying the configuration and templates to make new
355 # configurations from the class-attribute defaults. This will be
356 # private, but with a public read-only view. This mapping is what the
357 # descriptor interface of the class-level attributes will return when
358 # they are accessed on an instance. This is better than just assigning
359 # regular instance attributes as it makes it so removed connections
360 # cannot be accessed on instances, instead of having access to them
361 # silent fall through to the not-removed class connection instance.
362 instance._allConnections = {}
363 instance.allConnections = MappingProxyType(instance._allConnections)
364 for internal_name, connection in cls.allConnections.items():
365 dataset_type_name = getattr(config.connections, internal_name).format(**templateValues)
366 instance_connection = dataclasses.replace(
367 connection,
368 name=dataset_type_name,
369 doc=(
370 connection.doc
371 if connection.deprecated is None
372 else f"{connection.doc}\n{connection.deprecated}"
373 ),
374 _deprecation_context=connection._deprecation_context,
375 )
376 instance._allConnections[internal_name] = instance_connection
378 # Finally call __init__. The base class implementation does nothing;
379 # we could have left some of the above implementation there (where it
380 # originated), but putting it here instead makes it hard for derived
381 # class implementors to get things into a weird state by delegating to
382 # super().__init__ in the wrong place, or by forgetting to do that
383 # entirely.
384 instance.__init__(config=config) # type: ignore
386 # Derived-class implementations may have changed the contents of the
387 # various kinds-of-connection sets; update allConnections to have keys
388 # that are a union of all those. We get values for the new
389 # allConnections from the attributes, since any dynamically added new
390 # ones will not be present in the old allConnections. Typically those
391 # getattrs will invoke the descriptors and get things from the old
392 # allConnections anyway. After processing each set we replace it with
393 # a frozenset.
394 updated_all_connections = {}
395 for attrName in ("initInputs", "prerequisiteInputs", "inputs", "initOutputs", "outputs"):
396 updated_connection_names = getattr(instance, attrName)
397 updated_all_connections.update(
398 {name: getattr(instance, name) for name in updated_connection_names}
399 )
400 # Setting these to frozenset is at odds with the type annotation,
401 # but MyPy can't tell because we're using setattr, and we want to
402 # lie to it anyway to get runtime guards against post-__init__
403 # mutation.
404 setattr(instance, attrName, frozenset(updated_connection_names))
405 # Update the existing dict in place, since we already have a view of
406 # that.
407 instance._allConnections.clear()
408 instance._allConnections.update(updated_all_connections)
410 for connection_name, obj in instance._allConnections.items():
411 if obj.deprecated is not None:
412 warnings.warn(
413 f"Connection {connection_name} with datasetType {obj.name} "
414 f"(from {obj._deprecation_context}): {obj.deprecated}",
415 FutureWarning,
416 stacklevel=1, # Report from this location.
417 )
419 # Freeze the connection instance dimensions now. This at odds with the
420 # type annotation, which says [mutable] `set`, just like the connection
421 # type attributes (e.g. `inputs`, `outputs`, etc.), though MyPy can't
422 # tell with those since we're using setattr for them.
423 instance.dimensions = frozenset(instance.dimensions) # type: ignore
425 return instance
428class QuantizedConnection(SimpleNamespace):
429 r"""A Namespace to map defined variable names of connections to the
430 associated `lsst.daf.butler.DatasetRef` objects.
432 This class maps the names used to define a connection on a
433 `PipelineTaskConnections` class to the corresponding
434 `~lsst.daf.butler.DatasetRef`\s provided by a `~lsst.daf.butler.Quantum`
435 instance. This will be a quantum of execution based on the graph created
436 by examining all the connections defined on the
437 `PipelineTaskConnections` class.
439 Parameters
440 ----------
441 **kwargs : `~typing.Any`
442 Not used.
443 """
445 def __init__(self, **kwargs):
446 # Create a variable to track what attributes are added. This is used
447 # later when iterating over this QuantizedConnection instance
448 object.__setattr__(self, "_attributes", set())
450 def __setattr__(self, name: str, value: DatasetRef | list[DatasetRef]) -> None:
451 # Capture the attribute name as it is added to this object
452 self._attributes.add(name)
453 super().__setattr__(name, value)
455 def __delattr__(self, name):
456 object.__delattr__(self, name)
457 self._attributes.remove(name)
459 def __len__(self) -> int:
460 return len(self._attributes)
462 def __iter__(
463 self,
464 ) -> Generator[tuple[str, DatasetRef | list[DatasetRef]], None, None]:
465 """Make an iterator for this `QuantizedConnection`.
467 Iterating over a `QuantizedConnection` will yield a tuple with the name
468 of an attribute and the value associated with that name. This is
469 similar to dict.items() but is on the namespace attributes rather than
470 dict keys.
471 """
472 yield from ((name, getattr(self, name)) for name in self._attributes)
474 def keys(self) -> Generator[str, None, None]:
475 """Return an iterator over all the attributes added to a
476 `QuantizedConnection` class.
477 """
478 yield from self._attributes
481class InputQuantizedConnection(QuantizedConnection):
482 """Input variant of a `QuantizedConnection`."""
484 pass
487class OutputQuantizedConnection(QuantizedConnection):
488 """Output variant of a `QuantizedConnection`."""
490 pass
493@dataclass(frozen=True)
494class DeferredDatasetRef:
495 """A wrapper class for `~lsst.daf.butler.DatasetRef` that indicates that a
496 `PipelineTask` should receive a `~lsst.daf.butler.DeferredDatasetHandle`
497 instead of an in-memory dataset.
499 Attributes
500 ----------
501 datasetRef : `lsst.daf.butler.DatasetRef`
502 The `lsst.daf.butler.DatasetRef` that will be eventually used to
503 resolve a dataset.
504 """
506 datasetRef: DatasetRef
508 def __getattr__(self, name: str) -> Any:
509 return getattr(self.datasetRef, name)
512class PipelineTaskConnections(metaclass=PipelineTaskConnectionsMetaclass):
513 """PipelineTaskConnections is a class used to declare desired IO when a
514 PipelineTask is run by an activator.
516 Parameters
517 ----------
518 config : `PipelineTaskConfig`
519 A `PipelineTaskConfig` class instance whose class has been configured
520 to use this `PipelineTaskConnections` class.
522 See Also
523 --------
524 iterConnections : Iterator over selected connections.
526 Notes
527 -----
528 ``PipelineTaskConnection`` classes are created by declaring class
529 attributes of types defined in `lsst.pipe.base.connectionTypes` and are
530 listed as follows:
532 * ``InitInput`` - Defines connections in a quantum graph which are used as
533 inputs to the ``__init__`` function of the `PipelineTask` corresponding
534 to this class
535 * ``InitOuput`` - Defines connections in a quantum graph which are to be
536 persisted using a butler at the end of the ``__init__`` function of the
537 `PipelineTask` corresponding to this class. The variable name used to
538 define this connection should be the same as an attribute name on the
539 `PipelineTask` instance. E.g. if an ``InitOutput`` is declared with
540 the name ``outputSchema`` in a ``PipelineTaskConnections`` class, then
541 a `PipelineTask` instance should have an attribute
542 ``self.outputSchema`` defined. Its value is what will be saved by the
543 activator framework.
544 * ``PrerequisiteInput`` - An input connection type that defines a
545 `lsst.daf.butler.DatasetType` that must be present at execution time,
546 but that will not be used during the course of creating the quantum
547 graph to be executed. These most often are things produced outside the
548 processing pipeline, such as reference catalogs.
549 * ``Input`` - Input `lsst.daf.butler.DatasetType` objects that will be used
550 in the ``run`` method of a `PipelineTask`. The name used to declare
551 class attribute must match a function argument name in the ``run``
552 method of a `PipelineTask`. E.g. If the ``PipelineTaskConnections``
553 defines an ``Input`` with the name ``calexp``, then the corresponding
554 signature should be ``PipelineTask.run(calexp, ...)``
555 * ``Output`` - A `lsst.daf.butler.DatasetType` that will be produced by an
556 execution of a `PipelineTask`. The name used to declare the connection
557 must correspond to an attribute of a `Struct` that is returned by a
558 `PipelineTask` ``run`` method. E.g. if an output connection is
559 defined with the name ``measCat``, then the corresponding
560 ``PipelineTask.run`` method must return ``Struct(measCat=X,..)`` where
561 X matches the ``storageClass`` type defined on the output connection.
563 Attributes of these types can also be created, replaced, or deleted on the
564 `PipelineTaskConnections` instance in the ``__init__`` method, if more than
565 just the name depends on the configuration. It is preferred to define them
566 in the class when possible (even if configuration may cause the connection
567 to be removed from the instance).
569 The process of declaring a ``PipelineTaskConnection`` class involves
570 parameters passed in the declaration statement.
572 The first parameter is ``dimensions`` which is an iterable of strings which
573 defines the unit of processing the run method of a corresponding
574 `PipelineTask` will operate on. These dimensions must match dimensions that
575 exist in the butler registry which will be used in executing the
576 corresponding `PipelineTask`. The dimensions may be also modified in
577 subclass ``__init__`` methods if they need to depend on configuration.
579 The second parameter is labeled ``defaultTemplates`` and is conditionally
580 optional. The name attributes of connections can be specified as python
581 format strings, with named format arguments. If any of the name parameters
582 on connections defined in a `PipelineTaskConnections` class contain a
583 template, then a default template value must be specified in the
584 ``defaultTemplates`` argument. This is done by passing a dictionary with
585 keys corresponding to a template identifier, and values corresponding to
586 the value to use as a default when formatting the string. For example if
587 ``ConnectionsClass.calexp.name = '{input}Coadd_calexp'`` then
588 ``defaultTemplates`` = {'input': 'deep'}.
590 Once a `PipelineTaskConnections` class is created, it is used in the
591 creation of a `PipelineTaskConfig`. This is further documented in the
592 documentation of `PipelineTaskConfig`. For the purposes of this
593 documentation, the relevant information is that the config class allows
594 configuration of connection names by users when running a pipeline.
596 Instances of a `PipelineTaskConnections` class are used by the pipeline
597 task execution framework to introspect what a corresponding `PipelineTask`
598 will require, and what it will produce.
600 Examples
601 --------
602 >>> from lsst.pipe.base import connectionTypes as cT
603 >>> from lsst.pipe.base import PipelineTaskConnections
604 >>> from lsst.pipe.base import PipelineTaskConfig
605 >>> class ExampleConnections(
606 ... PipelineTaskConnections,
607 ... dimensions=("A", "B"),
608 ... defaultTemplates={"foo": "Example"},
609 ... ):
610 ... inputConnection = cT.Input(
611 ... doc="Example input",
612 ... dimensions=("A", "B"),
613 ... storageClass=Exposure,
614 ... name="{foo}Dataset",
615 ... )
616 ... outputConnection = cT.Output(
617 ... doc="Example output",
618 ... dimensions=("A", "B"),
619 ... storageClass=Exposure,
620 ... name="{foo}output",
621 ... )
622 >>> class ExampleConfig(
623 ... PipelineTaskConfig, pipelineConnections=ExampleConnections
624 ... ):
625 ... pass
626 >>> config = ExampleConfig()
627 >>> config.connections.foo = Modified
628 >>> config.connections.outputConnection = "TotallyDifferent"
629 >>> connections = ExampleConnections(config=config)
630 >>> assert connections.inputConnection.name == "ModifiedDataset"
631 >>> assert connections.outputConnection.name == "TotallyDifferent"
632 """
634 # We annotate these attributes as mutable sets because that's what they are
635 # inside derived ``__init__`` implementations and that's what matters most
636 # After that's done, the metaclass __call__ makes them into frozensets, but
637 # relatively little code interacts with them then, and that code knows not
638 # to try to modify them without having to be told that by mypy.
640 dimensions: set[str]
641 """Set of dimension names that define the unit of work for this task.
643 Required and implied dependencies will automatically be expanded later and
644 need not be provided.
646 This may be replaced or modified in ``__init__`` to change the dimensions
647 of the task. After ``__init__`` it will be a `frozenset` and may not be
648 replaced.
649 """
651 inputs: set[str]
652 """Set with the names of all `connectionTypes.Input` connection attributes.
654 This is updated automatically as class attributes are added, removed, or
655 replaced in ``__init__``. Removing entries from this set will cause those
656 connections to be removed after ``__init__`` completes, but this is
657 supported only for backwards compatibility; new code should instead just
658 delete the collection attributed directly. After ``__init__`` this will be
659 a `frozenset` and may not be replaced.
660 """
662 prerequisiteInputs: set[str]
663 """Set with the names of all `~connectionTypes.PrerequisiteInput`
664 connection attributes.
666 See `inputs` for additional information.
667 """
669 outputs: set[str]
670 """Set with the names of all `~connectionTypes.Output` connection
671 attributes.
673 See `inputs` for additional information.
674 """
676 initInputs: set[str]
677 """Set with the names of all `~connectionTypes.InitInput` connection
678 attributes.
680 See `inputs` for additional information.
681 """
683 initOutputs: set[str]
684 """Set with the names of all `~connectionTypes.InitOutput` connection
685 attributes.
687 See `inputs` for additional information.
688 """
690 allConnections: Mapping[str, BaseConnection]
691 """Mapping holding all connection attributes.
693 This is a read-only view that is automatically updated when connection
694 attributes are added, removed, or replaced in ``__init__``. It is also
695 updated after ``__init__`` completes to reflect changes in `inputs`,
696 `prerequisiteInputs`, `outputs`, `initInputs`, and `initOutputs`.
697 """
699 _allConnections: dict[str, BaseConnection]
701 def __init__(self, *, config: PipelineTaskConfig | None = None):
702 pass
704 def __setattr__(self, name: str, value: Any) -> None:
705 if isinstance(value, BaseConnection):
706 previous = self._allConnections.get(name)
707 try:
708 getattr(self, value._connection_type_set).add(name)
709 except AttributeError:
710 # Attempt to call add on a frozenset, which is what these sets
711 # are after __init__ is done.
712 raise TypeError("Connections objects are frozen after construction.") from None
713 if previous is not None and value._connection_type_set != previous._connection_type_set:
714 # Connection has changed type, e.g. Input to PrerequisiteInput;
715 # update the sets accordingly. To be extra defensive about
716 # multiple assignments we use the type of the previous instance
717 # instead of assuming that's the same as the type of the self,
718 # which is just the default. Use discard instead of remove so
719 # manually removing from these sets first is never an error.
720 getattr(self, previous._connection_type_set).discard(name)
721 self._allConnections[name] = value
722 if hasattr(self.__class__, name):
723 # Don't actually set the attribute if this was a connection
724 # declared in the class; in that case we let the descriptor
725 # return the value we just added to allConnections.
726 return
727 # Actually add the attribute.
728 super().__setattr__(name, value)
730 def __delattr__(self, name):
731 """Descriptor delete method."""
732 previous = self._allConnections.get(name)
733 if previous is not None:
734 # Delete this connection's name from the appropriate set, which we
735 # have to get from the previous instance instead of assuming it's
736 # the same set that was appropriate for the class-level default.
737 # Use discard instead of remove so manually removing from these
738 # sets first is never an error.
739 try:
740 getattr(self, previous._connection_type_set).discard(name)
741 except AttributeError:
742 # Attempt to call discard on a frozenset, which is what these
743 # sets are after __init__ is done.
744 raise TypeError("Connections objects are frozen after construction.") from None
745 del self._allConnections[name]
746 if hasattr(self.__class__, name):
747 # Don't actually delete the attribute if this was a connection
748 # declared in the class; in that case we let the descriptor
749 # see that it's no longer present in allConnections.
750 return
751 # Actually delete the attribute.
752 super().__delattr__(name)
754 def buildDatasetRefs(
755 self, quantum: Quantum
756 ) -> tuple[InputQuantizedConnection, OutputQuantizedConnection]:
757 """Build `QuantizedConnection` corresponding to input
758 `~lsst.daf.butler.Quantum`.
760 Parameters
761 ----------
762 quantum : `lsst.daf.butler.Quantum`
763 Quantum object which defines the inputs and outputs for a given
764 unit of processing.
766 Returns
767 -------
768 retVal : `tuple` of (`InputQuantizedConnection`, \
769 `OutputQuantizedConnection`)
770 Namespaces mapping attribute names
771 (identifiers of connections) to butler references defined in the
772 input `~lsst.daf.butler.Quantum`.
773 """
774 inputDatasetRefs = InputQuantizedConnection()
775 outputDatasetRefs = OutputQuantizedConnection()
777 # populate inputDatasetRefs from quantum inputs
778 for attributeName in itertools.chain(self.inputs, self.prerequisiteInputs):
779 # get the attribute identified by name
780 attribute = getattr(self, attributeName)
781 # if the dataset is marked to load deferred, wrap it in a
782 # DeferredDatasetRef
783 quantumInputRefs: list[DatasetRef] | list[DeferredDatasetRef]
784 if attribute.deferLoad:
785 quantumInputRefs = [
786 DeferredDatasetRef(datasetRef=ref) for ref in quantum.inputs[attribute.name]
787 ]
788 else:
789 quantumInputRefs = list(quantum.inputs[attribute.name])
790 # Unpack arguments that are not marked multiples (list of
791 # length one)
792 if not attribute.multiple:
793 if len(quantumInputRefs) > 1:
794 raise ScalarError(
795 "Received multiple datasets "
796 f"{', '.join(str(r.dataId) for r in quantumInputRefs)} "
797 f"for scalar connection {attributeName} "
798 f"({quantumInputRefs[0].datasetType.name}) "
799 f"of quantum for {quantum.taskName} with data ID {quantum.dataId}."
800 )
801 if len(quantumInputRefs) == 0:
802 continue
803 setattr(inputDatasetRefs, attributeName, quantumInputRefs[0])
804 else:
805 # Add to the QuantizedConnection identifier
806 setattr(inputDatasetRefs, attributeName, quantumInputRefs)
808 # populate outputDatasetRefs from quantum outputs
809 for attributeName in self.outputs:
810 # get the attribute identified by name
811 attribute = getattr(self, attributeName)
812 value = quantum.outputs[attribute.name]
813 # Unpack arguments that are not marked multiples (list of
814 # length one)
815 if not attribute.multiple:
816 setattr(outputDatasetRefs, attributeName, value[0])
817 else:
818 setattr(outputDatasetRefs, attributeName, value)
820 return inputDatasetRefs, outputDatasetRefs
822 def adjustQuantum(
823 self,
824 inputs: dict[str, tuple[BaseInput, Collection[DatasetRef]]],
825 outputs: dict[str, tuple[Output, Collection[DatasetRef]]],
826 label: str,
827 data_id: DataCoordinate,
828 ) -> tuple[
829 Mapping[str, tuple[BaseInput, Collection[DatasetRef]]],
830 Mapping[str, tuple[Output, Collection[DatasetRef]]],
831 ]:
832 """Override to make adjustments to `lsst.daf.butler.DatasetRef` objects
833 in the `lsst.daf.butler.Quantum` during the graph generation stage
834 of the activator.
836 Parameters
837 ----------
838 inputs : `dict`
839 Dictionary whose keys are an input (regular or prerequisite)
840 connection name and whose values are a tuple of the connection
841 instance and a collection of associated
842 `~lsst.daf.butler.DatasetRef` objects.
843 The exact type of the nested collections is unspecified; it can be
844 assumed to be multi-pass iterable and support `len` and ``in``, but
845 it should not be mutated in place. In contrast, the outer
846 dictionaries are guaranteed to be temporary copies that are true
847 `dict` instances, and hence may be modified and even returned; this
848 is especially useful for delegating to `super` (see notes below).
849 outputs : `~collections.abc.Mapping`
850 Mapping of output datasets, with the same structure as ``inputs``.
851 label : `str`
852 Label for this task in the pipeline (should be used in all
853 diagnostic messages).
854 data_id : `lsst.daf.butler.DataCoordinate`
855 Data ID for this quantum in the pipeline (should be used in all
856 diagnostic messages).
858 Returns
859 -------
860 adjusted_inputs : `~collections.abc.Mapping`
861 Mapping of the same form as ``inputs`` with updated containers of
862 input `~lsst.daf.butler.DatasetRef` objects. Connections that are
863 not changed should not be returned at all. Datasets may only be
864 removed, not added. Nested collections may be of any multi-pass
865 iterable type, and the order of iteration will set the order of
866 iteration within `PipelineTask.runQuantum`.
867 adjusted_outputs : `~collections.abc.Mapping`
868 Mapping of updated output datasets, with the same structure and
869 interpretation as ``adjusted_inputs``.
871 Raises
872 ------
873 ScalarError
874 Raised if any `Input` or `PrerequisiteInput` connection has
875 ``multiple`` set to `False`, but multiple datasets.
876 NoWorkFound
877 Raised to indicate that this quantum should not be run; not enough
878 datasets were found for a regular `Input` connection, and the
879 quantum should be pruned or skipped.
880 FileNotFoundError
881 Raised to cause QuantumGraph generation to fail (with the message
882 included in this exception); not enough datasets were found for a
883 `PrerequisiteInput` connection.
885 Notes
886 -----
887 The base class implementation performs important checks. It always
888 returns an empty mapping (i.e. makes no adjustments). It should
889 always called be via `super` by custom implementations, ideally at the
890 end of the custom implementation with already-adjusted mappings when
891 any datasets are actually dropped, e.g.:
893 .. code-block:: python
895 def adjustQuantum(self, inputs, outputs, label, data_id):
896 # Filter out some dataset refs for one connection.
897 connection, old_refs = inputs["my_input"]
898 new_refs = [ref for ref in old_refs if ...]
899 adjusted_inputs = {"my_input", (connection, new_refs)}
900 # Update the original inputs so we can pass them to super.
901 inputs.update(adjusted_inputs)
902 # Can ignore outputs from super because they are guaranteed
903 # to be empty.
904 super().adjustQuantum(inputs, outputs, label_data_id)
905 # Return only the connections we modified.
906 return adjusted_inputs, {}
908 Removing outputs here is guaranteed to affect what is actually
909 passed to `PipelineTask.runQuantum`, but its effect on the larger
910 graph may be deferred to execution, depending on the context in
911 which `adjustQuantum` is being run: if one quantum removes an output
912 that is needed by a second quantum as input, the second quantum may not
913 be adjusted (and hence pruned or skipped) until that output is actually
914 found to be missing at execution time.
916 Tasks that desire zip-iteration consistency between any combinations of
917 connections that have the same data ID should generally implement
918 `adjustQuantum` to achieve this, even if they could also run that
919 logic during execution; this allows the system to see outputs that will
920 not be produced because the corresponding input is missing as early as
921 possible.
922 """
923 for name, (input_connection, refs) in inputs.items():
924 dataset_type_name = input_connection.name
925 if not input_connection.multiple and len(refs) > 1:
926 raise ScalarError(
927 f"Found multiple datasets {', '.join(str(r.dataId) for r in refs)} "
928 f"for non-multiple input connection {label}.{name} ({dataset_type_name}) "
929 f"for quantum data ID {data_id}."
930 )
931 if len(refs) < input_connection.minimum:
932 if isinstance(input_connection, PrerequisiteInput):
933 # This branch should only be possible during QG generation,
934 # or if someone deleted the dataset between making the QG
935 # and trying to run it. Either one should be a hard error.
936 raise FileNotFoundError(
937 f"Not enough datasets ({len(refs)}) found for non-optional connection {label}.{name} "
938 f"({dataset_type_name}) with minimum={input_connection.minimum} for quantum data ID "
939 f"{data_id}."
940 )
941 else:
942 raise NoWorkFound(label, name, input_connection)
943 for name, (output_connection, refs) in outputs.items():
944 dataset_type_name = output_connection.name
945 if not output_connection.multiple and len(refs) > 1:
946 raise ScalarError(
947 f"Found multiple datasets {', '.join(str(r.dataId) for r in refs)} "
948 f"for non-multiple output connection {label}.{name} ({dataset_type_name}) "
949 f"for quantum data ID {data_id}."
950 )
951 return {}, {}
953 def getSpatialBoundsConnections(self) -> Iterable[str]:
954 """Return the names of regular input and output connections whose data
955 IDs should be used to compute the spatial bounds of this task's quanta.
957 The spatial bound for a quantum is defined as the union of the regions
958 of all data IDs of all connections returned here, along with the region
959 of the quantum data ID (if the task has spatial dimensions).
961 Returns
962 -------
963 connection_names : `collections.abc.Iterable` [ `str` ]
964 Names of collections with spatial dimensions. These are the
965 task-internal connection names, not butler dataset type names.
967 Notes
968 -----
969 The spatial bound is used to search for prerequisite inputs that have
970 skypix dimensions. The default implementation returns an empty
971 iterable, which is usually sufficient for tasks with spatial
972 dimensions, but if a task's inputs or outputs are associated with
973 spatial regions that extend beyond the quantum data ID's region, this
974 method may need to be overridden to expand the set of prerequisite
975 inputs found.
977 Tasks that do not have spatial dimensions that have skypix prerequisite
978 inputs should always override this method, as the default spatial
979 bounds otherwise cover the full sky.
980 """
981 return ()
983 def getTemporalBoundsConnections(self) -> Iterable[str]:
984 """Return the names of regular input and output connections whose data
985 IDs should be used to compute the temporal bounds of this task's
986 quanta.
988 The temporal bound for a quantum is defined as the union of the
989 timespans of all data IDs of all connections returned here, along with
990 the timespan of the quantum data ID (if the task has temporal
991 dimensions).
993 Returns
994 -------
995 connection_names : `collections.abc.Iterable` [ `str` ]
996 Names of collections with temporal dimensions. These are the
997 task-internal connection names, not butler dataset type names.
999 Notes
1000 -----
1001 The temporal bound is used to search for prerequisite inputs that are
1002 calibration datasets. The default implementation returns an empty
1003 iterable, which is usually sufficient for tasks with temporal
1004 dimensions, but if a task's inputs or outputs are associated with
1005 timespans that extend beyond the quantum data ID's timespan, this
1006 method may need to be overridden to expand the set of prerequisite
1007 inputs found.
1009 Tasks that do not have temporal dimensions that do not implement this
1010 method will use an infinite timespan for any calibration lookups.
1011 """
1012 return ()
1014 def adjust_all_quanta(self, adjuster: QuantaAdjuster) -> None:
1015 """Customize the set of quanta predicted for this task during quantum
1016 graph generation.
1018 Parameters
1019 ----------
1020 adjuster : `QuantaAdjuster`
1021 A helper object that implementations can use to modify the
1022 under-construction quantum graph.
1024 Notes
1025 -----
1026 This hook is called before `adjustQuantum`, which is where built-in
1027 checks for `NoWorkFound` cases and missing prerequisites are handled.
1028 This means that the set of preliminary quanta seen by this method could
1029 include some that would normally be dropped later.
1030 """
1031 pass
1034def iterConnections(
1035 connections: PipelineTaskConnections, connectionType: str | Iterable[str]
1036) -> Generator[BaseConnection, None, None]:
1037 """Create an iterator over the selected connections type which yields
1038 all the defined connections of that type.
1040 Parameters
1041 ----------
1042 connections : `PipelineTaskConnections`
1043 An instance of a `PipelineTaskConnections` object that will be iterated
1044 over.
1045 connectionType : `str`
1046 The type of connections to iterate over, valid values are inputs,
1047 outputs, prerequisiteInputs, initInputs, initOutputs.
1049 Yields
1050 ------
1051 connection: `~.connectionTypes.BaseConnection`
1052 A connection defined on the input connections object of the type
1053 supplied. The yielded value Will be an derived type of
1054 `~.connectionTypes.BaseConnection`.
1055 """
1056 if isinstance(connectionType, str):
1057 connectionType = (connectionType,)
1058 for name in itertools.chain.from_iterable(getattr(connections, ct) for ct in connectionType):
1059 yield getattr(connections, name)
1062@dataclass
1063class AdjustQuantumHelper:
1064 """Helper class for calling `PipelineTaskConnections.adjustQuantum`.
1066 This class holds `inputs` and `outputs` mappings in the form used by
1067 `lsst.daf.butler.Quantum` and execution harness code, i.e. with
1068 `~lsst.daf.butler.DatasetType` keys, translating them to and from the
1069 connection-oriented mappings used inside `PipelineTaskConnections`.
1070 """
1072 inputs: NamedKeyMapping[DatasetType, Sequence[DatasetRef]]
1073 """Mapping of regular input and prerequisite input datasets, grouped by
1074 `~lsst.daf.butler.DatasetType`.
1075 """
1077 outputs: NamedKeyMapping[DatasetType, Sequence[DatasetRef]]
1078 """Mapping of output datasets, grouped by `~lsst.daf.butler.DatasetType`.
1079 """
1081 inputs_adjusted: bool = False
1082 """Whether any inputs were removed in the last call to `adjust_in_place`.
1083 """
1085 outputs_adjusted: bool = False
1086 """Whether any outputs were removed in the last call to `adjust_in_place`.
1087 """
1089 def adjust_in_place(
1090 self,
1091 connections: PipelineTaskConnections,
1092 label: str,
1093 data_id: DataCoordinate,
1094 ) -> None:
1095 """Call `~PipelineTaskConnections.adjustQuantum` and update ``self``
1096 with its results.
1098 Parameters
1099 ----------
1100 connections : `PipelineTaskConnections`
1101 Instance on which to call `~PipelineTaskConnections.adjustQuantum`.
1102 label : `str`
1103 Label for this task in the pipeline (should be used in all
1104 diagnostic messages).
1105 data_id : `lsst.daf.butler.DataCoordinate`
1106 Data ID for this quantum in the pipeline (should be used in all
1107 diagnostic messages).
1108 """
1109 # Translate self's DatasetType-keyed, Quantum-oriented mappings into
1110 # connection-keyed, PipelineTask-oriented mappings.
1111 inputs_by_connection: dict[str, tuple[BaseInput, tuple[DatasetRef, ...]]] = {}
1112 outputs_by_connection: dict[str, tuple[Output, tuple[DatasetRef, ...]]] = {}
1113 for name in itertools.chain(connections.inputs, connections.prerequisiteInputs):
1114 connection = getattr(connections, name)
1115 dataset_type_name = connection.name
1116 inputs_by_connection[name] = (connection, tuple(self.inputs.get(dataset_type_name, ())))
1117 for name in itertools.chain(connections.outputs):
1118 connection = getattr(connections, name)
1119 dataset_type_name = connection.name
1120 outputs_by_connection[name] = (connection, tuple(self.outputs.get(dataset_type_name, ())))
1121 # Actually call adjustQuantum.
1122 # MyPy correctly complains that this call is not quite legal, but the
1123 # method docs explain exactly what's expected and it's the behavior we
1124 # want. It'd be nice to avoid this if we ever have to change the
1125 # interface anyway, but not an immediate problem.
1126 adjusted_inputs_by_connection, adjusted_outputs_by_connection = connections.adjustQuantum(
1127 inputs_by_connection, # type: ignore
1128 outputs_by_connection, # type: ignore
1129 label,
1130 data_id,
1131 )
1132 # Translate adjustments to DatasetType-keyed, Quantum-oriented form,
1133 # installing new mappings in self if necessary.
1134 if adjusted_inputs_by_connection:
1135 adjusted_inputs = NamedKeyDict[DatasetType, tuple[DatasetRef, ...]](self.inputs)
1136 for name, (connection, updated_refs) in adjusted_inputs_by_connection.items():
1137 dataset_type_name = connection.name
1138 if not set(updated_refs).issubset(self.inputs[dataset_type_name]):
1139 raise RuntimeError(
1140 f"adjustQuantum implementation for task with label {label} returned {name} "
1141 f"({dataset_type_name}) input datasets that are not a subset of those "
1142 f"it was given for data ID {data_id}."
1143 )
1144 adjusted_inputs[dataset_type_name] = tuple(updated_refs)
1145 self.inputs = adjusted_inputs.freeze()
1146 self.inputs_adjusted = True
1147 else:
1148 self.inputs_adjusted = False
1149 if adjusted_outputs_by_connection:
1150 adjusted_outputs = NamedKeyDict[DatasetType, tuple[DatasetRef, ...]](self.outputs)
1151 for name, (connection, updated_refs) in adjusted_outputs_by_connection.items():
1152 dataset_type_name = connection.name
1153 if not set(updated_refs).issubset(self.outputs[dataset_type_name]):
1154 raise RuntimeError(
1155 f"adjustQuantum implementation for task with label {label} returned {name} "
1156 f"({dataset_type_name}) output datasets that are not a subset of those "
1157 f"it was given for data ID {data_id}."
1158 )
1159 adjusted_outputs[dataset_type_name] = tuple(updated_refs)
1160 self.outputs = adjusted_outputs.freeze()
1161 self.outputs_adjusted = True
1162 else:
1163 self.outputs_adjusted = False
1166class QuantaAdjuster:
1167 """A helper class for the `PipelineTaskConnections.adjust_all_quanta` hook.
1169 Parameters
1170 ----------
1171 task_label : `str`
1172 Label of the task whose quanta are being adjusted.
1173 pipeline_graph : `pipeline_graph.PipelineGraph`
1174 Pipeline graph the quantum graph is being built from.
1175 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
1176 Under-construction quantum graph that will be modified in place.
1177 butler : `lsst.daf.butler.Butler`
1178 Read-only instance with its default collection search path set to the
1179 input collections passed to the quantum-graph builder.
1180 """
1182 def __init__(
1183 self, task_label: str, pipeline_graph: PipelineGraph, skeleton: QuantumGraphSkeleton, butler: Butler
1184 ):
1185 self._task_node = pipeline_graph.tasks[task_label]
1186 self._pipeline_graph = pipeline_graph
1187 self._skeleton = skeleton
1188 self._n_removed = 0
1189 self._butler = butler
1191 @property
1192 def task_label(self) -> str:
1193 """The label this task has been configured with."""
1194 return self._task_node.label
1196 @property
1197 def task_node(self) -> TaskNode:
1198 """The node for this task in the pipeline graph."""
1199 return self._task_node
1201 @property
1202 def butler(self) -> Butler:
1203 """Read-only instance with its default collection search path set to
1204 the input collections passed to the quantum-graph builder.
1205 """
1206 return self._butler
1208 def iter_data_ids(self) -> Iterator[DataCoordinate]:
1209 """Iterate over the data IDs of all quanta for this task.
1211 Returns
1212 -------
1213 data_ids : `~collections.abc.Iterator` [ \
1214 `~lsst.daf.butler.DataCoordinate` ]
1215 Data IDs. These are minimal data IDs without dimension records or
1216 implied values; use `expand_quantum_data_id` to get a full data ID
1217 when needed.
1218 """
1219 for key in self._skeleton.get_quanta(self._task_node.label):
1220 yield DataCoordinate.from_required_values(self._task_node.dimensions, key.data_id_values)
1222 def remove_quantum(self, data_id: DataCoordinate) -> None:
1223 """Remove a quantum from the graph.
1225 Parameters
1226 ----------
1227 data_id : `~lsst.daf.butler.DataCoordinate`
1228 Data ID of the quantum to remove. All outputs will be removed as
1229 well.
1230 """
1231 from .quantum_graph_skeleton import QuantumKey
1233 self._skeleton.remove_quantum_node(
1234 QuantumKey(self._task_node.label, data_id.required_values), remove_outputs=True
1235 )
1236 self._n_removed += 1
1238 def get_inputs(self, quantum_data_id: DataCoordinate) -> dict[str, list[DataCoordinate]]:
1239 """Return the data IDs of all regular inputs to a quantum.
1241 Parameters
1242 ----------
1243 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1244 Data ID of the quantum to get the inputs of.
1246 Returns
1247 -------
1248 inputs : `dict` [ `str`, `list` [ `~lsst.daf.butler.DataCoordinate` ] ]
1249 Data IDs of inputs, keyed by the connection name (the internal task
1250 name, not the dataset type name). This only contains regular
1251 inputs, not init-inputs or prerequisite inputs.
1253 Notes
1254 -----
1255 If two connections have the same dataset type, the current
1256 implementation assumes the set of datasets is the same for the two
1257 connections. This limitation may be removed in the future.
1258 """
1259 from .quantum_graph_skeleton import DatasetKey, QuantumKey
1261 by_dataset_type_name: defaultdict[str, list[DataCoordinate]] = defaultdict(list)
1262 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1263 for dataset_key in self._skeleton.iter_inputs_of(quantum_key):
1264 if not isinstance(dataset_key, DatasetKey):
1265 continue
1266 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
1267 by_dataset_type_name[dataset_key.parent_dataset_type_name].append(
1268 DataCoordinate.from_required_values(dataset_type_node.dimensions, dataset_key.data_id_values)
1269 )
1270 return {
1271 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name]
1272 for edge in self._task_node.iter_all_inputs()
1273 }
1275 def get_prerequisite_inputs(
1276 self,
1277 quantum_data_id: DataCoordinate,
1278 ) -> dict[str, dict[uuid.UUID, DataCoordinate]]:
1279 """Return the data IDs of all prerequisite inputs to a quantum.
1281 Parameters
1282 ----------
1283 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1284 Data ID of the quantum to get the inputs of.
1286 Returns
1287 -------
1288 inputs : `dict` [ `str`, \
1289 `dict` [ `uuid.UUID`, `~lsst.daf.butler.DataCoordinate` ] ]
1290 Dataset IDs and and data IDs of prerequisite inputs, keyed by the
1291 connection name (the internal task name, not the dataset type
1292 name). This only contains prerequisite inputs, not init-inputs or
1293 regular inputs.
1295 Notes
1296 -----
1297 If two connections have the same dataset type, the current
1298 implementation assumes the set of datasets is the same for the two
1299 connections. This limitation may be removed in the future.
1301 Unlike regular inputs, prerequisite inputs are not looked up from input
1302 collections or indexed by data ID. Instead, they are uniquely
1303 identified by dataset UUID and reused directly between quanta.
1304 """
1305 from .quantum_graph_skeleton import PrerequisiteDatasetKey, QuantumKey
1307 by_dataset_type_name: defaultdict[str, dict[uuid.UUID, DataCoordinate]] = defaultdict(dict)
1308 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1310 for dataset_key in self._skeleton.iter_inputs_of(quantum_key):
1311 if not isinstance(dataset_key, PrerequisiteDatasetKey):
1312 continue
1313 by_dataset_type_name[dataset_key.parent_dataset_type_name][
1314 uuid.UUID(bytes=dataset_key.dataset_id_bytes)
1315 ] = self._skeleton.get_data_id(dataset_key)
1316 return {
1317 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name]
1318 for edge in self._task_node.iter_all_inputs()
1319 }
1321 def get_outputs(self, quantum_data_id: DataCoordinate) -> dict[str, list[DataCoordinate]]:
1322 """Return the data IDs of all regular outputs to a quantum.
1324 Parameters
1325 ----------
1326 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1327 Data ID of the quantum to get the outputs of.
1329 Returns
1330 -------
1331 inputs : `dict` [ `str`, `list` [ `~lsst.daf.butler.DataCoordinate` ] ]
1332 Data IDs of inputs, keyed by the connection name (the internal task
1333 name, not the dataset type name). This only contains regular
1334 outputs, not init-outputs or log or metadata outputs.
1336 Notes
1337 -----
1338 If two connections have the same dataset type, the current
1339 implementation assumes the set of datasets is the same for the two
1340 connections. This limitation may be removed in the future.
1341 """
1342 from .quantum_graph_skeleton import QuantumKey
1344 by_dataset_type_name: defaultdict[str, list[DataCoordinate]] = defaultdict(list)
1345 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1346 for dataset_key in self._skeleton.iter_outputs_of(quantum_key):
1347 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
1348 by_dataset_type_name[dataset_key.parent_dataset_type_name].append(
1349 DataCoordinate.from_required_values(dataset_type_node.dimensions, dataset_key.data_id_values)
1350 )
1351 return {
1352 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name]
1353 for edge in self._task_node.outputs.values()
1354 }
1356 def add_input(
1357 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_data_id: DataCoordinate
1358 ) -> None:
1359 """Add a new input to a quantum.
1361 Parameters
1362 ----------
1363 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1364 Data ID of the quantum to add an input to.
1365 connection_name : `str`
1366 Name of the connection (the task-internal name, not the butler
1367 dataset type name).
1368 dataset_data_id : `~lsst.daf.butler.DataCoordinate`
1369 Data ID of the input dataset. Must already exist in the graph
1370 as an input to a different quantum of this task, and must be a
1371 regular input, not a prerequisite input or init-input.
1373 Notes
1374 -----
1375 If two connections have the same dataset type, the current
1376 implementation assumes the set of datasets is the same for the two
1377 connections. This limitation may be removed in the future.
1378 """
1379 from .quantum_graph_skeleton import DatasetKey, QuantumKey
1381 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1382 read_edge = self._task_node.inputs[connection_name]
1383 dataset_key = DatasetKey(read_edge.parent_dataset_type_name, dataset_data_id.required_values)
1384 if dataset_key not in self._skeleton:
1385 raise LookupError(
1386 f"Dataset {read_edge.parent_dataset_type_name}@{dataset_data_id} is not already in the graph."
1387 )
1388 self._skeleton.add_input_edge(quantum_key, dataset_key)
1390 def add_prerequisite_input(
1391 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_uuid: uuid.UUID
1392 ) -> None:
1393 """Add a new prerequisite input to a quantum.
1395 Parameters
1396 ----------
1397 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1398 Data ID of the quantum to add an input to.
1399 connection_name : `str`
1400 Name of the connection (the task-internal name, not the butler
1401 dataset type name).
1402 dataset_uuid : `uuid.UUID`
1403 UUID of the prerequisite input dataset. Must already exist in the
1404 graph as an input to a different quantum of this task, and must be
1405 a prerequisite input, not a regular input or init-input.
1407 Notes
1408 -----
1409 If two connections have the same dataset type, the current
1410 implementation assumes the set of datasets is the same for the two
1411 connections. This limitation may be removed in the future.
1413 Unlike regular inputs, prerequisite inputs are not looked up from input
1414 collections or indexed by data ID. Instead, they are uniquely
1415 identified by dataset UUID and reused directly between quanta.
1416 """
1417 from .quantum_graph_skeleton import PrerequisiteDatasetKey, QuantumKey
1419 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1420 read_edge = self._task_node.prerequisite_inputs[connection_name]
1421 dataset_key = PrerequisiteDatasetKey(read_edge.parent_dataset_type_name, dataset_uuid.bytes)
1422 if dataset_key not in self._skeleton:
1423 raise LookupError(
1424 f"Prerequisite Dataset {read_edge.parent_dataset_type_name}@{dataset_uuid} "
1425 "is not already in the graph."
1426 )
1427 self._skeleton.add_input_edge(quantum_key, dataset_key)
1429 def move_output(
1430 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_data_id: DataCoordinate
1431 ) -> None:
1432 """Remove an output of one quantum and make it a new output of another.
1434 Parameters
1435 ----------
1436 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1437 Data ID of the quantum to move the output to.
1438 connection_name : `str`
1439 Name of the connection (the task-internal name, not the butler
1440 dataset type name).
1441 dataset_data_id : `~lsst.daf.butler.DataCoordinate`
1442 Data ID of the output dataset. Must already exist in the graph
1443 as an output of a different quantum of this task.
1444 """
1445 from .quantum_graph_skeleton import DatasetKey, QuantumKey
1447 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1448 write_edge = self._task_node.outputs[connection_name]
1449 dataset_key = DatasetKey(write_edge.parent_dataset_type_name, dataset_data_id.required_values)
1450 if dataset_key not in self._skeleton:
1451 raise LookupError(
1452 f"Dataset {write_edge.parent_dataset_type_name}@{dataset_data_id} is "
1453 "not already in the graph."
1454 )
1455 self._skeleton.remove_output_edge(dataset_key)
1456 self._skeleton.add_output_edge(quantum_key, dataset_key)
1458 def expand_quantum_data_id(self, data_id: DataCoordinate) -> DataCoordinate:
1459 """Expand a quantum data ID to include implied values and dimension
1460 records.
1462 Parameters
1463 ----------
1464 data_id : `~lsst.daf.butler.DataCoordinate`
1465 A data ID of a quantum already in the graph.
1467 Returns
1468 -------
1469 expanded_data_id : `~lsst.daf.butler.DataCoordinate`
1470 The same data ID, with implied values included and dimension
1471 records attached.
1472 """
1473 from .quantum_graph_skeleton import QuantumKey
1475 return self._skeleton.get_data_id(QuantumKey(self._task_node.label, data_id.required_values))
1477 @property
1478 def n_removed(self) -> int:
1479 """The number of quanta that have been removed by this helper."""
1480 return self._n_removed