跟踪 GenAI 应用(笔记本)

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

你将实现的目标

在本教程结束时,你将拥有:

  • 一个与 MLflow 实验相关联的 Databricks Notebook,适用于您的 GenAI 应用程序
  • 借助于 MLflow 跟踪功能的简易 GenAI 应用
  • MLflow 试验中该应用的跟踪

trace

先决条件

  • Databricks 工作区:访问 Databricks 工作区。

步骤 1:创建 Databricks Notebook

注释

创建 Databricks Notebook 将创建一个 MLflow 试验,该试验是 GenAI 应用程序的容器。 详细了解试验及其包含在 数据模型 部分中的内容。

  1. 打开 Databricks 工作区

  2. 转到左侧边栏顶部的 “新建

  3. 单击 “笔记本”

笔记本创建

Databricks 运行时包括 MLflow。 但是,为了获得最佳 GenAI 功能体验,包括最全面的跟踪功能和可靠的支持,强烈建议使用最新版本的 MLflow。

通过运行以下命令更新笔记本中的 MLflow:

%pip install --upgrade "mlflow[databricks]>=3.1" openai
dbutils.library.restartPython()
  • mlflow[databricks]>=3.1:此命令可确保具有 MLflow 3.1 或更高版本,以及 databricks 用于在 Databricks 中实现无缝连接和功能的额外功能。
  • dbutils.library.restartPython():这一点对于确保 Python 内核使用新安装的版本至关重要。

警告

虽然跟踪功能在 MLflow 2.15.0+ 中可用,但强烈建议安装 MLflow 3(特别是 3.1 或更高版本(如果使用mlflow[databricks])以获取最新的 GenAI 功能,包括扩展的跟踪功能和可靠的支持。

步骤 3:给应用程序插装

Databricks 提供流行的前沿技术和开源基础 LLM 的开箱即用访问。 若要运行本快速入门,可以从以下模型托管选项中进行选择:

  • 直接从 LLM 提供程序(如 OpenAI 或 MLflow 支持的 20 多个其他 LLM SDK)使用自己的 API 密钥。
  • 创建外部模型以启用对 LLM 提供程序 API 密钥的受控制访问权限。

在笔记本单元中运行以下代码。 它结合使用 @mlflow.trace 修饰器OpenAI 自动化检测 来捕获 LLM 请求的详细信息。

OpenAI SDK

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


import mlflow
import os
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 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")

openai_client = openai.OpenAI()

# 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 = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    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)

步骤 4:在 MLflow 中查看追踪

“跟踪”将显示在 Notebook 单元格下方。

笔记本跟踪 UI

可选择进入 MLflow 实验 UI 查看跟踪。

  1. 单击屏幕右侧的“试验”图标。

  2. 单击“实验运行”旁边的打开按钮图标。

  3. 生成的跟踪将显示在“ 跟踪 ”选项卡中。

  4. 单击跟踪以查看其详细信息。

理解追踪技术

您刚刚创建的跟踪显示:

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

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

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

小窍门

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

后续步骤

继续您的旅程,并参考这些推荐的行动和教程。

参考指南

浏览本指南中提到的概念和功能的详细文档。