入门:本地 IDE 中 GenAI 的 MLflow 跟踪

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

本教程将指导你完成以下操作:

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

trace

先决条件

  • 访问Azure Databricks工作区。

步骤 1:安装 MLflow

使用 Databricks 连接安装 MLflow:

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

步骤 2:创建新的 MLflow 试验

MLflow 试验是 GenAI 应用程序的容器。 有关详细信息,请参阅 试验

  1. 打开 Azure Databricks 工作区。

  2. 在左侧边栏中的 “AI/ML”下,单击“ 试验”。

  3. 在“试验”页顶部,单击 GenAI 应用和代理

  4. 若要获取试验 ID 和路径,请单击左上角的信息 图标信息图标 。 你会在后续步骤中使用这些值。

    创建试验

步骤 3:将环境连接到 MLflow

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

使用环境变量

  1. 在 MLflow 试验中,单击 kebab 菜单图标 Kebab 菜单图标。>本地>日志跟踪生成 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_REGISTRY_URI=databricks-uc
    export MLFLOW_EXPERIMENT_ID=<experiment-id>
    

使用 .env 文件

  1. 在 MLflow 试验中,单击 kebab 菜单图标 Kebab 菜单图标。>本地>日志跟踪生成 API 密钥

  2. 将生成的代码复制到 .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>
    
  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:创建和检测应用程序

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

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

  2. 初始化 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"
    
  3. 定义并运行应用程序:

    @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(
            # Uses a Databricks-hosted LLM by default. To use an AI Gateway, Model Serving endpoint, or your own OpenAI credentials, replace `model_name` with a valid model such as `gpt-4o`.
            model=model_name,
            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
    

有关向应用添加跟踪的详细信息,请参阅 向应用程序添加跟踪:自动和手动跟踪MLflow 跟踪集成 (超过 20 个库集成)。

步骤 5:在 MLflow 中查看跟踪

  1. 返回到 MLflow 试验 UI。
  2. 生成的跟踪将显示在“ 跟踪 ”选项卡中。
  3. 点击跟踪以查看其详细信息。

跟踪详细信息

了解追踪

新的跟踪记录显示:

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

即使是此最小的跟踪图面也提供了有关应用程序行为的有用信息,包括:

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

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

指南和参考

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