跟踪 GenAI 应用(IDE)

本快速入门帮助你将 GenAI 应用与 MLflow 跟踪 集成(如果使用本地 IDE 作为开发环境)。 如果使用 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 应用程序的容器。 详细了解试验及其包含在 数据模型 部分中的内容。

  1. 打开 Databricks 工作区

  2. 转到机器学习下的左侧边栏中的试验

  3. 在“试验”页顶部,单击“新建 GenAI 试验

实验创建

步骤 3:将环境连接到 MLflow

注释

本快速入门介绍如何使用 Databricks 个人访问令牌。 MLflow 还适用于其他 Databricks 支持的身份验证方法

使用环境变量

  1. 单击“生成 API 密钥

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

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

使用 .env 文件

  1. 单击“生成 API 密钥

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

    DATABRICKS_TOKEN=<databricks-personal-access-token>
    DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com
    MLFLOW_TRACKING_URI=databricks
    MLFLOW_EXPERIMENT_ID=<experiment-id>
    
  3. 安装 python-dotenv

    pip install python-dotenv
    
  4. 在代码中加载环境变量

    # At the beginning of your Python script
    from dotenv import load_dotenv
    
    # Load environment variables from .env file
    load_dotenv()
    

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

以下示例快速入门使用 OpenAI SDK 连接到 Databricks 托管 LLM。 如果要使用自己的 OpenAI 密钥,请更新行 client = OpenAI(...)

如果希望使用 MLflow 支持的其他 20+ 个 LLM SDK(Anthropic、Bedrock 等)或 GenAI 创作框架(LangGraph 等),请按照上一步中的 MLflow 实验 UI 中的说明进行操作。

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

  2. 初始化 OpenAI 客户端以连接到 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"
    
  3. 定义并运行应用程序:

    @mlflow.trace使用修饰器可以轻松跟踪与 OpenAI 自动检测相结合的任何 Python 应用程序,以捕获对 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)
    
  4. 运行应用程序

    python app.py
    

步骤 5:在 MLflow 中查看追踪

  1. 返回到 MLflow 实验 UI

  2. 现在,你将在 “跟踪 ”选项卡中看到生成的跟踪

  3. 单击跟踪记录以查看其详细信息

跟踪详细信息

了解追踪

您刚刚创建的跟踪显示:

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

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

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

小窍门

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

后续步骤

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

参考指南

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