Repository: mlflow/mlflow

Issue:
[DOC-FIX] Figure Closed Before Saving in MLflow with Optuna Tutorial
### Willingness to contribute

No. I cannot contribute a documentation fix at this time.

### URL(s) with the issue

https://mlflow.org/docs/latest/traditional-ml/hyperparameter-tuning-with-child-runs/notebooks/hyperparameter-tuning-with-child-runs.html#:~:text=%23%20prevent%20matplotlib%20from%20displaying%20the%20chart%20every%20time%20we%20call%20this%20function%0A%20%20%20%20plt.close(fig)%0A%0A%20%20%20%20%23%20Save%20the%20plot%20if%20save_path%20is%20specified%0A%20%20%20%20if%20save_path%3A%0A%20%20%20%20%20%20%20%20plt.savefig(save_path%2C%20format%3D%22png%22%2C%20dpi%3D600)

### Description of proposal (what needs changing)

Move `plt.close(fig)` to be after `plt.savefig`

Discussion:

Patch:
diff --git a/docs/source/traditional-ml/hyperparameter-tuning-with-child-runs/notebooks/hyperparameter-tuning-with-child-runs.ipynb b/docs/source/traditional-ml/hyperparameter-tuning-with-child-runs/notebooks/hyperparameter-tuning-with-child-runs.ipynb
index 314e98ad1..0959b4042 100644
--- a/docs/source/traditional-ml/hyperparameter-tuning-with-child-runs/notebooks/hyperparameter-tuning-with-child-runs.ipynb
+++ b/docs/source/traditional-ml/hyperparameter-tuning-with-child-runs/notebooks/hyperparameter-tuning-with-child-runs.ipynb
@@ -564,13 +564,13 @@
     "\n",
     "    plt.tight_layout()\n",
     "\n",
-    "    # prevent matplotlib from displaying the chart every time we call this function\n",
-    "    plt.close(fig)\n",
-    "\n",
     "    # Save the plot if save_path is specified\n",
     "    if save_path:\n",
     "        plt.savefig(save_path, format=\"png\", dpi=600)\n",
     "\n",
+    "    # prevent matplotlib from displaying the chart every time we call this function\n",
+    "    plt.close(fig)\n",
+    "\n",
     "    return fig\n",
     "\n",
     "\n",
diff --git a/mlflow/data/digest_utils.py b/mlflow/data/digest_utils.py
index 37c5ef2a5..e6f8855bc 100644
--- a/mlflow/data/digest_utils.py
+++ b/mlflow/data/digest_utils.py
@@ -1,7 +1,5 @@
 from typing import Any, List
 
-from packaging.version import Version
-
 from mlflow.exceptions import MlflowException
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
 from mlflow.utils import insecure_hash
@@ -17,7 +15,6 @@ def compute_pandas_digest(df) -> str:
 
     Returns:
         A string digest.
-
     """
     import numpy as np
     import pandas as pd
@@ -50,7 +47,6 @@ def compute_numpy_digest(features, targets=None) -> str:
 
     Returns:
         A string digest.
-
     """
     import numpy as np
     import pandas as pd
@@ -84,82 +80,6 @@ def compute_numpy_digest(features, targets=None) -> str:
     return get_normalized_md5_digest(hashable_elements)
 
 
