Skip to main content

Instrument your app

Learn how to instrument your GenAI application with MLflow Tracing to capture and visualize the execution flow of your application. MLflow offers two main approaches to implementing tracing:

Prerequisites

This guide requires the following packages:

  • mlflow[databricks]>=3.1: Core MLflow functionality with GenAI features and Databricks connectivity.
  • openai>=1.0.0: Only required to run the Quick Start Examples below (if using other LLM providers, install their respective SDKs instead)

Install the required packages:

Python
%pip install --upgrade "mlflow[databricks]>=3.1" openai>=1.0.0
MLflow Version Recommendation

While tracing features are available in MLflow 2.15.0+, it is strongly recommended to install MLflow 3 (specifically 3.1 or newer if using mlflow[databricks]) for the latest GenAI capabilities, including expanded tracing features and robust support.

tip

Running in a Databricks notebook? MLflow is pre-installed in the Databricks runtime. You only need to install additional packages like openai if using external LLM providers.

Running locally? You'll need to install all packages listed above.

Environment Setup

Before running the examples below, configure your environment:

Configure Databricks Authentication

note

If you're running this code inside a Databricks notebook, you can skip this authentication setup. MLflow will automatically use your notebook's authentication context.

If you're running this code outside of Databricks (e.g., in your local IDE), you need to set up authentication:

Python
import os

# Set Databricks authentication (only needed when running outside Databricks)
os.environ["DATABRICKS_HOST"] = "https://your-workspace.databricks.com" # Your workspace URL
os.environ["DATABRICKS_TOKEN"] = "your-databricks-token" # Your access token

Configure API Keys

Set your LLM provider API keys as environment variables:

Python
import os

# Set your OpenAI API key (if using OpenAI models)
os.environ["OPENAI_API_KEY"] = "your-api-key-here"

# Set other provider keys as needed
# os.environ["ANTHROPIC_API_KEY"] = "your-anthropic-key"
# os.environ["GOOGLE_API_KEY"] = "your-google-key"
tip

You can also set these environment variables in your shell before running your script:

Bash
# For Databricks authentication (if running outside Databricks)
export DATABRICKS_HOST="https://your-workspace.databricks.com"
export DATABRICKS_TOKEN="your-databricks-token"

# For LLM providers
export OPENAI_API_KEY="your-api-key-here"

Tracing Approaches

  1. Automatic Tracing: Just add 1 line of code mlflow.<library>.autolog() to automatically capture your app's logic. Automatic tracing works with 20+ supported libraries and frameworks out of the box
  2. Manual Tracing: Designed for custom logic and complex workflows, manual tracing gives you full control over what gets traced and how using high-level APIs (decorators and fluent context managers) or low-level APIs.
note

Automatic and manual tracing can be used together. For example, you could use the auto-tracing for OpenAI's SDK and manual tracing to combine multiple LLM calls into a single trace that represents your application's end to end logic.

Quick Start Examples

Automatic Tracing Example

Enable automatic tracing for your favorite library with just one line of code:

Python
import mlflow
from openai import OpenAI
import os

# Set up environment (if not already configured)
os.environ["OPENAI_API_KEY"] = "your-api-key-here" # Replace with your actual API key

# Set up MLflow tracking
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/auto-tracing-demo")

# Enable automatic tracing for OpenAI
mlflow.openai.autolog()

# Your existing code works unchanged
client = OpenAI()
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "What is MLflow?"}],
max_tokens=100
)

print(response.choices[0].message.content)
# Traces are automatically captured and logged!

Manual Tracing Example

Use the @mlflow.trace decorator to instrument your custom functions:

Python
import mlflow
from mlflow.entities import SpanType
from openai import OpenAI
import os

# Set up environment (if not already configured)
os.environ["OPENAI_API_KEY"] = "your-api-key-here" # Replace with your actual API key

# Set up MLflow tracking
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/manual-tracing-demo")

# Enable automatic tracing for OpenAI
mlflow.openai.autolog()

@mlflow.trace(name="RAG Pipeline", span_type=SpanType.CHAIN)
def answer_question(question: str) -> str:
"""A simple RAG pipeline with manual tracing."""

# Step 1: Retrieve context (manually traced)
context = retrieve_context(question)

# Step 2: Generate answer (automatically traced by OpenAI autolog)
client = OpenAI()
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": f"Context: {context}"},
{"role": "user", "content": question}
],
max_tokens=150
)

return response.choices[0].message.content

@mlflow.trace(span_type=SpanType.RETRIEVER)
def retrieve_context(question: str) -> str:
"""Simulate context retrieval."""
# Simulate retrieval logic
return f"Relevant context for: {question}"

# Execute the traced pipeline
result = answer_question("What is MLflow Tracing?")
print(result)

What is the right approach for my use case?

Generally, we reccomend starting with automatic tracing and only moving to manual tracing if your application's logic is not accuratly captured or you need more control.

Determine the best tracing approach for your use case based on how you are writing your application's code:

  1. Using a single GenAI Authoring Library (e.g., LangGraph, CrewAI, OpenAI Agents, Bedrock Agents, etc)

    • Use auotomatic tracing by adding a one-line integration from our integrations page for your selected library: mlflow.<library>.autolog()
  2. Using an LLM provider's SDKs Directly (e.g., OpenAI SDK, Anthropic SDK, Bedrock SDK, etc)

    • Enable automatic tracing for the API library
    • Add manual tracing decorators to combine multiple LLM calls into a single trace
  3. Using Multiple Authoring Frameworks or combining an authoring framework with a LLM provider's SDK (e.g., LangGraph AND OpenAI SDK, etc)

    • Enable automatic tracing for each framework / SDK
    • Add manual tracing decorators to combine calls to multiple frameworks or SDKs into a single trace
  4. All other approaches or you have a need for fine-grained control

    • Use Manual Tracing
      • Start with the high-level APIs (@mlflow.trace decorator and fluent context managers) which provide a balance of control and ease of use
      • Use the low-level APIs only if the high-level APIs don't give you enough control

View logged traces

Trace Overview

tip

Start with automatic tracing for quick wins, then add manual tracing where you need more control. The combination of both approaches gives you comprehensive observability into your GenAI applications.

Next steps

Continue your journey with these recommended actions and tutorials.

Reference guides

Explore detailed documentation for concepts and features mentioned in this guide.