代理的 MLflow 跟踪

项目
02/10/2025

重要

此功能目前以公共预览版提供。

本文介绍 Databricks 上的 MLflow 跟踪以及如何使用它向生成式 AI 应用程序添加可观测性。

什么是 MLflow 跟踪？

MLflow 跟踪捕获有关执行生成式 AI 应用程序的详细信息。跟踪与请求的每个中间步骤关联的输入、输出和元数据，以便可以查明 bug 和意外行为的来源。例如，如果你的模型出现幻觉，你可以快速检查导致幻觉的每个步骤。

MLflow 跟踪与 Databricks 工具和基础结构集成，允许在 Databricks 笔记本或 MLflow 试验 UI 中存储和显示跟踪。

内联跟踪捕获生成式 AI 应用中每个步骤的详细信息中每个步骤的详细信息

为何使用 MLflow 跟踪？

MLflow 跟踪提供了以下几个优势：

查看交互式跟踪可视化效果，并使用调查工具来诊断问题。
验证提示模板和护栏是否产生合理的结果。
分析不同框架、模型和区块大小的延迟。
通过测量不同模型的令牌使用来估算应用程序成本。
建立基准“黄金”数据集以评估不同版本的性能。
存储来自生产模型终结点的日志，以调试问题并进行离线评估。

向代理添加跟踪

MLflow 跟踪支持向生成式 AI 应用程序添加跟踪的三种方法。有关 API 参考的详细信息，请参阅 MLflow 文档。

API	建议的用例	说明
MLflow 自动日志记录	使用集成的 GenAI 库进行开发	自动日志记录功能自动记录对支持的开源框架（如 LangChain、LlamaIndex 和 OpenAI）的跟踪。
Fluent API	使用 Pyfunc 的自定义代理	用于添加跟踪的低代码 API，无需担心管理跟踪的树结构。 MLflow 使用 Python 堆栈自动确定适当的父子跨度关系。
MLflow 客户端 API	高级用例，例如多线程	`MLflowClient` 为高级用例提供精细的线程安全 API。必须手动管理多个范围的父子关系。这可更好地控制跟踪生命周期，尤其是多线程用例。

安装 MLflow 跟踪

MLflow 跟踪在 MLflow 版本 2.13.0 及更高版本中提供，它预安装在 <DBR< 15.4 LTS ML 及更高版本中。如有必要，请使用以下代码安装 MLflow：

%pip install mlflow>=2.13.0 -qqqU
%restart_python

或者，可以安装最新版本的 databricks-agents，其中包括兼容的 MLflow 版本：

%pip install databricks-agents

使用自动日志记录向代理添加跟踪

如果 GenAI 库支持跟踪（如 LangChain 或 OpenAI），请通过向代码添加 mlflow.<library>.autolog() 来启用自动记录。例如：

mlflow.langchain.autolog()

注意

从 Databricks Runtime 15.4 LTS ML 开始，默认在笔记本中启用 MLflow 跟踪。若要禁用跟踪，例如使用 LangChain，可以在笔记本中执行 mlflow.langchain.autolog(log_traces=False)。

MLflow 支持用于跟踪自动日志记录的其他库。有关集成库的完整列表，请参阅 MLflow 追踪文档。

使用 Fluent API 手动将跟踪信息添加至代理

MLflow 中的 Fluent API 根据代码的执行流自动创建跟踪层次结构。

修饰函数

使用 @mlflow.trace 修饰器为修饰函数的作用域创建一个范围。

MLflow Span 对象组织跟踪步骤。范围捕获有关工作流中各个操作或步骤的信息，例如 API 调用或矢量存储查询。

该范围在调用函数时开始，并在返回时结束。 MLflow 记录函数的输入和输出以及从函数引发的任何异常。

例如，以下代码创建一个名为 my_function 的跨度，用于捕获输入参数 x 和 y 和输出。

@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
    return x + y

使用跟踪上下文管理器

如果要为任意代码块（而不仅仅是函数）创建跨度，可以使用 mlflow.start_span() 作为包装代码块的上下文管理器。范围在输入上下文时开始，并在退出上下文时结束。应使用上下文管理器生成的范围对象的 setter 方法手动提供范围输入和输出。

with mlflow.start_span("my_span") as span:
    span.set_inputs({"x": x, "y": y})
    result = x + y
    span.set_outputs(result)
    span.set_attribute("key", "value")

包装外部函数

若要跟踪外部库函数，请使用 mlflow.trace封装函数。

from sklearn.metrics import accuracy_score

y_pred = [0, 2, 1, 3]
y_true = [0, 1, 2, 3]

traced_accuracy_score = mlflow.trace(accuracy_score)
traced_accuracy_score(y_true, y_pred)

Fluent API 示例

以下示例演示如何使用 Fluent API mlflow.trace 和 mlflow.start_span 跟踪 quickstart-agent：

import mlflow
from mlflow.deployments import get_deploy_client

class QAChain(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.client = get_deploy_client("databricks")

    @mlflow.trace(name="quickstart-agent")
    def predict(self, model_input, system_prompt, params):
        messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content":  model_input[0]["query"]
                }
          ]

        traced_predict = mlflow.trace(self.client.predict)
        output = traced_predict(
            endpoint=params["model_name"],
            inputs={
                "temperature": params["temperature"],
                "max_tokens": params["max_tokens"],
                "messages": messages,
            },
        )

        with mlflow.start_span(name="_final_answer") as span:
          # Initiate another span generation
            span.set_inputs({"query": model_input[0]["query"]})

            answer = output["choices"][0]["message"]["content"]

            span.set_outputs({"generated_text": answer})
            # Attributes computed at runtime can be set using the set_attributes() method.
            span.set_attributes({
              "model_name": params["model_name"],
                        "prompt_tokens": output["usage"]["prompt_tokens"],
                        "completion_tokens": output["usage"]["completion_tokens"],
                        "total_tokens": output["usage"]["total_tokens"]
                    })
              return answer

