MLflow 追踪简介

备注

MLflow Tracing 目前处于 实验状态 ,可能会在没有弃用警告或通知的情况下发生变更。

MLflow Tracing 是一个功能,通过捕获应用程序服务的执行详细信息,增强了生成式AI(GenAI)应用程序中LLM的可观察性。Tracing提供了一种记录请求的每个中间步骤的输入、输出和元数据的方法,使您能够轻松地找出错误和意外行为的来源。

MLflow 提供了多种不同的选项来启用对您的 GenAI 应用程序的跟踪。

  • 自动追踪:MLflow 提供了与 LangChain、OpenAI、LlamaIndex 和 AutoGen 等集成库的完全自动化集成,只需启用 mlflow.<library>.autolog() 即可激活。

  • 使用高级流畅API进行手动跟踪检测:通过流畅API的装饰器、函数包装器和上下文管理器,允许您通过少量代码修改添加跟踪功能。

  • 用于追踪的低级客户端API: MLflow 客户端API提供了一种线程安全的方式来处理追踪实现,即使在异步操作模式下也是如此。

要了解更多关于追踪是什么的信息,请参阅我们的 追踪概念概述 指南。

要探索 MLflow Tracing 的结构和模式,请参阅 Tracing Schema 指南。

备注

MLflow 跟踪支持在 MLflow 2.14.0 版本中可用。此版本之前的 MLflow 版本不包含支持跟踪日志记录所需的完整功能集。

自动追踪

提示

你最喜欢的库不在列表中吗?考虑 为 MLflow Tracing 贡献提交一个功能请求 到我们的 Github 仓库。

开始使用 MLflow Tracing 的最简单方法是利用 MLflow 集成库的内置功能。MLflow 为一些集成库(如 LangChain、OpenAI、LlamaIndex 和 AutoGen)提供了自动追踪功能。对于这些库,您只需使用一个命令 mlflow.<library>.autolog() 来配置您的代码,MLflow 将自动记录模型/API 调用的追踪信息到当前活动的 MLflow 实验中。

LangChain Automatic Tracing


作为 LangChain 自动日志记录集成的一部分,当在链上调用调用 API 时,跟踪记录会记录到活动的 MLflow 实验中。您可以通过调用 mlflow.langchain.autolog() 函数来启用 LangChain 的跟踪。

import mlflow

mlflow.langchain.autolog()

在下面的完整示例中,模型及其关联的元数据将被记录为一个运行,而跟踪信息则单独记录到活动实验中。要了解更多信息,请访问 LangChain 自动记录文档

备注

此示例已确认在以下版本要求下正常工作:

pip install openai==1.30.5 langchain==0.2.1 langchain-openai==0.1.8 langchain-community==0.2.1 mlflow==2.14.0 tiktoken==0.7.0
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."

# Using a local MLflow tracking server
mlflow.set_tracking_uri("http://localhost:5000")

# Create a new experiment that the model and the traces will be logged to
mlflow.set_experiment("LangChain Tracing")

# Enable LangChain autologging
# Note that models and examples are not required to be logged in order to log traces.
# Simply enabling autolog for LangChain via mlflow.langchain.autolog() will enable trace logging.
mlflow.langchain.autolog(log_models=True, log_input_examples=True)

llm = OpenAI(temperature=0.7, max_tokens=1000)

prompt_template = (
    "Imagine that you are {person}, and you are embodying their manner of answering questions posed to them. "
    "While answering, attempt to mirror their conversational style, their wit, and the habits of their speech "
    "and prose. You will emulate them as best that you can, attempting to distill their quirks, personality, "
    "and habits of engagement to the best of your ability. Feel free to fully embrace their personality, whether "
    "aspects of it are not guaranteed to be productive or entirely constructive or inoffensive."
    "The question you are asked, to which you will reply as that person, is: {question}"
)

chain = prompt_template | llm

# Test the chain
chain.invoke(
    {
        "person": "Richard Feynman",
        "question": "Why should we colonize Mars instead of Venus?",
    }
)

# Let's test another call
chain.invoke(
    {
        "person": "Linus Torvalds",
        "question": "Can I just set everyone's access to sudo to make things easier?",
    }
)

如果我们导航到 MLflow UI,我们不仅可以查看自动记录的模型,还可以查看轨迹,如下面的视频所示:

