追踪 SDK 操作指南

在 Jupyter Notebook 中渲染追踪

注意

Jupyter 集成在 MLflow 2.20 及以上版本中可用

MLflow Trace UI in Jupyter Notebook

追踪 UI 也可在 Jupyter notebook 中使用！此功能需要使用 MLflow Tracking Server，因为 UI 资产是从此处获取的。要开始使用，只需确保 MLflow Tracking URI 已设置为您的追踪服务器（例如 mlflow.set_tracking_uri("https://:5000")）。

默认情况下，追踪 UI 将针对以下事件自动显示

当单元格代码生成追踪时（例如通过自动追踪，或通过运行手动追踪的函数时）
当调用 mlflow.search_traces() 时
当显示 mlflow.entities.Trace() 对象时（例如通过 IPython 的 display 函数，或当它是单元格中返回的最后一个值时）

要禁用显示，只需调用 mlflow.tracing.disable_notebook_display()，然后重新运行包含 UI 的单元格。要再次启用它，调用 mlflow.tracing.enable_notebook_display()。

要查看更完整的示例，请尝试运行此演示 notebook！

手动创建追踪和 Span

请参考手动追踪指南，了解如何手动创建追踪和 span。

设置追踪标签

可以向追踪添加标签，以在追踪级别提供附加元数据。例如，您可以将会话 ID 附加到追踪，以按对话会话对追踪进行分组。MLflow 提供 API 来设置和删除追踪上的标签。根据您是要在活动追踪上设置标签还是在已完成的追踪上设置标签，选择合适的 API。

API / 方法	用例
`mlflow.update_current_trace()` API。	在代码执行期间，为活动追踪设置标签。
`MlflowClient.set_trace_tag` API。	以编程方式为已完成的追踪设置标签。
MLflow UI	方便地为已完成的追踪设置标签。

在活动追踪上设置标签

如果您使用自动追踪或 fluent API 来创建追踪，并希望在追踪执行期间向其添加标签，您可以使用 mlflow.update_current_trace() 函数。

例如，以下代码示例为针对 my_func 函数创建的追踪添加了 "fruit": "apple" 标签

@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

注意

当键尚不存在时，mlflow.update_current_trace() 函数会将指定的标签添加到当前追踪。如果键已存在，它会使用新值更新该键。

在已完成的追踪上设置标签

要在后端存储中已完成并记录的追踪上设置标签，请使用 MlflowClient.set_trace_tag 方法设置追踪上的标签，并使用 MlflowClient.delete_trace_tag 方法从追踪中移除标签。

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

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

通过 MLflow UI 设置标签

或者，您可以通过 MLflow UI 更新或删除追踪上的标签。为此，导航到追踪选项卡，然后单击要更新的标签旁边的铅笔图标。

Traces tag update

删除追踪

您可以使用 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.tracing.disable() API 将停止从 MLflow 内部收集追踪数据，并且不会向 MLflow Tracking 服务记录任何关于追踪的数据。

要启用追踪（如果之前已暂时禁用），mlflow.tracing.enable() API 将重新启用对被调用的已插桩模型的追踪功能。

将追踪关联到 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 中导航到该运行。请参阅下面的视频，了解实际示例。

Tracing within a Run Context

您还可以使用 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)

异步记录追踪

默认情况下，MLflow 追踪是同步记录的。这可能会在记录追踪时引入性能开销，特别是当您的 MLflow Tracking Server 运行在远程服务器上时。如果性能开销是您关心的问题，您可以在 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%。

异步记录支持以下配置

环境变量	描述	默认值
`MLFLOW_ASYNC_TRACE_LOGGING_MAX_WORKERS`	用于异步追踪记录的最大工作线程数。	`10`
`MLFLOW_ASYNC_TRACE_LOGGING_MAX_QUEUE_SIZE`	可排队等待记录的最大追踪数。如果队列已满，追踪将被丢弃。	`1000`
`MLFLOW_ASYNC_TRACE_LOGGING_RETRY_TIMEOUT`	重试失败的追踪记录的超时时间（秒）。具体来说，失败的追踪将在此超时时间内进行带回退的重试，之后将被丢弃。	`60`

在 Jupyter Notebook 中渲染追踪​

手动创建追踪和 Span​

设置追踪标签​

在活动追踪上设置标签​

在已完成的追踪上设置标签​

通过 MLflow UI 设置标签​

删除追踪​

禁用追踪​

将追踪关联到 MLflow 运行​

异步记录追踪​