-def compute_tensorflow_dataset_digest(dataset, targets=None) -> str:
-    """Computes a digest for the given Tensorflow dataset.
-
-    Args:
-        dataset: A Tensorflow dataset.
-
-    Returns:
-        A string digest.
-
-    """
-    import numpy as np
-    import pandas as pd
-    import tensorflow as tf
-
-    hashable_elements = []
-
-    def hash_tf_dataset_iterator_element(element):
-        if element is None:
-            return
-        flat_element = tf.nest.flatten(element)
-        flattened_array = np.concatenate([x.flatten() for x in flat_element])
-        trimmed_array = flattened_array[0:MAX_ROWS]
-        try:
-            hashable_elements.append(pd.util.hash_array(trimmed_array))
-        except TypeError:
-            hashable_elements.append(np.int64(trimmed_array.size))
-
-    for element in dataset.as_numpy_iterator():
-        hash_tf_dataset_iterator_element(element)
-    if targets is not None:
-        for element in targets.as_numpy_iterator():
-            hash_tf_dataset_iterator_element(element)
-
-    return get_normalized_md5_digest(hashable_elements)
-
-
-def compute_tensor_digest(tensor_data, tensor_targets) -> str:
-    """Computes a digest for the given Tensorflow tensor.
-
-    Args:
-        tensor: A Tensorflow tensor.
-
-    Returns:
-        A string digest.
-
-    """
-    if tensor_targets is None:
-        return compute_numpy_digest(tensor_data.numpy())
-    else:
-        return compute_numpy_digest(tensor_data.numpy(), tensor_targets.numpy())
-
-
-def compute_spark_df_digest(df) -> str:
-    """Computes a digest for the given Spark DataFrame. Retrieve a semantic hash of the
-    DataFrame's logical plan, which is much more efficient and deterministic than hashing
-    DataFrame records
-
-    Args:
-        df: A Spark DataFrame.
-
-    Returns:
-        A string digest.
-
-    """
-
-    import numpy as np
-    import pyspark
-
-    # Spark 3.1.0+ has a semanticHash() method on DataFrame
-    if Version(pyspark.__version__) >= Version("3.1.0"):
-        semantic_hash = df.semanticHash()
-    else:
-        semantic_hash = df._jdf.queryExecution().analyzed().semanticHash()
-    return get_normalized_md5_digest([np.int64(semantic_hash)])
-
-
 def get_normalized_md5_digest(elements: List[Any]) -> str:
     """Computes a normalized digest for a list of hashable elements.
 
@@ -168,7 +88,6 @@ def get_normalized_md5_digest(elements: List[Any]) -> str:
 
     Returns:
         An 8-character, truncated md5 digest.
-
     """
 
     if not elements:
diff --git a/mlflow/data/spark_dataset.py b/mlflow/data/spark_dataset.py
index 8438b9bce..49d317c26 100644
--- a/mlflow/data/spark_dataset.py
+++ b/mlflow/data/spark_dataset.py
@@ -3,10 +3,12 @@ import logging
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
+from packaging.version import Version
+
 from mlflow.data.dataset import Dataset
 from mlflow.data.dataset_source import DatasetSource
 from mlflow.data.delta_dataset_source import DeltaDatasetSource
-from mlflow.data.digest_utils import compute_spark_df_digest
+from mlflow.data.digest_utils import get_normalized_md5_digest
 from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
 from mlflow.data.spark_dataset_source import SparkDatasetSource
 from mlflow.exceptions import MlflowException
@@ -53,7 +55,15 @@ class SparkDataset(Dataset, PyFuncConvertibleDatasetMixin):
         """
         # Retrieve a semantic hash of the DataFrame's logical plan, which is much more efficient
         # and deterministic than hashing DataFrame records
-        return compute_spark_df_digest(self._df)
+        import numpy as np
+        import pyspark
+
+        # Spark 3.1.0+ has a semanticHash() method on DataFrame
+        if Version(pyspark.__version__) >= Version("3.1.0"):
+            semantic_hash = self._df.semanticHash()
+        else:
+            semantic_hash = self._df._jdf.queryExecution().analyzed().semanticHash()
+        return get_normalized_md5_digest([np.int64(semantic_hash)])
 
     def to_dict(self) -> Dict[str, str]:
         """Create config dictionary for the dataset.
diff --git a/mlflow/data/tensorflow_dataset.py b/mlflow/data/tensorflow_dataset.py
index 03a1d3a07..b50fe401a 100644
--- a/mlflow/data/tensorflow_dataset.py
+++ b/mlflow/data/tensorflow_dataset.py
@@ -7,7 +7,11 @@ import numpy as np
 
 from mlflow.data.dataset import Dataset
 from mlflow.data.dataset_source import DatasetSource
-from mlflow.data.digest_utils import compute_tensor_digest, compute_tensorflow_dataset_digest
+from mlflow.data.digest_utils import (
+    MAX_ROWS,
+    compute_numpy_digest,
+    get_normalized_md5_digest,
+)
 from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
 from mlflow.data.schema import TensorDatasetSchema
 from mlflow.exceptions import MlflowException
@@ -73,6 +77,62 @@ class TensorFlowDataset(Dataset, PyFuncConvertibleDatasetMixin):
         self._targets = targets
         super().__init__(source=source, name=name, digest=digest)
 
