cancel
Showing results for 
Search instead for 
Did you mean: 
Data Engineering
Join discussions on data engineering best practices, architectures, and optimization strategies within the Databricks Community. Exchange insights and solutions with fellow data engineers.
cancel
Showing results for 
Search instead for 
Did you mean: 

Usage of MLFlow models inside Streamlit app in Databricks

pavelhym
New Contributor

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



1 REPLY 1

iyashk-DB
Databricks Employee
Databricks Employee

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.

Join Us as a Local Community Builder!

Passionate about hosting events and connecting people? Help us grow a vibrant local community—sign up today to get started!

Sign Up Now