Coverage for python/lsst/pipe/base/log_capture.py: 28%

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__ = ["LogCapture"] 

31 

32import logging 

33import os 

34import shutil 

35import tempfile 

36from collections.abc import Iterator 

37from contextlib import contextmanager, suppress 

38from logging import FileHandler 

39 

40from lsst.daf.butler import Butler, FileDataset, LimitedButler, Quantum 

41from lsst.daf.butler.logging import ButlerLogRecordHandler, ButlerLogRecords, ButlerMDC, JsonLogFormatter 

42 

43from ._status import InvalidQuantumError 

44from .automatic_connection_constants import METADATA_OUTPUT_TEMPLATE 

45from .pipeline_graph import TaskNode 

46 

47_LOG = logging.getLogger(__name__) 

48 

49 

50class _LogCaptureFlag: 

51 """Simple flag to enable/disable log-to-butler saving.""" 

52 

53 store: bool = True 

54 

55 

56class LogCapture: 

57 """Class handling capture of logging messages and their export to butler. 

58 

59 Parameters 

60 ---------- 

61 butler : `~lsst.daf.butler.LimitedButler` 

62 Data butler with limited API. 

63 full_butler : `~lsst.daf.butler.Butler` or `None` 

64 Data butler with full API, or `None` if full Butler is not available. 

65 If not none, then this must be the same instance as ``butler``. 

66 """ 

67 

68 stream_json_logs = True 

69 """If True each log record is written to a temporary file and ingested 

70 when quantum completes. If False the records are accumulated in memory 

71 and stored in butler on quantum completion. If full butler is not available 

72 then temporary file is not used.""" 

73 

74 def __init__( 

75 self, 

76 butler: LimitedButler, 

77 full_butler: Butler | None, 

78 ): 

79 self.butler = butler 

80 self.full_butler = full_butler 

81 

82 @classmethod 

83 def from_limited(cls, butler: LimitedButler) -> LogCapture: 

84 return cls(butler, None) 

85 

86 @classmethod 

87 def from_full(cls, butler: Butler) -> LogCapture: 

88 return cls(butler, butler) 

89 

90 @contextmanager 

91 def capture_logging(self, task_node: TaskNode, /, quantum: Quantum) -> Iterator[_LogCaptureFlag]: 

92 """Configure logging system to capture logs for execution of this task. 

93 

94 Parameters 

95 ---------- 

96 task_node : `~lsst.pipe.base.pipeline_graph.TaskNode` 

97 The task definition. 

98 quantum : `~lsst.daf.butler.Quantum` 

99 Single Quantum instance. 

100 

101 Notes 

102 ----- 

103 Expected to be used as a context manager to ensure that logging 

104 records are inserted into the butler once the quantum has been 

105 executed: 

106 

107 .. code-block:: py 

108 

109 with self.capture_logging(task_node, quantum): 

110 # Run quantum and capture logs. 

111 

112 Ths method can also setup logging to attach task- or 

113 quantum-specific information to log messages. Potentially this can 

114 take into account some info from task configuration as well. 

115 """ 

116 # include quantum dataId and task label into MDC 

117 mdc = {"LABEL": task_node.label, "RUN": ""} 

118 if quantum.dataId: 

119 mdc["LABEL"] += f":{quantum.dataId}" 

120 

121 metadata_ref = quantum.outputs[METADATA_OUTPUT_TEMPLATE.format(label=task_node.label)][0] 

122 mdc["RUN"] = metadata_ref.run 

123 

124 ctx = _LogCaptureFlag() 

125 log_dataset_name = ( 

126 task_node.log_output.dataset_type_name if task_node.log_output is not None else None 

127 ) 

128 

129 # Add a handler to the root logger to capture execution log output. 

130 if log_dataset_name is not None: 

131 # Either accumulate into ButlerLogRecords or stream JSON records to 

132 # file and ingest that (ingest is possible only with full butler). 

133 if self.stream_json_logs and self.full_butler is not None: 

134 # Create the log file in a temporary directory rather than 

