跟踪 GenAI 应用(笔记本)

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

你将实现的目标

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

  • 一个 Databricks Notebook,其中包含 GenAI 应用的链接 MLflow 试验
  • 简单的 GenAI 应用程序,集成了 MLflow 跟踪功能。
  • MLflow 试验中该应用的跟踪

跟踪

先决条件

  • 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 访问权限。 若要运行本快速入门,可以:

  1. 使用 Databricks 托管的 LLM模型
  2. 直接使用来自 LLM 提供商的自己的 API 密钥
  3. 创建 外部模型 以启用对 LLM 提供程序 API 密钥的受控制访问权限

如果您希望使用 MLflow 支持的另外 20 多个 LLM SDK(如 Anthropic、Bedrock 等)或 GenAI 创作框架(如 LangGraph 等),请按照创建新的 MLflow 实验时出现的说明进行操作。

为应用程序选择适当的集成:

  1. 使用以下代码在笔记本中创建单元格

    在这里,我们使用 @mlflow.trace 修饰器,这使我们能够轻松跟踪与 OpenAI 自动化监控工具 相结合的任何 Python 应用程序,从而捕获对 OpenAI SDK 调用的详细信息。

    下面的代码片段使用Anthropic的Claude Sonnet LLM。 可以从 受支持的基础模型列表中选择另一个 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
    mlflow_creds = mlflow.utils.databricks_utils.get_databricks_host_creds()
    client = OpenAI(
        api_key=mlflow_creds.token,
        base_url=f"{mlflow_creds.host}/serving-endpoints"
    )
    
    # 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="databricks-claude-sonnet-4",  # This example uses a Databricks hosted LLM - you can replace this with any AI Gateway or Model Serving endpoint. If you have an external model endpoint configured, you can use that model name here.
        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)
    
  2. 运行 Notebook 单元格以执行代码

步骤 4:在 MLflow 中查看轨迹

  1. 追踪信息会显示在笔记本单元格下方
  2. (可选)可以转到 MLflow 实验 UI 以查看跟踪
    1. 单击屏幕右侧的“试验”图标
    2. 单击试验运行旁边的打开图标
    3. 现在,你将在 “跟踪 ”选项卡中看到生成的跟踪
    4. 单击轨迹以查看其详细信息

了解轨迹

刚刚创建的追踪显示了:

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

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

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

小窍门

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

后续步骤

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

参考指南

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