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 if self.comms.config.incomplete: 

121 quantum_or_quanta = "quanta" if self.n_abandoned != 1 else "quantum" 

122 self.comms.progress.log.info( 

123 "%d %s incomplete/failed abandoned; re-run with incomplete=False to finish.", 

124 self.n_abandoned, 

125 quantum_or_quanta, 

126 ) 

127 self.comms.progress.log.info( 

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

129 self.comms.progress.elapsed_time, 

130 ) 

131 self.comms.wait_for_workers_to_finish() 

132 

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

134 """Handle a report from a scanner. 

135 

136 Parameters 

137 ---------- 

138 scan_report : `ScanReport` 

139 Information about the scan. 

140 """ 

141 match scan_report.status: 

142 case ProvenanceQuantumScanStatus.SUCCESSFUL | ProvenanceQuantumScanStatus.INIT: 

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

144 self.walker.finish(scan_report.quantum_id) 

145 case ProvenanceQuantumScanStatus.FAILED: 

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

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

148 for blocked_quantum_id in blocked_quanta: 

149 if self.comms.config.is_writing_provenance: 

150 self.comms.request_write( 

151 ProvenanceQuantumScanData( 

152 blocked_quantum_id, status=ProvenanceQuantumScanStatus.BLOCKED 

153 ) 

154 ) 

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

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

157 case ProvenanceQuantumScanStatus.ABANDONED: 

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

159 self.walker.fail(scan_report.quantum_id) 

160 self.n_abandoned += 1 

161 case unexpected: 

162 raise AssertionError( 

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

164 ) 

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

166 

167 

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

169 """Run the graph aggregator tool. 

170 

171 Parameters 

172 ---------- 

173 predicted_path : `str` 

174 Path to the predicted quantum graph. 

175 butler_path : `str` 

176 Path or alias to the central butler repository. 

177 config : `AggregatorConfig` 

178 Configuration for the aggregator. 

179 """ 

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

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

182 scanners: list[Worker] = [] 

183 ingester: Worker 

184 writer: Worker | None = None 

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

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

187 if config.is_writing_provenance: 

188 writer_comms = WriterCommunicator(comms) 

189 writer = ctx.make_worker( 

190 target=Writer.run, 

191 args=(predicted_path, writer_comms), 

192 name=writer_comms.name, 

193 ) 

194 writer.start() 

195 for scanner_id in range(config.n_processes): 

196 scanner_comms = ScannerCommunicator(comms, scanner_id) 

197 worker = ctx.make_worker( 

198 target=Scanner.run, 

199 args=(predicted_path, butler_path, scanner_comms), 

200 name=scanner_comms.name, 

201 ) 

202 worker.start() 

203 scanners.append(worker) 

204 ingester_comms = IngesterCommunicator(comms) 

205 ingester = ctx.make_worker( 

206 target=Ingester.run, 

207 args=(predicted_path, butler_path, ingester_comms), 

208 name=ingester_comms.name, 

209 ) 

210 ingester.start() 

211 supervisor = Supervisor(predicted_path, comms) 

212 supervisor.loop() 

213 for w in scanners: 

214 w.join() 

215 ingester.join() 

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

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

218 writer.join() 

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

220 parent_mem, child_mem = get_peak_mem_usage() 

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

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

223 # than spiky that's fine. 

224 total_mem: u.Quantity = parent_mem + child_mem 

225 log.info( 

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

227 comms.progress.elapsed_time, 

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

229 )