跳到主要内容

介绍 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() 即可激活跟踪。
  • 使用高级 Fluent API 手动跟踪仪器化:通过 Fluent 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

用于手动跟踪的 Fluent API

为了获得更多控制,您可以使用 MLflow 的 Fluent 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)

上下文处理器

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

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 社区。

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