mlflow.tracing

注意

The mlflow.tracing 命名空间仅包含几个用于管理追踪的实用函数。MLflow Tracing 的主入口点是直接在 mlflow 命名空间下定义的 Tracing Fluent APIs，或者底层的 Tracing Client APIs。

mlflow.tracing.configure(span_processors: list[typing.Callable[[ForwardRef('LiveSpan')], NoneType]] | None = None) → mlflow.tracing.config.TracingConfigContext[source]

注意

实验性功能：此函数可能会在未来版本中更改或移除，恕不另行通知。

配置 MLflow 追踪。可用作函数或上下文管理器。

仅更新显式提供的参数，其余参数保持不变。

参数

span_processors – 在导出之前处理 span 的函数列表。这对于过滤/屏蔽 span 中的特定属性以防止敏感数据被记录，或减小 span 的大小非常有用。每个函数必须接受一个 LiveSpan 类型的参数，并且不应返回任何值。当提供多个函数时，它们将按提供的顺序依次应用。

返回

用于临时配置更改的上下文管理器。: 当用作函数时，配置更改将持续存在。当用作上下文管理器时，更改将在退出时恢复。

返回类型

TracingConfigContext

示例

def pii_filter(span):
    """Example PII filter that masks sensitive data in span attributes."""
    # Mask sensitive inputs
    if inputs := span.inputs:
        for key, value in inputs.items():
            if "password" in key.lower() or "token" in key.lower():
                span.set_inputs({**inputs, key: "[REDACTED]"})

    # Mask sensitive outputs
    if outputs := span.outputs:
        if isinstance(outputs, dict):
            for key in outputs:
                if "secret" in key.lower():
                    outputs[key] = "[REDACTED]"
            span.set_outputs(outputs)

    # Mask sensitive attributes
    for attr_key in list(span.attributes.keys()):
        if "api_key" in attr_key.lower():
            span.set_attribute(attr_key, "[REDACTED]")

# Permanent configuration change
mlflow.tracing.configure(span_processors=[pii_filter])

# Temporary configuration change
with mlflow.tracing.configure(span_processors=[pii_filter]):
    # PII filtering enabled only in this block
    pass

mlflow.tracing.disable()[source]

禁用追踪。

注意

此函数将 OpenTelemetry 设置为使用 NoOpTracerProvider，从而有效地禁用所有追踪操作。

示例

import mlflow


@mlflow.trace
def f():
    return 0


# Tracing is enabled by default
f()
assert len(mlflow.search_traces()) == 1

# Disable tracing
mlflow.tracing.disable()
f()
assert len(mlflow.search_traces()) == 1

mlflow.tracing.disable_notebook_display()[source]: 禁用在 notebook 输出单元格中显示 MLflow Trace UI。请调用 mlflow.tracing.enable_notebook_display() 以重新启用显示。

mlflow.tracing.enable()[source]

启用追踪。

示例

import mlflow


@mlflow.trace
def f():
    return 0


# Tracing is enabled by default
f()
assert len(mlflow.search_traces()) == 1

# Disable tracing
mlflow.tracing.disable()
f()
assert len(mlflow.search_traces()) == 1

# Re-enable tracing
mlflow.tracing.enable()
f()
assert len(mlflow.search_traces()) == 2

mlflow.tracing.enable_notebook_display()[source]

在 notebook 输出单元格中启用 MLflow Trace UI。默认情况下显示是启用的，并且 Trace UI 会在执行以下任一操作时出现：

追踪完成时（即每次追踪导出时）
调用 mlflow.search_traces() 链式 API 时
调用 mlflow.client.MlflowClient.get_trace() 或 mlflow.client.MlflowClient.search_traces() 客户端 API 时

要禁用，请调用 mlflow.tracing.disable_notebook_display()。

mlflow.tracing.reset()[source]: 重置指示 MLflow Tracing Provider 是否已初始化的标志。这确保了在下一次执行追踪操作时 Tracing Provider 会被重新初始化。

mlflow.tracing.set_databricks_monitoring_sql_warehouse_id(sql_warehouse_id: str, experiment_id: Optional[str] = None) → None[source]

注意

实验性功能：此函数可能会在未来版本中更改或移除，恕不另行通知。

为记录到给定 MLflow 实验的追踪设置用于 Databricks 生产监控的 SQL warehouse ID。仅当实验的追踪位置使用 UC schema 时才有效。

参数

sql_warehouse_id – 用于监控的 SQL warehouse ID。
experiment_id – MLflow 实验 ID。如果未提供，将使用当前活动实验。

mlflow.tracing.set_destination(destination: mlflow.entities.trace_location.TraceLocationBase, *, context_local: bool = False)[source]

设置一个自定义 span 位置，MLflow 将追踪导出到该位置。

通过此函数指定的 destination 将优先于其他配置，例如 tracking URI、OTLP 环境变量。

