Coverage for python/lsst/ctrl/mpexec/execFixupDataId.py: 22%

41 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-18 09:41 +0000

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/>. 

27 

28__all__ = ["ExecutionGraphFixup"] 

29 

30import contextlib 

31import itertools 

32from collections import defaultdict 

33from collections.abc import Sequence 

34from typing import Any 

35 

36import networkx as nx 

37from lsst.pipe.base import QuantumGraph, QuantumNode 

38 

39from .executionGraphFixup import ExecutionGraphFixup 

40 

41 

42class ExecFixupDataId(ExecutionGraphFixup): 

43 """Implementation of ExecutionGraphFixup for ordering of tasks based 

44 on DataId values. 

45 

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.:: 

49 

50 # lsst/ap/verify/ci_fixup.py 

51 

52 from lsst.ctrl.mpexec.execFixupDataId import ExecFixupDataId 

53 

54 def assoc_fixup(): 

55 return ExecFixupDataId(taskLabel="ap_assoc", 

56 dimensions=("visit", "detector")) 

57 

58 

59 and then executing pipetask:: 

60 

61 pipetask run --graph-fixup=lsst.ap.verify.ci_fixup.assoc_fixup ... 

62 

63 This will add new dependencies between quanta executed by the task with 

64 label "ap_assoc". Quanta with higher visit number will depend on quanta 

65 with lower visit number and their execution will wait until lower visit 

66 number finishes. 

67 

68 Parameters 

69 ---------- 

70 taskLabel : `str` 

71 The label of the task for which to add dependencies. 

72 dimensions : `str` or sequence [`str`] 

73 One or more dimension names, quanta execution will be ordered 

74 according to values of these dimensions. 

75 reverse : `bool`, optional 

76 If `False` (default) then quanta with higher values of dimensions 

77 will be executed after quanta with lower values, otherwise the order 

78 is reversed. 

79 """ 

80 

81 def __init__(self, taskLabel: str, dimensions: str | Sequence[str], reverse: bool = False): 

82 self.taskLabel = taskLabel 

83 self.dimensions = dimensions 

84 self.reverse = reverse 

85 if isinstance(self.dimensions, str): 

86 self.dimensions = (self.dimensions,) 

87 else: 

88 self.dimensions = tuple(self.dimensions) 

89 

90 def _key(self, qnode: QuantumNode) -> tuple[Any, ...]: 

91 """Produce comparison key for quantum data. 

92 

93 Parameters 

94 ---------- 

95 qnode : `QuantumNode` 

96 An individual node in a `~lsst.pipe.base.QuantumGraph` 

97 

98 Returns 

99 ------- 

100 key : `tuple` 

101 """ 

102 dataId = qnode.quantum.dataId 

103 assert dataId is not None, "Quantum DataId cannot be None" 

104 key = tuple(dataId[dim] for dim in self.dimensions) 

105 return key 

106 

107 def fixupQuanta(self, graph: QuantumGraph) -> QuantumGraph: 

108 taskDef = graph.findTaskDefByLabel(self.taskLabel) 

109 if taskDef is None: 

110 raise ValueError(f"Cannot find task with label {self.taskLabel}") 

111 quanta = list(graph.getNodesForTask(taskDef)) 

112 keyQuanta = defaultdict(list) 

113 for q in quanta: 

114 key = self._key(q) 

115 keyQuanta[key].append(q) 

116 keys = sorted(keyQuanta.keys(), reverse=self.reverse) 

117 networkGraph = graph.graph 

118 

119 for prev_key, key in itertools.pairwise(keys): 

120 for prev_node in keyQuanta[prev_key]: 

121 for node in keyQuanta[key]: 

122 # remove any existing edges between the two nodes, but 

123 # don't fail if there are not any. Both directions need 

124 # tried because in a directed graph, order maters 

125 for edge in ((node, prev_node), (prev_node, node)): 

126 with contextlib.suppress(nx.NetworkXException): 

127 networkGraph.remove_edge(*edge) 

128 

129 networkGraph.add_edge(prev_node, node) 

130 return graph