MLflow OpenAI 自动记录

MLflow 的 OpenAI 风格支持自动日志记录,以确保您的实验、测试和验证想法可以在不使用日志记录样板代码包装代码的情况下动态捕获。

注意

Autologging 仅支持 1.17及以上版本的OpenAI SDK。

OpenAI SDK 的 MLflow autologging 支持以下接口:

  • 聊天补全 通过 client.chat.completions.create()

  • **完成**(遗留)通过 client.completions.create()

  • 嵌入 通过 client.embeddings.create()

其中 clientopenai.OpenAI() 的一个实例。

在本指南中,我们将讨论 autologging 功能中的一些关键特性。

快速入门

要开始使用 MLflow 的 OpenAI 自动记录功能,您只需在脚本或笔记本的开头调用 mlflow.openai.autolog()。启用无参数覆盖的自动记录将表现为下表中的 默认 配置。覆盖这些设置中的任何一项将允许您记录额外的元素。

小技巧

当自动记录功能激活时,唯一默认启用的元素是跟踪信息的记录。你可以在 这里 阅读更多关于 MLflow 跟踪的内容。

import os
import openai
import mlflow

# Enables trace logging by default
mlflow.openai.autolog()

openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What does turning something up to 11 refer to?",
    }
]

# The input messages and the response will be logged as a trace to the active experiment
answer = openai_client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    temperature=0.99,
)

备注

在使用 OpenAI SDK 时,请确保您的访问令牌已分配给环境变量 OPENAI_API_KEY

OpenAI 自动日志记录的配置

MLflow OpenAI 自动记录可以记录有关模型及其推理的各种信息。默认情况下,仅启用跟踪日志记录,但您可以通过在调用 mlflow.openai.autolog() 时设置相应的参数来启用其他信息的自动记录。

以下显示了可用选项及其默认值。要了解更多关于附加参数的信息,请参阅API文档。

目标

默认

参数

描述

跟踪

true

log_traces

是否为模型生成并记录跟踪。有关跟踪功能的更多详细信息,请参阅 MLflow 跟踪

模型工件

false

log_models

如果设置为 True ,OpenAI 模型在调用时将被记录。

模型签名

false

log_model_signatures

如果设置为 True,在推理期间会收集并记录描述模型输入和输出的 模型签名 以及 OpenAI 模型工件。此选项仅在 log_models 启用时可用。

输入示例

false

log_input_examples

如果设置为 True,推理数据中的输入示例将与OpenAI模型工件一起收集并记录在推理过程中。此选项仅在 log_models 启用时可用。

例如,要禁用跟踪日志记录,并改为启用模型日志记录,请运行以下代码:

import mlflow

mlflow.openai.autolog(
    log_traces=False,
    log_models=True,
)

使用 OpenAI Autologging 的示例

import os

import mlflow
import openai

API_KEY = os.environ.get("OPENAI_API_KEY")
EXPERIMENT_NAME = "OpenAI Autologging Demonstration"
REGISTERED_MODEL_NAME = "openai-auto"
MODEL_VERSION = 1

mlflow.openai.autolog(
    log_input_examples=True,
    log_model_signatures=True,
    log_models=True,
    log_traces=True,
    registered_model_name=REGISTERED_MODEL_NAME,
)

mlflow.set_experiment(EXPERIMENT_NAME)

openai_client = openai.OpenAI(api_key=API_KEY)

messages = [
    {
        "role": "user",
        "content": "State that you are responding to a test and that you are alive.",
    }
]

openai_client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    temperature=0.95,
)

在UI中调用OpenAI客户端时使用的已记录模型和跟踪信息可以在下面的图像中查看:

OpenAI 自动记录工件和追踪

可以通过使用 models uri 加载模型,该模型是通过记录和注册的,并通过 pyfunc API 进行接口操作,如下所示:

loaded_autologged_model = mlflow.pyfunc.load_model(
    f"models:/{REGISTERED_MODEL_NAME}/{MODEL_VERSION}"
)

loaded_autologged_model.predict(
    "How much relative time difference would occur between an astronaut travelling at 0.98c for 14 years "
    "as measured by an on-board clock on the spacecraft and humans on Earth, assuming constant speed?"
)

OpenAI Swarm 的自动追踪

