MLFlow Authentication from Databricks App for GenAI Tracing

pemidexx
New Contributor III

I am working on a Dash-based app that includes a call to a Databricks-hosted LLM endpoint. I am trying to track those calls with MLFlow. My code is (roughly) like this:

from openai import OpenAI
import mlflow

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Users/my-name@mycompany.com/my-experiment")
mlflow.openai.autolog()

client = OpenAI()

This works as expected when testing the app locally, but results in an authentication error when deployed to Databricks apps. Specifically, if I turn on debug logging I see this error before the app crashes:

DEBUG:urllib3.connectionpool:https://[MY WORKSPACE].databricks.com:443 "GET /api/2.0/mlflow/experiments/get-by-name?experiment_name=%2FUsers%2Fmy-name%40mycompany.com%2Fmy-experiment HTTP/1.1" 401 144

I have attempted the following (unsuccessfully):

  1. Manually setting the tracking URI to the specific URI for my workspace.
  2. Setting the DATABRICKS_TOKEN environment variable to DATABRICKS_CLIENT_SECRET from the app's environment.
  3. Setting the DATABRICKS_TOKEN to the value of a PAT I generated for the app.

What is the correct way to capture GenAI traces with MLFlow from within a Databricks App deployment?