通过 autolog 进行 LangChain 跟踪

备注

上面的例子故意简化(一个简单的聊天完成演示),以简洁为目的。在涉及复杂RAG链的实际场景中,MLflow记录的跟踪将会显著更加复杂和详细。

跟踪 Fluent API

MLflow 的 fluent APIs 提供了一种直接的方式来为你的函数和代码块添加追踪。通过使用装饰器、函数包装器和上下文管理器,你可以轻松地以最小的代码改动捕获详细的追踪数据。

作为流式API和客户端API之间追踪的比较,下图展示了两种API在复杂性上的差异,其中流式API更为简洁,如果你的追踪用例能够支持使用更高级别的API,那么这是推荐的方法。

Fluent vs 客户端 API

本节将介绍如何使用这些 fluent API 启动跟踪。

开始追踪

在本节中,我们将探讨使用 MLflow 的 fluent API 启动跟踪的不同方法。这些方法允许你通过最小的修改将跟踪功能添加到你的代码中,使你能够捕获有关函数和工作流执行的详细信息。

Trace Decorator

trace装饰器允许你通过简单地在函数定义中添加 @mlflow.trace 装饰器来自动捕获函数的输入和输出。这种方法非常适合在不显著改变现有代码的情况下快速为单个函数添加追踪。

import mlflow

# Create a new experiment to log the trace to
mlflow.set_experiment("Tracing Demo")


# Mark any function with the trace decorator to automatically capture input(s) and output(s)
@mlflow.trace
def some_function(x, y, z=2):
    return x + (y - z)


# Invoking the function will generate a trace that is logged to the active experiment
some_function(2, 4)

你可以按照以下方式向追踪装饰器添加额外的元数据:

@mlflow.trace(name="My Span", span_type="func", attributes={"a": 1, "b": 2})
def my_func(x, y):
    return x + y

当向 trace 装饰器构造函数添加额外元数据时,这些额外组件将与存储在活动 MLflow 实验中的 trace 内的 span 条目一起记录。

自 MLflow 2.16.0 起,trace 装饰器也支持异步函数:

from openai import AsyncOpenAI

client = AsyncOpenAI()


@mlflow.trace
async def async_func(message: str):
    return await client.chat.completion.create(
        model="gpt-4o", messages=[{"role": "user", "content": message}]
    )


await async_func("What is MLflow Tracing?")

捕获了什么?

如果我们导航到 MLflow UI,我们可以看到 trace 装饰器除了捕获与任何 span 相关的基本元数据(开始时间、结束时间、状态等)之外,还自动捕获了以下信息:

  • 输入: 在我们的装饰函数的情况下,这包括所有输入参数的状态(包括应用的默认 z 值)。

  • 响应: 函数的输出也被捕获,在这种情况下,是加法和减法操作的结果。

  • 跟踪名称: 装饰函数的名称。

Trace UI - 简单使用案例

错误处理与追踪

如果在处理跟踪检测的操作期间引发了 Exception,UI 中将显示调用未成功的指示,并且将提供部分数据捕获以帮助调试。此外,关于引发的异常的详细信息将包含在部分完成的跨度的 events 属性中,进一步帮助识别代码中问题发生的位置。

下面展示了一个从引发异常的代码中记录的跟踪示例:

# This will raise an AttributeError exception
do_math(3, 2, "multiply")
跟踪错误

如何处理父子关系

在使用 trace 装饰器时,每个被装饰的函数都会被视为 trace 中的一个单独的 span。依赖函数调用之间的关系通过 Python 中的原生调用执行顺序直接处理。例如,以下代码将在主父 span 中引入两个“子” span,所有这些都使用装饰器。

import mlflow


@mlflow.trace(span_type="func", attributes={"key": "value"})
def add_1(x):
    return x + 1


@mlflow.trace(span_type="func", attributes={"key1": "value1"})
def minus_1(x):
    return x - 1


@mlflow.trace(name="Trace Test")
def trace_test(x):
    step1 = add_1(x)
    return minus_1(step1)


trace_test(4)

如果我们在 MLflow UI 中查看这个跟踪,我们可以看到调用顺序在跟踪结构中显示的关系。

``Trace Decorator``

跨度类型

