Adding Model Data Init and Training Progress to MosaicMLLogger

composer/loggers/mosaicml_logger.py

@@ -18,6 +18,7 @@
 import mcli
 import torch
 
+from composer.core.time import TimeUnit
 from composer.loggers import Logger
 from composer.loggers.logger import Logger
 from composer.loggers.logger_destination import LoggerDestination
@@ -72,7 +73,6 @@ def __init__(
             self.time_last_logged = 0
             self.time_failed_count_adjusted = 0
             self.buffered_metadata: Dict[str, Any] = {}
-
             self.run_name = os.environ.get(RUN_NAME_ENV_VAR)
             if self.run_name is not None:
                 log.info(f'Logging to mosaic run {self.run_name}')
@@ -88,14 +88,58 @@ def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> No
         self._log_metadata(metrics)
 
     def after_load(self, state: State, logger: Logger) -> None:
+        # Log model data downloaded and initialized for run events
+        self._log_metadata({'model_initialized_time': time.time()})
         # Log WandB run URL if it exists. Must run on after_load as WandB is setup on event init
         for callback in state.callbacks:
             if isinstance(callback, WandBLogger):
                 run_url = callback.run_url
                 if run_url is not None:
                     self._log_metadata({'wandb/run_url': run_url})
 
+    def _get_training_progress_metrics(self, state: State) -> Dict[str, Any]:
+        """Calculates training progress metrics.
+
+        If user submits max duration:
+        - in tokens -> format: [token=x/xx]
+        - in batches -> format: [batch=x/xx]
+        - in epoch -> format: [epoch=x/xx] [batch=x/xx] (where batch refers to batches completed in current epoch)
+        If batches per epoch cannot be calculated, return [epoch=x/xx]
+
+        If no training duration given -> format: ''
+        """
+        assert state.max_duration is not None
+        if state.max_duration.unit == TimeUnit.TOKEN: 
+            return {
+                'training_progress': f'[token={state.timestamp.token.value}/{state.max_duration.value}]',
+            }
+        if state.max_duration.unit == TimeUnit.BATCH:
+            return {
+                'training_progress': f'[batch={state.timestamp.batch.value}/{state.max_duration.value}]',
+            }
+        training_progress_metrics = {}
+        if state.max_duration.unit == TimeUnit.EPOCH:
+            cur_batch = int(state.timestamp.batch_in_epoch)
+            cur_epoch = int(state.timestamp.epoch)
+            if int(state.timestamp.epoch) >= 1:
+                batches_per_epoch = int(
+                    (state.timestamp.batch - state.timestamp.batch_in_epoch).value / state.timestamp.epoch.value)
+                curr_progress = f'[batch={cur_batch}/{batches_per_epoch}]'
+            elif state.dataloader_len is None:
+                curr_progress = f'[batch={cur_batch}]'
+            else:
+                total = int(state.dataloader_len)
+                curr_progress = f'[batch={cur_batch}/{total}]'
+            if cur_epoch < state.max_duration.value:
+                cur_epoch += 1
+            training_progress_metrics = {
+                'training_progress': f'[epoch={cur_epoch}/{state.max_duration.value}]',
+            }
+            training_progress_metrics['training_sub_progress'] = curr_progress
+        return training_progress_metrics
+
     def batch_end(self, state: State, logger: Logger) -> None:
+        self._log_metadata(self._get_training_progress_metrics(state))
         self._flush_metadata()
 
     def epoch_end(self, state: State, logger: Logger) -> None:

tests/loggers/test_mosaicml_logger.py

@@ -3,13 +3,15 @@
 
 import json
 from typing import Type
+from unittest.mock import MagicMock
 
 import mcli
 import pytest
 import torch
 from torch.utils.data import DataLoader
 
 from composer.core import Callback
+from composer.core.time import Time, TimeUnit
 from composer.loggers import WandBLogger
 from composer.loggers.mosaicml_logger import (MOSAICML_ACCESS_TOKEN_ENV_VAR, MOSAICML_PLATFORM_ENV_VAR, MosaicMLLogger,
                                               format_data_to_json_serializable)
@@ -194,3 +196,93 @@ def test_auto_add_logger(monkeypatch, platform_env_var, access_token_env_var, lo
     # Otherwise, no logger
     else:
         assert logger_count == 0
+
+
+def test_model_initialized_time_logged(monkeypatch):
+    mock_mapi = MockMAPI()
+    monkeypatch.setattr(mcli, 'update_run_metadata', mock_mapi.update_run_metadata)
+    run_name = 'test-run-name'
+    monkeypatch.setenv('RUN_NAME', run_name)
+    trainer = Trainer(model=SimpleModel(),
+                      train_dataloader=DataLoader(RandomClassificationDataset()),
+                      train_subset_num_batches=1,
+                      max_duration='1ep',
+                      loggers=[MosaicMLLogger()])
+    trainer.fit()
+    assert isinstance(mock_mapi.run_metadata[run_name]['mosaicml/model_initialized_time'], float)
+
+
+def test_progress_logged(monkeypatch, tiny_bert_tokenizer):
+    mock_mapi = MockMAPI()
+    monkeypatch.setattr(mcli, 'update_run_metadata', mock_mapi.update_run_metadata)
+    run_name = 'test-run-name'
+    monkeypatch.setenv('RUN_NAME', run_name)
+    trainer = Trainer(model=SimpleModel(),
+                      train_dataloader=DataLoader(RandomClassificationDataset()),
+                      train_subset_num_batches=1,
+                      max_duration='4ba',
+                      loggers=[MosaicMLLogger()])
+    trainer.fit()
+    metadata = mock_mapi.run_metadata[run_name]
+    assert 'mosaicml/training_progress' in metadata
+    assert metadata['mosaicml/training_progress'] == '[batch=4/4]'
+    assert 'mosaicml/training_sub_progress' not in metadata
+
+
+def test_token_training_progress_logged():
+    logger = MosaicMLLogger()
+    state = MagicMock()
+    state.max_duration.unit = TimeUnit.TOKEN
+    state.max_duration.value = 64
+    state.timestamp.token.value = 50
+    training_progress = logger._get_training_progress_metrics(state)
+    assert 'training_progress' in training_progress
+    assert training_progress['training_progress'] == '[token=50/64]'
+    assert 'training_sub_progress' not in training_progress
+
+
+def test_epoch_training_progress_logged():
+    logger = MosaicMLLogger()
+    state = MagicMock()
+    state.max_duration.unit = TimeUnit.EPOCH
+    state.max_duration = Time(3, TimeUnit.EPOCH)
+    state.timestamp.epoch = Time(2, TimeUnit.EPOCH)
+    state.timestamp.batch = Time(11, TimeUnit.BATCH)
+    state.timestamp.batch_in_epoch = Time(1, TimeUnit.BATCH)
+    training_progress = logger._get_training_progress_metrics(state)
+    assert 'training_progress' in training_progress
+    assert training_progress['training_progress'] == '[epoch=3/3]'
+    assert 'training_sub_progress' in training_progress
+    assert training_progress['training_sub_progress'] == '[batch=1/5]'
+
+
+def test_epoch_zero_progress_logged():
+    logger = MosaicMLLogger()
+    state = MagicMock()
+    state.dataloader_len = 5
+    state.max_duration.unit = TimeUnit.EPOCH
+    state.max_duration = Time(3, TimeUnit.EPOCH)
+    state.timestamp.epoch = Time(0, TimeUnit.EPOCH)
+    state.timestamp.batch = Time(0, TimeUnit.BATCH)
+    state.timestamp.batch_in_epoch = Time(0, TimeUnit.BATCH)
+    training_progress = logger._get_training_progress_metrics(state)
+    assert 'training_progress' in training_progress
+    assert training_progress['training_progress'] == '[epoch=1/3]'
+    assert 'training_sub_progress' in training_progress
+    assert training_progress['training_sub_progress'] == '[batch=0/5]'
+
+
+def test_epoch_zero_no_dataloader_progress_logged():
+    logger = MosaicMLLogger()
+    state = MagicMock()
+    state.dataloader_len = None
+    state.max_duration.unit = TimeUnit.EPOCH
+    state.max_duration = Time(3, TimeUnit.EPOCH)
+    state.timestamp.epoch = Time(0, TimeUnit.EPOCH)
+    state.timestamp.batch = Time(1, TimeUnit.BATCH)
+    state.timestamp.batch_in_epoch = Time(1, TimeUnit.BATCH)
+    training_progress = logger._get_training_progress_metrics(state)
+    assert 'training_progress' in training_progress
+    assert training_progress['training_progress'] == '[epoch=1/3]'
+    assert 'training_sub_progress' in training_progress
+    assert training_progress['training_sub_progress'] == '[batch=1]'