分布式跟踪

当您的应用程序跨越多个服务时，您可能希望将这些服务的跨度（Span）连接成一个单一的追踪（Trace），以便在一个地方追踪端到端的执行情况。MLflow 通过分布式追踪（Distributed Tracing）支持此功能，它通过 HTTP 传播活跃的追踪上下文，从而将记录在不同服务中的跨度拼接在一起。

工作原理

MLflow 兼容 OpenTelemetry，因此追踪上下文会按照 W3C TraceContext 规范通过 HTTP 标头进行传播。MLflow 提供了两个 API 来简化客户端和服务端标头的处理。

• 在客户端（调用方服务）使用 get_tracing_context_headers_for_http_request 来获取包含当前追踪上下文的标头。
• 在服务端（被调用方服务）使用 set_tracing_context_from_http_request_headers 从传入的请求标头中提取追踪和跨度信息。

前提条件

为了使分布式追踪正常工作，两个服务必须将追踪记录到同一个 MLflow 追踪服务器和同一个实验（Experiment）中。

1. 启动 MLflow 追踪服务器（或使用共享的远程服务器）。
2. 在每个服务中，在创建跨度之前设置追踪 URI 和实验。

bash

export MLFLOW_TRACKING_URI="https://:5000"
export MLFLOW_EXPERIMENT_NAME="distributed-tracing-demo"

或

python

import mlflow

mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("distributed-tracing-demo")

如果服务使用不同的实验，它们的跨度最终将位于不同的追踪中，而不是合并为一个。

使用示例

以下是一个使用 FastAPI 的可运行的双服务示例。客户端创建根跨度并调用服务端，服务端在同一个追踪下添加一个嵌套的跨度。

服务端（被调用方服务）

创建一个包含以下内容的 server.py 文件

server.py

python

import mlflow
import uvicorn
from fastapi import FastAPI, Request
from mlflow.tracing import set_tracing_context_from_http_request_headers

# Set the tracking URI and experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("distributed-tracing-demo")

app = FastAPI()


@app.post("/handle")
async def handle(request: Request):
    headers = dict(request.headers)
    body = await request.json()
    with set_tracing_context_from_http_request_headers(headers):
        # This span will be added to the same trace as a child span
        with mlflow.start_span("server-handler") as span:
            span.set_inputs(body)
            result = {"response": f"Processed: {body.get('input', '')}"}
            span.set_outputs(result)
    return result


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=5001)

客户端（调用方服务）

创建一个包含以下内容的 client.py 文件

client.py

python

import requests
import mlflow
from mlflow.tracing import get_tracing_context_headers_for_http_request

# Connect to the same tracking server and experiment as the server
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("distributed-tracing-demo")

# Create a root span at client side
with mlflow.start_span("client-root") as span:
    span.set_inputs({"input": "hello"})
    headers = get_tracing_context_headers_for_http_request()
    response = requests.post(
        "https://:5001/handle", headers=headers, json={"input": "hello"}
    )
    span.set_outputs(response.json())

运行它

bash

# Terminal 1 — start the server
python server.py

# Terminal 2 — run the client
python client.py

在 https://:5000 打开 MLflow UI 并导航到 distributed-tracing-demo 实验。您将看到一个单一的追踪，其中包含来自客户端和服务端的跨度。

AI 网关集成

当代理（Agent）在启用了使用情况追踪的情况下调用 MLflow AI 网关端点时，网关会自动为每个请求创建一个追踪。如果代理还发送了 traceparent 标头，网关将在代理的追踪下创建一个轻量级跨度，该跨度链接到完整的网关追踪并包含令牌使用情况。

关于 AI 网关集成需要注意的一点是，网关追踪和代理追踪存储在不同的实验中，但它们通过跨度属性链接在一起。网关追踪是一个完整的追踪，包含请求/响应负载和令牌使用情况，而代理追踪包含一个轻量级跨度 gateway/<endpoint_name>，其中包含指向网关追踪的链接以及跨度属性中的令牌使用情况。这种设计允许在跨实验聚合网关使用情况的同时，避免存储重复的负载。

• 网关实验 — 包含请求/响应负载和令牌使用情况的完整追踪
• 代理实验 — 包含子跨度 gateway/<endpoint_name> 的代理追踪，该跨度包含指向网关追踪的链接和令牌使用情况（无重复负载）

要启用此功能，请在活跃跨度内调用网关时传递 traceparent 标头

python

import mlflow
import requests
from mlflow.tracing import get_tracing_context_headers_for_http_request
from openai import OpenAI

client = OpenAI(
    base_url="https://:5000/gateway/openai/v1",
    api_key="dummy",
)

with mlflow.start_span("my-agent"):
    headers = get_tracing_context_headers_for_http_request()
    response = client.chat.completions.create(
        model="my-endpoint",
        messages=[{"role": "user", "content": "Hello!"}],
        extra_headers=headers,
    )

注意

网关端点必须启用使用情况追踪，才能创建网关追踪和分布式跨度。

Databricks 中的限制

如果您将 MLflow 追踪设置到 Databricks，要使分布式追踪生效，追踪目的地必须设置为 Unity Catalog。详情请参考 将 MLflow 追踪存储在 Unity Catalog 中。