diff --git a/composer/loggers/mosaicml_logger.py b/composer/loggers/mosaicml_logger.py index 1fcb612ed5..701ddd78a0 100644 --- a/composer/loggers/mosaicml_logger.py +++ b/composer/loggers/mosaicml_logger.py @@ -148,6 +148,10 @@ def fit_start(self, state: State, logger: Logger) -> None: def eval_end(self, state: State, logger: Logger) -> None: self._flush_metadata(force_flush=True) + def eval_after_all(self, state: State, logger: Logger) -> None: + # Log model evaluation finished time for run events + self.log_metadata({'eval_after_all_time': time.time()}, force_flush=True) + def predict_end(self, state: State, logger: Logger) -> None: self._flush_metadata(force_flush=True) diff --git a/tests/loggers/test_mosaicml_logger.py b/tests/loggers/test_mosaicml_logger.py index e308dab122..9fb976202f 100644 --- a/tests/loggers/test_mosaicml_logger.py +++ b/tests/loggers/test_mosaicml_logger.py @@ -322,6 +322,7 @@ def test_run_events_logged(monkeypatch): assert 'mosaicml/training_sub_progress' not in metadata assert isinstance(metadata['mosaicml/train_finished_time'], float) assert isinstance(metadata['mosaicml/train_started_time'], float) + assert isinstance(metadata['mosaicml/eval_after_all_time'], float) def test_token_training_progress_metrics():