Add MLFlow as a logger option (#475)

llmfoundry/utils/builders.py

@@ -12,7 +12,7 @@
 from composer.core import Evaluator
 from composer.datasets.in_context_learning_evaluation import \
     get_icl_task_dataloader
-from composer.loggers import TensorboardLogger, WandBLogger
+from composer.loggers import MLFlowLogger, TensorboardLogger, WandBLogger
 from composer.optim import DecoupledAdamW
 from composer.optim.scheduler import (ConstantWithWarmupScheduler,
                                       CosineAnnealingWithWarmupScheduler,
@@ -67,6 +67,8 @@ def build_logger(name: str, kwargs: Dict[str, Any]):
         return WandBLogger(**kwargs)
     elif name == 'tensorboard':
         return TensorboardLogger(**kwargs)
+    elif name == 'mlflow':
+        return MLFlowLogger(**kwargs)
     else:
         raise ValueError(f'Not sure how to build logger: {name}')
 

llmfoundry/utils/config_utils.py

@@ -98,3 +98,11 @@ def log_config(cfg: DictConfig):
             raise e
         if wandb.run:
             wandb.config.update(om.to_container(cfg, resolve=True))
+
+    if 'mlflow' in cfg.get('loggers', {}):
+        try:
+            import mlflow
+        except ImportError as e:
+            raise e
+        if mlflow.active_run():
+            mlflow.log_params(params=om.to_container(cfg, resolve=True))

setup.py

@@ -47,7 +47,7 @@
 ]
 
 install_requires = [
-    'mosaicml[libcloud,nlp,wandb]>=0.15.0,<0.16',
+    'mosaicml[libcloud,nlp,wandb,mlflow]>=0.15.0,<0.16',
     'accelerate>=0.20,<0.21',  # for HF inference `device_map`
     'mosaicml-streaming>=0.5.1,<0.6',
     'torch>=1.13.1,<=2.0.1',