Coverage for python/lsst/pipe/base/quantum_graph/aggregator/

1# This file is part of pipe_base.

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.

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

28from __future__ import annotations

30__all__ = ("aggregate_graph",)

32import dataclasses

33import uuid

35import astropy.units as u

36import networkx

38from lsst.utils.logging import getLogger

39from lsst.utils.usage import get_peak_mem_usage

41from ...graph_walker import GraphWalker

42from ...pipeline_graph import TaskImportMode

43from .._predicted import PredictedQuantumGraphComponents, PredictedQuantumGraphReader

44from ._communicators import (

45 IngesterCommunicator,

46 ScannerCommunicator,

47 SpawnProcessContext,

48 SupervisorCommunicator,

49 ThreadingContext,

50 Worker,

51 WriterCommunicator,

52)

53from ._config import AggregatorConfig

54from ._ingester import Ingester

55from ._scanner import Scanner

56from ._structs import ScanReport, ScanResult, ScanStatus

57from ._writer import Writer

60@dataclasses.dataclass

61class Supervisor:

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

64 predicted_path: str

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

67 comms: SupervisorCommunicator

68 """Communicator object for the supervisor."""

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

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

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

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

76 n_abandoned: int = 0

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

78 we could not assume they had failed.

79 """

81 def __post_init__(self) -> None:

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

83 with PredictedQuantumGraphReader.open(

84 self.predicted_path, import_mode=TaskImportMode.DO_NOT_IMPORT

85 ) as reader:

86 reader.read_thin_graph()

87 reader.read_init_quanta()

88 self.predicted = reader.components

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

90 uuid_by_index = {

91 quantum_index: quantum_id for quantum_id, quantum_index in self.predicted.quantum_indices.items()

92 }

93 xgraph = networkx.DiGraph(

94 [(uuid_by_index[a], uuid_by_index[b]) for a, b in self.predicted.thin_graph.edges]

95 )

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

97 # quantum-only edges.

98 xgraph.add_nodes_from(uuid_by_index.values())

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

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

101 # care when we process them.

102 for init_quantum in self.predicted.init_quanta.root[1:]: # skip 'packages' producer

103 xgraph.add_node(init_quantum.quantum_id)

104 self.walker = GraphWalker(xgraph)

105

106 def loop(self) -> None:

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

108 ingest outputs.

109 """

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

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

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

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

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

115 for ready_quanta in self.walker:

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

117 ready_set.update(ready_quanta)

118 while ready_set:

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

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

121 self.handle_report(scan_return)

122

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

124 """Handle a report from a scanner.

125

126 Parameters

127 ----------

128 scan_report : `ScanReport`

129 Information about the scan.

130 """

131 match scan_report.status:

132 case ScanStatus.SUCCESSFUL | ScanStatus.INIT:

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

134 self.walker.finish(scan_report.quantum_id)

135 case ScanStatus.FAILED:

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

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

138 for blocked_quantum_id in blocked_quanta:

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

140 self.comms.request_write(ScanResult(blocked_quantum_id, status=ScanStatus.BLOCKED))

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

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

143 case ScanStatus.ABANDONED:

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

145 self.walker.fail(scan_report.quantum_id)

146 self.n_abandoned += 1

147 case unexpected:

148 raise AssertionError(

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

150 )

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

152

153

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

155 """Run the graph aggregator tool.

156

157 Parameters

158 ----------

159 predicted_path : `str`

160 Path to the predicted quantum graph.

161 butler_path : `str`

162 Path or alias to the central butler repository.

163 config : `AggregatorConfig`

164 Configuration for the aggregator.

165 """

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

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

168 scanners: list[Worker] = []

169 ingester: Worker

170 writer: Worker | None = None

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

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

173 if config.output_path is not None:

174 writer_comms = WriterCommunicator(comms)

175 writer = ctx.make_worker(

176 target=Writer.run,

177 args=(predicted_path, writer_comms),

178 name=writer_comms.name,

179 )

180 writer.start()

181 for scanner_id in range(config.n_processes):

182 scanner_comms = ScannerCommunicator(comms, scanner_id)

183 worker = ctx.make_worker(

184 target=Scanner.run,

185 args=(predicted_path, butler_path, scanner_comms),

186 name=scanner_comms.name,

187 )

188 worker.start()

189 scanners.append(worker)

190 ingester_comms = IngesterCommunicator(comms)

191 ingester = ctx.make_worker(

192 target=Ingester.run,

193 args=(predicted_path, butler_path, ingester_comms),

194 name=ingester_comms.name,

195 )

196 ingester.start()

197 supervisor = Supervisor(predicted_path, comms)

198 supervisor.loop()

199 log.info(

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

201 comms.progress.elapsed_time,

202 )

203 comms.wait_for_workers_to_finish()

204 if supervisor.n_abandoned:

205 raise RuntimeError(

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

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

208 "attempts have been exhausted."

209 )

210 for w in scanners:

211 w.join()

212 ingester.join()

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

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

215 writer.join()

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

217 parent_mem, child_mem = get_peak_mem_usage()

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

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

220 # than spiky that's fine.

221 total_mem: u.Quantity = parent_mem + child_mem

222 log.info(

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

224 comms.progress.elapsed_time,

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

226 )

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

108 statements