Lightweight Tracing SDK Optimized for Production Usage
MLflow offers a lightweight tracing SDK package called mlflow-tracing that includes only the essential functionality for tracing and monitoring of your GenAI applications. This package is designed for production environments where minimizing dependencies and deployment size is ideal.
Why Use the Lightweight SDK?
The package size and dependencies are significantly smaller than the full MLflow package, allowing for faster deployment times in dynamic environments such as Docker containers, serverless functions, and cloud-based applications.
A smaller set of dependencies means less work keeping up with dependency updates, security patches, and potential breaking changes from upstream libraries. It also reduces the chances of dependency collisions and incompatibilities.
With a fewer number of dependencies, MLflow Tracing can be seamlessly deployed across differing environments and platforms, without worrying about compatibility issues.
Each dependency potentially introduces security vulnerabilities. By reducing the number of dependencies, MLflow Tracing minimizes the attack surface and reduces the risk of security breaches.
Installation
Install the lightweight SDK using pip:
pip install mlflow-tracing
Do not install the full mlflow package together with the lightweight mlflow-tracing SDK, as this may cause version conflicts and namespace resolution issues.
Quickstart
Here's a simple example using the lightweight SDK with OpenAI for logging traces to an experiment on a remote MLflow server. If you haven't set up a remote MLflow server yet, please refer to the Choose Your Backend section below.
import mlflow
import openai
# Set the tracking URI (e.g., Databricks, SageMaker, or self-hosted server)
mlflow.set_tracking_uri("<your-tracking-uri>")
mlflow.set_experiment("<your-experiment-name>")
# Enable auto-tracing for OpenAI
mlflow.openai.autolog()
# Use OpenAI as usual
client = openai.OpenAI()
response = client.chat.completions.create(
model="gpt-4.1-mini", messages=[{"role": "user", "content": "What is MLflow?"}]
)
Choose Your Backend
The lightweight SDK works with various observability platforms. Choose your favorite one and follow the instructions to set up your tracing backend.
- Databricks
- Self-hosting
- OpenTelemetry
- Amazon SageMaker
- Nebius
Databricks Lakehouse Monitoring for GenAI is a go-to solution for monitoring your GenAI application with MLflow Tracing. It provides access to a robust, fully functional monitoring dashboard for operational excellence and quality analysis.
Lakehouse Monitoring for GenAI can be used regardless of whether your application is hosted on Databricks or not.
Sign up for free and get started in a minute to run your GenAI application with complete observability.
MLflow is a fully open-source project, allowing you to self-host your own MLflow server in your own infrastructure. This is a great option if you want to have full control over your data and are restricted in using cloud-based services.
In self-hosting mode, you will be responsible for running the tracking server instance and scaling it to your needs. We strongly recommend using a SQL-based tracking server on top of the performant database of your choice to minimize the operational overhead and high availability. Please refer to the tracking server setup guide for more guidance.
MLflow Tracing is built on the OpenTelemetry tracing spec, an industry-standard framework for observability, making it vendor-neutral solution for monitoring your GenAI applications.
If you are using OpenTelemetry as a part of observability tech stack in your team or organization, you can use MLflow Tracing without any additional service onboarding. Please refer to the OpenTelemetry Integration to connect MLflow Tracing to your OpenTelemetry backend.
MLflow on Amazon SageMaker is a fully managed service offered as part of the SageMaker platform by AWS, including tracing and other MLflow features such as model registry.
MLflow Tracing offers a one-line solution for tracing Amazon Bedrock, making it the best suitable solution for monitoring GenAI application powered by AWS and Amazon Bedrock.
Nebuis, a cutting-edge cloud platform for GenAI explorers, offers a fully managed MLflow server. Combining the powerful GPU infrastructure of Nebuis for training and hosting LLMs / foundation models with the observability capabilities of MLflow Tracing, Nebius serves as a comprehensive platform for AI/ML developers.
Refer to the Nebius documentation for more details about the managed MLflow service.
Supported Features
The lightweight SDK includes all essential tracing functionalities for monitoring your GenAI applications. Click the cards below to learn more about each supported feature.
MLflow Tracing SDK supports one-line integration with all of the most popular LLM/GenAI libraries including OpenAI, Anthropic, LangChain, LlamaIndex, Amazon Bedrock, DSPy, and any LLM provides that conforms to OpenAI API format. This automatic tracing capability allows you to monitor your GenAI application with MLflow Tracing with minimal effort and switch between different libraries easily.
MLflow Tracing SDK provides a simple and intuitive API for manually instrumenting your GenAI application. Manual instrumentation and automatic tracing can be used together, allowing you to trace advanced applications containing custom code and have fine-grained control over the tracing behavior.
By annotating traces with custom tags, you can add more context to your traces to group them and simplify the process of searching for them later. This is useful when you want to trace an application that runs across multiple request sessions.
Features Not Included
The following MLflow features are not available in the lightweight package:
- MLflow tracking server and UI
- Run management APIs (e.g.
mlflow.start_run) - Model logging and evaluation
- Model / prompt registry
- MLflow AI Gateway
- Other MLflow features unrelated to tracing
For these features, use the full MLflow package by installing pip install mlflow.