跨度类型是一种对跟踪中的跨度进行分类的方法。默认情况下,当使用跟踪装饰器时,跨度类型设置为 "UNKNOWN"。MLflow 提供了一组预定义的跨度类型用于常见用例,同时也允许你设置自定义跨度类型。

以下是可用的跨度类型:

跨度类型

描述

"LLM"

表示对LLM端点或本地模型的调用。

"CHAT_MODEL"

表示对聊天模型的查询。这是LLM交互的一个特例。

"CHAIN"

表示一系列操作的链。

"AGENT"

表示一个自主代理操作。

"工具"

表示一个工具执行(通常由代理执行),例如查询搜索引擎。

"EMBEDDING"

表示一个文本嵌入操作。

“检索器”

表示一个上下文检索操作,例如查询向量数据库。

"PARSER"

表示一个解析操作,将文本转换为结构化格式。

“RERANKER”

表示一个重新排序操作,根据相关性对检索到的上下文进行排序。

"UNKNOWN"

当没有指定其他跨度类型时使用的默认跨度类型。

要设置跨度类型,可以将 span_type 参数传递给 @mlflow.trace 装饰器或 mlflow.start_span 上下文管理器。当您使用 自动跟踪 时,跨度类型会由 MLflow 自动设置。

import mlflow
from mlflow.entities import SpanType


# Using a built-in span type
@mlflow.trace(span_type=SpanType.RETRIEVER)
def retrieve_documents(query: str):
    ...


# Setting a custom span type
with mlflow.start_span(name="add", span_type="MATH") as span:
    span.set_inputs({"x": z, "y": y})
    z = x + y
    span.set_outputs({"z": z})

    print(span.span_type)
    # Output: MATH

上下文处理程序

上下文处理器提供了一种创建嵌套跟踪或跨度的方式,这对于捕获代码中的复杂交互非常有用。通过使用 mlflow.start_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"):
    # Use the fluent API context handler to create a new span
    with mlflow.start_span(name="Math") as span:
        # Specify the inputs and attributes that will be associated with the span
        span.set_inputs({"a": a, "x": x})
        span.set_attributes({"mode": operation})

        # Both of these functions are decorated for tracing and will be associated
        # as 'children' of the parent 'span' defined with the context handler
        first = first_func(x)
        second = second_func(a)

        result = None

        if operation == "add":
            result = first + second
        elif operation == "subtract":
            result = first - second
        else:
            raise ValueError(f"Unsupported Operation Mode: {operation}")

        # Specify the output result to the span
        span.set_outputs({"result": result})

        return result

当调用 do_math 函数时,将生成一个跟踪,其根跨度(父)定义为上下文处理程序 with mlflow.start_span(): 调用。由于 first_funcsecond_func 调用都是装饰函数(在函数定义上有 @mlflow.trace 装饰),它们将被关联为此父跨度的子跨度。

运行以下代码将生成一个跟踪。

do_math(8, 3, "add")

这个跟踪可以在 MLflow UI 中查看:

在 MLflow UI 中追踪

函数包装

函数包装提供了一种灵活的方式,可以在不修改现有函数定义的情况下为其添加追踪功能。当你想要为第三方函数或不受你控制的函数添加追踪时,这尤其有用。通过使用 mlflow.trace() 包装外部函数,你可以捕获其输入、输出和执行上下文。

import math

import mlflow

mlflow.set_experiment("External Function Tracing")


def invocation(x, y=4, exp=2):
    # Initiate a context handler for parent logging
    with mlflow.start_span(name="Parent") as span:
        span.set_attributes({"level": "parent", "override": y == 4})
        span.set_inputs({"x": x, "y": y, "exp": exp})

        # Wrap an external function instead of modifying
        traced_pow = mlflow.trace(math.pow)

        # Call the wrapped function as you would call it directly
        raised = traced_pow(x, exp)

        # Wrap another external function
        traced_factorial = mlflow.trace(math.factorial)

        factorial = traced_factorial(int(raised))

        # Wrap another and call it directly
        response = mlflow.trace(math.sqrt)(factorial)

        # Set the outputs to the parent span prior to returning
        span.set_outputs({"result": response})

        return response


for i in range(8):
    invocation(i)

下面的视频展示了我们在 MLflow UI 中运行的外部函数包装。请注意

外部函数跟踪

跟踪客户端API

