wehub-resource-sync
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import TYPE_CHECKING, Any, Tuple, List, Dict, Optional
|
|
|
|
from .utils import check_comet_availability
|
|
from .monitor import Monitor
|
|
|
|
import deepspeed.comm as dist
|
|
|
|
if TYPE_CHECKING:
|
|
import comet_ml
|
|
from .config import CometConfig
|
|
|
|
Name = str
|
|
Value = Any
|
|
GlobalSamples = int
|
|
Event = Tuple[Name, Value, GlobalSamples]
|
|
|
|
|
|
class CometMonitor(Monitor):
|
|
|
|
def __init__(self, comet_config: "CometConfig"):
|
|
super().__init__(comet_config)
|
|
check_comet_availability()
|
|
import comet_ml
|
|
|
|
self.enabled = comet_config.enabled
|
|
self._samples_log_interval = comet_config.samples_log_interval
|
|
self._experiment: Optional["comet_ml.ExperimentBase"] = None
|
|
|
|
if self.enabled and dist.get_rank() == 0:
|
|
self._experiment = comet_ml.start(
|
|
api_key=comet_config.api_key,
|
|
project=comet_config.project,
|
|
workspace=comet_config.workspace,
|
|
experiment_key=comet_config.experiment_key,
|
|
mode=comet_config.mode,
|
|
online=comet_config.online,
|
|
)
|
|
|
|
if comet_config.experiment_name is not None:
|
|
self._experiment.set_name(comet_config.experiment_name)
|
|
|
|
self._events_log_scheduler = EventsLogScheduler(comet_config.samples_log_interval)
|
|
|
|
@property
|
|
def experiment(self) -> Optional["comet_ml.ExperimentBase"]:
|
|
return self._experiment
|
|
|
|
@property
|
|
def samples_log_interval(self) -> int:
|
|
return self._samples_log_interval
|
|
|
|
def write_events(self, event_list: List[Event]) -> None:
|
|
if not self.enabled or dist.get_rank() != 0:
|
|
return None
|
|
|
|
for event in event_list:
|
|
name = event[0]
|
|
value = event[1]
|
|
engine_global_samples = event[2]
|
|
|
|
if self._events_log_scheduler.needs_logging(name, engine_global_samples):
|
|
self._experiment.__internal_api__log_metric__(
|
|
name=name,
|
|
value=value,
|
|
step=engine_global_samples,
|
|
)
|
|
|
|
|
|
class EventsLogScheduler:
|
|
|
|
def __init__(self, samples_log_interval: int):
|
|
self._samples_log_interval = samples_log_interval
|
|
self._last_logged_events_samples: Dict[str, int] = {}
|
|
|
|
def needs_logging(self, name: str, current_sample: int) -> bool:
|
|
if name not in self._last_logged_events_samples:
|
|
self._last_logged_events_samples[name] = current_sample
|
|
return True
|
|
|
|
last_logged_sample = self._last_logged_events_samples[name]
|
|
samples_delta = current_sample - last_logged_sample
|
|
|
|
if samples_delta >= self._samples_log_interval:
|
|
self._last_logged_events_samples[name] = current_sample
|
|
return True
|
|
|
|
return False
|