入门：GenAI 的 MLflow 跟踪（本地 IDE）

2025-10-21

如果使用的是 IDE（VS Code、PyCharm、Cursor 或其他环境）或本地托管的笔记本环境（Jupyter 或其他环境），本快速入门可帮助你将 GenAI 应用与 MLflow 跟踪集成。如果使用 Databricks Notebook，请改用 Databricks Notebook 快速入门。

你将实现的目标

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

GenAI 应用的 MLflow 实验
连接到 MLflow 的本地开发环境
使用 MLflow 跟踪检测的简单 GenAI 应用程序
在你的 MLflow 实验中，该应用的跟踪记录

trace

先决条件

Databricks 工作区：访问 Databricks 工作区。

步骤 1：安装 MLflow

在本地 IDE 中工作时，需要使用 Databricks 连接安装 MLflow。

pip install --upgrade "mlflow[databricks]>=3.1" openai

步骤 2：创建新的 MLflow 试验

MLflow 试验是 GenAI 应用程序的容器。在概念部分中了解有关试验的详细信息。

打开 Databricks 工作区
在左侧边栏中的 “AI/ML”下，单击“ 试验”。
在“试验”页顶部，单击 GenAI 应用和代理
单击左上角的信息，获取试验 ID 和路径。稍后你会使用这些。

创建试验

步骤 3：将环境连接到 MLflow

下面的代码片段演示如何使用 Databricks 个人访问令牌设置身份验证。 MLflow 还适用于其他 Databricks 支持的身份验证方法。

使用环境变量

在 MLflow 试验中，单击 Kebab 菜单图标 >本地记录日志跟踪>，然后单击“生成 API 密钥”。

在终端中复制并运行生成的代码。

export DATABRICKS_TOKEN=<databricks-personal-access-token>
export DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com
export MLFLOW_TRACKING_URI=databricks
export MLFLOW_REGISTRY_URI=databricks-uc
export MLFLOW_EXPERIMENT_ID=<experiment-id>

使用 .env 文件

在 MLflow 试验中，单击 Kebab 菜单图标 >本地记录日志跟踪>，然后单击“生成 API 密钥”。

将生成的代码复制到 .env 项目根目录中的文件

DATABRICKS_TOKEN=<databricks-personal-access-token>
DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com
MLFLOW_TRACKING_URI=databricks
MLFLOW_REGISTRY_URI=databricks-uc
MLFLOW_EXPERIMENT_ID=<experiment-id>

安装 python-dotenv 包
```
pip install python-dotenv
```

在代码中加载环境变量

# At the beginning of your Python script
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

步骤 4：创建和检测应用程序

创建启用追踪功能的 GenAI 应用。

创建在项目目录中命名 app.py 的 Python 文件。

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

OpenAI 托管的 LLM

使用本地 OpenAI SDK 连接到由 OpenAI 托管的模型。从可用的 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/docs-demo")

# Create an OpenAI client connected to OpenAI SDKs
client = openai.OpenAI()

# Select an LLM
model_name = "gpt-4o-mini"

定义并运行应用程序：

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

# 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,  # This example uses a Databricks hosted LLM - you can replace this with any AI Gateway or Model Serving endpoint. If you provide your own OpenAI credentials, replace with a valid OpenAI model e.g., gpt-4o, etc.
        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)

运行应用程序
```
python app.py
```

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

步骤 5：在 MLflow 中查看跟踪

导航回 MLflow 试验 UI。
现在，你将在 “跟踪 ”选项卡中看到生成的跟踪。
单击追踪以查看其详细信息。

跟踪详细信息

了解追踪

您刚刚创建的跟踪显示：

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

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

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

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

指南和参考

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

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

通过