Add a helper function to get MlflowCredentialContext by run_id

[liangz1]

Related Issues/PRs

On Databricks, mlflow spark log_model and load_model use mlflowdbfs to write and read the model files. The read and write operations need credentials to be set properly. Currently, credentials are only set and unset within mlflow.spark.log_model(...) and mlflow.spark.load_model(...) via MlflowCredentialContext.

However, spark sometimes uses lazy loading, meaning that loaded_model = mlflow.spark.load_model(...) only loads model metadata. The actual model data loading will happen when loaded_model.transform() is called. But since loaded_model.transform() is not controlled by mlflow and is not wrapped by MlflowCredentialContext, the model data loading will fail. Users will see error:

...
Caused by: java.lang.IllegalStateException: No mlflowdbfs host found in local properties
...

The following wrokaround works:

from mlflow.utils import databricks_utils
from mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri

with mlflow.start_run(run_id="f59d8ac3911f4196a2c3c49ee185b719"):
    run_root_artifact_uri = mlflow.get_artifact_uri()

with databricks_utils.MlflowCredentialContext(
  get_databricks_profile_uri_from_artifact_uri(run_root_artifact_uri)
):
    loaded_model.transform(data)

However, it's too lengthy for the user to use the workaround. This PR proposes adding a helper function to simplify the workaround.

What changes are proposed in this pull request?

After the helper, we can update the error message to be:

No mlflowdbfs host found in local properties. Try to resolve this error in the following way:

import mlflow.utils import databricks_utils

run_id = ... # the run_id associated with the logged model
with databricks_utils.get_mlflow_credential_context_by_run_id(run_id):
    loaded_model.transform(data)  # code that causes the current error

How is this patch tested?

• Existing unit/integration tests
• New unit/integration tests
• Manual tests (describe details, including test results, below)

Does this PR change the documentation?

• No. You can skip the rest of this section.
• Yes. Make sure the changed pages / sections render correctly in the documentation preview.

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• 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/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• 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

Interface

• 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

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[mlflow-automation]

Documentation preview for 613b5e4 will be available here when this CircleCI job completes successfully.

More info

• Ignore this comment if this PR does not change the documentation.
• It takes a few minutes for the preview to be available.
• The preview is updated when a new commit is pushed to this PR.
• This comment was created by https://github.com/mlflow/mlflow/actions/runs/4808832481.