入门：GenAI 的 MLflow 跟踪（Databricks Notebook）

2025-10-21

本快速入门帮助你将 GenAI 应用与 MLflow 跟踪集成（如果使用 Databricks Notebook 作为开发环境）。如果使用本地 IDE，请改用 IDE 快速入门。

你将实现的目标

在本教程结束时，你将拥有：

一个与 MLflow 实验关联的 Databricks Notebook，用于你的 GenAI 应用。
借助于 MLflow 跟踪功能的简易 GenAI 应用
在你的 MLflow 实验中，该应用的跟踪记录

trace

先决条件

Databricks Notebook：在 Databricks 工作区中创建新笔记本。笔记本将具有默认的 MLflow 实验，该实验是 GenAI 应用程序的容器。在 MLflow 概念部分了解有关 MLflow 试验的详细信息。

步骤 1：更新 MLflow

将 MLflow 更新到最新版本，以获得最佳 GenAI 体验：

%pip install --upgrade "mlflow[databricks]>=3.1" openai
dbutils.library.restartPython()

步骤 2：为应用程序植入工具

下面的代码片段定义一个简单的 GenAI 应用，该应用使用 LLM 完成句子模板。

初始化 OpenAI 客户端以连接到由 Databricks 托管的 LLM 或者由 OpenAI 托管的 LLM。

Databricks 托管的 LLM

下面的代码片段使用 OpenAI 客户端连接到 Databricks 托管的 LLM。该应用使用人类学的 Claude Sonnet LLM，但可以从可用的基础模型中进行选择。

import mlflow
from databricks.sdk import WorkspaceClient

# Enable MLflow's autologging to instrument your application with Tracing
mlflow.openai.autolog()

# Create an OpenAI client that is connected to Databricks-hosted LLMs
w = WorkspaceClient()
client = w.serving_endpoints.get_open_ai_client()

model_name = "databricks-claude-sonnet-4"

OpenAI SDK

使用本地 OpenAI SDK 连接到由 OpenAI 托管的模型。下面的代码片段使用 gpt-4o-mini，但可以从可用的 OpenAI 模型中进行选择。

import mlflow
import openai

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured

# Enable MLflow's autologging to instrument your application with Tracing
mlflow.openai.autolog()

client = openai.OpenAI()

model_name = "gpt-4o-mini"

定义并运行应用程序：

@mlflow.trace使用修饰器可以轻松跟踪任何 Python 函数，并结合 OpenAI 自动检测来捕获对 OpenAI SDK 调用的详细信息。

import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

# Use the trace decorator to capture the application's entry point
@mlflow.trace
def my_app(input: str):
  # This call is automatically instrumented by `mlflow.openai.autolog()`
  response = client.chat.completions.create(
    model=model_name,
    temperature=0.1,
    max_tokens=100,
    messages=[
      {
        "role": "system",
        "content": "You are a helpful assistant.",
      },
      {
        "role": "user",
        "content": input,
      },
    ]
  )
  return response.choices[0].message.content

result = my_app(input="What is MLflow?")
print(result)

有关向应用添加跟踪的详细信息，请参阅跟踪检测指南和 20 多个库集成。

步骤 3：在 MLflow 中查看跟踪

追踪将显示在笔记本单元格下方。

笔记本跟踪 UI

（可选）您可以转到 MLflow 试验 UI 查看踪迹：

单击右侧栏中的
单击试验运行旁边的打开图标。
生成的跟踪将显示在“ 跟踪 ”选项卡中。
单击跟踪以查看其详细信息。

理解追踪技术

您刚刚创建的跟踪显示：

根范围：表示my_app(...)函数的输入
- 子范围：表示 OpenAI 完成请求
属性：包含模型名称、令牌计数和计时信息等元数据
输入：发送到模型的消息
输出：从模型接收的响应

此简单跟踪已提供对应用程序行为的宝贵见解，例如：

有人问什么
生成了什么响应
请求花费的时间
使用了多少个令牌（影响成本）

对于更复杂的应用程序（如 RAG 系统或多步骤代理），MLflow 跟踪通过揭示每个组件和步骤的内部工作来提供更多价值。

指南和参考

有关本指南中的概念和功能的详细信息，请参阅：

MLflow 跟踪指南 - 从此处开始详细了解 MLflow 跟踪
MLflow 跟踪集成 - 具有自动跟踪集成的 20 多个库
跟踪概念 - 了解 MLflow 跟踪的基础知识

通过