添加跟踪后，运行函数。以下内容继续讲解上一节中的 predict() 函数示例。运行调用方法 predict() 时会自动显示跟踪。

SYSTEM_PROMPT = """
You are an assistant for Databricks users. You answer Python, coding, SQL, data engineering, spark, data science, DW and platform, API, or infrastructure administration questions related to Databricks. If the question is unrelated to one of these topics, kindly decline to answer. If you don't know the answer, say that you don't know; don't try to make up an answer. Keep the answer as concise as possible. Use the following pieces of context to answer the question at the end:
"""

model = QAChain()

prediction = model.predict(
  [
      {"query": "What is in MLflow 5.0"},
  ],
  SYSTEM_PROMPT,
  {
    # Using Databricks Foundation Model for easier testing, feel free to replace it.
    "model_name": "databricks-dbrx-instruct",
    "temperature": 0.1,
    "max_tokens": 1000,
  }
)

MLflow 客户端 API

MlflowClient 公开精细的线程安全 API 来启动和结束跟踪、管理范围以及设置范围字段。它提供对跟踪生命周期和结构的完全控制。当 Fluent API 不足以满足要求（例如多线程应用程序和回调）时，这些 API 非常有用。

以下是使用 MLflow 客户端创建完整跟踪的步骤。

通过 client = MlflowClient()创建 MLflowClient 实例。
使用 client.start_trace() 方法启动跟踪。这会启动跟踪上下文，启动绝对根范围，并返回根范围对象。此方法必须在 start_span() API 之前运行。
1. 在 client.start_trace() 中设置跟踪的属性、输入和输出。
注意

Fluent API 中没有等效于 start_trace() 方法。这是因为，Fluent API 根据托管状态自动初始化跟踪上下文并确定它是否是根范围。
start_trace() API 返回一个跨度。使用 span.request_id 和 span.span_id获取请求 ID、跟踪的唯一标识符（也称为 trace_id）和返回范围的 ID。
使用 client.start_span(request_id, parent_id=span_id) 启动子范围来设置范围的属性、输入和输出。
1. 此方法需要 request_id 和 parent_id 将范围与跟踪层次结构中的正确位置相关联。它返回另一个 span 对象。
通过调用 client.end_span(request_id, span_id) 来结束子范围。
对于想要创建的任何子分段，重复步骤 3 - 5。
所有子范围结束后，调用 client.end_trace(request_id) 关闭跟踪并进行记录。

from mlflow.client import MlflowClient

mlflow_client = MlflowClient()

root_span = mlflow_client.start_trace(
  name="simple-rag-agent",
  inputs={
          "query": "Demo",
          "model_name": "DBRX",
          "temperature": 0,
          "max_tokens": 200
         }
  )

request_id = root_span.request_id

# Retrieve documents that are similar to the query
similarity_search_input = dict(query_text="demo", num_results=3)

span_ss = mlflow_client.start_span(
      "search",
      # Specify request_id and parent_id to create the span at the right position in the trace
        request_id=request_id,
        parent_id=root_span.span_id,
        inputs=similarity_search_input
  )
retrieved = ["Test Result"]

# You must explicitly end the span
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)

root_span.end_trace(request_id, outputs={"output": retrieved})

查看跟踪

若要在运行代理后查看跟踪，请使用以下选项之一：

跟踪可视化在单元格输出中以内联方式呈现。
跟踪将记录到 MLflow 试验。可以在试验页上的“跟踪”选项卡中查看和搜索历史跟踪的完整列表。当代理在活动的 MLflow 运行中运行时，跟踪会显示在“运行”页上。
使用 search_traces() API 以编程方式检索跟踪。

在生产环境中使用 MLflow 跟踪

MLflow 跟踪还与 Mosaic AI 模型服务集成，使你能够高效调试问题、监视性能，并创建用于离线评估的黄金数据集。为服务终结点启用 MLflow 跟踪后，跟踪将记录在列下的response中。

可以通过查询表，并在笔记本中显示结果来将记录到推理表中的跟踪进行可视化。在笔记本中使用 display(<the request logs table>)，然后选择要可视化的跟踪的各个行。

若要为服务终结点启用 MLflow 跟踪，必须将终结点配置中的 ENABLE_MLFLOW_TRACING 环境变量设置为 True。若要了解如何使用自定义环境变量部署终结点，请参阅添加纯文本环境变量。如果使用 deploy() API 部署了代理，则跟踪会自动记录到推理表中。请参阅为生成式 AI 应用程序部署代理。

注意

跟踪是以异步方式写入推理表的，因此在开发过程中，这不会产生与在笔记本环境中进行跟踪时的同样高的开销。但是，它仍可能会对终结点响应速度造成一定的开销，尤其是在每个推理请求的跟踪大小较大时。 Databricks 不保证任何服务级别协议 (SLA) 对模型终结点的实际延迟影响，因为它严重依赖于环境和模型实现。 Databricks 建议先测试终结点性能并深入了解跟踪开销，再部署到生产应用程序。

下表粗略地展示了不同跟踪大小对推理延迟的影响。

每个请求的跟踪大小	对延迟的影响 (ms)
~10 KB	~ 1 毫秒
~ 1 MB	50 ~ 100 毫秒
10 MB	150 毫秒 ~

限制

Databricks 笔记本、笔记本作业和模型服务中提供了 MLflow 跟踪。

LangChain 自动记录可能不支持所有 LangChain 预测 API。有关支持的 API 的完整列表，请参阅 MLflow 文档。

通过