135 # creating a temporary file. This is necessary because 

136 # temporary files are created with restrictive permissions 

137 # and during file ingest these permissions persist in the 

138 # datastore. Using a temp directory allows us to create 

139 # a file with umask default permissions. 

140 tmpdir = tempfile.mkdtemp(prefix="butler-temp-logs-") 

141 

142 # Construct a file to receive the log records and "touch" it. 

143 log_file = os.path.join(tmpdir, f"butler-log-{task_node.label}.json") 

144 with open(log_file, "w"): 

145 pass 

146 log_handler_file = FileHandler(log_file) 

147 log_handler_file.setFormatter(JsonLogFormatter()) 

148 logging.getLogger().addHandler(log_handler_file) 

149 

150 try: 

151 with ButlerMDC.set_mdc(mdc): 

152 yield ctx 

153 finally: 

154 # Ensure that the logs are stored in butler. 

155 logging.getLogger().removeHandler(log_handler_file) 

156 log_handler_file.close() 

157 if ctx.store: 

158 self._ingest_log_records(quantum, log_dataset_name, log_file) 

159 shutil.rmtree(tmpdir, ignore_errors=True) 

160 

161 else: 

162 log_handler_memory = ButlerLogRecordHandler() 

163 logging.getLogger().addHandler(log_handler_memory) 

164 

165 try: 

166 with ButlerMDC.set_mdc(mdc): 

167 yield ctx 

168 finally: 

169 # Ensure that the logs are stored in butler. 

170 logging.getLogger().removeHandler(log_handler_memory) 

171 if ctx.store: 

172 self._store_log_records(quantum, log_dataset_name, log_handler_memory) 

173 log_handler_memory.records.clear() 

174 

175 else: 

176 with ButlerMDC.set_mdc(mdc): 

177 yield ctx 

178 

179 def _store_log_records( 

180 self, quantum: Quantum, dataset_type: str, log_handler: ButlerLogRecordHandler 

181 ) -> None: 

182 # DatasetRef has to be in the Quantum outputs, can lookup by name. 

183 try: 

184 [ref] = quantum.outputs[dataset_type] 

185 except LookupError as exc: 

186 raise InvalidQuantumError( 

187 f"Quantum outputs is missing log output dataset type {dataset_type};" 

188 " this could happen due to inconsistent options between QuantumGraph generation" 

189 " and execution" 

190 ) from exc 

191 

192 self.butler.put(log_handler.records, ref) 

193 

194 def _ingest_log_records(self, quantum: Quantum, dataset_type: str, filename: str) -> None: 

195 # If we are logging to an external file we must always try to 

196 # close it. 

197 assert self.full_butler is not None, "Expected to have full butler for ingest" 

198 ingested = False 

199 try: 

200 # DatasetRef has to be in the Quantum outputs, can lookup by name. 

201 try: 

202 [ref] = quantum.outputs[dataset_type] 

203 except LookupError as exc: 

204 raise InvalidQuantumError( 

205 f"Quantum outputs is missing log output dataset type {dataset_type};" 

206 " this could happen due to inconsistent options between QuantumGraph generation" 

207 " and execution" 

208 ) from exc 

209 

210 # Need to ingest this file directly into butler. 

211 dataset = FileDataset(path=filename, refs=ref) 

212 try: 

213 self.full_butler.ingest(dataset, transfer="move") 

214 ingested = True 

215 except NotImplementedError: 

216 # Some datastores can't receive files (e.g. in-memory datastore 

217 # when testing), we store empty list for those just to have a 

218 # dataset. Alternative is to read the file as a 

219 # ButlerLogRecords object and put it. 

220 _LOG.info( 

221 "Log records could not be stored in this butler because the" 

222 " datastore can not ingest files, empty record list is stored instead." 

223 ) 

224 records = ButlerLogRecords.from_records([]) 

225 self.full_butler.put(records, ref) 

226 finally: 

227 # remove file if it is not ingested 

228 if not ingested: 

229 with suppress(OSError): 

230 os.remove(filename)