Skip to main content
Open In ColabOpen on GitHub

MLflow

MLflow 是一个多功能、开源的平台,用于管理机器学习生命周期中的工作流和工件。它与许多流行的机器学习库有内置的集成,但也可以与任何库、算法或部署工具一起使用。它被设计为可扩展的,因此您可以编写插件来支持新的工作流、库和工具。

在LangChain集成的背景下,MLflow提供了以下功能:

  • 实验跟踪: MLflow 跟踪并存储来自您的 LangChain 实验的工件,包括模型、代码、提示、指标等。
  • 依赖管理: MLflow 自动记录模型依赖,确保开发和生产环境之间的一致性。
  • 模型评估 MLflow 提供了评估 LangChain 应用程序的原生功能。
  • 追踪: MLflow 允许您通过 LangChain 链、代理、检索器或其他组件直观地追踪数据流。

注意: 跟踪功能仅在 MLflow 2.14.0 及更高版本中可用。

本笔记本演示了如何使用MLflow跟踪您的LangChain实验。有关此功能的更多信息以及探索使用LangChain与MLflow的教程和示例,请参阅MLflow文档中的LangChain集成部分

设置

安装 MLflow Python 包:

%pip install google-search-results num
%pip install mlflow -qU

此示例使用了OpenAI LLM。如果愿意,可以跳过下面的命令并继续使用不同的LLM。

%pip install langchain-openai -qU
import os

# Set MLflow tracking URI if you have MLflow Tracking Server running
os.environ["MLFLOW_TRACKING_URI"] = ""
os.environ["OPENAI_API_KEY"] = ""

首先,让我们创建一个专门的MLflow实验,以便跟踪我们的模型和工件。虽然您可以选择跳过此步骤并使用默认实验,但我们强烈建议将您的运行和工件组织到单独的实验中,以避免混乱并保持一个干净、结构化的工作流程。

import mlflow

mlflow.set_experiment("LangChain MLflow Integration")

概述

使用以下方法之一将MLflow与您的LangChain应用程序集成:

  1. 自动日志记录:通过mlflow.langchain.autolog()命令启用无缝跟踪,这是我们推荐的利用LangChain MLflow集成的首选选项。
  2. 手动日志记录: 使用MLflow API记录LangChain链和代理,提供对实验中要跟踪内容的细粒度控制。
  3. 自定义回调:在调用链时手动传递MLflow回调,允许半自动地自定义您的工作负载,例如跟踪特定的调用。

场景 1: MLFlow 自动日志记录

要开始使用自动日志记录,只需调用mlflow.langchain.autolog()。在这个例子中,我们将log_models参数设置为True,这使得链定义及其依赖库能够被记录为MLflow模型,提供全面的跟踪体验。

import mlflow

mlflow.langchain.autolog(
# These are optional configurations to control what information should be logged automatically (default: False)
# For the full list of the arguments, refer to https://mlflow.org/docs/latest/llms/langchain/autologging.html#id1
log_models=True,
log_input_examples=True,
log_model_signatures=True,
)

定义一个链

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"You are a helpful assistant that translates {input_language} to {output_language}.",
),
("human", "{input}"),
]
)
parser = StrOutputParser()

chain = prompt | llm | parser

调用链

请注意,此步骤可能比平常多花几秒钟,因为MLflow在后台运行多个后台任务,以将模型、跟踪和工件记录到跟踪服务器。

test_input = {
"input_language": "English",
"output_language": "German",
"input": "I love programming.",
}

chain.invoke(test_input)
'Ich liebe das Programmieren.'

花点时间探索MLflow跟踪用户界面,在这里你可以更深入地了解正在记录的信息。

  • 追踪 - 导航到实验中的“追踪”标签,并点击第一行的请求ID链接。显示的追踪树可视化您的链式调用的调用堆栈,为您提供深入了解链中每个组件如何执行的洞察。
  • MLflow 模型 - 当我们设置 log_model=True 时,MLflow 会自动创建一个 MLflow 运行来跟踪您的链定义。导航到最新的运行页面并打开“Artifacts”选项卡,其中列出了作为 MLflow 模型记录的文件工件,包括依赖项、输入示例、模型签名等。