MLflow 2.17.1 引入了对 OpenAI Swarm 的内置追踪能力,这是一个来自 OpenAI 的多代理编排框架。该框架提供了一个清晰的接口,用于在 OpenAI 的函数调用能力和 交接与例程模式 的概念之上构建多代理系统。

MLflow 的自动追踪功能提供了对代理之间、工具调用及其集体输出的无缝追踪。只需调用 mlflow.openai.autolog() 函数,即可为 OpenAI Swarm 启用自动追踪。

import mlflow
from swarm import Swarm, Agent

# Calling the autolog API will enable trace logging by default.
mlflow.openai.autolog()

mlflow.set_experiment("OpenAI Swarm")

client = Swarm()


def transfer_to_agent_b():
    return agent_b


agent_a = Agent(
    name="Agent A",
    instructions="You are a helpful agent.",
    functions=[transfer_to_agent_b],
)

agent_b = Agent(
    name="Agent B",
    instructions="Only speak in Haikus.",
)

response = client.run(
    agent=agent_a,
    messages=[{"role": "user", "content": "I want to talk to agent B."}],
)
print(response)

OpenAI Swarm 实验相关的日志跟踪可以在 MLflow UI 中看到,如下所示:

OpenAI 集群追踪

常见问题解答

如何使用 MLflow 手动记录 OpenAI SDK 的跟踪?

通过设置一个活跃的实验(不建议为此使用默认实验),你可以在处理模型接口时(无论你是否记录模型)使用高级跟踪流式API,方法是利用MLflow跟踪流式API。

你可以在这里发现如何使用 fluent API

如果我正在为我的 OpenAI 模型使用流式传输,自动日志记录会正确记录跟踪数据吗?

是的。对于每个支持从OpenAI流式传输响应的MLflow支持的客户端接口类型,自动记录将在输出中记录迭代器响应块。

作为一个例子:

import openai
import mlflow

mlflow.set_experiment("OpenAI")

# Enable trace logging
mlflow.openai.autolog()

client = openai.OpenAI()

stream = client.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "user", "content": "How fast would a glass of water freeze on Titan?"}
    ],
    stream=True,  # Stream response
)
for chunk in stream:
    print(chunk.choices[0].delta.content or "", end="")

在 MLflow UI 中,流式模型的跟踪将显示如下:

OpenAI 自动日志记录流跟踪

备注

指定流式响应的 OpenAI 配置 尚未支持 在 MLflow 中使用 predict_stream() pyfunc 调用 API。不过,您仍然可以记录流式跟踪。当通过 mlflow.pyfunc.load_model() 将记录的 openai 模型作为 pyfunc 加载时,推理的唯一可用接口是同步阻塞的 predict() API。

autologging 中是否支持异步 API?

MLflow OpenAI 的自动记录功能 不支持用于记录模型或追踪的异步API

保存你的异步实现最好通过使用 代码模型功能 来完成。

如果你想为异步 OpenAI API 记录跟踪事件,下面是一个简化示例,用于记录流式异步请求的跟踪:

import openai
import mlflow
import asyncio

# Activate an experiment for logging traces to
mlflow.set_experiment("OpenAI")


async def fetch_openai_response(messages, model="gpt-4o", temperature=0.99):
    """
    Asynchronously gets a response from the OpenAI API using the provided messages and streams the response.

    Args:
        messages (list): List of message dictionaries for the OpenAI API.
        model (str): The model to use for the OpenAI API. Default is "gpt-4o".
        temperature (float): The temperature to use for the OpenAI API. Default is 0.99.

    Returns:
        None
    """
    client = openai.AsyncOpenAI()

    # Create the response stream
    response_stream = await client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
        stream=True,
    )

    # Manually log traces using the tracing fluent API
    with mlflow.start_span() as trace:
        trace.set_inputs(messages)
        full_response = []

        async for chunk in response_stream:
            content = chunk.choices[0].delta.content
            if content is not None:
                print(content, end="")
                full_response.append(content)

        trace.set_outputs("".join(full_response))


messages = [
    {
        "role": "user",
        "content": "How much additional hydrogen mass would Jupiter require to ignite a sustainable fusion cycle?",
    }
]

await fetch_openai_response(messages)