Coverage for python / lsst / pipe / base / connections.py: 39%
388 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-18 01:04 -0800
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-18 01:04 -0800
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28"""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 # make sure reduce is called on DeferredDatasetRef and not on
510 # the DatasetRef
511 if name in ("__reduce__", "datasetRef", "__deepcopy__"):
512 object.__getattribute__(self, name)
513 return getattr(self.datasetRef, name)
515 def __deepcopy__(self, memo: dict) -> DeferredDatasetRef:
516 # dataset refs should be immutable deferred version should be too
517 return self
519 def __reduce__(self) -> tuple:
520 return (self.__class__, (self.datasetRef,))
523class PipelineTaskConnections(metaclass=PipelineTaskConnectionsMetaclass):
524 """PipelineTaskConnections is a class used to declare desired IO when a
525 PipelineTask is run by an activator.
527 Parameters
528 ----------
529 config : `PipelineTaskConfig`
530 A `PipelineTaskConfig` class instance whose class has been configured
531 to use this `PipelineTaskConnections` class.
533 See Also
534 --------
535 iterConnections : Iterator over selected connections.
537 Notes
538 -----
539 ``PipelineTaskConnection`` classes are created by declaring class
540 attributes of types defined in `lsst.pipe.base.connectionTypes` and are
541 listed as follows:
543 * ``InitInput`` - Defines connections in a quantum graph which are used as
544 inputs to the ``__init__`` function of the `PipelineTask` corresponding
545 to this class
546 * ``InitOuput`` - Defines connections in a quantum graph which are to be
547 persisted using a butler at the end of the ``__init__`` function of the
548 `PipelineTask` corresponding to this class. The variable name used to
549 define this connection should be the same as an attribute name on the
550 `PipelineTask` instance. E.g. if an ``InitOutput`` is declared with
551 the name ``outputSchema`` in a ``PipelineTaskConnections`` class, then
552 a `PipelineTask` instance should have an attribute
553 ``self.outputSchema`` defined. Its value is what will be saved by the
554 activator framework.
555 * ``PrerequisiteInput`` - An input connection type that defines a
556 `lsst.daf.butler.DatasetType` that must be present at execution time,
557 but that will not be used during the course of creating the quantum
558 graph to be executed. These most often are things produced outside the
559 processing pipeline, such as reference catalogs.
560 * ``Input`` - Input `lsst.daf.butler.DatasetType` objects that will be used
561 in the ``run`` method of a `PipelineTask`. The name used to declare
562 class attribute must match a function argument name in the ``run``
563 method of a `PipelineTask`. E.g. If the ``PipelineTaskConnections``
564 defines an ``Input`` with the name ``calexp``, then the corresponding
565 signature should be ``PipelineTask.run(calexp, ...)``
566 * ``Output`` - A `lsst.daf.butler.DatasetType` that will be produced by an
567 execution of a `PipelineTask`. The name used to declare the connection
568 must correspond to an attribute of a `Struct` that is returned by a
569 `PipelineTask` ``run`` method. E.g. if an output connection is
570 defined with the name ``measCat``, then the corresponding
571 ``PipelineTask.run`` method must return ``Struct(measCat=X,..)`` where
572 X matches the ``storageClass`` type defined on the output connection.
574 Attributes of these types can also be created, replaced, or deleted on the
575 `PipelineTaskConnections` instance in the ``__init__`` method, if more than
576 just the name depends on the configuration. It is preferred to define them
577 in the class when possible (even if configuration may cause the connection
578 to be removed from the instance).
580 The process of declaring a ``PipelineTaskConnection`` class involves
581 parameters passed in the declaration statement.
583 The first parameter is ``dimensions`` which is an iterable of strings which
584 defines the unit of processing the run method of a corresponding
585 `PipelineTask` will operate on. These dimensions must match dimensions that
586 exist in the butler registry which will be used in executing the
587 corresponding `PipelineTask`. The dimensions may be also modified in
588 subclass ``__init__`` methods if they need to depend on configuration.
590 The second parameter is labeled ``defaultTemplates`` and is conditionally
591 optional. The name attributes of connections can be specified as python
592 format strings, with named format arguments. If any of the name parameters
593 on connections defined in a `PipelineTaskConnections` class contain a
594 template, then a default template value must be specified in the
595 ``defaultTemplates`` argument. This is done by passing a dictionary with
596 keys corresponding to a template identifier, and values corresponding to
597 the value to use as a default when formatting the string. For example if
598 ``ConnectionsClass.calexp.name = '{input}Coadd_calexp'`` then
599 ``defaultTemplates`` = {'input': 'deep'}.
601 Once a `PipelineTaskConnections` class is created, it is used in the
602 creation of a `PipelineTaskConfig`. This is further documented in the
603 documentation of `PipelineTaskConfig`. For the purposes of this
604 documentation, the relevant information is that the config class allows
605 configuration of connection names by users when running a pipeline.
607 Instances of a `PipelineTaskConnections` class are used by the pipeline
608 task execution framework to introspect what a corresponding `PipelineTask`
609 will require, and what it will produce.
611 Examples
612 --------
613 >>> from lsst.pipe.base import connectionTypes as cT
614 >>> from lsst.pipe.base import PipelineTaskConnections
615 >>> from lsst.pipe.base import PipelineTaskConfig
616 >>> class ExampleConnections(
617 ... PipelineTaskConnections,
618 ... dimensions=("A", "B"),
619 ... defaultTemplates={"foo": "Example"},
620 ... ):
621 ... inputConnection = cT.Input(
622 ... doc="Example input",
623 ... dimensions=("A", "B"),
624 ... storageClass=Exposure,
625 ... name="{foo}Dataset",
626 ... )
627 ... outputConnection = cT.Output(
628 ... doc="Example output",
629 ... dimensions=("A", "B"),
630 ... storageClass=Exposure,
631 ... name="{foo}output",
632 ... )
633 >>> class ExampleConfig(
634 ... PipelineTaskConfig, pipelineConnections=ExampleConnections
635 ... ):
636 ... pass
637 >>> config = ExampleConfig()
638 >>> config.connections.foo = Modified
639 >>> config.connections.outputConnection = "TotallyDifferent"
640 >>> connections = ExampleConnections(config=config)
641 >>> assert connections.inputConnection.name == "ModifiedDataset"
642 >>> assert connections.outputConnection.name == "TotallyDifferent"
643 """
645 # We annotate these attributes as mutable sets because that's what they are
646 # inside derived ``__init__`` implementations and that's what matters most
647 # After that's done, the metaclass __call__ makes them into frozensets, but
648 # relatively little code interacts with them then, and that code knows not
649 # to try to modify them without having to be told that by mypy.
651 dimensions: set[str]
652 """Set of dimension names that define the unit of work for this task.
654 Required and implied dependencies will automatically be expanded later and
655 need not be provided.
657 This may be replaced or modified in ``__init__`` to change the dimensions
658 of the task. After ``__init__`` it will be a `frozenset` and may not be
659 replaced.
660 """
662 inputs: set[str]
663 """Set with the names of all `connectionTypes.Input` connection attributes.
665 This is updated automatically as class attributes are added, removed, or
666 replaced in ``__init__``. Removing entries from this set will cause those
667 connections to be removed after ``__init__`` completes, but this is
668 supported only for backwards compatibility; new code should instead just
669 delete the collection attributed directly. After ``__init__`` this will be
670 a `frozenset` and may not be replaced.
671 """
673 prerequisiteInputs: set[str]
674 """Set with the names of all `~connectionTypes.PrerequisiteInput`
675 connection attributes.
677 See `inputs` for additional information.
678 """
680 outputs: set[str]
681 """Set with the names of all `~connectionTypes.Output` connection
682 attributes.
684 See `inputs` for additional information.
685 """
687 initInputs: set[str]
688 """Set with the names of all `~connectionTypes.InitInput` connection
689 attributes.
691 See `inputs` for additional information.
692 """
694 initOutputs: set[str]
695 """Set with the names of all `~connectionTypes.InitOutput` connection
696 attributes.
698 See `inputs` for additional information.
699 """
701 allConnections: Mapping[str, BaseConnection]
702 """Mapping holding all connection attributes.
704 This is a read-only view that is automatically updated when connection
705 attributes are added, removed, or replaced in ``__init__``. It is also
706 updated after ``__init__`` completes to reflect changes in `inputs`,
707 `prerequisiteInputs`, `outputs`, `initInputs`, and `initOutputs`.
708 """
710 _allConnections: dict[str, BaseConnection]
712 def __init__(self, *, config: PipelineTaskConfig | None = None):
713 pass
715 def __setattr__(self, name: str, value: Any) -> None:
716 if isinstance(value, BaseConnection):
717 previous = self._allConnections.get(name)
718 try:
719 getattr(self, value._connection_type_set).add(name)
720 except AttributeError:
721 # Attempt to call add on a frozenset, which is what these sets
722 # are after __init__ is done.
723 raise TypeError("Connections objects are frozen after construction.") from None
724 if previous is not None and value._connection_type_set != previous._connection_type_set:
725 # Connection has changed type, e.g. Input to PrerequisiteInput;
726 # update the sets accordingly. To be extra defensive about
727 # multiple assignments we use the type of the previous instance
728 # instead of assuming that's the same as the type of the self,
729 # which is just the default. Use discard instead of remove so
730 # manually removing from these sets first is never an error.
731 getattr(self, previous._connection_type_set).discard(name)
732 self._allConnections[name] = value
733 if hasattr(self.__class__, name):
734 # Don't actually set the attribute if this was a connection
735 # declared in the class; in that case we let the descriptor
736 # return the value we just added to allConnections.
737 return
738 # Actually add the attribute.
739 super().__setattr__(name, value)
741 def __delattr__(self, name):
742 """Descriptor delete method."""
743 previous = self._allConnections.get(name)
744 if previous is not None:
745 # Delete this connection's name from the appropriate set, which we
746 # have to get from the previous instance instead of assuming it's
747 # the same set that was appropriate for the class-level default.
748 # Use discard instead of remove so manually removing from these
749 # sets first is never an error.
750 try:
751 getattr(self, previous._connection_type_set).discard(name)
752 except AttributeError:
753 # Attempt to call discard on a frozenset, which is what these
754 # sets are after __init__ is done.
755 raise TypeError("Connections objects are frozen after construction.") from None
756 del self._allConnections[name]
757 if hasattr(self.__class__, name):
758 # Don't actually delete the attribute if this was a connection
759 # declared in the class; in that case we let the descriptor
760 # see that it's no longer present in allConnections.
761 return
762 # Actually delete the attribute.
763 super().__delattr__(name)
765 def buildDatasetRefs(
766 self, quantum: Quantum
767 ) -> tuple[InputQuantizedConnection, OutputQuantizedConnection]:
768 """Build `QuantizedConnection` corresponding to input
769 `~lsst.daf.butler.Quantum`.
771 Parameters
772 ----------
773 quantum : `lsst.daf.butler.Quantum`
774 Quantum object which defines the inputs and outputs for a given
775 unit of processing.
777 Returns
778 -------
779 retVal : `tuple` of (`InputQuantizedConnection`, \
780 `OutputQuantizedConnection`)
781 Namespaces mapping attribute names
782 (identifiers of connections) to butler references defined in the
783 input `~lsst.daf.butler.Quantum`.
784 """
785 inputDatasetRefs = InputQuantizedConnection()
786 outputDatasetRefs = OutputQuantizedConnection()
788 # populate inputDatasetRefs from quantum inputs
789 for attributeName in itertools.chain(self.inputs, self.prerequisiteInputs):
790 # get the attribute identified by name
791 attribute = getattr(self, attributeName)
792 # if the dataset is marked to load deferred, wrap it in a
793 # DeferredDatasetRef
794 quantumInputRefs: list[DatasetRef] | list[DeferredDatasetRef]
795 if attribute.deferLoad:
796 quantumInputRefs = [
797 DeferredDatasetRef(datasetRef=ref) for ref in quantum.inputs[attribute.name]
798 ]
799 else:
800 quantumInputRefs = list(quantum.inputs[attribute.name])
801 # Unpack arguments that are not marked multiples (list of
802 # length one)
803 if not attribute.multiple:
804 if len(quantumInputRefs) > 1:
805 raise ScalarError(
806 "Received multiple datasets "
807 f"{', '.join(str(r.dataId) for r in quantumInputRefs)} "
808 f"for scalar connection {attributeName} "
809 f"({quantumInputRefs[0].datasetType.name}) "
810 f"of quantum for {quantum.taskName} with data ID {quantum.dataId}."
811 )
812 if len(quantumInputRefs) == 0:
813 continue
814 setattr(inputDatasetRefs, attributeName, quantumInputRefs[0])
815 else:
816 # Add to the QuantizedConnection identifier
817 setattr(inputDatasetRefs, attributeName, quantumInputRefs)
819 # populate outputDatasetRefs from quantum outputs
820 for attributeName in self.outputs:
821 # get the attribute identified by name
822 attribute = getattr(self, attributeName)
823 value = quantum.outputs[attribute.name]
824 # Unpack arguments that are not marked multiples (list of
825 # length one)
826 if not attribute.multiple:
827 setattr(outputDatasetRefs, attributeName, value[0])
828 else:
829 setattr(outputDatasetRefs, attributeName, value)
831 return inputDatasetRefs, outputDatasetRefs
833 def adjustQuantum(
834 self,
835 inputs: dict[str, tuple[BaseInput, Collection[DatasetRef]]],
836 outputs: dict[str, tuple[Output, Collection[DatasetRef]]],
837 label: str,
838 data_id: DataCoordinate,
839 ) -> tuple[
840 Mapping[str, tuple[BaseInput, Collection[DatasetRef]]],
841 Mapping[str, tuple[Output, Collection[DatasetRef]]],
842 ]:
843 """Override to make adjustments to `lsst.daf.butler.DatasetRef` objects
844 in the `lsst.daf.butler.Quantum` during the graph generation stage
845 of the activator.
847 Parameters
848 ----------
849 inputs : `dict`
850 Dictionary whose keys are an input (regular or prerequisite)
851 connection name and whose values are a tuple of the connection
852 instance and a collection of associated
853 `~lsst.daf.butler.DatasetRef` objects.
854 The exact type of the nested collections is unspecified; it can be
855 assumed to be multi-pass iterable and support `len` and ``in``, but
856 it should not be mutated in place. In contrast, the outer
857 dictionaries are guaranteed to be temporary copies that are true
858 `dict` instances, and hence may be modified and even returned; this
859 is especially useful for delegating to `super` (see notes below).
860 outputs : `~collections.abc.Mapping`
861 Mapping of output datasets, with the same structure as ``inputs``.
862 label : `str`
863 Label for this task in the pipeline (should be used in all
864 diagnostic messages).
865 data_id : `lsst.daf.butler.DataCoordinate`
866 Data ID for this quantum in the pipeline (should be used in all
867 diagnostic messages).
869 Returns
870 -------
871 adjusted_inputs : `~collections.abc.Mapping`
872 Mapping of the same form as ``inputs`` with updated containers of
873 input `~lsst.daf.butler.DatasetRef` objects. Connections that are
874 not changed should not be returned at all. Datasets may only be
875 removed, not added. Nested collections may be of any multi-pass
876 iterable type, and the order of iteration will set the order of
877 iteration within `PipelineTask.runQuantum`.
878 adjusted_outputs : `~collections.abc.Mapping`
879 Mapping of updated output datasets, with the same structure and
880 interpretation as ``adjusted_inputs``.
882 Raises
883 ------
884 ScalarError
885 Raised if any `Input` or `PrerequisiteInput` connection has
886 ``multiple`` set to `False`, but multiple datasets.
887 NoWorkFound
888 Raised to indicate that this quantum should not be run; not enough
889 datasets were found for a regular `Input` connection, and the
890 quantum should be pruned or skipped.
891 FileNotFoundError
892 Raised to cause QuantumGraph generation to fail (with the message
893 included in this exception); not enough datasets were found for a
894 `PrerequisiteInput` connection.
896 Notes
897 -----
898 The base class implementation performs important checks. It always
899 returns an empty mapping (i.e. makes no adjustments). It should
900 always called be via `super` by custom implementations, ideally at the
901 end of the custom implementation with already-adjusted mappings when
902 any datasets are actually dropped, e.g.:
904 .. code-block:: python
906 def adjustQuantum(self, inputs, outputs, label, data_id):
907 # Filter out some dataset refs for one connection.
908 connection, old_refs = inputs["my_input"]
909 new_refs = [ref for ref in old_refs if ...]
910 adjusted_inputs = {"my_input", (connection, new_refs)}
911 # Update the original inputs so we can pass them to super.
912 inputs.update(adjusted_inputs)
913 # Can ignore outputs from super because they are guaranteed
914 # to be empty.
915 super().adjustQuantum(inputs, outputs, label_data_id)
916 # Return only the connections we modified.
917 return adjusted_inputs, {}
919 Removing outputs here is guaranteed to affect what is actually
920 passed to `PipelineTask.runQuantum`, but its effect on the larger
921 graph may be deferred to execution, depending on the context in
922 which `adjustQuantum` is being run: if one quantum removes an output
923 that is needed by a second quantum as input, the second quantum may not
924 be adjusted (and hence pruned or skipped) until that output is actually
925 found to be missing at execution time.
927 Tasks that desire zip-iteration consistency between any combinations of
928 connections that have the same data ID should generally implement
929 `adjustQuantum` to achieve this, even if they could also run that
930 logic during execution; this allows the system to see outputs that will
931 not be produced because the corresponding input is missing as early as
932 possible.
933 """
934 for name, (input_connection, refs) in inputs.items():
935 dataset_type_name = input_connection.name
936 if not input_connection.multiple and len(refs) > 1:
937 raise ScalarError(
938 f"Found multiple datasets {', '.join(str(r.dataId) for r in refs)} "
939 f"for non-multiple input connection {label}.{name} ({dataset_type_name}) "
940 f"for quantum data ID {data_id}."
941 )
942 if len(refs) < input_connection.minimum:
943 if isinstance(input_connection, PrerequisiteInput):
944 # This branch should only be possible during QG generation,
945 # or if someone deleted the dataset between making the QG
946 # and trying to run it. Either one should be a hard error.
947 raise FileNotFoundError(
948 f"Not enough datasets ({len(refs)}) found for non-optional connection {label}.{name} "
949 f"({dataset_type_name}) with minimum={input_connection.minimum} for quantum data ID "
950 f"{data_id}."
951 )
952 else:
953 raise NoWorkFound(label, name, input_connection)
954 for name, (output_connection, refs) in outputs.items():
955 dataset_type_name = output_connection.name
956 if not output_connection.multiple and len(refs) > 1:
957 raise ScalarError(
958 f"Found multiple datasets {', '.join(str(r.dataId) for r in refs)} "
959 f"for non-multiple output connection {label}.{name} ({dataset_type_name}) "
960 f"for quantum data ID {data_id}."
961 )
962 return {}, {}
964 def getSpatialBoundsConnections(self) -> Iterable[str]:
965 """Return the names of regular input and output connections whose data
966 IDs should be used to compute the spatial bounds of this task's quanta.
968 The spatial bound for a quantum is defined as the union of the regions
969 of all data IDs of all connections returned here, along with the region
970 of the quantum data ID (if the task has spatial dimensions).
972 Returns
973 -------
974 connection_names : `collections.abc.Iterable` [ `str` ]
975 Names of collections with spatial dimensions. These are the
976 task-internal connection names, not butler dataset type names.
978 Notes
979 -----
980 The spatial bound is used to search for prerequisite inputs that have
981 skypix dimensions. The default implementation returns an empty
982 iterable, which is usually sufficient for tasks with spatial
983 dimensions, but if a task's inputs or outputs are associated with
984 spatial regions that extend beyond the quantum data ID's region, this
985 method may need to be overridden to expand the set of prerequisite
986 inputs found.
988 Tasks that do not have spatial dimensions that have skypix prerequisite
989 inputs should always override this method, as the default spatial
990 bounds otherwise cover the full sky.
991 """
992 return ()
994 def getTemporalBoundsConnections(self) -> Iterable[str]:
995 """Return the names of regular input and output connections whose data
996 IDs should be used to compute the temporal bounds of this task's
997 quanta.
999 The temporal bound for a quantum is defined as the union of the
1000 timespans of all data IDs of all connections returned here, along with
1001 the timespan of the quantum data ID (if the task has temporal
1002 dimensions).
1004 Returns
1005 -------
1006 connection_names : `collections.abc.Iterable` [ `str` ]
1007 Names of collections with temporal dimensions. These are the
1008 task-internal connection names, not butler dataset type names.
1010 Notes
1011 -----
1012 The temporal bound is used to search for prerequisite inputs that are
1013 calibration datasets. The default implementation returns an empty
1014 iterable, which is usually sufficient for tasks with temporal
1015 dimensions, but if a task's inputs or outputs are associated with
1016 timespans that extend beyond the quantum data ID's timespan, this
1017 method may need to be overridden to expand the set of prerequisite
1018 inputs found.
1020 Tasks that do not have temporal dimensions that do not implement this
1021 method will use an infinite timespan for any calibration lookups.
1022 """
1023 return ()
1025 def adjust_all_quanta(self, adjuster: QuantaAdjuster) -> None:
1026 """Customize the set of quanta predicted for this task during quantum
1027 graph generation.
1029 Parameters
1030 ----------
1031 adjuster : `QuantaAdjuster`
1032 A helper object that implementations can use to modify the
1033 under-construction quantum graph.
1035 Notes
1036 -----
1037 This hook is called before `adjustQuantum`, which is where built-in
1038 checks for `NoWorkFound` cases and missing prerequisites are handled.
1039 This means that the set of preliminary quanta seen by this method could
1040 include some that would normally be dropped later.
1041 """
1042 pass
1045def iterConnections(
1046 connections: PipelineTaskConnections, connectionType: str | Iterable[str]
1047) -> Generator[BaseConnection, None, None]:
1048 """Create an iterator over the selected connections type which yields
1049 all the defined connections of that type.
1051 Parameters
1052 ----------
1053 connections : `PipelineTaskConnections`
1054 An instance of a `PipelineTaskConnections` object that will be iterated
1055 over.
1056 connectionType : `str`
1057 The type of connections to iterate over, valid values are inputs,
1058 outputs, prerequisiteInputs, initInputs, initOutputs.
1060 Yields
1061 ------
1062 connection: `~.connectionTypes.BaseConnection`
1063 A connection defined on the input connections object of the type
1064 supplied. The yielded value Will be an derived type of
1065 `~.connectionTypes.BaseConnection`.
1066 """
1067 if isinstance(connectionType, str):
1068 connectionType = (connectionType,)
1069 for name in itertools.chain.from_iterable(getattr(connections, ct) for ct in connectionType):
1070 yield getattr(connections, name)
1073@dataclass
1074class AdjustQuantumHelper:
1075 """Helper class for calling `PipelineTaskConnections.adjustQuantum`.
1077 This class holds `inputs` and `outputs` mappings in the form used by
1078 `lsst.daf.butler.Quantum` and execution harness code, i.e. with
1079 `~lsst.daf.butler.DatasetType` keys, translating them to and from the
1080 connection-oriented mappings used inside `PipelineTaskConnections`.
1081 """
1083 inputs: NamedKeyMapping[DatasetType, Sequence[DatasetRef]]
1084 """Mapping of regular input and prerequisite input datasets, grouped by
1085 `~lsst.daf.butler.DatasetType`.
1086 """
1088 outputs: NamedKeyMapping[DatasetType, Sequence[DatasetRef]]
1089 """Mapping of output datasets, grouped by `~lsst.daf.butler.DatasetType`.
1090 """
1092 inputs_adjusted: bool = False
1093 """Whether any inputs were removed in the last call to `adjust_in_place`.
1094 """
1096 outputs_adjusted: bool = False
1097 """Whether any outputs were removed in the last call to `adjust_in_place`.
1098 """
1100 def adjust_in_place(
1101 self,
1102 connections: PipelineTaskConnections,
1103 label: str,
1104 data_id: DataCoordinate,
1105 ) -> None:
1106 """Call `~PipelineTaskConnections.adjustQuantum` and update ``self``
1107 with its results.
1109 Parameters
1110 ----------
1111 connections : `PipelineTaskConnections`
1112 Instance on which to call `~PipelineTaskConnections.adjustQuantum`.
1113 label : `str`
1114 Label for this task in the pipeline (should be used in all
1115 diagnostic messages).
1116 data_id : `lsst.daf.butler.DataCoordinate`
1117 Data ID for this quantum in the pipeline (should be used in all
1118 diagnostic messages).
1119 """
1120 # Translate self's DatasetType-keyed, Quantum-oriented mappings into
1121 # connection-keyed, PipelineTask-oriented mappings.
1122 inputs_by_connection: dict[str, tuple[BaseInput, tuple[DatasetRef, ...]]] = {}
1123 outputs_by_connection: dict[str, tuple[Output, tuple[DatasetRef, ...]]] = {}
1124 for name in itertools.chain(connections.inputs, connections.prerequisiteInputs):
1125 connection = getattr(connections, name)
1126 dataset_type_name = connection.name
1127 inputs_by_connection[name] = (connection, tuple(self.inputs.get(dataset_type_name, ())))
1128 for name in itertools.chain(connections.outputs):
1129 connection = getattr(connections, name)
1130 dataset_type_name = connection.name
1131 outputs_by_connection[name] = (connection, tuple(self.outputs.get(dataset_type_name, ())))
1132 # Actually call adjustQuantum.
1133 # MyPy correctly complains that this call is not quite legal, but the
1134 # method docs explain exactly what's expected and it's the behavior we
1135 # want. It'd be nice to avoid this if we ever have to change the
1136 # interface anyway, but not an immediate problem.
1137 adjusted_inputs_by_connection, adjusted_outputs_by_connection = connections.adjustQuantum(
1138 inputs_by_connection, # type: ignore
1139 outputs_by_connection, # type: ignore
1140 label,
1141 data_id,
1142 )
1143 # Translate adjustments to DatasetType-keyed, Quantum-oriented form,
1144 # installing new mappings in self if necessary.
1145 if adjusted_inputs_by_connection:
1146 adjusted_inputs = NamedKeyDict[DatasetType, tuple[DatasetRef, ...]](self.inputs)
1147 for name, (connection, updated_refs) in adjusted_inputs_by_connection.items():
1148 dataset_type_name = connection.name
1149 if not set(updated_refs).issubset(self.inputs[dataset_type_name]):
1150 raise RuntimeError(
1151 f"adjustQuantum implementation for task with label {label} returned {name} "
1152 f"({dataset_type_name}) input datasets that are not a subset of those "
1153 f"it was given for data ID {data_id}."
1154 )
1155 adjusted_inputs[dataset_type_name] = tuple(updated_refs)
1156 self.inputs = adjusted_inputs.freeze()
1157 self.inputs_adjusted = True
1158 else:
1159 self.inputs_adjusted = False
1160 if adjusted_outputs_by_connection:
1161 adjusted_outputs = NamedKeyDict[DatasetType, tuple[DatasetRef, ...]](self.outputs)
1162 for name, (connection, updated_refs) in adjusted_outputs_by_connection.items():
1163 dataset_type_name = connection.name
1164 if not set(updated_refs).issubset(self.outputs[dataset_type_name]):
1165 raise RuntimeError(
1166 f"adjustQuantum implementation for task with label {label} returned {name} "
1167 f"({dataset_type_name}) output datasets that are not a subset of those "
1168 f"it was given for data ID {data_id}."
1169 )
1170 adjusted_outputs[dataset_type_name] = tuple(updated_refs)
1171 self.outputs = adjusted_outputs.freeze()
1172 self.outputs_adjusted = True
1173 else:
1174 self.outputs_adjusted = False
1177class QuantaAdjuster:
1178 """A helper class for the `PipelineTaskConnections.adjust_all_quanta` hook.
1180 Parameters
1181 ----------
1182 task_label : `str`
1183 Label of the task whose quanta are being adjusted.
1184 pipeline_graph : `pipeline_graph.PipelineGraph`
1185 Pipeline graph the quantum graph is being built from.
1186 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton`
1187 Under-construction quantum graph that will be modified in place.
1188 butler : `lsst.daf.butler.Butler`
1189 Read-only instance with its default collection search path set to the
1190 input collections passed to the quantum-graph builder.
1191 """
1193 def __init__(
1194 self, task_label: str, pipeline_graph: PipelineGraph, skeleton: QuantumGraphSkeleton, butler: Butler
1195 ):
1196 self._task_node = pipeline_graph.tasks[task_label]
1197 self._pipeline_graph = pipeline_graph
1198 self._skeleton = skeleton
1199 self._n_removed = 0
1200 self._butler = butler
1202 @property
1203 def task_label(self) -> str:
1204 """The label this task has been configured with."""
1205 return self._task_node.label
1207 @property
1208 def task_node(self) -> TaskNode:
1209 """The node for this task in the pipeline graph."""
1210 return self._task_node
1212 @property
1213 def butler(self) -> Butler:
1214 """Read-only instance with its default collection search path set to
1215 the input collections passed to the quantum-graph builder.
1216 """
1217 return self._butler
1219 def iter_data_ids(self) -> Iterator[DataCoordinate]:
1220 """Iterate over the data IDs of all quanta for this task.
1222 Returns
1223 -------
1224 data_ids : `~collections.abc.Iterator` [ \
1225 `~lsst.daf.butler.DataCoordinate` ]
1226 Data IDs. These are minimal data IDs without dimension records or
1227 implied values; use `expand_quantum_data_id` to get a full data ID
1228 when needed.
1229 """
1230 for key in self._skeleton.get_quanta(self._task_node.label):
1231 yield DataCoordinate.from_required_values(self._task_node.dimensions, key.data_id_values)
1233 def remove_quantum(self, data_id: DataCoordinate) -> None:
1234 """Remove a quantum from the graph.
1236 Parameters
1237 ----------
1238 data_id : `~lsst.daf.butler.DataCoordinate`
1239 Data ID of the quantum to remove. All outputs will be removed as
1240 well.
1241 """
1242 from .quantum_graph_skeleton import QuantumKey
1244 self._skeleton.remove_quantum_node(
1245 QuantumKey(self._task_node.label, data_id.required_values), remove_outputs=True
1246 )
1247 self._n_removed += 1
1249 def get_inputs(self, quantum_data_id: DataCoordinate) -> dict[str, list[DataCoordinate]]:
1250 """Return the data IDs of all regular inputs to a quantum.
1252 Parameters
1253 ----------
1254 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1255 Data ID of the quantum to get the inputs of.
1257 Returns
1258 -------
1259 inputs : `dict` [ `str`, `list` [ `~lsst.daf.butler.DataCoordinate` ] ]
1260 Data IDs of inputs, keyed by the connection name (the internal task
1261 name, not the dataset type name). This only contains regular
1262 inputs, not init-inputs or prerequisite inputs.
1264 Notes
1265 -----
1266 If two connections have the same dataset type, the current
1267 implementation assumes the set of datasets is the same for the two
1268 connections. This limitation may be removed in the future.
1269 """
1270 from .quantum_graph_skeleton import DatasetKey, QuantumKey
1272 by_dataset_type_name: defaultdict[str, list[DataCoordinate]] = defaultdict(list)
1273 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1274 for dataset_key in self._skeleton.iter_inputs_of(quantum_key):
1275 if not isinstance(dataset_key, DatasetKey):
1276 continue
1277 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
1278 by_dataset_type_name[dataset_key.parent_dataset_type_name].append(
1279 DataCoordinate.from_required_values(dataset_type_node.dimensions, dataset_key.data_id_values)
1280 )
1281 return {
1282 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name]
1283 for edge in self._task_node.iter_all_inputs()
1284 }
1286 def get_prerequisite_inputs(
1287 self,
1288 quantum_data_id: DataCoordinate,
1289 ) -> dict[str, dict[uuid.UUID, DataCoordinate]]:
1290 """Return the data IDs of all prerequisite inputs to a quantum.
1292 Parameters
1293 ----------
1294 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1295 Data ID of the quantum to get the inputs of.
1297 Returns
1298 -------
1299 inputs : `dict` [ `str`, \
1300 `dict` [ `uuid.UUID`, `~lsst.daf.butler.DataCoordinate` ] ]
1301 Dataset IDs and and data IDs of prerequisite inputs, keyed by the
1302 connection name (the internal task name, not the dataset type
1303 name). This only contains prerequisite inputs, not init-inputs or
1304 regular inputs.
1306 Notes
1307 -----
1308 If two connections have the same dataset type, the current
1309 implementation assumes the set of datasets is the same for the two
1310 connections. This limitation may be removed in the future.
1312 Unlike regular inputs, prerequisite inputs are not looked up from input
1313 collections or indexed by data ID. Instead, they are uniquely
1314 identified by dataset UUID and reused directly between quanta.
1315 """
1316 from .quantum_graph_skeleton import PrerequisiteDatasetKey, QuantumKey
1318 by_dataset_type_name: defaultdict[str, dict[uuid.UUID, DataCoordinate]] = defaultdict(dict)
1319 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1321 for dataset_key in self._skeleton.iter_inputs_of(quantum_key):
1322 if not isinstance(dataset_key, PrerequisiteDatasetKey):
1323 continue
1324 by_dataset_type_name[dataset_key.parent_dataset_type_name][
1325 uuid.UUID(bytes=dataset_key.dataset_id_bytes)
1326 ] = self._skeleton.get_data_id(dataset_key)
1327 return {
1328 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name]
1329 for edge in self._task_node.iter_all_inputs()
1330 }
1332 def get_outputs(self, quantum_data_id: DataCoordinate) -> dict[str, list[DataCoordinate]]:
1333 """Return the data IDs of all regular outputs to a quantum.
1335 Parameters
1336 ----------
1337 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1338 Data ID of the quantum to get the outputs of.
1340 Returns
1341 -------
1342 inputs : `dict` [ `str`, `list` [ `~lsst.daf.butler.DataCoordinate` ] ]
1343 Data IDs of inputs, keyed by the connection name (the internal task
1344 name, not the dataset type name). This only contains regular
1345 outputs, not init-outputs or log or metadata outputs.
1347 Notes
1348 -----
1349 If two connections have the same dataset type, the current
1350 implementation assumes the set of datasets is the same for the two
1351 connections. This limitation may be removed in the future.
1352 """
1353 from .quantum_graph_skeleton import QuantumKey
1355 by_dataset_type_name: defaultdict[str, list[DataCoordinate]] = defaultdict(list)
1356 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1357 for dataset_key in self._skeleton.iter_outputs_of(quantum_key):
1358 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name]
1359 by_dataset_type_name[dataset_key.parent_dataset_type_name].append(
1360 DataCoordinate.from_required_values(dataset_type_node.dimensions, dataset_key.data_id_values)
1361 )
1362 return {
1363 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name]
1364 for edge in self._task_node.outputs.values()
1365 }
1367 def add_input(
1368 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_data_id: DataCoordinate
1369 ) -> None:
1370 """Add a new input to a quantum.
1372 Parameters
1373 ----------
1374 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1375 Data ID of the quantum to add an input to.
1376 connection_name : `str`
1377 Name of the connection (the task-internal name, not the butler
1378 dataset type name).
1379 dataset_data_id : `~lsst.daf.butler.DataCoordinate`
1380 Data ID of the input dataset. Must already exist in the graph
1381 as an input to a different quantum of this task, and must be a
1382 regular input, not a prerequisite input or init-input.
1384 Notes
1385 -----
1386 If two connections have the same dataset type, the current
1387 implementation assumes the set of datasets is the same for the two
1388 connections. This limitation may be removed in the future.
1389 """
1390 from .quantum_graph_skeleton import DatasetKey, QuantumKey
1392 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1393 read_edge = self._task_node.inputs[connection_name]
1394 dataset_key = DatasetKey(read_edge.parent_dataset_type_name, dataset_data_id.required_values)
1395 if dataset_key not in self._skeleton:
1396 raise LookupError(
1397 f"Dataset {read_edge.parent_dataset_type_name}@{dataset_data_id} is not already in the graph."
1398 )
1399 self._skeleton.add_input_edge(quantum_key, dataset_key)
1401 def add_prerequisite_input(
1402 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_uuid: uuid.UUID
1403 ) -> None:
1404 """Add a new prerequisite input to a quantum.
1406 Parameters
1407 ----------
1408 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1409 Data ID of the quantum to add an input to.
1410 connection_name : `str`
1411 Name of the connection (the task-internal name, not the butler
1412 dataset type name).
1413 dataset_uuid : `uuid.UUID`
1414 UUID of the prerequisite input dataset. Must already exist in the
1415 graph as an input to a different quantum of this task, and must be
1416 a prerequisite input, not a regular input or init-input.
1418 Notes
1419 -----
1420 If two connections have the same dataset type, the current
1421 implementation assumes the set of datasets is the same for the two
1422 connections. This limitation may be removed in the future.
1424 Unlike regular inputs, prerequisite inputs are not looked up from input
1425 collections or indexed by data ID. Instead, they are uniquely
1426 identified by dataset UUID and reused directly between quanta.
1427 """
1428 from .quantum_graph_skeleton import PrerequisiteDatasetKey, QuantumKey
1430 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1431 read_edge = self._task_node.prerequisite_inputs[connection_name]
1432 dataset_key = PrerequisiteDatasetKey(read_edge.parent_dataset_type_name, dataset_uuid.bytes)
1433 if dataset_key not in self._skeleton:
1434 raise LookupError(
1435 f"Prerequisite Dataset {read_edge.parent_dataset_type_name}@{dataset_uuid} "
1436 "is not already in the graph."
1437 )
1438 self._skeleton.add_input_edge(quantum_key, dataset_key)
1440 def move_output(
1441 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_data_id: DataCoordinate
1442 ) -> None:
1443 """Remove an output of one quantum and make it a new output of another.
1445 Parameters
1446 ----------
1447 quantum_data_id : `~lsst.daf.butler.DataCoordinate`
1448 Data ID of the quantum to move the output to.
1449 connection_name : `str`
1450 Name of the connection (the task-internal name, not the butler
1451 dataset type name).
1452 dataset_data_id : `~lsst.daf.butler.DataCoordinate`
1453 Data ID of the output dataset. Must already exist in the graph
1454 as an output of a different quantum of this task.
1455 """
1456 from .quantum_graph_skeleton import DatasetKey, QuantumKey
1458 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values)
1459 write_edge = self._task_node.outputs[connection_name]
1460 dataset_key = DatasetKey(write_edge.parent_dataset_type_name, dataset_data_id.required_values)
1461 if dataset_key not in self._skeleton:
1462 raise LookupError(
1463 f"Dataset {write_edge.parent_dataset_type_name}@{dataset_data_id} is "
1464 "not already in the graph."
1465 )
1466 self._skeleton.remove_output_edge(dataset_key)
1467 self._skeleton.add_output_edge(quantum_key, dataset_key)
1469 def expand_quantum_data_id(self, data_id: DataCoordinate) -> DataCoordinate:
1470 """Expand a quantum data ID to include implied values and dimension
1471 records.
1473 Parameters
1474 ----------
1475 data_id : `~lsst.daf.butler.DataCoordinate`
1476 A data ID of a quantum already in the graph.
1478 Returns
1479 -------
1480 expanded_data_id : `~lsst.daf.butler.DataCoordinate`
1481 The same data ID, with implied values included and dimension
1482 records attached.
1483 """
1484 from .quantum_graph_skeleton import QuantumKey
1486 return self._skeleton.get_data_id(QuantumKey(self._task_node.label, data_id.required_values))
1488 @property
1489 def n_removed(self) -> int:
1490 """The number of quanta that have been removed by this helper."""
1491 return self._n_removed