Cherry picks for mlflow 1.12.1 (mlflow#3717)

* Do not disable existing logger when init alembic (mlflow#3653)

Signed-off-by: ruoxuany <ruoxuan.yin@citrix.com>

Co-authored-by: ruoxuany <ruoxuan.yin@citrix.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* Only import matplotlib where necessary for artifact plotting (mlflow#3703)

* Only import matplotlib where necessary for artifact plotting

Signed-off-by: Corey Zumar <corey.zumar@databricks.com>

* Silly change for formatter

Signed-off-by: Corey Zumar <corey.zumar@databricks.com>
Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

* Fix run_link for cross-workspace model versions (mlflow#3681)

Fix run_link for cross-workspace model versions

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>

Co-authored-by: ruoxuany <ruoxuan.yin@citrix.com>

Author: harupy

mlflow/sklearn/utils.py

@@ -4,7 +4,6 @@
 import inspect
 import logging
 from numbers import Number
-import matplotlib.pyplot as plt
 import numpy as np
 import time
 
@@ -500,6 +499,8 @@ def _log_specialized_estimator_content(fitted_estimator, run_id, fit_args, fit_k
                     display.ax_.set_title(artifact.title)
                     filepath = tmp_dir.path("{}.png".format(artifact.name))
                     display.figure_.savefig(filepath)
+                    import matplotlib.pyplot as plt
+
                     plt.close(display.figure_)
                 except Exception as e:  # pylint: disable=broad-except
                     _log_warning_for_artifacts(artifact.name, artifact.function, e)

mlflow/store/db_migrations/env.py

@@ -11,7 +11,7 @@
 
 # Interpret the config file for Python logging.
 # This line sets up loggers basically.
-fileConfig(config.config_file_name)
+fileConfig(config.config_file_name, disable_existing_loggers=False)
 
 # add your model's MetaData object here
 # for 'autogenerate' support

mlflow/utils/databricks_utils.py

@@ -164,7 +164,7 @@ def get_job_type():
 
 
 def get_webapp_url():
-    """Should only be called if is_in_databricks_notebook is true"""
+    """Should only be called if is_in_databricks_notebook or is_in_databricks_jobs is true"""
     url = _get_property_from_spark_context("spark.databricks.api.url")
     if url is not None:
         return url
@@ -191,7 +191,8 @@ def get_browser_hostname():
 def get_workspace_info_from_dbutils():
     dbutils = _get_dbutils()
     if dbutils:
-        workspace_host = get_browser_hostname()
+        browser_hostname = get_browser_hostname()
+        workspace_host = "https://" + browser_hostname if browser_hostname else get_webapp_url()
         workspace_id = get_workspace_id()
         return workspace_host, workspace_id
     return None, None

tests/utils/test_databricks_utils.py

@@ -104,6 +104,23 @@ def test_get_workspace_info_from_databricks_secrets():
 def test_get_workspace_info_from_dbutils():
     mock_dbutils = mock.MagicMock()
     mock_dbutils.notebook.entry_point.getDbutils.return_value.notebook.return_value.getContext.return_value.browserHostName.return_value.get.return_value = (  # noqa
+        "mlflow.databricks.com"
+    )
+    mock_dbutils.notebook.entry_point.getDbutils.return_value.notebook.return_value.getContext.return_value.workspaceId.return_value.get.return_value = (  # noqa
+        "1111"
+    )
+    with mock.patch("mlflow.utils.databricks_utils._get_dbutils", return_value=mock_dbutils):
+        workspace_host, workspace_id = get_workspace_info_from_dbutils()
+        assert workspace_host == "https://mlflow.databricks.com"
+        assert workspace_id == "1111"
+
+
+def test_get_workspace_info_from_dbutils_no_browser_host_name():
+    mock_dbutils = mock.MagicMock()
+    mock_dbutils.notebook.entry_point.getDbutils.return_value.notebook.return_value.getContext.return_value.browserHostName.return_value.get.return_value = (  # noqa
+        None
+    )
+    mock_dbutils.notebook.entry_point.getDbutils.return_value.notebook.return_value.getContext.return_value.apiUrl.return_value.get.return_value = (  # noqa
         "https://mlflow.databricks.com"
     )
     mock_dbutils.notebook.entry_point.getDbutils.return_value.notebook.return_value.getContext.return_value.workspaceId.return_value.get.return_value = (  # noqa
@@ -121,7 +138,7 @@ def test_get_workspace_info_from_dbutils_old_runtimes():
         '{"tags": {"orgId" : "1111", "browserHostName": "mlflow.databricks.com"}}'
     )
     mock_dbutils.notebook.entry_point.getDbutils.return_value.notebook.return_value.getContext.return_value.browserHostName.return_value.get.return_value = (  # noqa
-        "https://mlflow.databricks.com"
+        "mlflow.databricks.com"
     )
     # Mock out workspace ID tag
     mock_workspace_id_tag_opt = mock.MagicMock()