Repository: mlflow/mlflow

Issue:
[BUG] sklearn.autolog fails without matplotlib installed
### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

1.26.1.dev0

### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: MacOS (with M1 chip)
- **Python version**: 3.8.13
- **yarn version, if running the dev UI**:


### Describe the problem

Since release 1.26.0, mlflow.sklearn.autolog fails to log experiments if matplotlib is not installed.

The issue is that you import matplotlib outside of the try block on line 489 in mlflow/mlflow/sklearn/utils.py.

Proposed fix: move matplotlib code into the try/except block (only place where it is used anyways).



### Tracking information

_No response_

### Code to reproduce issue

def tree_trainer_with_mlflow(
        X_train: np.ndarray,
        y_train: np.ndarray,
    ):
    mlflow.sklearn.autolog()  # log all model hparams and metrics to MLflow
    model = DecisionTreeClassifier()
    model.fit(X_train, y_train)
    return model

### Other info / logs

_No response_

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/sklearn/utils.py b/mlflow/sklearn/utils.py
--- a/mlflow/sklearn/utils.py
+++ b/mlflow/sklearn/utils.py
@@ -486,7 +486,6 @@ def _log_specialized_estimator_content(
     autologging_client, fitted_estimator, run_id, prefix, X, y_true, sample_weight, pos_label
 ):
     import sklearn
-    import matplotlib
 
     metrics = dict()
 
@@ -522,21 +521,28 @@ def _log_specialized_estimator_content(
                 + str(e)
             )
             _logger.warning(msg)
-            return
+            return metrics
 
-        _matplotlib_config = {"savefig.dpi": 175, "figure.autolayout": True, "font.size": 8}
+        try:
+            import matplotlib
+            import matplotlib.pyplot as plt
+        except ImportError as ie:
+            _logger.warning(
+                f"Failed to import matplotlib (error: {repr(ie)}). Skipping artifact logging."
+            )
+            return metrics
 
-        with TempDir() as tmp_dir, matplotlib.rc_context(_matplotlib_config):
+        _matplotlib_config = {"savefig.dpi": 175, "figure.autolayout": True, "font.size": 8}
+        with TempDir() as tmp_dir:
             for artifact in artifacts:
                 try:
-                    display = artifact.function(**artifact.arguments)
-                    display.ax_.set_title(artifact.title)
-                    artifact_path = "{}.png".format(artifact.name)
-                    filepath = tmp_dir.path(artifact_path)
-                    display.figure_.savefig(fname=filepath, format="png")
-                    import matplotlib.pyplot as plt
-
-                    plt.close(display.figure_)
+                    with matplotlib.rc_context(_matplotlib_config):
+                        display = artifact.function(**artifact.arguments)
+                        display.ax_.set_title(artifact.title)
+                        artifact_path = "{}.png".format(artifact.name)
+                        filepath = tmp_dir.path(artifact_path)
+                        display.figure_.savefig(fname=filepath, format="png")
+                        plt.close(display.figure_)
                 except Exception as e:
                     _log_warning_for_artifacts(artifact.name, artifact.function, e)
 


Test patch:
diff --git a/tests/sklearn/test_sklearn_autolog_without_matplotlib.py b/tests/sklearn/test_sklearn_autolog_without_matplotlib.py
new file mode 100644
--- /dev/null
+++ b/tests/sklearn/test_sklearn_autolog_without_matplotlib.py
@@ -0,0 +1,42 @@
+import pytest
+from sklearn.datasets import load_breast_cancer
+from sklearn.ensemble import RandomForestClassifier
+from unittest import mock
+
+import mlflow
+from mlflow.tracking import MlflowClient
+from tests.helper_functions import AnyStringWith
+
+
+def is_matplotlib_installed():
+    try:
+        import matplotlib  # pylint: disable=unused-import
+
+        return True
+    except ImportError:
+        return False
+
+
+@pytest.mark.skipif(
+    is_matplotlib_installed(), reason="matplotlib must be uninstalled to run this test"
+)
+def test_sklearn_autolog_works_without_matplotlib():
+    mlflow.sklearn.autolog()
+    model = RandomForestClassifier(max_depth=2, random_state=0, n_estimators=10)
+    X, y = load_breast_cancer(return_X_y=True)
+    with mlflow.start_run() as run, mock.patch(
+        "mlflow.sklearn.utils._logger.warning"
+    ) as mock_warning:
+        model.fit(X, y)
+        mock_warning.assert_called_once_with(AnyStringWith("Failed to import matplotlib"))
+
+    run = MlflowClient().get_run(run.info.run_id)
+    expected_metric_keys = {
+        "training_score",
+        "training_accuracy_score",
+        "training_precision_score",
+        "training_recall_score",
+        "training_f1_score",
+        "training_log_loss",
+    }
+    assert set(run.data.metrics).issuperset(expected_metric_keys)