MLflow 客户端 API 提供了一套全面的线程安全方法,用于手动管理跟踪。这些 API 允许对跟踪进行细粒度控制,使您能够以编程方式创建、操作和检索跟踪。本节将介绍如何使用这些 API 手动跟踪模型,提供逐步说明和示例。

开始追踪

与流利的API不同,MLflow Trace Client API要求您在添加子跨度之前显式启动一个跟踪。此初始API调用启动跟踪的根跨度,提供一个用于将后续跨度与根跨度关联的上下文request_id。

要开始一个新的跟踪,请使用 mlflow.client.MlflowClient.start_trace() 方法。此方法创建一个新的跟踪并返回根跨度对象。

from mlflow import MlflowClient

client = MlflowClient()

# Start a new trace
root_span = client.start_trace("my_trace")

# The request_id is used for creating additional spans that have a hierarchical association to this root span
request_id = root_span.request_id

添加子跨度

一旦开始跟踪,您可以使用 mlflow.client.MlflowClient.start_span() API 为其添加子跨度。子跨度允许您将跟踪分解为更小、更易管理的段,每个段代表整个过程中的特定操作或步骤。

# 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"},
)

结束一个跨度

在执行与一个跨度相关的操作后,您必须使用 mlflow.client.MlflowClient.end_span() 方法显式结束该跨度。请注意API签名中的两个必需字段:

  • request_id: 与根跨度关联的标识符

  • span_id:与正在结束的跨度相关联的标识符

为了有效地结束一个特定的跨度,在调用 end_span API 时,需要识别根跨度(从调用 start_trace 返回)和目标跨度(从调用 start_span 返回)。发起的 request_id 可以从任何父跨度对象的属性中访问。

备注

通过客户端API创建的跨度需要手动终止。确保所有使用 start_span API 启动的跨度都已使用 end_span API 结束。

# 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"},
)

结束跟踪

要完成跟踪,请使用 mlflow.client.MlflowClient.end_trace() 方法结束根跨度。这还将确保所有相关子跨度都正确结束。

# End the root span (trace)
client.end_trace(
    request_id=request_id,
    outputs={"final_output_key": "final_output_value"},
    attributes={"token_usage": "1174"},
)

搜索和检索追踪

搜索痕迹

你可以使用 mlflow.client.MlflowClient.search_traces() 方法根据各种标准搜索跟踪。此方法允许你按实验ID、过滤字符串和其他参数过滤跟踪。

# Search for traces in specific experiments
traces = client.search_traces(
    experiment_ids=["1", "2"],
    filter_string="attributes.status = 'OK'",
    max_results=5,
)

或者,您可以使用 fluent API mlflow.search_traces() 来搜索跟踪,这将返回一个包含每行一个跟踪的 pandas DataFrame。此方法允许您使用格式 "span_name.[inputs|outputs]""span_name.[inputs|outputs].field_name" 指定要从跟踪中提取的字段。提取的字段作为额外的列包含在 pandas DataFrame 中。此功能可用于构建评估数据集,以进一步提高模型和代理的性能。

import mlflow

with mlflow.start_span(name="span1") as span:
    span.set_inputs({"a": 1, "b": 2})
    span.set_outputs({"c": 3, "d": 4})

# Search for traces with specific fields extracted
traces = mlflow.search_traces(
    extract_fields=["span1.inputs", "span1.outputs.c"],
)

print(traces)

这输出:

    request_id                              ...     span1.inputs        span1.outputs.c
0   tr-97c4ef97c21f4348a5698f069c1320f1     ...     {'a': 1, 'b': 2}    3.0
1   tr-4dc3cd5567764499b5532e3af61b9f78     ...     {'a': 1, 'b': 2}    3.0

检索特定跟踪

要通过请求ID检索特定跟踪,请使用 mlflow.client.MlflowClient.get_trace() 方法。此方法返回与给定请求ID对应的跟踪对象。

# Retrieve a trace by request ID
trace = client.get_trace(request_id="12345678")

管理跟踪数据

删除痕迹

你可以使用 mlflow.client.MlflowClient.delete_traces() 方法根据特定标准删除跟踪。此方法允许你通过 实验ID最大时间戳请求ID 删除跟踪。

小技巧

删除跟踪是一个不可逆的过程。确保 delete_traces API 中的设置符合预期的删除范围。

import time

# Get the current timestamp in milliseconds
current_time = int(time.time() * 1000)

