Coverage for python/lsst/ctrl/mpexec/execFixupDataId.py: 22%
41 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 02:56 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 02:56 -0700
1# This file is part of ctrl_mpexec.
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__all__ = ["ExecutionGraphFixup"]
30import contextlib
31import itertools
32from collections import defaultdict
33from collections.abc import Sequence
34from typing import Any
36import networkx as nx
37from lsst.pipe.base import QuantumGraph, QuantumNode
39from .executionGraphFixup import ExecutionGraphFixup
42class ExecFixupDataId(ExecutionGraphFixup):
43 """Implementation of ExecutionGraphFixup for ordering of tasks based
44 on DataId values.
46 This class is a trivial implementation mostly useful as an example,
47 though it can be used to make actual fixup instances by defining
48 a method that instantiates it, e.g.::
50 # lsst/ap/verify/ci_fixup.py
52 from lsst.ctrl.mpexec.execFixupDataId import ExecFixupDataId
54 def assoc_fixup():
55 return ExecFixupDataId(taskLabel="ap_assoc",
56 dimensions=("visit", "detector"))
58 and then executing pipetask::
60 pipetask run --graph-fixup=lsst.ap.verify.ci_fixup.assoc_fixup ...
62 This will add new dependencies between quanta executed by the task with
63 label "ap_assoc". Quanta with higher visit number will depend on quanta
64 with lower visit number and their execution will wait until lower visit
65 number finishes.
67 Parameters
68 ----------
69 taskLabel : `str`
70 The label of the task for which to add dependencies.
71 dimensions : `str` or sequence [`str`]
72 One or more dimension names, quanta execution will be ordered
73 according to values of these dimensions.
74 reverse : `bool`, optional
75 If `False` (default) then quanta with higher values of dimensions
76 will be executed after quanta with lower values, otherwise the order
77 is reversed.
78 """
80 def __init__(self, taskLabel: str, dimensions: str | Sequence[str], reverse: bool = False):
81 self.taskLabel = taskLabel
82 self.dimensions = dimensions
83 self.reverse = reverse
84 if isinstance(self.dimensions, str):
85 self.dimensions = (self.dimensions,)
86 else:
87 self.dimensions = tuple(self.dimensions)
89 def _key(self, qnode: QuantumNode) -> tuple[Any, ...]:
90 """Produce comparison key for quantum data.
92 Parameters
93 ----------
94 qnode : `QuantumNode`
95 An individual node in a `~lsst.pipe.base.QuantumGraph`
97 Returns
98 -------
99 key : `tuple`
100 """
101 dataId = qnode.quantum.dataId
102 assert dataId is not None, "Quantum DataId cannot be None"
103 key = tuple(dataId[dim] for dim in self.dimensions)
104 return key
106 def fixupQuanta(self, graph: QuantumGraph) -> QuantumGraph:
107 taskDef = graph.findTaskDefByLabel(self.taskLabel)
108 if taskDef is None:
109 raise ValueError(f"Cannot find task with label {self.taskLabel}")
110 quanta = list(graph.getNodesForTask(taskDef))
111 keyQuanta = defaultdict(list)
112 for q in quanta:
113 key = self._key(q)
114 keyQuanta[key].append(q)
115 keys = sorted(keyQuanta.keys(), reverse=self.reverse)
116 networkGraph = graph.graph
118 for prev_key, key in itertools.pairwise(keys):
119 for prev_node in keyQuanta[prev_key]:
120 for node in keyQuanta[key]:
121 # remove any existing edges between the two nodes, but
122 # don't fail if there are not any. Both directions need
123 # tried because in a directed graph, order maters
124 for edge in ((node, prev_node), (prev_node, node)):
125 with contextlib.suppress(nx.NetworkXException):
126 networkGraph.remove_edge(*edge)
128 networkGraph.add_edge(prev_node, node)
129 return graph