Get started: MLflow Tracing for GenAI (Databricks Notebook)

This quickstart helps you integrate your GenAI app with MLflow Tracing if you use a Databricks Notebook as your development environment. If you use a local IDE, please use the IDE quickstart instead.

What you will achieve

By the end of this tutorial, you will have:

• A Databricks Notebook with a linked MLflow experiment for your GenAI app
• A simple GenAI application instrumented with MLflow Tracing
• A trace from that app in your MLflow experiment

Prerequisites

• Databricks Notebook: Create a new notebook in your Databricks workspace. The notebook will have a default MLflow experiment that is the container for your GenAI application. Learn more about MLflow experiments in the MLflow concepts section.

Step 1: Update MLflow

Update MLflow to the latest version for the best GenAI experience:

Python

%pip install --upgrade "mlflow[databricks]>=3.1" openai
dbutils.library.restartPython()

Step 2: Instrument your application

The code snippets below define a simple GenAI app that completes sentence templates using an LLM.

1. Initialize an OpenAI client to connect to either Databricks-hosted LLMs or LLMs hosted by OpenAI.

   Databricks-hosted LLMs

   The code snippet below uses the OpenAI client to connect to Databricks-hosted LLMs. The app uses Anthropic's Claude Sonnet LLM, but you can choose from the available foundation models.

   Python

   import mlflow
   from databricks.sdk import WorkspaceClient
   
   # Enable MLflow's autologging to instrument your application with Tracing
   mlflow.openai.autolog()
   
   # Create an OpenAI client that is connected to Databricks-hosted LLMs
   w = WorkspaceClient()
   client = w.serving_endpoints.get_open_ai_client()
   
   model_name = "databricks-claude-sonnet-4"

   OpenAI SDK

   Use the native OpenAI SDK to connect to OpenAI-hosted models. The code snippet below uses gpt-4o-mini, but you can select from available OpenAI models.

   Python

   import mlflow
   import openai
   
   # Ensure your OPENAI_API_KEY is set in your environment
   # os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured
   
   # Enable MLflow's autologging to instrument your application with Tracing
   mlflow.openai.autolog()
   
   client = openai.OpenAI()
   
   model_name = "gpt-4o-mini"

2. Define and run your application:

   Use the @mlflow.trace decorator, which makes it easy to trace any Python function, combined with OpenAI automatic instrumentation to capture the details of the call to the OpenAI SDK.

   Python

   import mlflow
   import os
   
   # Ensure your OPENAI_API_KEY is set in your environment
   # os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured
   
   # Enable auto-tracing for OpenAI
   mlflow.openai.autolog()
   
   # Set up MLflow tracking to Databricks
   mlflow.set_tracking_uri("databricks")
   mlflow.set_experiment("/Shared/openai-tracing-demo")
   
   # Use the trace decorator to capture the application's entry point
   @mlflow.trace
   def my_app(input: str):
     # This call is automatically instrumented by `mlflow.openai.autolog()`
     response = client.chat.completions.create(
       model=model_name,
       temperature=0.1,
       max_tokens=100,
       messages=[
         {
           "role": "system",
           "content": "You are a helpful assistant.",
         },
         {
           "role": "user",
           "content": input,
         },
       ]
     )
     return response.choices[0].message.content
   
   result = my_app(input="What is MLflow?")
   print(result)

For details on adding tracing to apps, see the tracing instrumentation guide and the 20+ library integrations.

Step 3: View the trace in MLflow

The trace will appear below the notebook cell.

Optionally, you can go to the MLflow experiment UI to see the trace:

1. Click the experiments icon on the right sidebar.
2. Click the open icon next to experiment runs.
3. The generated trace appears in the Traces tab.
4. Click the trace to view its details.

Understanding the trace

The trace you've just created shows:

• Root span: Represents the inputs to the my_app(...) function
  • Child span: Represents the OpenAI completion request
• Attributes: Contains metadata like model name, token counts, and timing information
• Inputs: The messages sent to the model
• Outputs: The response received from the model

This simple trace already provides valuable insights into your application's behavior, such as:

• What was asked
• What response was generated
• How long the request took
• How many tokens were used (affecting cost)

For more complex applications like RAG systems or multi-step agents, MLflow Tracing provides even more value by revealing the inner workings of each component and step.

Guides and references

For details on concepts and features in this guide, see:

• MLflow Tracing guide - Start here for more in-depth learning about MLflow Tracing
• MLflow Tracing integrations - 20+ libraries with automatic tracing integrations
• Tracing concepts - Understand the fundamentals of MLflow Tracing