Coverage for python / lsst / pipe / base / quantum_graph / aggregator / _supervisor.py: 24%

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

27 

28from __future__ import annotations 

29 

30__all__ = ("aggregate_graph",) 

31 

32import dataclasses 

33import itertools 

34import uuid 

35 

36import astropy.units as u 

37import networkx 

38 

39from lsst.utils.logging import getLogger 

40from lsst.utils.usage import get_peak_mem_usage 

41 

42from ...graph_walker import GraphWalker 

43from ...pipeline_graph import TaskImportMode 

44from .._predicted import PredictedQuantumGraphComponents, PredictedQuantumGraphReader 

45from ._communicators import ( 

46 IngesterCommunicator, 

47 ScannerCommunicator, 

48 SpawnProcessContext, 

49 SupervisorCommunicator, 

50 ThreadingContext, 

51 Worker, 

52 WriterCommunicator, 

53) 

54from ._config import AggregatorConfig 

55from ._ingester import Ingester 

56from ._scanner import Scanner 

57from ._structs import ScanReport, ScanStatus, WriteRequest 

58from ._writer import Writer 

59 

60 

61@dataclasses.dataclass 

62class Supervisor: 

63 """The main process/thread for the provenance aggregator.""" 

64 

65 predicted_path: str 

66 """Path to the predicted quantum graph.""" 

67 

68 comms: SupervisorCommunicator 

69 """Communicator object for the supervisor.""" 

70 

71 predicted: PredictedQuantumGraphComponents = dataclasses.field(init=False) 

72 """Components of the predicted quantum graph.""" 

73 

74 walker: GraphWalker[uuid.UUID] = dataclasses.field(init=False) 

75 """Iterator that traverses the quantum graph.""" 

76 

77 n_abandoned: int = 0 

78 """Number of quanta we abandoned because they did not complete in time and 

79 we could not assume they had failed. 

80 """ 

81 

82 def __post_init__(self) -> None: 

83 self.comms.progress.log.info("Reading predicted quantum graph.") 

84 with PredictedQuantumGraphReader.open( 

85 self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT 

86 ) as reader: 

87 reader.read_thin_graph() 

88 reader.read_init_quanta() 

89 self.predicted = reader.components 

90 self.comms.progress.log.info("Analyzing predicted graph.") 

91 xgraph = networkx.DiGraph(self.predicted.thin_graph.edges) 

92 # Make sure all quanta are in the graph, even if they don't have any 

93 # quantum-only edges. 

94 for thin_quantum in itertools.chain.from_iterable(self.predicted.thin_graph.quanta.values()): 

95 xgraph.add_node(thin_quantum.quantum_id) 

96 # Add init quanta as nodes without edges, because the scanner should 

97 # only be run after init outputs are all written and hence we don't 

98 # care when we process them. 

99 for init_quantum in self.predicted.init_quanta.root: 

100 xgraph.add_node(init_quantum.quantum_id) 

101 self.walker = GraphWalker(xgraph) 

102 

103 def loop(self) -> None: 

104 """Scan the outputs of the quantum graph to gather provenance and 

105 ingest outputs. 

106 """ 

107 n_quanta = self.predicted.header.n_quanta + len(self.predicted.init_quanta.root) 

108 self.comms.progress.scans.total = n_quanta 

109 self.comms.progress.writes.total = n_quanta 

110 self.comms.progress.quantum_ingests.total = n_quanta 

111 ready_set: set[uuid.UUID] = set() 

112 for ready_quanta in self.walker: 

113 self.comms.log.debug("Sending %d new quanta to scan queue.", len(ready_quanta)) 

114 ready_set.update(ready_quanta) 

115 while ready_set: 

116 self.comms.request_scan(ready_set.pop()) 

117 for scan_return in self.comms.poll(): 

118 self.handle_report(scan_return) 

119 

120 def handle_report(self, scan_report: ScanReport) -> None: 

121 """Handle a report from a scanner. 

122 

123 Parameters 

124 ---------- 

125 scan_report : `ScanReport` 

126 Information about the scan. 

127 """ 

128 match scan_report.status: 

129 case ScanStatus.SUCCESSFUL | ScanStatus.INIT: 

130 self.comms.log.debug("Scan complete for %s: quantum succeeded.", scan_report.quantum_id) 