调用已记录的链

接下来,让我们重新加载模型并验证我们是否可以重现相同的预测,确保一致性和可靠性。

有两种加载模型的方法

  1. mlflow.langchain.load_model(MODEL_URI) - 这将模型加载为原始的LangChain对象。
  2. mlflow.pyfunc.load_model(MODEL_URI) - 这将在PythonModel包装器中加载模型,并使用predict() API封装预测逻辑,该API包含诸如模式强制执行等附加逻辑。
# Replace YOUR_RUN_ID with the Run ID displayed on the MLflow UI
loaded_model = mlflow.langchain.load_model("runs:/{YOUR_RUN_ID}/model")
loaded_model.invoke(test_input)
'Ich liebe Programmieren.'
pyfunc_model = mlflow.pyfunc.load_model("runs:/{YOUR_RUN_ID}/model")
pyfunc_model.predict(test_input)
['Ich liebe das Programmieren.']

配置自动日志记录

mlflow.langchain.autolog() 函数提供了多个参数,允许对记录到 MLflow 的工件进行细粒度控制。有关可用配置的完整列表,请参阅最新的 MLflow LangChain 自动记录文档

场景 2: 从代码中手动记录代理

先决条件

本示例使用SerpAPI,一个搜索引擎API,作为代理检索Google搜索结果的工具。LangChain原生集成了SerpAPI,只需一行代码即可为您的代理配置该工具。

开始使用:

  • 通过 pip 安装所需的 Python 包:pip install google-search-results numexpr
  • SerpAPI的官方网站上创建一个账户并获取一个API密钥。
  • 在环境变量中设置API密钥:os.environ["SERPAPI_API_KEY"] = "YOUR_API_KEY"

定义一个代理

在这个例子中,我们将以代码的形式记录代理定义,而不是直接提供Python对象并以序列化格式保存。这种方法有几个好处:

  1. 无需序列化: 通过将模型保存为代码,我们避免了序列化的需求,这在处理不支持序列化的组件时可能会出现问题。这种方法还消除了在不同环境中反序列化模型时可能出现的兼容性问题。
  2. 更好的透明度:通过检查保存的代码文件,您可以获得有关模型行为的宝贵见解。这与像pickle这样的序列化格式形成对比,在pickle中,模型的行为在加载回来之前仍然是不透明的,可能会暴露诸如远程代码执行之类的安全风险。

首先,创建一个单独的.py文件来定义代理实例。

为了节省时间,您可以运行以下单元格来生成一个Python文件agent.py,其中包含代理定义代码。在实际开发场景中,您可能会在另一个笔记本或手工编写的Python脚本中定义它。

script_content = """
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain_openai import ChatOpenAI
import mlflow

llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
tools = load_tools(["serpapi", "llm-math"], llm=llm)
agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)

# IMPORTANT: call set_model() to register the instance to be logged.
mlflow.models.set_model(agent)
"""

with open("agent.py", "w") as f:
f.write(script_content)

记录代理

返回原始笔记本并运行以下单元格以记录您在agent.py文件中定义的代理。

question = "How long would it take to drive to the Moon with F1 racing cars?"

with mlflow.start_run(run_name="search-math-agent") as run:
info = mlflow.langchain.log_model(
lc_model="agent.py", # Specify the relative code path to the agent definition
artifact_path="model",
input_example=question,
)

print("The agent is successfully logged to MLflow!")
The agent is successfully logged to MLflow!

现在,打开MLflow UI并导航到运行详情页面的“Artifacts”选项卡。您应该看到agent.py文件已成功记录,以及其他模型工件,如依赖项、输入示例等。

调用已记录的代理

现在加载代理并调用它。有两种加载模型的方式

# Let's turn on the autologging with default configuration, so we can see the trace for the agent invocation.
mlflow.langchain.autolog()

# Load the model back
agent = mlflow.pyfunc.load_model(info.model_uri)

