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

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 .._provenance import ProvenanceQuantumScanData, ProvenanceQuantumScanStatus 

46from ._communicators import ( 

47 IngesterCommunicator, 

48 ScannerCommunicator, 

49 SpawnProcessContext, 

50 SupervisorCommunicator, 

51 ThreadingContext, 

52 Worker, 

53 WriterCommunicator, 

54) 

55from ._config import AggregatorConfig 

56from ._ingester import Ingester 

57from ._scanner import Scanner 

58from ._structs import ScanReport 

59from ._writer import Writer 

60 

61 

62@dataclasses.dataclass 

63class Supervisor: 

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

65 

66 predicted_path: str 

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

68 

69 comms: SupervisorCommunicator 

70 """Communicator object for the supervisor.""" 

71 

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

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

74 

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

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

77 

78 n_abandoned: int = 0 

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

80 we could not assume they had failed. 

81 """ 

82 

83 def __post_init__(self) -> None: 

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

85 with PredictedQuantumGraphReader.open( 

86 self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT 

87 ) as reader: 

88 reader.read_thin_graph() 

89 reader.read_init_quanta() 

90 self.predicted = reader.components 

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

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

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

94 # quantum-only edges. 

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

96 xgraph.add_node(thin_quantum.quantum_id) 

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

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

99 # care when we process them. 

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

101 xgraph.add_node(init_quantum.quantum_id) 

102 self.walker = GraphWalker(xgraph) 

103 

104 def loop(self) -> None: 

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

106 ingest outputs. 

107 """ 

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

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

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

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

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

113 for ready_quanta in self.walker: 

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

115 ready_set.update(ready_quanta) 

116 while ready_set: 

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

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

119 self.handle_report(scan_return) 

120 

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

122 """Handle a report from a scanner. 

123 

124 Parameters 

125 ---------- 

126 scan_report : `ScanReport` 

127 Information about the scan. 

128 """ 

129 match scan_report.status: 

130 case ProvenanceQuantumScanStatus.SUCCESSFUL | ProvenanceQuantumScanStatus.INIT: 

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

132 self.walker.finish(scan_report.quantum_id) 

133 case ProvenanceQuantumScanStatus.FAILED: 

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

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

136 for blocked_quantum_id in blocked_quanta: 

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

138 self.comms.request_write( 

139 ProvenanceQuantumScanData( 

140 blocked_quantum_id, status=ProvenanceQuantumScanStatus.BLOCKED 

141 ) 

142 ) 

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

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

145 case ProvenanceQuantumScanStatus.ABANDONED: 

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

147 self.walker.fail(scan_report.quantum_id) 

148 self.n_abandoned += 1 

149 case unexpected: 

150 raise AssertionError( 

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

152 ) 

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

154 

155 

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

157 """Run the graph aggregator tool. 

158 

159 Parameters 

160 ---------- 

161 predicted_path : `str` 

162 Path to the predicted quantum graph. 

163 butler_path : `str` 

164 Path or alias to the central butler repository. 

165 config : `AggregatorConfig` 

166 Configuration for the aggregator. 

167 """ 

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

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

170 scanners: list[Worker] = [] 

171 ingester: Worker 

172 writer: Worker | None = None 

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

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

175 if config.output_path is not None: 

176 writer_comms = WriterCommunicator(comms) 

177 writer = ctx.make_worker( 

178 target=Writer.run, 

179 args=(predicted_path, writer_comms), 

180 name=writer_comms.name, 

181 ) 

182 writer.start() 

183 for scanner_id in range(config.n_processes): 

184 scanner_comms = ScannerCommunicator(comms, scanner_id) 

185 worker = ctx.make_worker( 

186 target=Scanner.run, 

187 args=(predicted_path, butler_path, scanner_comms), 

188 name=scanner_comms.name, 

189 ) 

190 worker.start() 

191 scanners.append(worker) 

192 ingester_comms = IngesterCommunicator(comms) 

193 ingester = ctx.make_worker( 

194 target=Ingester.run, 

195 args=(predicted_path, butler_path, ingester_comms), 

196 name=ingester_comms.name, 

197 ) 

198 ingester.start() 

199 supervisor = Supervisor(predicted_path, comms) 

200 supervisor.loop() 

201 log.info( 

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

203 comms.progress.elapsed_time, 

204 ) 

205 comms.wait_for_workers_to_finish() 

206 if supervisor.n_abandoned: 

207 raise RuntimeError( 

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

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

210 "attempts have been exhausted." 

211 ) 

212 for w in scanners: 

213 w.join() 

214 ingester.join() 

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

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

217 writer.join() 

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

219 parent_mem, child_mem = get_peak_mem_usage() 

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

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

222 # than spiky that's fine. 

223 total_mem: u.Quantity = parent_mem + child_mem 

224 log.info( 

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

226 comms.progress.elapsed_time, 

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

228 )