跳到主要内容

Jupyter Notebook 中的 MLflow 追踪

·3 分钟阅读
Daniel Lok
Databricks 软件工程师

Thumbnail

🚀 我们很高兴地宣布,MLflow Tracing 体验迎来了一次重大升级!

如果您不熟悉 MLflow Tracing,它是一款可观测性工具,允许您记录中间函数执行的输入和输出。它在调试生成式 AI(GenAI)应用程序时特别有用,并且 MLflow 提供了十多种与流行 GenAI 框架的集成,可以自动生成追踪信息,而无需您更改现有代码。

MLflow 2.20 版本开始,您现在可以直接在 Jupyter Notebook 中查看 MLflow Trace UI,从而无需切换出开发环境即可调试您的应用程序。上下文切换常常会打断工作流程,而此功能让您更容易保持专注,同时还能可视化生成的追踪数据。

MLflow Trace UI in Jupyter Notebook

JupyterLab 中的 UI 示例

开始使用

要开始使用,您需要使用一个 MLflow 跟踪服务器。在底层,MLflow 客户端需要发出网络请求来获取 UI 资源和追踪数据。

如果您不使用远程服务器,您随时可以通过运行 mlflow server CLI 命令在本地启动一个。默认情况下,服务器将在 https://:5000 启动。

在您的 notebook 中,只需确保 MLflow Tracking URI 设置为您的跟踪服务器,一切就准备就绪了!

import mlflow

# replace this with your own URI, if it differs
tracking_uri = "https://:5000"
mlflow.set_tracking_uri(tracking_uri)

# create a new experiment to avoid cluttering the default experiment
experiment = mlflow.set_experiment("mlflow-trace-ui-demo")

# the trace UI should now show up whenever traces are generated,
# for example:
@mlflow.trace
def add(a, b):
return a + b

# running the traced function triggers the UI display
add(1, 2)

当发生以下任何一种情况时,追踪 UI 就会显示出来

  1. 在单元格中生成了追踪(通过自动追踪,或在运行手动追踪的函数时)
  2. 当显式显示一个追踪对象时(例如,通过 IPython 的 display() 函数)
  3. 当调用 mlflow.search_traces() API 时

要亲身体验此功能,请尝试运行我们的演示 notebook!该 notebook 包含了上述所有三种场景的详细示例,以及一个简短的 LangChain RAG 演示,让您更真实地感受此功能在开发循环中的体验。

禁用和重新启用显示

此功能默认启用,但可以随时通过调用 mlflow.tracing.disable_notebook_display() 来关闭。要移除已经渲染的显示内容,您需要重新运行单元格(或简单地清除单元格输出)。

如果您想重新启用显示,可以调用 mlflow.tracing.enable_notebook_display()

错误报告和反馈

要报告错误或提供反馈,请在 MLflow GitHub 仓库中提交一个 issue。我们期待您的反馈!