+    def _compute_tensorflow_dataset_digest(
+        self,
+        dataset,
+        targets=None,
+    ) -> str:
+        """Computes a digest for the given Tensorflow dataset.
+
+        Args:
+            dataset: A Tensorflow dataset.
+
+        Returns:
+            A string digest.
+        """
+        import pandas as pd
+        import tensorflow as tf
+
+        hashable_elements = []
+
+        def hash_tf_dataset_iterator_element(element):
+            if element is None:
+                return
+            flat_element = tf.nest.flatten(element)
+            flattened_array = np.concatenate([x.flatten() for x in flat_element])
+            trimmed_array = flattened_array[0:MAX_ROWS]
+            try:
+                hashable_elements.append(pd.util.hash_array(trimmed_array))
+            except TypeError:
+                hashable_elements.append(np.int64(trimmed_array.size))
+
+        for element in dataset.as_numpy_iterator():
+            hash_tf_dataset_iterator_element(element)
+        if targets is not None:
+            for element in targets.as_numpy_iterator():
+                hash_tf_dataset_iterator_element(element)
+
+        return get_normalized_md5_digest(hashable_elements)
+
+    def _compute_tensor_digest(
+        self,
+        tensor_data,
+        tensor_targets,
+    ) -> str:
+        """Computes a digest for the given Tensorflow tensor.
+
+        Args:
+            tensor_data: A Tensorflow tensor, representing the features.
+            tensor_targets: A Tensorflow tensor, representing the targets. Optional.
+
+        Returns:
+            A string digest.
+        """
+        if tensor_targets is None:
+            return compute_numpy_digest(tensor_data.numpy())
+        else:
+            return compute_numpy_digest(tensor_data.numpy(), tensor_targets.numpy())
+
     def _compute_digest(self) -> str:
         """
         Computes a digest for the dataset. Called if the user doesn't supply
@@ -80,11 +140,9 @@ class TensorFlowDataset(Dataset, PyFuncConvertibleDatasetMixin):
         """
         import tensorflow as tf
 
-        return (
-            compute_tensorflow_dataset_digest(self._features, self._targets)
-            if isinstance(self._features, tf.data.Dataset)
-            else compute_tensor_digest(self._features, self._targets)
-        )
+        if isinstance(self._features, tf.data.Dataset):
+            return self._compute_tensorflow_dataset_digest(self._features, self._targets)
+        return self._compute_tensor_digest(self._features, self._targets)
 
     def to_dict(self) -> Dict[str, str]:
         """Create config dictionary for the dataset.
diff --git a/mlflow/utils/databricks_utils.py b/mlflow/utils/databricks_utils.py
index 76de63d3f..d36fb1c14 100644
--- a/mlflow/utils/databricks_utils.py
+++ b/mlflow/utils/databricks_utils.py
@@ -1,7 +1,9 @@
 import functools
+import json
 import logging
 import os
 import subprocess
+import time
 from sys import stderr
 from typing import Optional, TypeVar
 
@@ -162,6 +164,10 @@ def is_in_databricks_job():
         return False
 
 
+def is_in_databricks_model_serving_environment():
+    return "DATABRICKS_MODEL_SERVING_ENV" in os.environ
+
+
 def is_in_databricks_repo():
     try:
         return get_git_repo_relative_path() is not None
@@ -419,28 +425,33 @@ def _fail_malformed_databricks_auth(profile):
     )
 
 
-def get_databricks_host_creds(server_uri=None):
-    """
-    Reads in configuration necessary to make HTTP requests to a Databricks server. This
-    uses the Databricks CLI's ConfigProvider interface to load the DatabricksConfig object.
-    If no Databricks CLI profile is found corresponding to the server URI, this function
-    will attempt to retrieve these credentials from the Databricks Secret Manager. For that to work,
-    the server URI will need to be of the following format: "databricks://scope:prefix". In the
-    Databricks Secret Manager, we will query for a secret in the scope "<scope>" for secrets with
-    keys of the form "<prefix>-host" and "<prefix>-token". Note that this prefix *cannot* be empty
-    if trying to authenticate with this method. If found, those host credentials will be used. This
-    method will throw an exception if sufficient auth cannot be found.
+# constant defined outside function for testing override
+_MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH = "/var/credentials-secret/model-dependencies-oauth-token"
 
-    Args:
-        server_uri: A URI that specifies the Databricks profile you want to use for making
-            requests.
 
