Usage of MLFlow models inside Streamlit app in Databricks

pavelhym
New Contributor

‎12-24-2025 09:50 AM

I have an issue with loading registered MLflow model into streamlit app inside the Databricks

This is the sample code used for model load:

import mlflow
from mlflow.tracking import MlflowClient

mlflow.set_tracking_uri("databricks")
mlflow.set_registry_uri("databricks-uc")
client = MlflowClient()
model_uri = "models:/workspace.default.xgboost_units/2"
model = mlflow.pyfunc.load_model(model_uri=model_uri)
The streamlit app served with the databricks app is failing silently loading this model, why the same code inside the databricks python notebook is working without any issues.

Any thoughts what can be the reason? Or maybe there is sample app which uses MLflow models successfully? 

mlflow==2.22.0 in both envs



0 Kudos

iyashk-DB
Databricks Employee
Databricks Employee

‎12-25-2025 07:31 AM

Authentication context isn’t automatically available in Apps. Notebooks automatically inject workspace host and token for mlflow when you use mlflow.set_tracking_uri("databricks") and mlflow.set_registry_uri("databricks-uc"). In Databricks Apps, you typically must supply credentials explicitly via environment variables (for example, DATABRICKS_HOST and DATABRICKS_TOKEN) or by using a Databricks CLI profile in the URIs (for example, databricks-uc://<profile>).

The app runs under an app service principal or (if configured) on-behalf-of the user. That identity must have UC privileges on the model’s catalog/schema (USE CATALOG, USE SCHEMA) and read access to the registered model to download artifacts through MLflow APIs.

put your mlflow.pyfunc.load_model(...) in try/except and display exceptions with st.exception(e) to diagnose the root cause instead of failing silently.