MLflow Tracing for LLM Observability

MLflow Tracing is a feature that enhances LLM observability in your Generative AI (GenAI) applications by capturing detailed information about the execution of your application's services. Tracing provides a way to record the inputs, outputs, and metadata associated with each intermediate step of a request, enabling you to easily pinpoint the source of bugs and unexpected behaviors.

Why Choose MLflow?

• 🪽 Free and Open - MLflow is open source and 100% FREE. You don’t need to pay additional SaaS costs to add observability to your GenAI stack. Your trace data is hosted on your own infrastructure.

• 🥇 Standard - MLflow Tracing is compatible with OpenTelemetry, an industry-standard observability spec. You can export your trace data to various services in your existing observability stack, such as Grafana, Prometheus, Datadog, New Relic, and more.

• 🤝 Framework Support - MLflow Tracing integrates with 15+ GenAI libraries, including OpenAI, LangChain, LlamaIndex, DSPy, and others. See the Automatic Tracing section for the full list of supported libraries.

• 🔄 End-to-End - MLflow is designed for managing the end-to-end machine learning lifecycle. With its model tracking and evaluation capabilities, MLflow empowers you to leverage your trace data fully.

• 👥 Community - MLflow boasts a vibrant Open Source community as a part of the Linux Foundation. With 19,000+ GitHub Stars and 15MM+ monthly downloads, MLflow is a trusted standard in the MLOps/LLMOps ecosystem.

Introduction to Observability and Traces

tip

If you are new to the tracing or observability concepts, we recommend starting with the Tracing 101 page.

Traces empowers your throughout the end-to-end lifecycle of a machine learning project. Click on the tabs below to see how it helps you at each step of the workflow in more details.

Debugging

Complete Debugging Experience in Your IDE or Notebook

MLflow's tracing capabilities provide deep insights into what happens beneath the abstractions of GenAI libraries, helping you precisely identify where issues occur.

You can navigate traces seamlessly within your preferred IDE or notebook, eliminating the hassle of switching between multiple tabs or searching through an overwhelming list of traces.

Evaluation Dataset

Building a High-Quality Evaluation Dataset

Evaluating the performance of your GenAI application is crucial, but creating a reliable evaluation dataset can be challenging. Traces serve as a rich data source, helping you build high-quality datasets with precise metrics for internal components like retrievers and tools.

When combined with MLflow LLM Evaluation, MLflow offers a seamless experience for assessing and improving your application’s performance.

Inspect Quality

Root Cause Analysis for Improved Quality

After evaluating your model using MLflow LLM Evaluation, you can explore auto-generated traces during the evaluation run to identify root causes of quality issues — for instance, insufficiently retrieved documents.

Traces empower you to analyze issues in detail and iterate quickly to enhance the quality of your application.

Production Monitoring

Monitor Applications with Your Favorite Observability Stack

Machine learning projects don’t end with the first launch. Continuous monitoring and incremental improvement are critical to long-term success.

Integrated with various observability platforms such as Databricks, Datadog, Grafana, and Prometheus, MLflow Tracing provides a comprehensive solution for monitoring your GenAI applications in production. Refer to Monitoring GenAI Application in Production for more details.

Getting Started

The following is the simplest example of MLflow Tracing with OpenAI Python SDK:

import mlflow
import openai

# Enable MLflow automatic tracing for OpenAI with one line of code!
mlflow.openai.autolog()

# Use OpenAI Python SDK as usual
openai.OpenAI().chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "system", "content": "You are a chatbot."},
        {"role": "user", "content": "What is the weather like today?"},
    ],
)

Then go to MLflow UI (if not started, run mlflow ui in your terminal) to see the trace!

Automatic Tracing

MLflow Tracing is integrated with various GenAI libraries and provide one-line automatic tracing experience for each library (and the combination of them!). Click on the icon below to see detailed examples to integrate MLflow with your favorite library.

Hint

Is your favorite library missing from the list? Consider contributing to MLflow Tracing or submitting a feature request to our Github repository.

Manual Tracing

In addition to the one-line auto tracing experience, MLflow offers Python SDK for manually instrumenting you code and manipulating traces.

1. Instrument a function with @mlflow.trace decorator.
2. Instrument any block of code using mlflow.start_span context manager.
3. Grouping or annotating traces using a tag.
4. Disabling trace globally.

Refer to the Tracing How To Guide for more details about the SDK.

Reviewing Traces

MLflow Traces can be reviewed in several ways:

1. MLflow UI: The MLflow UI provides a rich interface for exploring traces. You can view traces for a specific experiment, run, search and filter traces based on various criteria. You can start the UI by running mlflow ui in your terminal and navigating to http://localhost:5000.

2. Jupyter Notebook: The trace UI is also available within Jupyter notebooks! The trace UI will automatically be displayed when a cell generates a trace, eliminating the need to switch between the notebook and web browser. See the Jupyter Notebook Integration for more details.

Query Traces

Trace data are useful for various downstream tasks, such as creating an evaluation dataset for offline evaluation and production monitoring. MLflow provides several APIs to search and retrieve recorded traces programmatically. See Searching and Retrieving Traces for more details.

Production Monitoring

MLflow Tracing is production ready. Read Production Tracing for the guidance to use it for monitoring models in production and various backend options.

tip

When using MLflow Tracing in a production environment, we recommend using the Lightweight Tracing SDK (mlflow-tracing) that is optimized for reducing the total installation size and minimizing the set of dependencies.

note

MLflow Tracing support is available with the MLflow 2.14.0 release.