131 self.walker.finish(scan_report.quantum_id) 

132 case ScanStatus.FAILED: 

133 self.comms.log.debug("Scan complete for %s: quantum failed.", scan_report.quantum_id) 

134 blocked_quanta = self.walker.fail(scan_report.quantum_id) 

135 for blocked_quantum_id in blocked_quanta: 

136 if self.comms.config.output_path is not None: 

137 self.comms.request_write(WriteRequest(blocked_quantum_id, status=ScanStatus.BLOCKED)) 

138 self.comms.progress.scans.update(1) 

139 self.comms.progress.quantum_ingests.update(len(blocked_quanta)) 

140 case ScanStatus.ABANDONED: 

141 self.comms.log.debug("Abandoning scan for %s: quantum has not succeeded (yet).") 

142 self.walker.fail(scan_report.quantum_id) 

143 self.n_abandoned += 1 

144 case unexpected: 

145 raise AssertionError( 

146 f"Unexpected status {unexpected!r} in scanner loop for {scan_report.quantum_id}." 

147 ) 

148 self.comms.progress.scans.update(1) 

149 

150 

151def aggregate_graph(predicted_path: str, butler_path: str, config: AggregatorConfig) -> None: 

152 """Run the graph aggregator tool. 

153 

154 Parameters 

155 ---------- 

156 predicted_path : `str` 

157 Path to the predicted quantum graph. 

158 butler_path : `str` 

159 Path or alias to the central butler repository. 

160 config : `AggregatorConfig` 

161 Configuration for the aggregator. 

162 """ 

163 log = getLogger("lsst.pipe.base.quantum_graph.aggregator") 

164 ctx = ThreadingContext() if config.n_processes == 1 else SpawnProcessContext() 

165 scanners: list[Worker] = [] 

166 ingester: Worker 

167 writer: Worker | None = None 

168 with SupervisorCommunicator(log, config.n_processes, ctx, config) as comms: 

169 comms.progress.log.verbose("Starting workers.") 

170 if config.output_path is not None: 

171 writer_comms = WriterCommunicator(comms) 

172 writer = ctx.make_worker( 

173 target=Writer.run, 

174 args=(predicted_path, writer_comms), 

175 name=writer_comms.name, 

176 ) 

177 writer.start() 

178 for scanner_id in range(config.n_processes): 

179 scanner_comms = ScannerCommunicator(comms, scanner_id) 

180 worker = ctx.make_worker( 

181 target=Scanner.run, 

182 args=(predicted_path, butler_path, scanner_comms), 

183 name=scanner_comms.name, 

184 ) 

185 worker.start() 

186 scanners.append(worker) 

187 ingester_comms = IngesterCommunicator(comms) 

188 ingester = ctx.make_worker( 

189 target=Ingester.run, 

190 args=(predicted_path, butler_path, ingester_comms), 

191 name=ingester_comms.name, 

192 ) 

193 ingester.start() 

194 supervisor = Supervisor(predicted_path, comms) 

195 supervisor.loop() 

196 log.info( 

197 "Scanning complete after %0.1fs; waiting for workers to finish.", 

198 comms.progress.elapsed_time, 

199 ) 

200 comms.wait_for_workers_to_finish() 

201 if supervisor.n_abandoned: 

202 raise RuntimeError( 

203 f"{supervisor.n_abandoned} {'quanta' if supervisor.n_abandoned > 1 else 'quantum'} " 

204 "abandoned because they did not succeed. Re-run with assume_complete=True after all retry " 

205 "attempts have been exhausted." 

206 ) 

207 for w in scanners: 

208 w.join() 

209 ingester.join() 

210 if writer is not None and writer.is_alive(): 

211 log.info("Waiting for writer process to close (garbage collecting can be very slow).") 

212 writer.join() 

213 # We can't get memory usage for children until they've joined. 

214 parent_mem, child_mem = get_peak_mem_usage() 

215 # This is actually an upper bound on the peak (since the peaks could be 

216 # at different times), but since we expect memory usage to be more smooth 

217 # than spiky that's fine. 

218 total_mem: u.Quantity = parent_mem + child_mem 

219 log.info( 

220 "All aggregation tasks complete after %0.1fs; peak memory usage ≤ %0.1f MB.", 

221 comms.progress.elapsed_time, 

222 total_mem.to(u.MB).value, 

223 )