Add analytics logging to MosaicMLLogger

composer/core/engine.py

@@ -83,6 +83,7 @@ def run_last(algorithms: Sequence[Algorithm], event: Event) -> Sequence[Algorith
 from composer.core.event import Event
 from composer.core.state import State
 from composer.loggers import Logger, LoggerDestination
+from composer.loggers.mosaicml_logger import log_run_analytics
 from composer.profiler import ProfilerAction
 from composer.utils import ensure_tuple
 
@@ -293,6 +294,10 @@ def run_event(
             self._run_loggers(event)
             self._run_nonlogger_callbacks(event)
             traces = self._run_algorithms(event)
+
+            # If a MosaicMLLogger is present, log analytics for the run to metadata.
+            log_run_analytics(self.logger.destinations)
+
         else:
             traces = self._run_algorithms(event)
             # Run callbacks first, so any log calls from a callback that are executed lazily

composer/loggers/mosaicml_logger.py

@@ -15,16 +15,20 @@
 import warnings
 from concurrent.futures import wait
 from functools import reduce
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 import mcli
 import torch
+import torch.utils.data
 
-from composer.core.time import TimeUnit
+from composer.core.event import Event
+from composer.core.time import Time, TimeUnit
 from composer.loggers import Logger
 from composer.loggers.logger_destination import LoggerDestination
 from composer.loggers.wandb_logger import WandBLogger
 from composer.utils import dist
+from composer.utils.analytics_helpers import get_logger_type
+from composer.utils.file_helpers import parse_uri
 
 if TYPE_CHECKING:
     from composer.core import State
@@ -69,10 +73,12 @@ def __init__(
         log_interval: int = 60,
         ignore_keys: Optional[List[str]] = None,
         ignore_exceptions: bool = False,
+        analytics_data: Optional[Dict[str, Any]] = None,
     ) -> None:
         self.log_interval = log_interval
         self.ignore_keys = ignore_keys
         self.ignore_exceptions = ignore_exceptions
+        self.analytics_data = analytics_data
         self._enabled = dist.get_global_rank() == 0
         if self._enabled:
             self.time_last_logged = 0
@@ -96,6 +102,90 @@ def log_hyperparameters(self, hyperparameters: Dict[str, Any]):
     def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
         self._log_metadata(metrics)
 
+    def log_analytics(self,) -> None:
+        if self.analytics_data is None:
+            return
+
+        # Fetch / cast metrics that we want to log from self.analytics_data
+        autoresume: bool = self.analytics_data['autoresume']
+        trainer_state: State = self.analytics_data['state']
+        save_interval: Union[str, int, Time, Callable[[State, Event], bool]] = self.analytics_data['save_interval']
+        loggers: List[LoggerDestination] = self.analytics_data['loggers']
+        load_path: Union[str, None] = self.analytics_data['load_path']
+        save_folder: Union[str, None] = self.analytics_data['save_folder']
+
+        metrics: Dict[str, Any] = {'composer/autoresume': autoresume, 'composer/precision': trainer_state.precision}
+
+        train_dataloader = trainer_state.train_dataloader
+        if train_dataloader is not None and isinstance(train_dataloader, torch.utils.data.DataLoader):
+            metrics['composer/train_loader_workers'] = train_dataloader.num_workers
+
+        metrics['composer/eval_loaders'] = []
+        for evaluator in trainer_state.evaluators:
+            dataloader = evaluator.dataloader.dataloader
+            if isinstance(dataloader, torch.utils.data.DataLoader):
+                metrics['composer/eval_loaders'].append(
+                    json.dumps({
+                        'label': evaluator.label,
+                        'num_workers': dataloader.num_workers,
+                    }),
+                )
+
+        metrics['composer/optimizers'] = [
+            json.dumps(optimizer.state_dict(), sort_keys=True) for optimizer in trainer_state.optimizers
+        ]
+        metrics['composer/algorithms'] = [
+            json.dumps(algorithm.state_dict(), sort_keys=True) for algorithm in trainer_state.algorithms
+        ]
+        metrics['composer/loggers'] = [
+            get_logger_type(logger) if not isinstance(logger, MosaicMLLogger) else 'MosaicMLLogger'
+            for logger in loggers
+        ]
+
+        # Take the service provider out of the URI and log it to metadata. If no service provider
+        # is found (i.e. backend = ''), then we assume 'local' for the cloud provider.
+        if load_path is not None:
+            backend, _, _ = parse_uri(load_path)
+            metrics['composer/cloud_provider_data'] = backend if backend else 'local'
+        if save_folder is not None:
+            backend, _, _ = parse_uri(save_folder)
+            metrics['composer/cloud_provider_checkpoints'] = backend if backend else 'local'
+
+        # Save interval can be passed in w/ multiple types. If the type is a function, then
+        # we log 'callable' as the save_interval value for analytics.
+        if isinstance(save_interval, Union[str, int]):
+            save_interval_str = str(save_interval)
+        elif isinstance(save_interval, Time):
+            save_interval_str = f'{save_interval._value}{save_interval._unit}'
+        else:
+            save_interval_str = 'callable'
+        metrics['composer/save_interval'] = save_interval_str
+
+        if trainer_state.fsdp_config:
+            metrics['composer/sharding_strategy'] = trainer_state.fsdp_config.get('sharding_strategy', None)
+            metrics['composer/activation_checkpointing'] = trainer_state.fsdp_config.get(
+                'activation_checkpointing',
+                False,
+            )
+            metrics['composer/forward_prefetch'] = trainer_state.fsdp_config.get('forward_prefetch', False)
+            metrics['composer/backward_prefetch'] = trainer_state.fsdp_config.get('backward_prefetch', None)
+
+            # Get device_mesh from config so it is in list form and JSON parsable
+            metrics['composer/device_mesh'] = trainer_state.fsdp_config.get('device_mesh', [])
+
+            mixed_precision = trainer_state.fsdp_config.get('mixed_precision', None)
+            if mixed_precision is not None and isinstance(mixed_precision, dict):
+                # Sorting the keys allows us to parse this dict value as JSON in a SQL query if needed
+                metrics['composer/mixed_precision'] = json.dumps(mixed_precision, sort_keys=True)
+            else:
+                metrics['composer/mixed_precision'] = mixed_precision
+
+        if trainer_state.fsdp_state_dict_type is not None:
+            metrics['composer/state_dict_type'] = trainer_state.fsdp_state_dict_type
+
+        self.log_metrics(metrics)
+        self._flush_metadata(force_flush=True)
+
     def log_exception(self, exception: Exception):
         self._log_metadata({'exception': exception_to_json_serializable_dict(exception)})
         self._flush_metadata(force_flush=True)
@@ -284,3 +374,14 @@ def exception_to_json_serializable_dict(exc: Exception):
             except AttributeError:
                 pass
     return exc_data
+
+
+def log_run_analytics(loggers: Tuple[LoggerDestination, ...]):
+    """Log run analytics to metadata if a MosaicMLLogger is available in the list."""
+    # Avoids a casting bug during testing
+    if not isinstance(loggers, Iterable):
+        return
+
+    mosaicml_logger = next((logger for logger in loggers if isinstance(logger, MosaicMLLogger)), None)
+    if mosaicml_logger is not None:
+        mosaicml_logger.log_analytics()

composer/trainer/trainer.py

@@ -1284,7 +1284,16 @@ def __init__(
             MOSAICML_ACCESS_TOKEN_ENV_VAR,
         ) is not None and not any(isinstance(x, MosaicMLLogger) for x in loggers):
             log.info('Detected run on MosaicML platform. Adding MosaicMLLogger to loggers.')
-            mosaicml_logger = MosaicMLLogger()
+            mosaicml_logger = MosaicMLLogger(
+                analytics_data={
+                    'autoresume': autoresume,
+                    'state': self.state,
+                    'save_interval': save_interval,
+                    'loggers': loggers,
+                    'load_path': load_path,
+                    'save_folder': save_folder,
+                },
+            )
             loggers.append(mosaicml_logger)
 
         # Remote Uploader Downloader

composer/utils/analytics_helpers.py

@@ -0,0 +1,42 @@
+# Copyright 2024 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Helpers for logging analytics with the MosaicMLLogger."""
+
+from typing import Any
+
+from composer.loggers.cometml_logger import CometMLLogger
+from composer.loggers.console_logger import ConsoleLogger
+from composer.loggers.file_logger import FileLogger
+from composer.loggers.in_memory_logger import InMemoryLogger
+from composer.loggers.logger_destination import LoggerDestination
+from composer.loggers.mlflow_logger import MLFlowLogger
+from composer.loggers.neptune_logger import NeptuneLogger
+from composer.loggers.progress_bar_logger import ProgressBarLogger
+from composer.loggers.remote_uploader_downloader import RemoteUploaderDownloader
+from composer.loggers.slack_logger import SlackLogger
+from composer.loggers.tensorboard_logger import TensorboardLogger
+from composer.loggers.wandb_logger import WandBLogger
+
+LOGGER_TYPES = [
+    FileLogger,
+    SlackLogger,
+    WandBLogger,
+    MLFlowLogger,
+    NeptuneLogger,
+    ConsoleLogger,
+    CometMLLogger,
+    InMemoryLogger,
+    TensorboardLogger,
+    ProgressBarLogger,
+    RemoteUploaderDownloader,
+    LoggerDestination,
+]
+
+
+def get_logger_type(logger: Any) -> str:
+    """Returns the type of a logger as a string. If the logger is not a known type, returns 'Custom'."""
+    for logger_type in LOGGER_TYPES:
+        if isinstance(logger, logger_type):
+            return logger_type.__name__
+    return 'Custom'