Add CometML Support

composer/loggers/__init__.py

@@ -12,6 +12,7 @@
 define a custom logger and use it when training.
 """
 
+from composer.loggers.cometml_logger import CometMLLogger
 from composer.loggers.file_logger import FileLogger
 from composer.loggers.in_memory_logger import InMemoryLogger
 from composer.loggers.logger import Logger, LogLevel
@@ -32,4 +33,5 @@
     'WandBLogger',
     'ObjectStoreLogger',
     'TensorboardLogger',
+    'CometMLLogger',
 ]

composer/loggers/cometml_logger.py

@@ -0,0 +1,98 @@
+# Copyright 2022 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Log to `Comet <https://www.comet.com/docs/v2/>`_."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from composer.core.state import State
+from composer.loggers.logger import Logger
+from composer.loggers.logger_destination import LoggerDestination
+from composer.utils import dist
+from composer.utils.import_helpers import MissingConditionalImportError
+
+__all__ = ['CometMLLogger']
+
+
+class CometMLLogger(LoggerDestination):
+    """Log to `Comet <https://www.comet.com/docs/v2/>`_.
+
+    Args:
+        workspace (str, optional): The name of the workspace which contains the project
+            you want to attach your experiment to. If nothing specified will default to your
+            default workspace as configured in your comet account settings.
+        project_name (str, optional): The name of the project to categorize your experiment in.
+            A new project with this name will be created under the Comet workspace if one
+            with this name does not exist. If no project name specified, the experiment will go
+            under Uncategorized Experiments.
+        log_code (bool): Whether to log your code in your experiment (default: ``False``).
+        log_graph (bool): Whether to log your computational graph in your experiment
+            (default: ``False``).
+        name (str, optional): The name of your experiment. If not specified, it will be set
+            to :attr:`.State.run_name`.
+        rank_zero_only (bool, optional): Whether to log only on the rank-zero process.
+            (default: ``False``).
+        exp_kwargs (Dict[str, Any], optional): Any additional kwargs to
+            comet_ml.Experiment(see
+            `Comet documentation <https://www.comet.com/docs/v2/api-and-sdk/python-sdk/reference/Experiment/>`_).
+    """
+
+    def __init__(
+        self,
+        workspace: Optional[str] = None,
+        project_name: Optional[str] = None,
+        log_code: bool = False,
+        log_graph: bool = False,
+        name: Optional[str] = None,
+        rank_zero_only: bool = True,
+        exp_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        try:
+            from comet_ml import Experiment
+        except ImportError as e:
+            raise MissingConditionalImportError(extra_deps_group='comet_ml',
+                                                conda_package='comet_ml',
+                                                conda_channel='conda-forge') from e
+
+        self._enabled = (not rank_zero_only) or dist.get_global_rank() == 0
+
+        if exp_kwargs is None:
+            exp_kwargs = {}
+
+        if workspace is not None:
+            exp_kwargs['workspace'] = workspace
+
+        if project_name is not None:
+            exp_kwargs['project_name'] = project_name
+
+        exp_kwargs['log_code'] = log_code
+        exp_kwargs['log_graph'] = log_graph
+
+        self.name = name
+        self._rank_zero_only = rank_zero_only
+        self._exp_kwargs = exp_kwargs
+        self.experiment = Experiment(**self._exp_kwargs)
+
+    def init(self, state: State, logger: Logger) -> None:
+        del logger  # unused
+
+        # Use the logger run name if the name is not set.
+        if self.name is None:
+            self.name = state.run_name
+
+        # Adjust name and group based on `rank_zero_only`.
+        if not self._rank_zero_only:
+            self.name += f'-rank{dist.get_global_rank()}'
+
+        if self._enabled:
+            self.experiment.set_name(self.name)
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        if self._enabled:
+            self.experiment.log_metrics(dic=metrics, step=step)
+
+    def log_hyperparameters(self, hyperparameters: Dict[str, Any]):
+        if self._enabled:
+            self.experiment.log_parameters(hyperparameters)

composer/loggers/logger_hparams_registry.py

@@ -14,6 +14,7 @@
 
 import yahp as hp
 
+from composer.loggers.cometml_logger import CometMLLogger
 from composer.loggers.file_logger import FileLogger
 from composer.loggers.in_memory_logger import InMemoryLogger
 from composer.loggers.logger_destination import LoggerDestination
@@ -86,4 +87,5 @@ def initialize_object(self) -> ObjectStoreLogger:
     'tensorboard': TensorboardLogger,
     'in_memory': InMemoryLogger,
     'object_store': ObjectStoreLoggerHparams,
+    'comet_ml': CometMLLogger,
 }

docs/source/getting_started/installation.rst

@@ -23,6 +23,7 @@ the following installation targets are available:
 * ``pip install 'mosaicml[unet]'``: Installs Composer with support for :doc:`Unet </model_cards/unet>`.
 * ``pip install 'mosaicml[timm]'``: Installs Composer with support for :mod:`timm`.
 * ``pip install 'mosaicml[wandb]'``: Installs Composer with support for :mod:`wandb`.
+* ``pip install 'mosaicml[comet_ml]'``: Installs Composer with support for :mod:`comet_ml`.
 * ``pip install 'mosaicml[all]'``: Install all optional dependencies.
 
 For a developer install, clone directly:

docs/source/notes/run_name.md

@@ -53,13 +53,11 @@ See {class}`~.CheckpointSaver` for more information on specifying the arguments
 
 In addition to checkpointing, loggers also use the `run_name` for default logging.
 
-#### Tensorboard Logger
+#### Experiment Tracking Loggers
 
-The {class}`~.TensorboardLogger` will save all the logs for a run to a folder called `run_name` and the name of each run in the Tensorboard GUI will be `run_name`.
-
-#### Weights and Biases Logger
-
-The `run_name` you specify will be used by the {class}`~.WandBLogger` as the run name for Weights and Biases.
+* The {class}`~.TensorboardLogger` will save all the logs for a run to a folder called `run_name` and the name of each run in the Tensorboard GUI will be `run_name`.
+* The `run_name` you specify will be used by the {class}`~.WandBLogger` as the run name for Weights and Biases.
+* The `run_name` you specify will be used by the {class}`~.CometMLLogger` as the run name for your Comet experiment.
 
 #### Object Store Logger
 

docs/source/trainer/logging.rst

@@ -45,6 +45,7 @@ Available Loggers
 
     ~file_logger.FileLogger
     ~wandb_logger.WandBLogger
+    ~cometml_logger.CometMLLogger
     ~progress_bar_logger.ProgressBarLogger
     ~tensorboard_logger.TensorboardLogger
     ~in_memory_logger.InMemoryLogger

examples/run_composer_trainer.py

@@ -44,14 +44,8 @@ def _main():
 
     trainer = hparams.initialize_object()
 
-    # if using wandb, store the config inside the wandb run
-    try:
-        import wandb
-    except ImportError:
-        pass
-    else:
-        if wandb.run is not None:
-            wandb.config.update(hparams.to_dict())
+    # Log all hyperparameters.
+    trainer.logger.log_hyperparameters(hparams.to_dict())
 
     # Only log the config once, since it should be the same on all ranks.
     if dist.get_global_rank() == 0:

setup.py

@@ -144,6 +144,8 @@ def package_files(prefix: str, directory: str, extension: str):
     'wandb>=0.13.2,<0.14',
 ]
 
+extra_deps['comet_ml'] = ['comet_ml>=3.31.12,<4.0.0']
+
 extra_deps['tensorboard'] = [
     'tensorboard>=2.9.1,<3.0.0',
     'tensorflow-io>=0.26.0,<0.27',

tests/callbacks/callback_settings.py

@@ -1,6 +1,7 @@
 # Copyright 2022 MosaicML Composer authors
 # SPDX-License-Identifier: Apache-2.0
 
+import os
 from typing import Any, Dict, List, Type, Union
 
 import pytest
@@ -14,7 +15,7 @@
 from composer.callbacks.callback_hparams_registry import callback_registry
 from composer.callbacks.export_for_inference import ExportForInferenceCallback
 from composer.callbacks.mlperf import MLPerfCallback
-from composer.loggers import ObjectStoreLogger, TensorboardLogger, WandBLogger
+from composer.loggers import CometMLLogger, ObjectStoreLogger, TensorboardLogger, WandBLogger
 from composer.loggers.logger_destination import LoggerDestination
 from composer.loggers.logger_hparams_registry import ObjectStoreLoggerHparams, logger_registry
 from composer.loggers.progress_bar_logger import ProgressBarLogger
@@ -35,6 +36,17 @@
 except ImportError:
     _TENSORBOARD_INSTALLED = False
 
+try:
+    import comet_ml
+    _COMETML_INSTALLED = True
+    os.environ['COMET_API_KEY']
+    del comet_ml  # unused
+except ImportError:
+    _COMETML_INSTALLED = False
+# If COMET_API_KEY not set.
+except KeyError:
+    _COMETML_INSTALLED = False
+
 try:
     import mlperf_logging
     _MLPERF_INSTALLED = True
@@ -115,6 +127,7 @@
         pytest.mark.filterwarnings(
             r'ignore:Specifying the ProgressBarLogger via `loggers` is deprecated:DeprecationWarning')
     ],
+    CometMLLogger: [pytest.mark.skipif(not _COMETML_INSTALLED, reason='comet_ml is optional'),],
     TensorboardLogger: [pytest.mark.skipif(not _TENSORBOARD_INSTALLED, reason='Tensorboard is optional'),],
     ObjectStoreLoggerHparams: [
         pytest.mark.filterwarnings(

tests/loggers/test_cometml_logger.py

@@ -0,0 +1,63 @@
+# Copyright 2022 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import zipfile
+from json import JSONDecoder
+from pathlib import Path
+
+import pytest
+
+
+def test_comet_ml_logging(monkeypatch, tmp_path):
+    """Check metrics logged with CometMLLogger are properly written to offline dump."""
+    pytest.importorskip('comet_ml', reason='comet_ml is optional')
+    import comet_ml
+
+    # Set some dummy log values.
+    steps = [0, 1, 2]
+    metric_values = [0.1, 0.4, 0.7]
+    metric_name = 'my_test_metric'
+    param_names = ['my_cool_parameter1', 'my_cool_parameter2']
+    param_values = [10, 3]
+
+    # Set offline directory.
+    offline_directory = str(tmp_path / Path('.my_cometml_runs'))
+    os.environ['COMET_OFFLINE_DIRECTORY'] = offline_directory
+
+    # Monkeypatch Experiment with OfflineExperiment to avoid uploading to CometML and
+    # avoid needing an API+KEY.
+    monkeypatch.setattr(comet_ml, 'Experiment', comet_ml.OfflineExperiment)
+    from composer.loggers import CometMLLogger
+
+    # Log dummy values with CometMLLogger.
+    comet_logger = CometMLLogger()
+    comet_logger.log_hyperparameters(dict(zip(param_names, param_values)))
+    for step, metric_value in zip(steps, metric_values):
+        comet_logger.log_metrics({'my_test_metric': metric_value}, step=step)
+    comet_logger.experiment.end()
+
+    # Open, decompress, decode, and extract offline dump of metrics.
+    comet_exp_dump_filepath = str(Path(offline_directory) / Path(comet_logger.experiment.id).with_suffix('.zip'))
+    zf = zipfile.ZipFile(comet_exp_dump_filepath)
+    comet_logs_path = zf.extract('messages.json', path=offline_directory)
+    jd = JSONDecoder()
+    metric_msgs = []
+    param_msgs = []
+    with open(comet_logs_path) as f:
+        for line in f.readlines():
+            comet_msg = jd.decode(line)
+            if (comet_msg['type'] == 'metric_msg') and (comet_msg['payload']['metric']['metricName']
+                                                        == 'my_test_metric'):
+                metric_msgs.append(comet_msg['payload']['metric'])
+            if comet_msg['type'] == 'parameter_msg' and (
+                    comet_msg['payload']['param']['paramName'].startswith('my_cool')):
+                param_msgs.append(comet_msg['payload']['param'])
+
+    # Assert dummy metrics input to log_metrics are the same as
+    # those written to offline dump.
+    assert [msg['metricValue'] for msg in metric_msgs] == metric_values
+    assert [msg['step'] for msg in metric_msgs] == steps
+    assert all([msg['metricName'] == metric_name for msg in metric_msgs])
+    assert [msg['paramValue'] for msg in param_msgs] == param_values
+    assert [msg['paramName'] for msg in param_msgs] == param_names