[DO NOT MERGE] prototyping

.gitignore

@@ -3,6 +3,8 @@ my-copy-c4*/
 my-copy-arxiv*/
 *.jsonl*
 
+ygong/notebook/*
+
 # WandB
 wandb/
 
@@ -156,3 +158,5 @@ notebooks/
 **/mlruns/*
 **/tokenizer-save-dir-*/**
 **/.downloaded_finetuning/
+
+.databricks

llmfoundry/composerpatch/MLFlowLogger.py

@@ -0,0 +1,31 @@
+from composer.loggers import MLFlowLogger as ComposerMLFlowLogger
+from composer.utils import dist
+import json
+import os
+from composer.core.state import State
+from composer.loggers.logger import Logger
+
+
+
+CONFIG_FILE = "/tmp/mlflow_config.yaml"
+EXPERIMENT_ID_FIELD = "experiment_id"
+RUN_ID_FIELD = "run_id"
+TRACKING_URI_FIELD = "tracking_uri"
+
+
+class MLFlowLogger(ComposerMLFlowLogger):
+
+    def init(self, state: State, logger: Logger) -> None:
+        super().init(state, logger)
+
+        if self._enabled and dist.get_local_rank() == 0:
+           if os.path.exists(CONFIG_FILE):
+                os.remove(CONFIG_FILE)
+
+           with open(CONFIG_FILE, "w") as f:
+               data = {
+                     EXPERIMENT_ID_FIELD: self._experiment_id,
+                     RUN_ID_FIELD: self._run_id,
+                     TRACKING_URI_FIELD : self.tracking_uri,
+               }
+               json.dump(data, f)

llmfoundry/utils/builders.py

@@ -14,6 +14,8 @@
 from composer.core import Algorithm, Callback, Evaluator
 from composer.datasets.in_context_learning_evaluation import \
     get_icl_task_dataloader
+from composer.loggers import (InMemoryLogger, LoggerDestination,
+                              TensorboardLogger, WandBLogger)
 from composer.loggers import LoggerDestination
 from composer.models import ComposerModel
 from composer.optim.scheduler import ComposerScheduler
@@ -24,6 +26,8 @@
 from torchmetrics import Metric
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
+from llmfoundry.composerpatch import MLFlowLogger
+
 from llmfoundry import registry
 from llmfoundry.callbacks import EvalGauntlet
 from llmfoundry.data.dataloader import build_dataloader
@@ -236,15 +240,19 @@ def build_callback(
                                    kwargs=kwargs)
 
 
-def build_logger(name: str,
-                 kwargs: Optional[Dict[str, Any]] = None) -> LoggerDestination:
-    """Builds a logger from the registry."""
-    return construct_from_registry(name=name,
-                                   registry=registry.loggers,
-                                   partial_function=True,
-                                   pre_validation_function=LoggerDestination,
-                                   post_validation_function=None,
-                                   kwargs=kwargs)
+def build_logger(name: str, kwargs: Dict[str, Any]) -> LoggerDestination:
+    if name == 'wandb':
+        return WandBLogger(**kwargs)
+    elif name == 'tensorboard':
+        return TensorboardLogger(**kwargs)
+    elif name == 'in_memory_logger':
+        return InMemoryLogger(**kwargs)
+    elif name == 'mlflow':
+        return MLFlowLogger.MLFlowLogger(**kwargs)
+    elif name == 'inmemory':
+        return InMemoryLogger(**kwargs)
+    else:
+        raise ValueError(f'Not sure how to build logger: {name}')
 
 
 def build_algorithm(name: str,