MLflow Tracing SDK

3.10.1 · active · verified Sun Apr 05

MLflow Tracing SDK (mlflow-tracing) is an open-source, lightweight Python package that provides a minimum set of dependencies and functionality to instrument your code, models, or agents with MLflow Tracing. It is designed for production environments to enable faster deployment, simplified dependency management, enhanced portability, and reduced security risks compared to the full MLflow package. It supports LLM and AI agent observability, capturing inputs, outputs, and metadata for each step of a request.

Warnings

breaking Do NOT co-install `mlflow-tracing` with the full `mlflow` package.
Fix: If you need the full MLflow features, install `mlflow`. If you only need tracing in production, use `mlflow-tracing`. Uninstall one before installing the other (`pip uninstall mlflow && pip install mlflow-tracing` or vice-versa).
gotcha Calling `mlflow.set_tracking_uri()` after OpenTelemetry auto-instrumentation can reset the global `TracerProvider`.
Fix: To avoid silent data loss or breaking other telemetry, ensure `mlflow.set_tracking_uri()` is called early in your application's lifecycle, preferably before other OpenTelemetry instrumentation is initialized. This issue mainly applies when combining MLflow tracing with external OTel auto-instrumentation.
gotcha Using a file-based backend store for the MLflow server can lead to poor UI/SDK performance.
Fix: For better performance in production or with significant trace volume, configure your MLflow server to use a database-based backend store (e.g., PostgreSQL, MySQL, SQLite with `--backend-store-uri sqlite:///mlflow.db`).
gotcha Traces might get stuck in 'in progress' and not be viewable if a process hangs or runs too long.
Fix: Set the `MLFLOW_TRACE_TIMEOUT_SECONDS` environment variable to automatically halt and export traces that exceed a specified duration. This allows analysis even for stuck traces.
gotcha MLflow client and server version compatibility is important for new features.
Fix: While `mlflow-tracing` SDK and the MLflow server within the same major version are generally compatible, it's recommended to keep both client and server up-to-date to ensure new tracing features (introduced in MLflow 2.14.0 and enhanced in 3.x) are available and function correctly.

Install

pip install mlflow-tracing Install MLflow Tracing SDK

Imports

mlflow
```
import mlflow
```
While `mlflow.tracing` contains utility functions, the main entry points for autologging, manual tracing (e.g., `@mlflow.trace`), and configuration are directly under the `mlflow` namespace.
mlflow.tracing.configure
```
from mlflow import tracing
tracing.configure(...)
```
Used for advanced configuration like span processors.
mlflow.tracing.disable
```
from mlflow import tracing
tracing.disable()
```
Used to temporarily disable tracing.
mlflow.tracing.enable
```
from mlflow import tracing
tracing.enable()
```
Used to re-enable tracing if previously disabled.

@mlflow.trace

import mlflow

@mlflow.trace
def my_traced_function():
    pass

Decorator for manual function instrumentation.

mlflow.start_span

import mlflow

with mlflow.start_span('my_span_name'):
    # ... code to trace ...

Context manager for manual code block instrumentation.

Quickstart

This quickstart demonstrates how to set up MLflow Tracing for an OpenAI call. It configures the MLflow tracking URI, sets an experiment, enables autologging for OpenAI, and then performs a simple API call. The trace for this call will be automatically logged and viewable in the MLflow UI. Make sure an MLflow server is running and `OPENAI_API_KEY` is set in your environment.

import os
import mlflow
from openai import OpenAI

# Set your MLflow Tracking URI (replace with your server, e.g., 'http://localhost:5000')
# For Databricks, use 'databricks' and ensure DATABRICKS_HOST/TOKEN are set.
mlflow.set_tracking_uri(os.environ.get('MLFLOW_TRACKING_URI', 'http://127.0.0.1:5000'))

# Set a new MLflow experiment to log traces to
mlflow.set_experiment("my_genai_app_traces")

# Ensure OpenAI API key is set for the example
if not os.environ.get("OPENAI_API_KEY"): 
    # In a real app, use a secure way to load keys (e.g., environment variable, secret manager)
    # For quick testing, you can set it directly here, but it's not recommended for production
    print("WARNING: OPENAI_API_KEY environment variable not set. Skipping OpenAI example.")
    openai_client = None
else:
    # Enable auto-tracing for OpenAI calls
    mlflow.openai.autolog()
    
    # Initialize OpenAI client
    openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    # Make an OpenAI call - this will be automatically traced
    print("Invoking OpenAI completion...")
    response = openai_client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": "Tell me a fun fact about Python programming."}
        ],
        max_tokens=50
    )
    print("OpenAI Response:", response.choices[0].message.content)
    print("Trace should now be visible in MLflow UI under 'my_genai_app_traces' experiment.")

view raw JSON →