Coverage for python / lsst / pipe / base / log_capture.py: 36%
123 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:14 -0800
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/>.
28from __future__ import annotations
30__all__ = ["LogCapture"]
32import dataclasses
33import logging
34import os
35import shutil
36import tempfile
37import uuid
38from collections.abc import Iterator
39from contextlib import contextmanager, suppress
40from logging import FileHandler
42import pydantic
44from lsst.daf.butler import Butler, FileDataset, LimitedButler, Quantum
45from lsst.daf.butler.logging import (
46 ButlerLogRecord,
47 ButlerLogRecordHandler,
48 ButlerLogRecords,
49 ButlerMDC,
50 JsonLogFormatter,
51)
53from ._status import ExceptionInfo, InvalidQuantumError
54from ._task_metadata import TaskMetadata
55from .automatic_connection_constants import METADATA_OUTPUT_TEMPLATE
56from .pipeline_graph import TaskNode
58_LOG = logging.getLogger(__name__)
61class _ExecutionLogRecordsExtra(pydantic.BaseModel):
62 """Extra information about a quantum's execution stored with logs.
64 This middleware-private model includes information that is not directly
65 available via any public interface, as it is used exclusively for
66 provenance extraction and then made available through the provenance
67 quantum graph.
68 """
70 exception: ExceptionInfo | None = None
71 """Exception information for this quantum, if it failed.
72 """
74 metadata: TaskMetadata | None = None
75 """Metadata for this quantum, if it failed.
77 Metadata datasets are written if and only if a quantum succeeds, but we
78 still want to capture metadata from failed attempts, so we store it in the
79 log dataset. This field is always `None` when the quantum succeeds,
80 because in that case the metadata is already stored separately.
81 """
83 previous_process_quanta: list[uuid.UUID] = pydantic.Field(default_factory=list)
84 """The IDs of other quanta previously executed in the same process as this
85 one.
86 """
88 logs: list[ButlerLogRecord] = pydantic.Field(default_factory=list)
89 """Logs for this attempt.
91 This is always empty for the most recent attempt, because that stores logs
92 in the main section of the butler log records.
93 """
95 previous_attempts: list[_ExecutionLogRecordsExtra] = pydantic.Field(default_factory=list)
96 """Information about previous attempts to run this task within the same
97 `~lsst.daf.butler.CollectionType.RUN` collection.
99 This is always empty for any attempt other than the most recent one,
100 as all previous attempts are flattened into one list.
101 """
103 def attach_previous_attempt(self, log_records: ButlerLogRecords) -> None:
104 """Attach logs from a previous attempt to this struct.
106 Parameters
107 ----------
108 log_records : `ButlerLogRecords`
109 Logs from a past attempt to run a quantum.
110 """
111 previous = self.model_validate(log_records.extra)
112 previous.logs.extend(log_records)
113 self.previous_attempts.extend(previous.previous_attempts)
114 self.previous_attempts.append(previous)
115 previous.previous_attempts.clear()
118@dataclasses.dataclass
119class _LogCaptureContext:
120 """Controls for log capture returned by the `LogCapture.capture_logging`
121 context manager.
122 """
124 store: bool = True
125 """Whether to store logs at all."""
127 extra: _ExecutionLogRecordsExtra = dataclasses.field(default_factory=_ExecutionLogRecordsExtra)
128 """Extra information about the quantum's execution to store for provenance
129 extraction.
130 """
133class LogCapture:
134 """Class handling capture of logging messages and their export to butler.
136 Parameters
137 ----------
138 butler : `~lsst.daf.butler.LimitedButler`
139 Data butler with limited API.
140 full_butler : `~lsst.daf.butler.Butler` or `None`
141 Data butler with full API, or `None` if full Butler is not available.
142 If not none, then this must be the same instance as ``butler``.
143 """
145 stream_json_logs = True
146 """If True each log record is written to a temporary file and ingested
147 when quantum completes. If False the records are accumulated in memory
148 and stored in butler on quantum completion. If full butler is not available
149 then temporary file is not used."""
151 def __init__(
152 self,
153 butler: LimitedButler,
154 full_butler: Butler | None,
155 ):
156 self.butler = butler
157 self.full_butler = full_butler
159 @classmethod
160 def from_limited(cls, butler: LimitedButler) -> LogCapture:
161 return cls(butler, None)
163 @classmethod
164 def from_full(cls, butler: Butler) -> LogCapture:
165 return cls(butler, butler)
167 @contextmanager
168 def capture_logging(self, task_node: TaskNode, /, quantum: Quantum) -> Iterator[_LogCaptureContext]:
169 """Configure logging system to capture logs for execution of this task.
171 Parameters
172 ----------
173 task_node : `~lsst.pipe.base.pipeline_graph.TaskNode`
174 The task definition.
175 quantum : `~lsst.daf.butler.Quantum`
176 Single Quantum instance.
178 Notes
179 -----
180 Expected to be used as a context manager to ensure that logging
181 records are inserted into the butler once the quantum has been
182 executed:
184 .. code-block:: py
186 with self.capture_logging(task_node, quantum):
187 # Run quantum and capture logs.
189 Ths method can also setup logging to attach task- or
190 quantum-specific information to log messages. Potentially this can
191 take into account some info from task configuration as well.
192 """
193 # include quantum dataId and task label into MDC
194 mdc = {"LABEL": task_node.label, "RUN": ""}
195 if quantum.dataId:
196 mdc["LABEL"] += f":{quantum.dataId}"
198 metadata_ref = quantum.outputs[METADATA_OUTPUT_TEMPLATE.format(label=task_node.label)][0]
199 mdc["RUN"] = metadata_ref.run
201 ctx = _LogCaptureContext()
202 log_dataset_name = (
203 task_node.log_output.dataset_type_name if task_node.log_output is not None else None
204 )
206 # Add a handler to the root logger to capture execution log output.
207 if log_dataset_name is not None:
208 # Either accumulate into ButlerLogRecords or stream JSON records to
209 # file and ingest that (ingest is possible only with full butler).
210 if self.stream_json_logs and self.full_butler is not None:
211 # Create the log file in a temporary directory rather than
212 # creating a temporary file. This is necessary because
213 # temporary files are created with restrictive permissions
214 # and during file ingest these permissions persist in the
215 # datastore. Using a temp directory allows us to create
216 # a file with umask default permissions.
217 tmpdir = tempfile.mkdtemp(prefix="butler-temp-logs-")
219 # Construct a file to receive the log records and "touch" it.
220 log_file = os.path.join(tmpdir, f"butler-log-{task_node.label}.json")
221 with open(log_file, "w"):
222 pass
223 log_handler_file = FileHandler(log_file)
224 log_handler_file.setFormatter(JsonLogFormatter())
225 logging.getLogger().addHandler(log_handler_file)
227 try:
228 with ButlerMDC.set_mdc(mdc):
229 yield ctx
230 finally:
231 # Ensure that the logs are stored in butler.
232 logging.getLogger().removeHandler(log_handler_file)
233 log_handler_file.close()
234 if ctx.extra:
235 with open(log_file, "a") as log_stream:
236 ButlerLogRecords.write_streaming_extra(
237 log_stream,
238 ctx.extra.model_dump_json(exclude_unset=True, exclude_defaults=True),
239 )
240 if ctx.store:
241 self._ingest_log_records(quantum, log_dataset_name, log_file)
242 shutil.rmtree(tmpdir, ignore_errors=True)
244 else:
245 log_handler_memory = ButlerLogRecordHandler()
246 logging.getLogger().addHandler(log_handler_memory)
248 try:
249 with ButlerMDC.set_mdc(mdc):
250 yield ctx
251 except:
252 raise
253 else:
254 # If the quantum succeeded, we don't need to save the
255 # metadata in the logs, because we'll have saved them in
256 # the metadata.
257 ctx.extra.metadata = None
258 finally:
259 log_handler_memory.records.extra = ctx.extra.model_dump()
260 # Ensure that the logs are stored in butler.
261 logging.getLogger().removeHandler(log_handler_memory)
262 if ctx.store:
263 self._store_log_records(quantum, log_dataset_name, log_handler_memory)
264 log_handler_memory.records.clear()
266 else:
267 with ButlerMDC.set_mdc(mdc):
268 yield ctx
270 def _store_log_records(
271 self, quantum: Quantum, dataset_type: str, log_handler: ButlerLogRecordHandler
272 ) -> None:
273 # DatasetRef has to be in the Quantum outputs, can lookup by name.
274 try:
275 [ref] = quantum.outputs[dataset_type]
276 except LookupError as exc:
277 raise InvalidQuantumError(
278 f"Quantum outputs is missing log output dataset type {dataset_type};"
279 " this could happen due to inconsistent options between QuantumGraph generation"
280 " and execution"
281 ) from exc
283 self.butler.put(log_handler.records, ref)
285 def _ingest_log_records(self, quantum: Quantum, dataset_type: str, filename: str) -> None:
286 # If we are logging to an external file we must always try to
287 # close it.
288 assert self.full_butler is not None, "Expected to have full butler for ingest"
289 ingested = False
290 try:
291 # DatasetRef has to be in the Quantum outputs, can lookup by name.
292 try:
293 [ref] = quantum.outputs[dataset_type]
294 except LookupError as exc:
295 raise InvalidQuantumError(
296 f"Quantum outputs is missing log output dataset type {dataset_type};"
297 " this could happen due to inconsistent options between QuantumGraph generation"
298 " and execution"
299 ) from exc
301 # Need to ingest this file directly into butler.
302 dataset = FileDataset(path=filename, refs=ref)
303 try:
304 self.full_butler.ingest(dataset, transfer="move")
305 ingested = True
306 except NotImplementedError:
307 # Some datastores can't receive files (e.g. in-memory datastore
308 # when testing), we store empty list for those just to have a
309 # dataset. Alternative is to read the file as a
310 # ButlerLogRecords object and put it.
311 _LOG.info(
312 "Log records could not be stored in this butler because the"
313 " datastore can not ingest files, empty record list is stored instead."
314 )
315 records = ButlerLogRecords.from_records([])
316 self.full_butler.put(records, ref)
317 finally:
318 # remove file if it is not ingested
319 if not ingested:
320 with suppress(OSError):
321 os.remove(filename)