通过

使用上下文管理器进行范围跟踪

上下文 mlflow.start_span 管理器允许为任意代码块创建范围。 当仅捕获单个函数的执行边界无法提供足够细节时，此功能便可用于更精细地记录代码内部的复杂交互。

使用上下文管理器进行跨度跟踪可让你对跟踪的代码进行精细控制：

  • 任意代码块：跟踪任何代码块，而不仅仅是整个函数
  • 灵活的边界：定义跨度的确切起点和终点
  • 自动上下文管理：MLflow 处理父子关系和清理
  • 适用于函数修饰器：与 @mlflow.trace 灵活组合以实现混合方法
  • 异常处理：自动错误捕获（如修饰器）

先决条件

MLflow 3

此页面需要以下包：

  • mlflow[databricks] 3.1 及更高版本：具有 GenAI 功能和 Databricks 连接的核心 MLflow 功能。
  • openai 1.0.0 及更高版本：（可选）仅当自定义代码与 OpenAI 交互时;如果需要，请替换为其他 SDK。

安装基本组件：

%pip install --upgrade "mlflow[databricks]>=3.1"
# %pip install --upgrade openai>=1.0.0 # Install if needed

MLflow 2.x

本指南需要以下包：

  • mlflow[databricks] 2.15.0 及更高版本：具有 Databricks 连接的核心 MLflow 功能。
  • openai 1.0.0 及更高版本：（可选）仅当自定义代码与 OpenAI 交互时。

注释

Databricks 强烈建议安装 MLflow 3.1 或更高版本（如果使用 mlflow[databricks]）。

安装基本组件：

%pip install --upgrade "mlflow[databricks]>=2.15.0,<3.0.0"
# pip install --upgrade openai>=1.0.0 # Install if needed

上下文管理器 API

与修饰器类似，上下文管理器会自动捕获父子关系、异常、执行时间，并支持自动追踪。 但是，必须手动提供范围的名称、输入和输出。 您可以使用从上下文管理器返回的 mlflow.entities.Span 对象来配置它们。

with mlflow.start_span(name="my_span") as span:
    span.set_inputs({"x": 1, "y": 2})
    z = x + y
    span.set_outputs(z)

下面是一个稍微复杂一些的示例，该示例将 mlflow.start_span 上下文管理器与 OpenAI 的修饰器和自动跟踪结合使用。

import mlflow
import openai
from mlflow.entities import SpanType

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Create OpenAI client
client = openai.OpenAI()

@mlflow.trace(span_type=SpanType.CHAIN)
def start_session():
    messages = [{"role": "system", "content": "You are a friendly chat bot"}]
    while True:
        with mlflow.start_span(name="User") as span:
            span.set_inputs(messages)
            user_input = input(">> ")
            span.set_outputs(user_input)

        if user_input == "BYE":
            break

        messages.append({"role": "user", "content": user_input})

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            max_tokens=100,
            messages=messages,
        )
        answer = response.choices[0].message.content
        print(f"Assistant: {answer}")

        messages.append({"role": "assistant", "content": answer})

start_session()

后续步骤

  • Last updated on