Fix MLFlow Tag Name for Resumption

[KuuCi]

https://databricks.atlassian.net/browse/GRT-2836

[jerrychen109]

Can you also update test_mlflow_logger.py to unit test the behavior when the env var is set differently from the state.run_name? Would also be good to verify the fallback behavior if the env var isn't set.

composer/tests/loggers/test_mlflow_logger.py

Lines 173 to 210 in 83cb5c1

	def test_mlflow_experiment_init_existing_composer_run(monkeypatch):
	""" Test that an existing MLFlow run is used if one tagged with `run_name` exists in the experiment for the Composer run.
	"""
	mlflow = pytest.importorskip('mlflow')
	monkeypatch.setattr(mlflow, 'set_tracking_uri', MagicMock())
	monkeypatch.setattr(mlflow, 'start_run', MagicMock())
	mock_state = MagicMock()
	mock_state.run_name = 'dummy-run-name'
	existing_id = 'dummy-id'
	mock_search_runs = MagicMock(return_value=[MagicMock(info=MagicMock(run_id=existing_id))])
	monkeypatch.setattr(mlflow, 'search_runs', mock_search_runs)
	test_logger = MLFlowLogger()
	test_logger.init(state=mock_state, logger=MagicMock())
	assert test_logger._run_id == existing_id
	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.
	"""
	mlflow = pytest.importorskip('mlflow')
	monkeypatch.setattr(mlflow, 'set_tracking_uri', MagicMock())
	monkeypatch.setattr(mlflow, 'start_run', MagicMock())
	mock_state = MagicMock()
	mock_state.composer_run_name = 'dummy-run-name'
	existing_id = 'dummy-id'
	mock_search_runs = MagicMock(return_value=[MagicMock(info=MagicMock(run_id=existing_id))])
	monkeypatch.setattr(mlflow, 'search_runs', mock_search_runs)
	test_logger = MLFlowLogger()
	test_logger.init(state=mock_state, logger=MagicMock())
	assert test_logger._run_id == existing_id

[jerrychen109]

lgtm once you get the code quality/linting to pass - thanks for adding the tests and handling the mocking! (but you'll need to get @dakinggg's approval as code owner)