Coverage for python / lsst / pipe / base / log_capture.py: 40%
102 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-19 01:13 -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 uuid
35from collections.abc import Iterator
36from contextlib import contextmanager
37from logging import FileHandler
39import pydantic
41from lsst.daf.butler import Butler, LimitedButler, Quantum
42from lsst.daf.butler._rubin.temporary_for_ingest import TemporaryForIngest
43from lsst.daf.butler.logging import (
44 ButlerLogRecord,
45 ButlerLogRecordHandler,
46 ButlerLogRecords,
47 ButlerMDC,
48 JsonLogFormatter,
49)
51from ._status import ExceptionInfo, InvalidQuantumError
52from ._task_metadata import TaskMetadata
53from .automatic_connection_constants import METADATA_OUTPUT_TEMPLATE
54from .pipeline_graph import TaskNode
56_LOG = logging.getLogger(__name__)
59class _ExecutionLogRecordsExtra(pydantic.BaseModel):
60 """Extra information about a quantum's execution stored with logs.
62 This middleware-private model includes information that is not directly
63 available via any public interface, as it is used exclusively for
64 provenance extraction and then made available through the provenance
65 quantum graph.
66 """
68 exception: ExceptionInfo | None = None
69 """Exception information for this quantum, if it failed.
70 """
72 metadata: TaskMetadata | None = None
73 """Metadata for this quantum, if it failed.
75 Metadata datasets are written if and only if a quantum succeeds, but we
76 still want to capture metadata from failed attempts, so we store it in the
77 log dataset. This field is always `None` when the quantum succeeds,
78 because in that case the metadata is already stored separately.
79 """
81 previous_process_quanta: list[uuid.UUID] = pydantic.Field(default_factory=list)
82 """The IDs of other quanta previously executed in the same process as this
83 one.
84 """
86 logs: list[ButlerLogRecord] = pydantic.Field(default_factory=list)
87 """Logs for this attempt.
89 This is always empty for the most recent attempt, because that stores logs
90 in the main section of the butler log records.
91 """
93 previous_attempts: list[_ExecutionLogRecordsExtra] = pydantic.Field(default_factory=list)
94 """Information about previous attempts to run this task within the same
95 `~lsst.daf.butler.CollectionType.RUN` collection.
97 This is always empty for any attempt other than the most recent one,
98 as all previous attempts are flattened into one list.
99 """
101 def attach_previous_attempt(self, log_records: ButlerLogRecords) -> None:
102 """Attach logs from a previous attempt to this struct.
104 Parameters
105 ----------
106 log_records : `ButlerLogRecords`
107 Logs from a past attempt to run a quantum.
108 """
109 previous = self.model_validate(log_records.extra)
110 previous.logs.extend(log_records)
111 self.previous_attempts.extend(previous.previous_attempts)
112 self.previous_attempts.append(previous)
113 previous.previous_attempts.clear()
116@dataclasses.dataclass
117class _LogCaptureContext:
118 """Controls for log capture returned by the `LogCapture.capture_logging`
119 context manager.
120 """
122 store: bool = True
123 """Whether to store logs at all."""
125 extra: _ExecutionLogRecordsExtra = dataclasses.field(default_factory=_ExecutionLogRecordsExtra)
126 """Extra information about the quantum's execution to store for provenance
127 extraction.
128 """
131class LogCapture:
132 """Class handling capture of logging messages and their export to butler.
134 Parameters
135 ----------
136 butler : `~lsst.daf.butler.LimitedButler`
137 Data butler with limited API.
138 full_butler : `~lsst.daf.butler.Butler` or `None`
139 Data butler with full API, or `None` if full Butler is not available.
140 If not none, then this must be the same instance as ``butler``.
141 """
143 stream_json_logs = True
144 """If True each log record is written to a temporary file and ingested
145 when quantum completes. If False the records are accumulated in memory
146 and stored in butler on quantum completion. If full butler is not available
147 then temporary file is not used."""
149 def __init__(
150 self,
151 butler: LimitedButler,
152 full_butler: Butler | None,
153 ):
154 self.butler = butler
155 self.full_butler = full_butler
157 @classmethod
158 def from_limited(cls, butler: LimitedButler) -> LogCapture:
159 return cls(butler, None)
161 @classmethod
162 def from_full(cls, butler: Butler) -> LogCapture:
163 return cls(butler, butler)
165 @contextmanager
166 def capture_logging(self, task_node: TaskNode, /, quantum: Quantum) -> Iterator[_LogCaptureContext]:
167 """Configure logging system to capture logs for execution of this task.
169 Parameters
170 ----------
171 task_node : `~lsst.pipe.base.pipeline_graph.TaskNode`
172 The task definition.
173 quantum : `~lsst.daf.butler.Quantum`
174 Single Quantum instance.
176 Notes
177 -----
178 Expected to be used as a context manager to ensure that logging
179 records are inserted into the butler once the quantum has been
180 executed:
182 .. code-block:: py
184 with self.capture_logging(task_node, quantum):
185 # Run quantum and capture logs.
187 Ths method can also setup logging to attach task- or
188 quantum-specific information to log messages. Potentially this can
189 take into account some info from task configuration as well.
190 """
191 # include quantum dataId and task label into MDC
192 mdc = {"LABEL": task_node.label, "RUN": ""}
193 if quantum.dataId:
194 mdc["LABEL"] += f":{quantum.dataId}"
196 metadata_ref = quantum.outputs[METADATA_OUTPUT_TEMPLATE.format(label=task_node.label)][0]
197 mdc["RUN"] = metadata_ref.run
199 ctx = _LogCaptureContext()
200 log_dataset_name = (
201 task_node.log_output.dataset_type_name if task_node.log_output is not None else None
202 )
204 # Add a handler to the root logger to capture execution log output.
205 if log_dataset_name is not None:
206 try:
207 [ref] = quantum.outputs[log_dataset_name]
208 except LookupError as exc:
209 raise InvalidQuantumError(
210 f"Quantum outputs is missing log output dataset type {log_dataset_name};"
211 " this could happen due to inconsistent options between QuantumGraph generation"
212 " and execution"
213 ) from exc
214 # Either accumulate into ButlerLogRecords or stream JSON records to
215 # file and ingest that (ingest is possible only with full butler).
216 if self.stream_json_logs and self.full_butler is not None:
217 with TemporaryForIngest(self.full_butler, ref) as temporary:
218 log_handler_file = FileHandler(temporary.ospath)
219 log_handler_file.setFormatter(JsonLogFormatter())
220 logging.getLogger().addHandler(log_handler_file)
222 try:
223 with ButlerMDC.set_mdc(mdc):
224 yield ctx
225 finally:
226 # Ensure that the logs are stored in butler.
227 logging.getLogger().removeHandler(log_handler_file)
228 log_handler_file.close()
229 if ctx.extra:
230 with open(temporary.ospath, "a") as log_stream:
231 ButlerLogRecords.write_streaming_extra(
232 log_stream,
233 ctx.extra.model_dump_json(exclude_unset=True, exclude_defaults=True),
234 )
235 if ctx.store:
236 temporary.ingest()
238 else:
239 log_handler_memory = ButlerLogRecordHandler()
240 logging.getLogger().addHandler(log_handler_memory)
242 try:
243 with ButlerMDC.set_mdc(mdc):
244 yield ctx
245 except:
246 raise
247 else:
248 # If the quantum succeeded, we don't need to save the
249 # metadata in the logs, because we'll have saved them in
250 # the metadata.
251 ctx.extra.metadata = None
252 finally:
253 log_handler_memory.records.extra = ctx.extra.model_dump()
254 # Ensure that the logs are stored in butler.
255 logging.getLogger().removeHandler(log_handler_memory)
256 if ctx.store:
257 self._store_log_records(quantum, log_dataset_name, log_handler_memory)
258 log_handler_memory.records.clear()
260 else:
261 with ButlerMDC.set_mdc(mdc):
262 yield ctx
264 def _store_log_records(
265 self, quantum: Quantum, dataset_type: str, log_handler: ButlerLogRecordHandler
266 ) -> None:
267 # DatasetRef has to be in the Quantum outputs, can lookup by name.
268 try:
269 [ref] = quantum.outputs[dataset_type]
270 except LookupError as exc:
271 raise InvalidQuantumError(
272 f"Quantum outputs is missing log output dataset type {dataset_type};"
273 " this could happen due to inconsistent options between QuantumGraph generation"
274 " and execution"
275 ) from exc
277 self.butler.put(log_handler.records, ref)