DeepSpeed/deepspeed/monitor/monitor.py at f33ad0a6aa0ab660970c10b02d48bc049c063c6d · deepspeedai/DeepSpeed · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team
"""
Support different forms of monitoring such as wandb and tensorboard
"""

from abc import ABC, abstractmethod
import deepspeed.comm as dist


class Monitor(ABC):

    @abstractmethod
    def __init__(self, monitor_config):
        self.monitor_config = monitor_config

    @abstractmethod
    def write_events(self, event_list):
        pass


from .wandb import WandbMonitor
from .tensorboard import TensorBoardMonitor
from .csv_monitor import csvMonitor
from .comet import CometMonitor


class MonitorMaster(Monitor):

    def __init__(self, monitor_config):
        super().__init__(monitor_config)
        self.tb_monitor = None
        self.wandb_monitor = None
        self.csv_monitor = None
        self.comet_monitor = None
        self.enabled = monitor_config.enabled

        if dist.get_rank() == 0:
            if monitor_config.tensorboard.enabled:
                self.tb_monitor = TensorBoardMonitor(monitor_config.tensorboard)
            if monitor_config.wandb.enabled:
                self.wandb_monitor = WandbMonitor(monitor_config.wandb)
            if monitor_config.csv_monitor.enabled:
                self.csv_monitor = csvMonitor(monitor_config.csv_monitor)
            if monitor_config.comet.enabled:
                self.comet_monitor = CometMonitor(monitor_config.comet)

    def write_events(self, event_list):
        if dist.get_rank() == 0:
            if self.tb_monitor is not None:
                self.tb_monitor.write_events(event_list)
            if self.wandb_monitor is not None:
                self.wandb_monitor.write_events(event_list)
            if self.csv_monitor is not None:
                self.csv_monitor.write_events(event_list)
            if self.comet_monitor is not None:
                self.comet_monitor.write_events(event_list)

    def update_config(self, config_dict):
        if dist.get_rank() == 0:
            if self.tb_monitor is not None:
                self.tb_monitor.update_config(config_dict)
            if self.wandb_monitor is not None:
                self.wandb_monitor.update_config(config_dict)
            if self.csv_monitor is not None:
                self.csv_monitor.update_config(config_dict)
            if self.comet_monitor is not None:
                self.comet_monitor.update_config(config_dict)