Rename composer_run_name tag to run_name when logging to MLflow

composer/loggers/mlflow_logger.py

@@ -144,7 +144,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['composer_run_name'] = state.run_name
+        self.tags['run_name'] = state.run_name
 
         # Adjust name and group based on `rank_zero_only`.
         if not self._rank_zero_only:
@@ -161,9 +161,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
-                existing_runs = mlflow.search_runs(experiment_ids=[self._experiment_id],
-                                                   filter_string=f'tags.composer_run_name = "{state.run_name}"',
-                                                   output_format='list')
+                existing_runs = mlflow.search_runs(
+                    experiment_ids=[self._experiment_id],
+                    filter_string=f'tags.run_name = "{state.run_name}" OR 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

@@ -231,7 +231,7 @@ def test_mlflow_experiment_set_up(tmp_path):
     assert actual_run_name == expected_run_name
 
     # Check run tagged with Composer run name.
-    assert tags['composer_run_name'] == mock_state.run_name
+    assert tags['run_name'] == mock_state.run_name
 
     # Check run ended.
     test_mlflow_logger.post_close()