跳到主要内容

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

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

MLflow 追踪:生成式 AI 的可观测性

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

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

本博文将向您展示如何(在大约五分钟内)开始使用 MLflow 追踪。它假定您对生成式 AI API(例如 OpenAI API)有一定的了解,但不假定您事先了解 MLflow。

Alt text
使用 MLflow 追踪来追踪 LangChain 应用程序

快速开始

我们将从展示如何使用 MLflow 自动日志记录来自动追踪对 OpenAI 模型的调用开始,尽管 MLflow 支持对数量不断增加的提供商和框架的自动追踪,其中包括 Anthropic、Ollama、Langchain、LlamaIndex 等等。要开始使用,请使用以下命令安装 MLflow 和 OpenAI Python 包

pip install mlflow openai

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

在 Python 脚本或笔记本中,导入 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)”选项卡。这将列出所有收集到的追踪。点击追踪的追踪 ID 即可打开一个显示更多详细信息的新窗格。

Traces in the MLflow UI

提示

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

使用 mlflow ui 启动 MLflow 跟踪服务器后,您还可以直接在 Jupyter 笔记本中查看追踪!您只需要将跟踪 URI 设置为上述指定的地址

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

然后,当您调用启用了追踪的 AI 模型/框架时,生成的追踪将直接显示在笔记本输出中。

Tracing in Jupyter Notebooks

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

组织您的追踪

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

组织追踪最简单的方法是将它们分离到实验 (experiments)中。每个实验都有自己的追踪选项卡,其中显示该实验的追踪。

您可以在 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”实验。

您还可以使用标签 (tags)运行 (runs)来组织您的追踪。

追踪其他提供商

我们的快速开始示例侧重于 OpenAI,但 MLflow 支持对许多不同的 AI 提供商和框架的自动追踪。方法是相同的:只需将 mlflow.<provider>.autolog 这一行添加到您的笔记本或脚本中即可。

以下是一些示例。有关支持的提供商的完整列表,请此处查看。

使用 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 应用程序的工具,而不仅仅是追踪。