跳到主要内容

实用的 AI 可观测性：MLflow 追踪入门

·7分钟阅读
Daniel Liden
Daniel Liden
Databricks 开发者布道师

MLflow Tracing：面向生成式AI的可观测性​

生成式AI提供商和框架通常会返回复杂难读的数据结构，或者返回隐藏中间步骤的简单响应。此外，随着时间的推移，跟踪和比较生成式AI模型/框架调用可能会很困难，特别是当您在不同框架和脚本之间切换时。

MLflow 的 LLM 追踪通过记录您的所有生成式AI调用（包括单独的LLM调用和多步代理工作流），并提供易于阅读的界面以供浏览和比较，从而解决了这些问题。您只需一行代码即可为大多数生成式AI提供商启用此功能：mlflow.<provider>.autolog()。

本博客将展示如何在大约五分钟内开始使用 MLflow 追踪。它假定您对生成式AI API（例如 OpenAI API）有一定了解，但不需要任何 MLflow 的先验知识。

Alt text
使用 MLflow Tracing 追踪 LangChain 应用

快速入门​

我们将首先展示如何使用 MLflow 自动日志记录来自动追踪对 OpenAI 模型的调用，尽管 MLflow 支持对Anthropic、Ollama、Langchain、LlamaIndex 等越来越多的提供商和框架进行自动追踪。要开始使用，请安装 MLflow 和 OpenAI Python 包，命令如下：

pip install mlflow openai

使用自动日志记录收集追踪​

在 Python 脚本或 notebook 中，导入 MLflow 和您正在使用的生成式AI提供商，并使用 mlflow.<provider>.autolog 启用追踪。以下是为 OpenAI 设置自动追踪的方法：

import mlflow
from openai import OpenAI

mlflow.openai.autolog()

请务必创建并设置您的 OpenAI API 密钥！您可以在您的环境中通过以下方式设置它：

export OPENAI_API_KEY="your_api_key_here"

现在，当您使用 OpenAI 库时，MLflow 将捕获您的模型调用的*追踪*。例如，由于我们启用了自动日志记录，MLflow 将记录以下 OpenAI 调用的追踪。

client = OpenAI()

completion = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {
            "role": "user",
            "content": "What is an MLflow tracking server?"
        }
    ]
)

查看您的 LLM 追踪​

MLflow UI 提供了一个 AI 可观测性仪表板，用于查看您的追踪。从您的终端启动 MLflow UI，命令如下：

mlflow ui

导航到 UI。`mlflow ui` 命令的输出将告诉您访问地址（默认为 `https://:5000`）。在 UI 中，导航到“Traces”选项卡。这将列出所有已收集的追踪。点击追踪的“Trace ID”以打开一个新窗格，其中包含更多详细信息。

Traces in the MLflow UI

提示

默认情况下，MLflow 服务器将侦听 `https://:5000`。您可以使用 `mlflow ui -p <port>` 选择不同的端口。例如，要侦听端口 5001，请使用 `mlflow ui -p 5001`。

使用 `mlflow ui` 启动 MLflow 追踪服务器还可以让您直接在 Jupyter notebook 中查看追踪！您只需将追踪 URI 设置为上面指定的位置。

mlflow.set_tracking_uri("https://:5000")

然后，当您调用启用了追踪功能的 AI 模型/框架时，生成的追踪将直接显示在 notebook 输出中。

Tracing in Jupyter Notebooks

您可以使用 `mlflow.tracing.disable_notebook_display()` 禁用此功能。

组织您的追踪​

如果您在多个不同项目和任务中使用追踪，您可能希望将追踪组织成单独的组。

组织追踪最简单的方法是将其分成实验。每个实验都有自己的追踪选项卡，显示该实验的追踪。

您可以在 UI 中（通过“Experiments”旁边的“+”按钮）、使用 MLflow CLI 或通过 Python 创建实验。我们来创建一个名为“Quickstart”的新实验并记录一个追踪。

mlflow.set_experiment("quickstart")

completion = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {
            "role": "user",
            "content": "What is an MLflow tracking server?"
        }
    ]
)

现在我们可以在“quickstart”实验的“Traces”选项卡中找到此追踪。

Trace in Experiment

`set_experiment` 函数指定追踪应记录到哪个实验中，如果该实验不存在则创建它，因此上面的代码片段创建了一个新的“quickstart”实验。

您还可以使用标签和运行来组织您的追踪。

追踪其他提供商​

我们的快速入门示例侧重于 OpenAI，但 MLflow 支持对许多不同的 AI 提供商和框架进行自动追踪。方法是相同的：只需将 `mlflow.<provider>.autolog` 行添加到您的 notebook 或脚本中。

这里有一些示例。有关支持的提供商的完整列表，请参见此处。

使用 `mlflow.anthropic.autolog()` 启用 Anthropic 模型调用的自动追踪。

import anthropic
import mlflow

mlflow.anthropic.autolog()

client = anthropic.Anthropic()

message = client.messages.create(
    model="claude-3-7-sonnet-20250219",
    max_tokens=1000,
    temperature=1,
    messages=[
        {
            "role": "user",
            "content": "What is an MLflow tracking server?"
        }
    ]
)

这将在 MLflow UI 中返回以下内容：

Anthropic tracing

结论：一行代码实现高效 LLM 追踪​

在本指南中，您学习了如何使用 MLflow 的自动日志记录功能，仅用一行代码即可获得完整的 AI 可观测性解决方案。如果您正在使用 MLflow 提供自动追踪支持的众多生成式AI框架/提供商之一（包括任何具有 OpenAI 兼容端点的提供商），那么自动日志记录是可视化和调试您的 AI 应用行为的最简单方法。您所需要的只是 `mlflow.<provider>.autolog()`。

后续步骤​

自动日志记录是开始使用 MLflow 追踪的绝佳起点，但随着您开发更复杂的生成式AI应用，您可能需要更灵活地收集和使用追踪。此外，MLflow 除了追踪之外，还包含许多用于处理生成式AI应用的工具。