# Delete traces older than a specific timestamp
deleted_count = client.delete_traces(
    experiment_id="1", max_timestamp_millis=current_time, max_traces=10
)

设置和删除跟踪标签

可以向跟踪添加标签以提供额外的元数据。使用 mlflow.client.MlflowClient.set_trace_tag() 方法在跟踪上设置标签,并使用 mlflow.client.MlflowClient.delete_trace_tag() 方法从跟踪中删除标签。

# Set a tag on a trace
client.set_trace_tag(request_id="12345678", key="tag_key", value="tag_value")

# Delete a tag from a trace
client.delete_trace_tag(request_id="12345678", key="tag_key")

异步日志记录

默认情况下,MLflow 跟踪是同步记录的。这可能会在记录跟踪时引入性能开销,尤其是在您的 MLflow 跟踪服务器运行在远程服务器上时。如果性能开销对您来说是一个问题,您可以在 MLflow 2.16.0 及更高版本中启用 异步记录 进行跟踪。

要为跟踪启用异步日志记录,请在代码中调用 mlflow.config.enable_async_logging()。这将使跟踪日志记录操作变为非阻塞,并减少性能开销。

import mlflow

mlflow.config.enable_async_logging()

# Traces will be logged asynchronously
with mlflow.start_span(name="foo") as span:
    span.set_inputs({"a": 1})
    span.set_outputs({"b": 2})

# If you don't see the traces in the UI after waiting for a while, you can manually flush the traces
# mlflow.flush_trace_async_logging()

请注意,异步日志记录并不能完全消除性能开销。一些后端调用仍然需要同步进行,并且还有其他因素,如数据序列化。然而,异步日志记录可以显著减少日志记录跟踪的整体开销,对于典型工作负载,经验上大约减少了~80%。

使用 OpenTelemetry Collector 导出跟踪

MLflow 生成的追踪记录与 OpenTelemetry 追踪规范 兼容。因此,MLflow 追踪支持将追踪记录导出到 OpenTelemetry 收集器,然后可以用于将追踪记录导出到各种后端,如 Jaeger、Zipkin 和 AWS X-Ray。

默认情况下,MLflow 将跟踪导出到 MLflow 跟踪服务器。要启用将跟踪导出到 OpenTelemetry 收集器,请在启动任何跟踪之前,将 OTEL_EXPORTER_OTLP_ENDPOINT 环境变量(或 OTEL_EXPORTER_OTLP_TRACES_ENDPOINT)设置为 OpenTelemetry 收集器的目标 URL。

import mlflow
import os

# Set the endpoint of the OpenTelemetry Collector
os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = "http://localhost:4317/v1/traces"
# Optionally, set the service name to group traces
os.environ["OTEL_SERVICE_NAME"] = "<your-service-name>"

# Trace will be exported to the OTel collector at http://localhost:4317/v1/traces
with mlflow.start_span(name="foo") as span:
    span.set_inputs({"a": 1})
    span.set_outputs({"b": 2})

警告

MLflow 仅将跟踪导出到一个目的地。当配置了 OTEL_EXPORTER_OTLP_ENDPOINT 环境变量时,MLflow 不会 将跟踪导出到 MLflow 跟踪服务器,您将不会在 MLflow UI 中看到跟踪。

同样地,如果你将模型部署到启用了跟踪的 Databricks Model Serving,使用 OpenTelemetry Collector 将导致跟踪未记录在推理表中。

配置

MLflow 使用标准的 OTLP 导出器将跟踪数据导出到 OpenTelemetry Collector 实例。因此,您可以使用 OpenTelemetry 支持的所有配置 。以下示例配置 OTLP 导出器使用 HTTP 协议而不是默认的 gRPC,并设置自定义头:

export OTEL_EXPORTER_OTLP_TRACES_ENDPOINT="http://localhost:4317/v1/traces"
export OTEL_EXPORTER_OTLP_TRACES_PROTOCOL="http/protobuf"
export OTEL_EXPORTER_OTLP_TRACES_HEADERS="api_key=12345"

常见问题解答

Q: 我可以全局禁用和重新启用跟踪吗?

是的。

有两种流式API用于全局启用或禁用MLflow跟踪功能,以支持那些可能不希望在短时间内记录与启用跟踪的模型的交互,或者如果他们对长期存储与请求负载一起发送到交互模式模型的数据有顾虑的用户。

