wehub-resource-sync
60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
# 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)
|