MLflow Trace UI 在 Jupyter Notebook 中的演示

下载此 Notebook

本 Notebook 快速展示了 MLflow Trace UI 在 Jupyter Notebook 中的使用。

我们首先通过一些简单的示例来解释显示功能，最后构建一个简单的 RAG 演示来展示更多 UI 特性。

先决条件

请确保您已为此演示安装了以下软件包。

mlflow >= 2.20
openai

另外，对于最后的 RAG 演示，您还需要安装

langchain
langchain-community
beautifulsoup4

您可以运行下面的单元格来安装所有这些软件包（请确保之后重启内核）

%pip install mlflow>=2.20 openai langchain langchain-community beautifulsoup4

何时显示 MLflow Trace UI？

只有当 MLflow Tracking URI 设置为 HTTP 跟踪服务器时，UI 才会显示，因为 UI 资源是从那里提供的。如果您不使用远程跟踪服务器，可以通过运行 mlflow server CLI 命令在本地启动一个。默认情况下，跟踪服务器将在 http://localhost:5000 运行。

对于本教程，请确保您的跟踪 URI 设置正确！

import mlflow

# replace with your own URI
tracking_uri = "http://localhost:5000"
mlflow.set_tracking_uri(tracking_uri)

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

设置完成后，跟踪 UI 应在以下事件发生时自动显示。下面将提供每个事件的示例

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

示例 1：在单元格中生成跟踪

跟踪可以通过自动跟踪集成（例如使用 mlflow.openai.autolog()）生成，或者当您运行手动跟踪的函数时生成。例如

# Simple manual tracing example
import mlflow


@mlflow.trace
def foo(input):
  return input + 1


# running foo() generates a trace
foo(1)

# Automatic tracing with OpenAI
import os
from getpass import getpass

if "OPENAI_API_KEY" not in os.environ:
  os.environ["OPENAI_API_KEY"] = getpass("Enter your OpenAI API key: ")

from openai import OpenAI

import mlflow

mlflow.openai.autolog()

client = OpenAI()

# creating a chat completion will generate a trace
client.chat.completions.create(
  model="gpt-4o-mini",
  messages=[{"role": "user", "content": "hello!"}],
)

示例 2：显示 Trace 对象

当显示 MLflow Trace 实体时，跟踪 UI 也会显示。这可以通过两种方式实现

使用 IPython 的 display() 函数显式显示 trace 对象
当 trace 对象恰好是单元格中最后一个被评估的表达式时

# Explicitly calling `display()`
trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id)
display(trace)

# Even if the last expression does not result in a trace,
# display(trace) will still trigger the UI display
print("Test")

# Displaying as a result of the trace being the last expression
trace

示例 3：调用 `mlflow.search_traces()`

MLflow 提供了 mlflow.search_traces() API，以便方便地搜索实验中的所有跟踪。在 Jupyter Notebook 中调用此 API 时，跟踪 UI 将以分页视图渲染所有跟踪。单个单元格输出中可以渲染的跟踪数量是有限的。默认最大值为 10，但可以通过设置 MLFLOW_MAX_TRACES_TO_DISPLAY_IN_NOTEBOOK 环境变量进行配置。

mlflow.search_traces(experiment_ids=[experiment.experiment_id])

禁用 UI

默认情况下显示是启用的，但如果您希望不显示，可以运行 mlflow.tracing.disable_notebook_display() 来禁用它。您需要重新运行单元格（或简单地清除单元格输出），以移除已渲染的显示内容。

如果您想重新启用自动显示功能，只需调用 mlflow.tracing.enable_notebook_display()。

mlflow.tracing.disable_notebook_display()

# no UI will be rendered
trace

mlflow.tracing.enable_notebook_display()

# re-enable display
trace

结论

这就是基础知识！希望您觉得 Jupyter 集成很有用。一如既往，如果您发现任何问题或想留下任何反馈，请在 https://github.com/mlflow/mlflow/issues 提交 issue。

在接下来的几个单元格中，我们将提供一个简短的 RAG 演示，它将创建包含更真实数据的跟踪，以便您更好地了解使用此 UI 的感受。

from langchain_core.vectorstores import InMemoryVectorStore
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

# define necessary RAG entities
llm = ChatOpenAI(model="gpt-4o-mini")
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
vector_store = InMemoryVectorStore(embeddings)

import bs4
from langchain import hub
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.runnables import RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter

# generate sample doc chunks from the MLflow documentation
loader = WebBaseLoader(
  web_paths=("https://mlflow.org.cn/docs/latest/llms/tracing/index.html",),
  bs_kwargs={"parse_only": bs4.SoupStrainer(class_=("document"))},
)
docs = loader.load()

# add documents to the vector store
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
all_splits = text_splitter.split_documents(docs)
vector_store.add_documents(documents=all_splits)

retriever = vector_store.as_retriever(
  search_type="similarity",
  search_kwargs={"k": 3},
)

# Define prompt for question-answering
prompt = hub.pull("rlm/rag-prompt")

# define our chain
chain = {"context": retriever, "question": RunnablePassthrough()} | prompt | llm

import mlflow

# call the langchain autolog function so that traces will be generated
mlflow.langchain.autolog()

response = chain.invoke("What is MLflow Tracing?")

先决条件​

何时显示 MLflow Trace UI？​

示例 1：在单元格中生成跟踪​

示例 2：显示 Trace 对象​

示例 3：调用 mlflow.search_traces()​

禁用 UI​

结论​