Fix MLFlow Tag Name for Resumption

composer/loggers/mlflow_logger.py

@@ -148,7 +148,7 @@ def init(self, state: State, logger: Logger) -> None:
 
         # Store the Composer run name in the MLFlow run tags so it can be retrieved for autoresume.
         self.tags = self.tags or {}
-        self.tags['run_name'] = state.run_name
+        self.tags['run_name'] = os.environ['RUN_NAME'] if 'RUN_NAME' in os.environ else state.run_name
 
         # Adjust name and group based on `rank_zero_only`.
         if not self._rank_zero_only:
@@ -171,16 +171,6 @@ def init(self, state: State, logger: Logger) -> None:
                     output_format='list',
                 )
 
-                # Check for the old tag (`composer_run_name`) For backwards compatibility in case a run using the old
-                # tag fails and the run is resumed with a newer version of Composer that uses `run_name` instead of
-                # `composer_run_name`.
-                if len(existing_runs) == 0:
-                    existing_runs = mlflow.search_runs(
-                        experiment_ids=[self._experiment_id],
-                        filter_string=f'tags.composer_run_name = "{state.run_name}"',
-                        output_format='list',
-                    )
-
                 if len(existing_runs) > 0:
                     self._run_id = existing_runs[0].info.run_id
                 else:

tests/loggers/test_mlflow_logger.py

@@ -190,6 +190,37 @@ def test_mlflow_experiment_init_existing_composer_run(monkeypatch):
     assert test_logger._run_id == existing_id
 
 
+def test_mlflow_logger_uses_env_var_run_name(monkeypatch):
+    """Test that MLFlowLogger uses the 'RUN_NAME' environment variable if set."""
+    mlflow = pytest.importorskip('mlflow')
+    monkeypatch.setattr(mlflow, 'set_tracking_uri', MagicMock())
+    monkeypatch.setattr(mlflow, 'start_run', MagicMock())
+
+    mock_state = MagicMock()
+    mock_state.run_name = 'state-run-name'
+    env_var_run_name = 'env-run-name'
+    monkeypatch.setenv('RUN_NAME', env_var_run_name)
+
+    logger = MLFlowLogger()
+    logger.init(state=mock_state, logger=MagicMock())
+    assert logger.tags['run_name'] == env_var_run_name, "Logger should use the run name from the environment variable."
+    monkeypatch.delenv('RUN_NAME')
+
+
+def test_mlflow_logger_uses_state_run_name_if_no_env_var_set(monkeypatch):
+    """Test that MLFlowLogger uses the state's run name if no 'RUN_NAME' environment variable is set."""
+    mlflow = pytest.importorskip('mlflow')
+    monkeypatch.setattr(mlflow, 'set_tracking_uri', MagicMock())
+    monkeypatch.setattr(mlflow, 'start_run', MagicMock())
+
+    mock_state = MagicMock()
+    mock_state.run_name = 'state-run-name'
+
+    logger = MLFlowLogger()
+    logger.init(state=mock_state, logger=MagicMock())
+    assert logger.tags['run_name'] == 'state-run-name', "Logger should fallback to the state's run name when no environment variable is set."
+
+
 def test_mlflow_experiment_init_existing_composer_run_with_old_tag(monkeypatch):
     """ Test that an existing MLFlow run is used if one exists with the old `composer_run_name` tag.
     """