Skip to main content

Instrument Your App with MLflow Tracing

tip

New to MLflow Tracing? Checkout the Quick Start Guide to get started.

This page provides a comprehensive guide for instrumenting your GenAI applications with MLflow Tracing.

1. Installation

Add mlflow-tracing to your Python environment to enable tracing features with minimum set of dependencies.

pip install mlflow-tracing
info

The mlflow-tracing Python package does not include other MLflow features, such as the self-host tracking server, evaluation capabilities, prompt management, etc. If you want to use the full set of MLflow's features, please install the full MLflow package. We recommend using the mlflow-tracing package for production applications to keep the dependencies and package footprint minimal.

pip install mlflow

2. Connect Your Application with MLflow

For local development or personal use, locally hosting MLflow is a good option. While this is not recommended for production use, it is the fastest way to get started. For logging traces, we recommend using a SQL Backend for better performance.

To start a local MLflow server, run the following command in your terminal. Alternatively, you can use the official Docker image to run the server.

# Replace sqlite:/// part with your preferred database URI.
mlflow server --host 127.0.0.1 --port 5000 --backend-store-uri sqlite:///mlruns.db

Then, in your application, configure the tracking URI and set an active experiment using the following code:

Python:

import mlflow

mlflow.set_tracking_uri("http://127.0.0.1:5000")
mlflow.set_experiment("<your-experiment-name>")

JS/TS:

import * as mlflow from "mlflow-tracing";

mlflow.init(
trackingUri: "http://127.0.0.1:5000",
experimentId: "<your-experiment-id>", // If no experiment is created yet, create it on UI first.
);

3. Instrumenting Your Application Logic

MLflow offers different ways to instrument your application logic. Follow the links below to learn more about each approach to instrument your application:

⚙️ Common Patterns

Production Considerations

MLflow Tracing is production ready, but in order to ensure the scalability and reliability of the tracing system, we recommend the following best practices:

  1. Enable Async Logging and set up appropriate queue size and timeout.
  2. Use the lightweight mlflow-tracing package for minimizing the package footprint and dependencies.
  3. Use managed MLflow services for reducing the operational overhead and ensure the scalability of the tracing system.
  4. When using self-hosted MLflow, make sure to use the SQL Backend with a scalable database like PostgreSQL. The default file-based backend has scalability limitations and is not recommended for production use.

Async Applications

Async programming is an effective tool for improving the throughput of your application, particularly for LLM-based applications that are typically I/O bound. MLflow Tracing natively supports instrumenting async applications.

Multi-Threaded Applications

Multi-threading is a common strategy for parallelizing IO-bound operations in applications. MLflow Tracing supports multi-threaded applications using context propagation.

Managing User sessions

Many LLM-based applications are deployed as chat-based applications, where each user session is a separate thread. Grouping traces by user session is a common practice. MLflow Tracing supports managing user sessions.

Redacting PII Data

Traces can contain sensitive data such as raw user inputs, internal document contents, etc. MLflow Tracing supports redacting PII data using flexible masking rules, custom functions, and integration with external PII masking libraries.

Collecting User Feedbacks

User feedback is a valuable source of information for improving the user experience of your application. MLflow Tracing supports collecting user feedback on traces to track and analyze the feedbacks effectively.