参数

destination –
一个 trace location 对象，指定 trace 数据的存储位置。当前支持以下位置：
- MlflowExperimentLocation：将追踪记录到
  一个 MLflow 实验。
- UCSchemaLocation：将追踪记录到
  Databricks Unity Catalog schema。仅在 Databricks 中可用。
context_local – 如果为 False（默认值），则 destination 是全局设置的。如果为 True，则 destination 会按异步任务或线程隔离，为并发应用程序提供隔离。

示例

将追踪记录到 MLflow 实验

from mlflow.entities.trace_location import MlflowExperimentLocation

mlflow.tracing.set_destination(MlflowExperimentLocation(experiment_id="123"))

注意：这与通过 MLFLOW_EXPERIMENT_ID 环境变量或 mlflow.set_experiment API 设置活动 MLflow 实验效果相同，但范围更窄。

将追踪记录到 Databricks Unity Catalog

from mlflow.entities.trace_location import UCSchemaLocation

mlflow.tracing.set_destination(
    UCSchemaLocation(catalog_name="catalog", schema_name="schema")
)

隔离异步任务或线程之间的 destination

from mlflow.tracing.destination import Databricks

mlflow.tracing.set_destination(
    MlflowExperimentLocation(experiment_id="123"),
    context_local=True,
)

使用 context_local 标志设置的 destination 将仅在当前异步任务或线程中生效。这在您希望从多租户应用程序将追踪发送到不同 destination 时特别有用。

**重置 destination：**

mlflow.tracing.reset()

mlflow.tracing.set_experiment_trace_location(location: UCSchemaLocation, experiment_id: Optional[str] = None, sql_warehouse_id: Optional[str] = None) → UCSchemaLocation[source]

注意

实验性功能：此函数可能会在未来版本中更改或移除，恕不另行通知。

配置实验追踪的存储位置。

用于存储追踪数据的 Unity Catalog 表将在指定的 schema 中创建。启用追踪后，指定实验的所有追踪都将存储在提供的 Unity Catalog schema 中。

注意

如果实验已链接到存储位置，这将引发错误。请先使用 mlflow.tracing.unset_experiment_trace_location 移除现有存储位置，然后再设置新的。

参数

location – 实验追踪在 Unity Catalog 中的存储位置。
experiment_id – 要为其设置存储位置的 MLflow 实验 ID。如果未指定，将使用当前活动实验。
sql_warehouse_id – 用于创建视图和查询的 SQL warehouse ID。如果未指定，则使用 MLFLOW_TRACING_SQL_WAREHOUSE_ID 的值，如果环境变量未设置，则回退到默认 SQL warehouse。

返回

表示配置的存储位置的 UCSchemaLocation 对象，包括 spans 和 logs 表的表名。

示例

import mlflow
from mlflow.entities import UCSchemaLocation

location = UCSchemaLocation(catalog_name="my_catalog", schema_name="my_schema")

result = mlflow.tracing.set_experiment_trace_location(
    location=location,
    experiment_id="12345",
)
print(result.full_otel_spans_table_name)  # my_catalog.my_schema.otel_spans_table


@mlflow.trace
def add(x):
    return x + 1


add(1)  # this writes the trace to the storage location set above

mlflow.tracing.set_span_chat_tools(span: LiveSpan, tools: list[ChatTool])[source]

在指定的 span 上设置 mlflow.chat.tools 属性。此属性在 UI 中使用，也被下游应用程序（如 MLflow evaluate）使用。

参数

span – 要添加属性的 LiveSpan。
tools – 标准化聊天工具定义列表（有关详细信息，请参阅规范）。

示例

import mlflow
from mlflow.tracing import set_span_chat_tools

tools = [
    {
        "type": "function",
        "function": {
            "name": "add",
            "description": "Add two numbers",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "number"},
                    "b": {"type": "number"},
                },
                "required": ["a", "b"],
            },
        },
    }
]


@mlflow.trace
def f():
    span = mlflow.get_current_active_span()
    set_span_chat_tools(span, tools)
    return 0


f()

mlflow.tracing.unset_experiment_trace_location(location: UCSchemaLocation, experiment_id: Optional[str] = None) → None[source]

注意

实验性功能：此函数可能会在未来版本中更改或移除，恕不另行通知。

取消设置实验追踪的存储位置。

此函数将移除实验存储位置的配置，包括视图和实验标签。

参数

location – 要取消设置的存储位置。
experiment_id – 要为其取消设置存储位置的 MLflow 实验 ID。如果未提供，将使用当前活动实验。

示例

import mlflow
from mlflow.entities import UCSchemaLocation

mlflow.tracing.unset_experiment_trace_location(
    location=UCSchemaLocation(catalog_name="my_catalog", schema_name="my_schema"),
    experiment_id="12345",
)