Coverage for python/lsst/pipe/base/quantum_graph/aggregator/_progress.py: 22%

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__ = ("Progress", "make_worker_log") 

31 

32import logging 

33import os 

34import time 

35from types import TracebackType 

36from typing import Self 

37 

38from lsst.utils.logging import TRACE, VERBOSE, LsstLogAdapter, PeriodicLogger, getLogger 

39 

40from ._config import AggregatorConfig 

41 

42 

43class Progress: 

44 """A helper class for the provenance aggregator that handles reporting 

45 progress to the user. 

46 

47 This includes both logging (including periodic logging) and optional 

48 progress bars. 

49 

50 Parameters 

51 ---------- 

52 log : `lsst.utils.logging.LsstLogAdapter` 

53 LSST-customized logger. 

54 config : `AggregatorConfig` 

55 Configuration for the aggregator. 

56 

57 Notes 

58 ----- 

59 This class is a context manager in order to manage the redirection of 

60 logging when progress bars for interactive display are in use. The context 

61 manager does nothing otherwise. 

62 """ 

63 

64 def __init__(self, log: LsstLogAdapter, config: AggregatorConfig): 

65 self.start = time.time() 

66 self.log = log 

67 self.config = config 

68 self._periodic_log = PeriodicLogger(self.log, config.log_status_interval) 

69 self._n_scanned: int = 0 

70 self._n_ingested: int = 0 

71 self._n_written: int = 0 

72 self._n_quanta: int | None = None 

73 self.interactive = config.interactive_status 

74 

75 def __enter__(self) -> Self: 

76 if self.interactive: 

77 from tqdm.contrib.logging import logging_redirect_tqdm 

78 

79 self._logging_redirect = logging_redirect_tqdm() 

80 self._logging_redirect.__enter__() 

81 return self 

82 

83 def __exit__( 

84 self, 

85 exc_type: type[BaseException] | None, 

86 exc_value: BaseException | None, 

87 traceback: TracebackType | None, 

88 ) -> bool | None: 

89 if self.interactive: 

90 self._logging_redirect.__exit__(exc_type, exc_value, traceback) 

91 return None 

92 

93 def set_n_quanta(self, n_quanta: int) -> None: 

94 """Set the total number of quanta. 

95 

96 Parameters 

97 ---------- 

98 n_quanta : `int` 

99 Total number of quanta, including special "init" quanta. 

100 

101 Notes 

102 ----- 

103 This method must be called before any of the ``report_*`` methods. 

104 """ 

105 self._n_quanta = n_quanta 

106 if self.interactive: 

107 from tqdm import tqdm 

108 

109 self._scan_progress = tqdm(desc="Scanning", total=n_quanta, leave=False, unit="quanta") 

110 self._ingest_progress = tqdm( 

111 desc="Ingesting", total=n_quanta, leave=False, smoothing=0.1, unit="quanta" 

112 ) 

113 if self.config.output_path is not None: 

114 self._write_progress = tqdm(desc="Writing", total=n_quanta, leave=False, unit="quanta") 

115 

116 @property 

117 def elapsed_time(self) -> float: 

118 """The time in seconds since the start of the aggregator.""" 

119 return time.time() - self.start 

120 

121 def _log_status(self) -> None: 

122 """Invoke the periodic logger with the current status.""" 

123 self._periodic_log.log( 

124 "%s quanta scanned, %s quantum outputs ingested, " 

125 "%s provenance quanta written (of %s) after %0.1fs.", 

126 self._n_scanned, 

127 self._n_ingested, 

128 self._n_written, 

129 self._n_quanta, 

130 self.elapsed_time, 

131 ) 

132 

133 def report_scan(self) -> None: 

134 """Report that a quantum was scanned.""" 

135 self._n_scanned += 1 

136 if self.interactive: 

137 self._scan_progress.update(1) 

138 else: 

139 self._log_status() 

140 

141 def finish_scans(self) -> None: 

142 """Report that all scanning is done.""" 

143 if self.interactive: 

144 self._scan_progress.close() 

145 

146 def report_ingests(self, n_quanta: int) -> None: 

147 """Report that ingests for multiple quanta were completed. 

148 

149 Parameters 

150 ---------- 

151 n_quanta : `int` 

152 Number of quanta whose outputs were ingested. 

153 """ 

154 self._n_ingested += n_quanta 

155 if self.interactive: 

156 self._ingest_progress.update(n_quanta) 

157 else: 

158 self._log_status() 

159 

160 def finish_ingests(self) -> None: 

161 """Report that all ingests are done.""" 

162 if self.interactive: 

163 self._ingest_progress.close() 

164 

165 def report_write(self) -> None: 

166 """Report that a quantum's provenance was written.""" 

167 self._n_written += 1 

168 if self.interactive: 

169 self._write_progress.update() 

170 else: 

171 self._log_status() 

172 

173 def finish_writes(self) -> None: 

174 """Report that all writes are done.""" 

175 if self.interactive: 

176 self._write_progress.close() 

177 

178 

179def make_worker_log(name: str, config: AggregatorConfig) -> LsstLogAdapter: 

180 """Make a logger for a worker. 

181 

182 Parameters 

183 ---------- 

184 name : `str` 

185 Name of the worker, to be used as part of the name for the logger. 

186 config : `AggregatorConfig` 

187 Configuration for the aggregator. 

188 """ 

189 base_log = logging.getLogger(f"lsst.pipe.base.quantum_graph.aggregator.{name}") 

190 base_log.propagate = False 

191 log = getLogger(logger=base_log) 

192 if config.worker_log_dir is not None: 

193 os.makedirs(config.worker_log_dir, exist_ok=True) 

194 match config.worker_log_level.upper(): 

195 case "VERBOSE": 

196 log.setLevel(VERBOSE) 

197 case "TRACE": 

198 log.setLevel(TRACE) 

199 case std: 

200 log.setLevel(getattr(logging, std)) 

201 handler = logging.FileHandler(os.path.join(config.worker_log_dir, f"{name}.log")) 

202 handler.setFormatter( 

203 logging.Formatter("%(levelname)s %(asctime)s.%(msecs)03d %(message)s", "%Y-%m-%dT%H:%M:%S") 

204 ) 

205 log.addHandler(handler) 

206 else: 

207 log.addHandler(logging.NullHandler()) 

208 return log