禁用 跟踪,mlflow.tracing.disable() API 将停止从 MLflow 内部收集跟踪数据,并且不会将任何跟踪数据记录到 MLflow Tracking 服务中。

启用 跟踪(如果它已被暂时禁用),mlflow.tracing.enable() API 将重新启用被调用模型的跟踪功能。

Q: 如何将一个追踪与MLflow运行关联?

如果在运行上下文中生成跟踪,记录的跟踪将与活动的实验关联到活动的运行。

例如,在以下代码中,跟踪是在 start_run 上下文中生成的。

import mlflow

# Create and activate an Experiment
mlflow.set_experiment("Run Associated Tracing")

# Start a new MLflow Run
with mlflow.start_run() as run:
    # Initiate a trace by starting a Span context from within the Run context
    with mlflow.start_span(name="Run Span") as parent_span:
        parent_span.set_inputs({"input": "a"})
        parent_span.set_outputs({"response": "b"})
        parent_span.set_attribute("a", "b")
        # Initiate a child span from within the parent Span's context
        with mlflow.start_span(name="Child Span") as child_span:
            child_span.set_inputs({"input": "b"})
            child_span.set_outputs({"response": "c"})
            child_span.set_attributes({"b": "c", "c": "d"})

当导航到 MLflow UI 并选择活动实验时,跟踪显示视图将显示与跟踪关联的运行,并提供一个链接以在 MLflow UI 中导航到该运行。请参见下面的视频以查看此操作的示例。

在运行上下文中的跟踪

您还可以通过使用 mlflow.client.MlflowClient.search_traces() 方法以编程方式检索与特定运行相关联的跟踪。

from mlflow import MlflowClient

client = MlflowClient()

# Retrieve traces associated with a specific Run
traces = client.search_traces(run_id=run.info.run_id)

print(traces)

Q: 我可以同时使用 fluent API 和客户端 API 吗?

你当然可以。然而,客户端API比流畅API更为冗长,并且是为更复杂的用例设计的,在这些用例中,你需要控制异步任务,而上下文管理器将无法处理对上下文的适当关闭。

虽然完全可能将两者混合,但通常不推荐这样做。

例如,以下内容将有效:

import mlflow

# Initiate a fluent span creation context
with mlflow.start_span(name="Testing!") as span:
    # Use the client API to start a child span
    child_span = client.start_span(
        name="Child Span From Client",
        request_id=span.request_id,
        parent_id=span.span_id,
        inputs={"request": "test input"},
        attributes={"attribute1": "value1"},
    )

    # End the child span
    client.end_span(
        request_id=span.request_id,
        span_id=child_span.span_id,
        outputs={"response": "test output"},
        attributes={"attribute2": "value2"},
    )
在 fluent 上下文中使用客户端 API

警告

使用 fluent API 来管理客户端发起的根跨度或子跨度的子跨度是不可能的。尝试在使用客户端 API 的同时启动 start_span 上下文处理程序将导致创建两个跟踪,一个用于 fluent API,另一个用于客户端 API。

Q: 如何为 span 添加自定义元数据?

有几种方法。

Fluent API

  1. mlflow.start_span() 构造函数本身中。

with mlflow.start_span(
    name="Parent", attributes={"attribute1": "value1", "attribute2": "value2"}
) as span:
    span.set_inputs({"input1": "value1", "input2": "value2"})
    span.set_outputs({"output1": "value1", "output2": "value2"})
  1. 使用从 start_span 返回的对象返回的 span 对象上的 set_attributeset_attributes 方法。

with mlflow.start_span(name="Parent") as span:
    # Set multiple attributes
    span.set_attributes({"attribute1": "value1", "attribute2": "value2"})
    # Set a single attribute
    span.set_attribute("attribute3", "value3")

客户端 API

  1. 在开始一个跨度时,你可以将属性作为 start_tracestart_span 方法调用的一部分传递。

parent_span = client.start_trace(
    name="Parent Span",
    attributes={"attribute1": "value1", "attribute2": "value2"}
)

child_span = client.start_span(
    name="Child Span",
    request_id=parent_span.request_id,
    parent_id=parent_span.span_id,
    attributes={"attribute1": "value1", "attribute2": "value2"}
)
  1. 直接在 Span 对象上使用 set_attributeset_attributes API。

