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
- MLflow 3
- MLflow 2.x
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:
%pip install --upgrade "mlflow[databricks]>=3.1" openai>=1.0.0
This guide requires the following packages:
- mlflow[databricks]>=2.15.0,<3.0.0: Core MLflow functionality with 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:
%pip install --upgrade "mlflow[databricks]>=2.15.0,<3.0.0" openai>=1.0.0
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.
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
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:
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:
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"
You can also set these environment variables in your shell before running your script:
# 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
- 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 - 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.
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:
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:
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:
-
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()
- Use auotomatic tracing by adding a one-line integration from our integrations page for your selected library:
-
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
-
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
-
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
- Start with the high-level APIs (
- Use Manual Tracing
View logged traces
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.
- Automatic Tracing - Explore all supported libraries and frameworks for one-line integration
- Manual Tracing - Learn advanced techniques for custom instrumentation
- Debug & observe your app - Use traces for debugging and monitoring
Reference guides
Explore detailed documentation for concepts and features mentioned in this guide.
- Tracing concepts - Understand the fundamentals of MLflow Tracing
- Tracing data model - Learn about traces, spans, and attributes
- Query traces - Programmatically access and analyze your traces