跳到主要内容

介绍 MLflow 追踪

·4 分钟阅读
MLflow maintainers
MLflow 维护者

我们很高兴地宣布 MLflow 中一项强大新功能的发布:MLflow Tracing。此功能为您的 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 Tracing”实验并调用链后,您将看到以下内容:

Traces in UI

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

敬请关注更多更新,一如既往,祝您编码愉快!