介绍 MLflow 追踪

我们很高兴地宣布 MLflow 发布了一项强大的新功能：MLflow Tracing。此功能为您的 GenAI 应用程序带来了全面的仪器化能力，使您能够深入了解模型和工作流的执行情况，从简单的聊天界面到复杂的多阶段检索增强生成 (RAG) 应用程序。

注意：MLflow Tracing 已在 MLflow 2.14.0 中发布，不适用于以前的版本。

MLflow Tracing 简介

Tracing 是理解和优化复杂应用程序的关键方面，尤其是在机器学习和人工智能领域。随着 MLflow Tracing 的发布，您现在可以轻松捕获、可视化和分析 GenAI 应用程序的详细执行轨迹。这项新功能旨在提高应用程序性能和行为的可见性和控制力，有助于从微调到调试的各个方面。

什么是 MLflow Tracing？

MLflow Tracing 提供了多种方法来启用应用程序中的追踪

• 通过 LangChain 实现自动化追踪：与 LangChain 的完全自动化集成允许您通过简单地启用 mlflow.langchain.autolog() 来激活追踪。
• 使用高级流式 API 进行手动追踪仪器化：通过流式 API 使用装饰器、函数包装器和上下文管理器，以最少的代码修改添加追踪功能。
• 用于追踪的低级客户端 API：MLflow 客户端 API 提供了一种线程安全的方式来处理追踪实现，从而对记录的数据和时间进行精细控制。

MLflow Tracing 入门

LangChain 自动追踪

开始使用 MLflow Tracing 最简单的方法是通过与 LangChain 的内置集成。通过启用自动日志记录，当在链上调用调用 API 时，追踪将自动记录到活动的 MLflow 实验中。下面是一个快速示例：

import os
from langchain.prompts import PromptTemplate
from langchain_openai import OpenAI
import mlflow

assert "OPENAI_API_KEY" in os.environ, "Please set your OPENAI_API_KEY environment variable."

mlflow.set_experiment("LangChain Tracing")
mlflow.langchain.autolog(log_models=True, log_input_examples=True)

llm = OpenAI(temperature=0.7, max_tokens=1000)
prompt_template = "Imagine you are {person}, and you are answering a question: {question}"
chain = prompt_template | llm

chain.invoke({"person": "Richard Feynman", "question": "Why should we colonize Mars?"})
chain.invoke({"person": "Linus Torvalds", "question": "Can I set everyone's access to sudo?"})

当您在 MLflow UI 中导航到LangChain Tracing实验时，调用链后您将看到以下内容：

用于手动追踪的流式 API

为了获得更多控制，您可以使用 MLflow 的流式 API 手动检测您的代码。这种方法允许您以最少的现有代码更改捕获详细的追踪数据。

追踪装饰器

追踪装饰器捕获函数的输入和输出

import mlflow

mlflow.set_experiment("Tracing Demo")

@mlflow.trace
def some_function(x, y, z=2):
    return x + (y - z)

some_function(2, 4)

上下文处理器

上下文处理器非常适合在信息生成时用额外数据补充跨度信息

import mlflow

@mlflow.trace
def first_func(x, y=2):
    return x + y

@mlflow.trace
def second_func(a, b=3):
    return a * b

def do_math(a, x, operation="add"):
    with mlflow.start_span(name="Math") as span:
        span.set_inputs({"a": a, "x": x})
        span.set_attributes({"mode": operation})
        first = first_func(x)
        second = second_func(a)
        result = first + second if operation == "add" else first - second
        span.set_outputs({"result": result})
        return result

do_math(8, 3, "add")

通过客户端 API 实现全面追踪

对于高级用例，MLflow 客户端 API 提供了对追踪管理的细粒度控制。这些 API 允许您以编程方式创建、操作和检索追踪，尽管在整个实现过程中会增加复杂性。

使用客户端 API 启动和管理追踪

from mlflow import MlflowClient

client = MlflowClient()

# Start a new trace
root_span = client.start_trace("my_trace")
request_id = root_span.request_id

# Create a child span
child_span = client.start_span(
    name="child_span",
    request_id=request_id,
    parent_id=root_span.span_id,
    inputs={"input_key": "input_value"},
    attributes={"attribute_key": "attribute_value"},
)

# End the child span
client.end_span(
    request_id=child_span.request_id,
    span_id=child_span.span_id,
    outputs={"output_key": "output_value"},
    attributes={"custom_attribute": "value"},
)

# End the root span (trace)
client.end_trace(
    request_id=request_id,
    outputs={"final_output_key": "final_output_value"},
    attributes={"token_usage": "1174"},
)

深入了解追踪

MLflow Tracing 旨在灵活而强大，支持从简单函数追踪到复杂异步工作流的各种用例。

要了解有关此功能的更多信息，请阅读指南，查看 API 文档，并立即开始使用 LangChain 集成！

与我们一同踏上这段旅程

MLflow Tracing 的推出标志着我们提供管理机器学习工作流的全面工具的使命取得了重大里程碑。我们对这项新功能所开启的可能性感到兴奋，并期待您的反馈和贡献。

对于我们社区中热衷于分享知识的人，我们邀请您协作。无论是编写教程、分享用例还是提供反馈，每一份贡献都丰富了 MLflow 社区。

请继续关注更多更新，一如既往，编程愉快！