nano_gpt.log

 1"""Logging utilities.
 2
 3This module provides a logging utility that can be used to log messages to a file or
 4standard output. This uses a structured log format that can be parsed by tools
 5for parsing and metrics gathering.
 6"""
 7
 8import logging
 9import pathlib
10from abc import ABC, abstractmethod
11from dataclasses import dataclass
12from typing import Any
13
14_LOGGER = logging.getLogger(__name__)
15
16
17@dataclass(frozen=True, kw_only=True)
18class LogRecord(ABC):
19    """A structured log record."""
20
21    log_type: str
22    data: dict[str, Any]
23
24    @property
25    def message(self) -> str:
26        """String representation of the data."""
27        return " | ".join(f"{key}: {value}" for key, value in self.data.items())
28
29    def __str__(self) -> str:
30        """String representation."""
31        return f"{self.log_type}: {self.message}"
32
33
34class TrainLog(ABC):
35    """Train logger abstract base class."""
36
37    @abstractmethod
38    def log(self, message: LogRecord) -> None:
39        """Log a message."""
40
41
42class LogFile(TrainLog):
43    """Train logger that logs to a file."""
44
45    def __init__(self, log_file: pathlib.Path) -> None:
46        """Initialize the train logger."""
47        self.log_file = log_file
48        self.log_file.parent.mkdir(parents=True, exist_ok=True)
49        self.log_file.unlink(missing_ok=True)
50
51    def log(self, message: LogRecord) -> None:
52        """Log a message."""
53        with self.log_file.open("a") as f:
54            f.write(str(message))
55            f.write("\n")
56            f.flush()
57
58
59class LogStdout(TrainLog):
60    """Train logger that logs to standard output."""
61
62    def log(self, message: LogRecord) -> None:
63        """Log a message."""
64        print(str(message))
65
66
67class LogMulti(TrainLog):
68    """Train logger that logs to multiple loggers."""
69
70    def __init__(self, logs: list[TrainLog]) -> None:
71        """Initialize the train logger."""
72        self.logs = logs
73
74    def log(self, message: LogRecord) -> None:
75        """Log a message."""
76        for log in self.logs:
77            log.log(message)
78
79
80def create_log(
81    log_file: pathlib.Path | None = None, log_stdout: bool = True
82) -> TrainLog:
83    """Create a train logger that logs to a file and standard output."""
84    logs: list[TrainLog] = []
85    if log_file:
86        logs.append(LogFile(log_file))
87    if log_stdout:
88        logs.append(LogStdout())
89    return LogMulti(logs)