Search for proper run_name and not state name

composer/loggers/mlflow_logger.py

@@ -9,6 +9,7 @@
 import logging
 import os
 import pathlib
+import posixpath
 import textwrap
 import time
 import warnings
@@ -110,6 +111,7 @@ def __init__(
 
         self._experiment_id: Optional[str] = None
         self._run_id = None
+        self.run_url = None
 
         if self._enabled:
             self.tracking_uri = str(tracking_uri or mlflow.get_tracking_uri())
@@ -165,9 +167,10 @@ def init(self, state: State, logger: Logger) -> None:
             else:
                 # Search for an existing run tagged with this Composer run.
                 assert self._experiment_id is not None
+                run_name = self.tags['run_name']
                 existing_runs = mlflow.search_runs(
                     experiment_ids=[self._experiment_id],
-                    filter_string=f'tags.run_name = "{state.run_name}"',
+                    filter_string=f'tags.run_name = "{run_name}"',
                     output_format='list',
                 )
 
@@ -194,6 +197,14 @@ def init(self, state: State, logger: Logger) -> None:
 
     def after_load(self, state: State, logger: Logger) -> None:
         logger.log_hyperparameters({'mlflow_experiment_id': self._experiment_id, 'mlflow_run_id': self._run_id})
+        self.run_url = posixpath.join(
+            os.environ.get('DATABRICKS_HOST', ''),
+            'ml',
+            'experiments',
+            str(self._experiment_id),
+            'runs',
+            str(self._run_id),
+        )
 
     def log_table(
         self,

composer/loggers/mosaicml_logger.py

@@ -23,6 +23,7 @@
 from composer.core.time import TimeUnit
 from composer.loggers import Logger
 from composer.loggers.logger_destination import LoggerDestination
+from composer.loggers.mlflow_logger import MLFlowLogger
 from composer.loggers.wandb_logger import WandBLogger
 from composer.utils import dist
 
@@ -113,6 +114,9 @@ def after_load(self, state: State, logger: Logger) -> None:
                     log.debug(f'Logging WandB run URL to metadata: {run_url}')
                 else:
                     log.debug('WandB run URL not found, not logging to metadata')
+            if isinstance(callback, MLFlowLogger) and callback._enabled:
+                self._log_metadata({'mlflow/run_url': callback.run_url})
+                log.debug(f'Logging MLFlow run URL to metadata: {callback.run_url}')
         self._flush_metadata(force_flush=True)
 
     def batch_start(self, state: State, logger: Logger) -> None: