| |
| |
|
|
| |
|
|
| 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 |
|
|
