介绍 MLflow 追踪

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

注意：MLflow Tracing 已在 MLflow 2.14.0 版本中发布，旧版本中不提供此功能。

推出 MLflow 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 跟踪实验时，调用链后您将看到如下内容

用于手动跟踪的流畅 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 社区。

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