1 from __future__ import annotations

2

3 import json

4 import sys

5 import time

6 import typing

7 from typing import TextIO, Union

8

9 from muutils.json_serialize import JSONitem, json_serialize

10

11

12 class NullIO:

13 """null IO class"""

14

15 def __init__(self) -> None:

16 pass

17

18 def write(self, msg: str) -> int:

19 """write to nothing! this throws away the message"""

20 return len(msg)

21

22 def flush(self) -> None:

23 """flush nothing! this is a no-op"""

24 pass

25

26 def close(self) -> None:

27 """close nothing! this is a no-op"""

28 pass

29

30

31 AnyIO = Union[TextIO, NullIO]

32

33

34 class SimpleLogger:

35 """logs training data to a jsonl file"""

36

37 def __init__(

38 self,

39 log_path: str | None = None,

40 log_file: AnyIO | None = None,

41 timestamp: bool = True,

42 ):

43 self._timestamp: bool = timestamp

44 self._log_path: str | None = log_path

45

46 self._log_file_handle: AnyIO

47

48 if (log_path is None) and (log_file is None):

49 print(

50 "[logger_internal] # no log file specified, will only write to console",

51 sys.stderr,

52 )

53 self._log_file_handle = sys.stdout

54

55 elif (log_path is not None) and (log_file is not None):

56 raise ValueError(

57 "cannot specify both log_path and log_file, use streams in `SimpleLogger`"

58 )

59 else:

60 # now exactly one of the two is None

61 if log_file is not None:

62 self._log_file_handle = log_file

63 else:

64 assert log_path is not None

65 self._log_file_handle = open(log_path, "w", encoding="utf-8")

66

67 def log(self, msg: JSONitem, console_print: bool = False, **kwargs):

68 """log a message to the log file, and optionally to the console"""

69 if console_print:

70 print(msg)

71

72 if not isinstance(msg, typing.Mapping):

73 msg = {"_msg": msg}

74

75 if self._timestamp:

76 msg["_timestamp"] = time.time()

77

78 if len(kwargs) > 0:

79 msg["_kwargs"] = kwargs

80

81 self._log_file_handle.write(json.dumps(json_serialize(msg)) + "\n")