-    Returns:
-        MlflowHostCreds which includes the hostname and authentication information necessary to
-        talk to the Databricks server.
-    """
+# Helper function to attempt to read OAuth Token from
+# mounted file in Databricks Model Serving environment
+def _get_model_dependency_oauth_token(should_retry=True):
+    try:
+        with open(_MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH) as f:
+            oauth_dict = json.load(f)
+            return oauth_dict["OAUTH_TOKEN"][0]["oauthTokenValue"]
+    except Exception as e:
+        # sleep and retry in case of any race conditions with OAuth refreshing
+        if should_retry:
+            time.sleep(0.5)
+            return _get_model_dependency_oauth_token(should_retry=False)
+        else:
+            raise MlflowException(
+                "Unable to read Oauth credentials from file mount for Databricks "
+                "Model Serving dependency failed"
+            ) from e
+
+
+def _default_databricks_host_creds(server_uri):
     profile, path = get_db_info_from_uri(server_uri)
     config = ProfileConfigProvider(profile).get_config() if profile else get_config()
+    insecure = getattr(config, "insecure", False)
     # if a path is specified, that implies a Databricks tracking URI of the form:
     # databricks://profile-name/path-specifier
     if (not config or not config.host) and path:
@@ -454,9 +465,6 @@ def get_databricks_host_creds(server_uri=None):
                 config = DatabricksConfig.from_token(host=host, token=token, insecure=False)
     if not config or not config.host:
         _fail_malformed_databricks_auth(profile)
-
-    insecure = hasattr(config, "insecure") and config.insecure
-
     if config.username is not None and config.password is not None:
         return MlflowHostCreds(
             config.host,
@@ -469,6 +477,67 @@ def get_databricks_host_creds(server_uri=None):
     _fail_malformed_databricks_auth(profile)
 
 
+def _model_serving_env_databricks_host_creds():
+    # Since we do not record OAuth expiration time in OAuth file, perform periodic refresh
+    # of OAuth environment variable cache here. As currently configured (02/24) OAuth token
+    # in model serving environment is guaranteed to have at least 30 min remaining on TTL
+    # at any point in time but refresh at higher rate of every 5 min here to be safe
+    # and conform with refresh logic for Brickstore tables.
+    OAUTH_CACHE_REFRESH_DURATION_SEC = 5 * 60
+    OAUTH_CACHE_ENV_VAR = "DATABRICKS_DEPENDENCY_OAUTH_CACHE"
+    OAUTH_CACHE_EXPIRATION_ENV_VAR = "DATABRICKS_DEPENDENCY_OAUTH_CACHE_EXIRY_TS"
+    MODEL_SERVING_HOST_ENV_VAR = "DATABRICKS_MODEL_SERVING_HOST_URL"
+
+    # check if dependency is cached in env var before reading from file
+    oauth_token = ""
+    if (
+        OAUTH_CACHE_ENV_VAR in os.environ
+        and OAUTH_CACHE_EXPIRATION_ENV_VAR in os.environ
+        and float(os.environ[OAUTH_CACHE_EXPIRATION_ENV_VAR]) > time.time()
+    ):
+        oauth_token = os.environ[OAUTH_CACHE_ENV_VAR]
+    else:
+        oauth_token = _get_model_dependency_oauth_token()
+        os.environ[OAUTH_CACHE_ENV_VAR] = oauth_token
+        os.environ[OAUTH_CACHE_EXPIRATION_ENV_VAR] = str(
+            time.time() + OAUTH_CACHE_REFRESH_DURATION_SEC
+        )
+
+    return MlflowHostCreds(
+        os.environ[MODEL_SERVING_HOST_ENV_VAR], token=oauth_token, ignore_tls_verification=True
+    )
+
+
+def get_databricks_host_creds(server_uri=None):
+    """
+    Reads in configuration necessary to make HTTP requests to a Databricks server. This
+    uses the Databricks CLI's ConfigProvider interface to load the DatabricksConfig object.
+    If no Databricks CLI profile is found corresponding to the server URI, this function
+    will attempt to retrieve these credentials from the Databricks Secret Manager. For that to work,
+    the server URI will need to be of the following format: "databricks://scope:prefix". In the
+    Databricks Secret Manager, we will query for a secret in the scope "<scope>" for secrets with
+    keys of the form "<prefix>-host" and "<prefix>-token". Note that this prefix *cannot* be empty
+    if trying to authenticate with this method. If found, those host credentials will be used. This
+    method will throw an exception if sufficient auth cannot be found. If called, within Databricks
+    Model Serving environment, this method will attempt to read the OAuth token from where it would
+    expect to find it in the serving environment.
+
+    Args:
+        server_uri: A URI that specifies the Databricks profile you want to use for making
+            requests.
+
+    Returns:
+        MlflowHostCreds which includes the hostname and authentication information necessary to
+        talk to the Databricks server.
+    """
+    # if in model serving environment databricks attempt to fetch model dependency OAuth token
+    if is_in_databricks_model_serving_environment():
+        return _model_serving_env_databricks_host_creds()
+    # default host creds behavior if not fetching OAuth token for model serving dependency
+    else:
+        return _default_databricks_host_creds(server_uri)
+
+
 @_use_repl_context_if_available("mlflowGitRepoUrl")
 def get_git_repo_url():
     try:


Test patch:
diff --git a/tests/utils/test_databricks_utils.py b/tests/utils/test_databricks_utils.py
index adb066fea..9ef1f903b 100644
--- a/tests/utils/test_databricks_utils.py
+++ b/tests/utils/test_databricks_utils.py
@@ -1,5 +1,8 @@
 import builtins
+import json
+import os
 import sys
+import time
 from unittest import mock
 
 import pytest
@@ -98,6 +101,62 @@ def test_databricks_single_slash_in_uri_scheme_throws(get_config):
         databricks_utils.get_databricks_host_creds("databricks:/profile:path")
 
 
+def test_databricks_model_serving_throws(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_ENV", "true")
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_HOST_URL", "host")
+    with pytest.raises(MlflowException, match="Unable to read Oauth credentials"):
+        databricks_utils.get_databricks_host_creds()
+
+
+def test_databricks_params_model_serving_oauth_cache(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_ENV", "true")
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_HOST_URL", "host")
+    monkeypatch.setenv("DATABRICKS_DEPENDENCY_OAUTH_CACHE", "token")
+    monkeypatch.setenv("DATABRICKS_DEPENDENCY_OAUTH_CACHE_EXIRY_TS", str(time.time() + 5))
+    params = databricks_utils.get_databricks_host_creds()
+    assert params.host == "host"
+    assert params.token == "token"
+
+
+@pytest.fixture
+def oauth_file(tmp_path):
+    token_contents = {"OAUTH_TOKEN": [{"oauthTokenValue": "token2"}]}
+    oauth_file = tmp_path.joinpath("model-dependencies-oauth-token")
+    with open(oauth_file, "w") as f:
+        json.dump(token_contents, f)
+    return oauth_file
+
+
+def test_databricks_params_model_serving_oauth_cache_expired(monkeypatch, oauth_file):
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_ENV", "true")
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_HOST_URL", "host")
+    monkeypatch.setenv("DATABRICKS_DEPENDENCY_OAUTH_CACHE", "token")
+    monkeypatch.setenv("DATABRICKS_DEPENDENCY_OAUTH_CACHE_EXIRY_TS", str(time.time() - 5))
+    with mock.patch(
+        "mlflow.utils.databricks_utils._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH", str(oauth_file)
+    ):
+        params = databricks_utils.get_databricks_host_creds()
+        # cache should get updated with new token
+        assert os.environ["DATABRICKS_DEPENDENCY_OAUTH_CACHE"] == "token2"
+        assert float(os.environ["DATABRICKS_DEPENDENCY_OAUTH_CACHE_EXIRY_TS"]) > time.time()
+        assert params.host == "host"
+        # should use token2 from oauthfile, rather than token from cache
+        assert params.token == "token2"
+
+
+def test_databricks_params_model_serving_read_oauth(monkeypatch, oauth_file):
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_ENV", "true")
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_HOST_URL", "host")
+    with mock.patch(
+        "mlflow.utils.databricks_utils._MODEL_DEPENDENCY_OAUTH_TOKEN_FILE_PATH", str(oauth_file)
+    ):
+        params = databricks_utils.get_databricks_host_creds()
+        assert os.environ["DATABRICKS_DEPENDENCY_OAUTH_CACHE"] == "token2"
+        assert float(os.environ["DATABRICKS_DEPENDENCY_OAUTH_CACHE_EXIRY_TS"]) > time.time()
+        assert params.host == "host"
+        assert params.token == "token2"
+
+
 def test_get_workspace_info_from_databricks_secrets():
     mock_dbutils = mock.MagicMock()
     mock_dbutils.secrets.get.return_value = "workspace-placeholder-info"
@@ -224,6 +283,14 @@ def test_is_in_databricks_runtime(monkeypatch):
     assert not databricks_utils.is_in_databricks_runtime()
 
 
+def test_is_in_databricks_model_serving_environment(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_MODEL_SERVING_ENV", "true")
+    assert databricks_utils.is_in_databricks_model_serving_environment()
+
+    monkeypatch.delenv("DATABRICKS_MODEL_SERVING_ENV")
+    assert not databricks_utils.is_in_databricks_model_serving_environment()
+
+
 def test_get_repl_id():
     # Outside of Databricks environments, the Databricks REPL ID should be absent
     assert databricks_utils.get_repl_id() is None