向代理添加跟踪

项目
09/25/2024

重要

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

本文介绍如何使用 MLflow 跟踪提供的 Fluent 和 MLflowClient API 向代理添加跟踪。

注意

有关 MLflow 跟踪的详细 API 参考和代码示例，请参阅 MLflow 文档。

要求

MLflow 2.13.1

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

如果使用支持跟踪的 GenAI 库（如 LangChain、LlamaIndex 或 OpenAI），你可以为库集成启用 MLflow 自动日志记录以启用跟踪。例如，使用 mlflow.langchain.autolog() 自动将跟踪添加到基于 LangChain 的代理。

注意

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

mlflow.langchain.autolog()

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

使用 Fluent API 手动向代理添加跟踪

下面是一个使用 Fluent API: mlflow.trace 和 mlflow.start_span 将跟踪添加到 quickstart-agent 的简单示例。建议对 PyFunc 模型执行此操作。


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 are answering python, coding, SQL, data engineering, spark, data science, DW and platform, API or infrastructure administration question related to Databricks. If the question is not related to one of these topics, kindly decline to answer. If you don't know the answer, just 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,
  }
)

Fluent API

MLflow 中的 Fluent API 会根据代码的执行位置和时间自动生成跟踪层次结构。以下部分介绍了使用 MLflow 跟踪 Fluent API 支持的任务。

修饰函数

可以使用 @mlflow.trace 修饰器来修饰函数，从而为修饰函数的作用域创建一个范围。该范围在函数调用时开始，在函数返回时结束。 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)

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 返回范围。获取请求 ID、跟踪的唯一标识符（也称为 trace_id）以及使用 span.request_id 和 span.span_id 返回范围的 ID。
使用 client.start_span(request_id, parent_id=span_id) 启动子范围来设置范围的属性、输入和输出。
1. 此方法需要通过 request_id 和 parent_id 将范围与跟踪层次结构中的正确位置相关联。它返回另一个范围对象。
通过调用 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"]

# Span has to be ended explicitly
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)

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

通过