# Invoke
agent.predict(question)
Downloading artifacts: 100%|██████████| 10/10 [00:00<00:00, 331.57it/s]
['It would take approximately 1194.5 hours to drive to the Moon with an F1 racing car.']

导航到实验中的"追踪"标签,并点击第一行的请求ID链接。追踪可视化展示了代理在单个预测调用中如何操作多个任务:

  1. 确定回答这些问题需要哪些子任务。
  2. 搜索F1赛车的速度。
  3. 搜索地球到月球的距离。
  4. 使用LLM计算除法。

场景 3. 使用 MLflow 回调

MLflow 回调 提供了一种半自动化的方式来在 MLflow 中跟踪您的 LangChain 应用程序。主要有两种回调可用:

  1. MlflowLangchainTracer: 主要用于生成跟踪,适用于 mlflow >= 2.14.0
  2. MLflowCallbackHandler: 将指标和工件记录到MLflow跟踪服务器。

MlflowLangchainTracer

当使用MlflowLangchainTracer回调调用链或代理时,MLflow将自动为调用堆栈生成跟踪并将其记录到MLflow跟踪服务器。结果与mlflow.langchain.autolog()完全相同,但当你只想跟踪特定的调用时,这特别有用。另一方面,自动日志记录应用于同一笔记本/脚本中的所有调用。

from mlflow.langchain.langchain_tracer import MlflowLangchainTracer

mlflow_tracer = MlflowLangchainTracer()

# This call generates a trace
chain.invoke(test_input, config={"callbacks": [mlflow_tracer]})

# This call does not generate a trace
chain.invoke(test_input)

在哪里传递回调函数

LangChain 支持两种传递回调实例的方式:(1) 请求时回调 - 将它们传递给 invoke 方法或与 with_config() 绑定 (2) 构造函数回调 - 在链的构造函数中设置它们。当使用 MlflowLangchainTracer 作为回调时,必须使用请求时回调。在构造函数中设置它只会将回调应用于顶级对象,阻止其传播到子组件,从而导致跟踪不完整。有关此行为的更多信息,请参阅 回调文档 以获取更多详细信息。

# OK
chain.invoke(test_input, config={"callbacks": [mlflow_tracer]})
chain.with_config(callbacks=[mlflow_tracer])
# NG
chain = TheNameOfSomeChain(callbacks=[mlflow_tracer])

支持的方法

MlflowLangchainTracer 支持来自 Runnable Interfaces 的以下调用方法。

  • 标准接口:invoke, stream, batch
  • 异步接口:astream, ainvoke, abatch, astream_log, astream_events

其他方法不能保证完全兼容。

MlflowCallbackHandler

MlflowCallbackHandler 是一个回调处理程序,位于 LangChain 社区代码库中。

此回调可以传递给链/代理调用,但必须通过调用flush_tracker()方法显式完成。

当使用回调调用链时,它会执行以下操作:

  1. 创建一个新的MLflow运行,或者在活动的MLflow实验中检索一个可用的活动运行。
  2. 记录诸如LLM调用次数、令牌使用量等相关指标。如果链/代理包含LLM调用并且您安装了spacy库,它还会记录文本复杂度指标,如flesch_kincaid_grade
  3. 将内部步骤记录为JSON文件(这是traces的旧版本)。
  4. 将链的输入和输出记录为Pandas Dataframe。
  5. 调用flush_tracker()方法,使用链/代理实例,将链/代理记录为MLflow模型。
from langchain_community.callbacks import MlflowCallbackHandler

mlflow_callback = MlflowCallbackHandler()

chain.invoke("What is LangChain callback?", config={"callbacks": [mlflow_callback]})

mlflow_callback.flush_tracker()
API Reference:MlflowCallbackHandler

参考文献

要了解更多关于该功能的信息,并访问使用LangChain与MLflow的教程和示例,请参阅MLflow文档中的LangChain集成部分

MLflow 还提供了几个关于 LangChain 集成的 教程示例


这个页面有帮助吗?