1 2
import json
2 2
import logging
3 2
import os
4 2
from collections import defaultdict
5 2
from datetime import datetime
6 2
from typing import Any, DefaultDict, List, Mapping, Optional
7

8 2
from snorkel.types import Config
9

10

11 2
class LogWriterConfig(Config):
12
    """Manager for checkpointing model.
13

14
    Parameters
15
    ----------
16
    log_dir
17
        The root directory where logs should be saved
18
    run_name
19
        The name of this particular run (defaults to date-time combination if None)
20
    """
21

22 2
    log_dir: str = "logs"
23 2
    run_name: Optional[str] = None
24

25

26 2
class LogWriter:
27
    """A class for writing logs.
28

29
    Parameters
30
    ----------
31
    kwargs
32
        Settings to merge into LogWriterConfig
33

34
    Attributes
35
    ----------
36
    config
37
        Merged configuration
38
    run_name
39
        Name of run if provided, otherwise date-time combination
40
    log_dir
41
        The root directory where logs should be saved
42
    run_log
43
        Dictionary of scalar values to log, keyed by value name
44
    """
45

46 2
    def __init__(self, **kwargs: Any) -> None:
47 2
        self.config = LogWriterConfig(**kwargs)
48

49 2
        self.run_name = self.config.run_name
50 2
        if self.run_name is None:
51 2
            date = datetime.now().strftime("%Y_%m_%d")
52 2
            time = datetime.now().strftime("%H_%M_%S")
53 2
            self.run_name = f"{date}/{time}/"
54

55 2
        self.log_dir = os.path.join(self.config.log_dir, self.run_name)
56 2
        if not os.path.exists(self.log_dir):
57 2
            os.makedirs(self.log_dir)
58

59 2
        self.run_log: DefaultDict[str, List[List[float]]] = defaultdict(list)
60

61 2
    def add_scalar(self, name: str, value: float, step: float) -> None:
62
        """Log a scalar variable.
63

64
        Parameters
65
        ----------
66
        name
67
            Name of the scalar collection
68
        value
69
            Value of scalar
70
        step
71
            Step axis value
72
        """
73
        # Note: storing as list for JSON roundtripping
74 2
        self.run_log[name].append([step, value])
75

76 2
    def write_config(
77
        self, config: Config, config_filename: str = "config.json"
78
    ) -> None:
79
        """Dump the config to file.
80

81
        Parameters
82
        ----------
83
        config
84
            JSON-compatible config to write to file
85
        config_filename
86
            Name of file in logging directory to write to
87
        """
88 2
        self.write_json(config._asdict(), config_filename)
89

90 2
    def write_log(self, log_filename: str) -> None:
91
        """Dump the scalar value log to file.
92

93
        Parameters
94
        ----------
95
        log_filename
96
            Name of file in logging directory to write to
97
        """
98 2
        self.write_json(self.run_log, log_filename)
99

100 2
    def write_text(self, text: str, filename: str) -> None:
101
        """Dump user-provided text to filename (e.g., the launch command).
102

103
        Parameters
104
        ----------
105
        text
106
            Text to write
107
        filename
108
            Name of file in logging directory to write to
109
        """
110 2
        text_path = os.path.join(self.log_dir, filename)
111 2
        with open(text_path, "w") as f:
112 2
            f.write(text)
113

114 2
    def write_json(self, dict_to_write: Mapping[str, Any], filename: str) -> None:
115
        """Dump a JSON-compatbile object to root log directory.
116

117
        Parameters
118
        ----------
119
        dict_to_write
120
            JSON-compatbile object to log
121
        filename
122
            Name of file in logging directory to write to
123
        """
124
        if not filename.endswith(".json"):  # pragma: no cover
125
            logging.warning(
126
                f"Using write_json() method with a filename without a .json extension: {filename}"
127
            )
128 2
        log_path = os.path.join(self.log_dir, filename)
129 2
        with open(log_path, "w") as f:
130 2
            json.dump(dict_to_write, f)
131

132 2
    def cleanup(self) -> None:
133
        """Perform final operations and close writer if necessary."""
134 2
        self.write_log("log.json")

Read our documentation on viewing source code .

Loading