parent_span = client.start_trace(
    name="Parent Span", attributes={"attribute1": "value1", "attribute2": "value2"}
)

# Set a single attribute
parent_span.set_attribute("attribute3", "value3")
# Set multiple attributes
parent_span.set_attributes({"attribute4": "value4", "attribute5": "value5"})
  1. 在结束一个跨度或整个跟踪时设置属性。

client.end_span(
    request_id=parent_span.request_id,
    span_id=child_span.span_id,
    attributes={"attribute1": "value1", "attribute2": "value2"},
)

client.end_trace(
    request_id=parent_span.request_id,
    attributes={"attribute3": "value3", "attribute4": "value4"},
)

Q: 如何查看捕获异常的Span的堆栈跟踪?

MLflow UI 在记录跟踪时如果发生故障,不会显示异常类型、消息或堆栈跟踪。然而,跟踪确实包含了这些关键的调试信息,作为构成跟踪的 Span 对象的一部分。

从经历了异常的跨度中检索特定堆栈跟踪信息的最简单方法是直接在交互式环境中(如 Jupyter Notebook)检索跟踪。

这是一个在收集跟踪时故意抛出异常的示例,以及查看异常详细信息的简单方法:

import mlflow

experiment = mlflow.set_experiment("Intentional Exception")

with mlflow.start_span(name="A Problematic Span") as span:
    span.set_inputs({"input": "Exception should log as event"})
    span.set_attribute("a", "b")
    raise Exception("Intentionally throwing!")
    span.set_outputs({"This": "should not be recorded"})

运行此代码时,如预期般会抛出一个异常。然而,跟踪信息仍会被记录到当前实验中,并且可以按如下方式检索:

from pprint import pprint

trace = mlflow.get_trace(span.request_id)
trace_data = trace.data
pprint(trace_data.to_dict(), indent=1)  # Minimum indent due to depth of Span object

在交互式环境中,例如 Jupyter Notebook,stdout 返回将呈现如下输出:

{'spans': [{'name': 'A Span',
    'context': {'span_id': '0x896ff177c0942903',
        'trace_id': '0xcae9cb08ec0a273f4c0aab36c484fe87'},
    'parent_id': None,
    'start_time': 1718063629190062000,
    'end_time': 1718063629190595000,
    'status_code': 'ERROR',
    'status_message': 'Exception: Intentionally throwing!',
    'attributes': {'mlflow.traceRequestId': '"7d418211df5945fa94e5e39b8009039e"',
        'mlflow.spanType': '"UNKNOWN"',
        'mlflow.spanInputs': '{"input": "Exception should log as event"}',
        'a': '"b"'},
    'events': [{'name': 'exception',
        'timestamp': 1718063629190527000,
        'attributes': {'exception.type': 'Exception',
        'exception.message': 'Intentionally throwing!',
        'exception.stacktrace': 'Traceback (most recent call last):\n
                                 File "/usr/local/lib/python3.8/site-packages/opentelemetry/trace/__init__.py",
                                 line 573, in use_span\n
                                    yield span\n  File "/usr/local/mlflow/mlflow/tracing/fluent.py",
                                 line 241, in start_span\n
                                    yield mlflow_span\n  File "/var/folders/cd/n8n0rm2x53l_s0xv_j_xklb00000gp/T/ipykernel_9875/4089093747.py",
                                 line 4, in <cell line: 1>\n
                                    raise Exception("Intentionally throwing!")\nException: Intentionally throwing!\n',
        'exception.escaped': 'False'}}]}],
 'request': '{"input": "Exception should log as event"}',
 'response': None
}

exception.stacktrace 属性包含在 span 执行期间引发的异常的完整堆栈跟踪。

或者,如果你使用 MLflowClient API 来搜索跟踪,从失败中检索跨度的事件数据的访问方式会有所不同(由于返回值是一个 pandas DataFrame)。使用 search_traces API 访问相同的异常数据如下所示:

import mlflow

client = mlflow.MlflowClient()

traces = client.search_traces(
    experiment_ids=[experiment.experiment_id]
)  # This returns a pandas DataFrame
pprint(traces["trace"][0].data.spans[0].to_dict(), indent=1)

从此调用渲染的 stdout 值与上面示例跨度数据中的值相同。