本快速入门帮助你将 GenAI 应用与 MLflow 跟踪 集成(如果使用本地 IDE 作为开发环境)。 如果使用 Databricks Notebook,请改用 Databricks Notebook 快速入门 。
你将实现的目标
在本教程结束时,你将拥有:
- GenAI 应用的 MLflow 实验
- 连接到 MLflow 的本地开发环境
- 使用 MLflow 跟踪检测的简单 GenAI 应用程序
- MLflow 实验中该应用的追踪记录
先决条件
- Databricks 工作区:访问 Databricks 工作区。
步骤 1:安装 MLflow
在本地 IDE 中工作时,需要使用 Databricks 连接安装 MLflow。
pip install --upgrade "mlflow[databricks]>=3.1" openai
步骤 2:创建新的 MLflow 试验
MLflow 试验是 GenAI 应用程序的容器。 详细了解试验及其包含在 数据模型 部分中的内容。
打开 Databricks 工作区
转到机器学习下的左侧边栏中的试验
在“试验”页顶部,单击“新建 GenAI 试验”
步骤 3:将环境连接到 MLflow
注释
本快速入门介绍如何使用 Databricks 个人访问令牌。 MLflow 还适用于其他 Databricks 支持的身份验证方法。
使用环境变量
单击“生成 API 密钥”
在终端中复制并运行生成的代码。
export DATABRICKS_TOKEN=<databricks-personal-access-token> export DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com export MLFLOW_TRACKING_URI=databricks export MLFLOW_EXPERIMENT_ID=<experiment-id>
使用 .env 文件
单击“生成 API 密钥”
将生成的代码复制到
.env
项目根目录中的文件DATABRICKS_TOKEN=<databricks-personal-access-token> DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com MLFLOW_TRACKING_URI=databricks MLFLOW_EXPERIMENT_ID=<experiment-id>
安装
python-dotenv
包pip install python-dotenv
在代码中加载环境变量
# At the beginning of your Python script from dotenv import load_dotenv # Load environment variables from .env file load_dotenv()
步骤 4:创建和检测应用程序
以下示例快速入门使用 OpenAI SDK 连接到 Databricks 托管 LLM。 如果要使用自己的 OpenAI 密钥,请更新行 client = OpenAI(...)
。
如果希望使用 MLflow 支持的其他 20+ 个 LLM SDK(Anthropic、Bedrock 等)或 GenAI 创作框架(LangGraph 等),请按照上一步中的 MLflow 实验 UI 中的说明进行操作。
创建在项目目录中命名
app.py
的 Python 文件。初始化 OpenAI 客户端以连接到 OpenAI 托管的 LLM。
OpenAI 托管的 LLM
使用本地 OpenAI SDK 连接到由 OpenAI 托管的模型。 从 可用的 OpenAI 模型中选择一个模型。
import mlflow import os import openai # Ensure your OPENAI_API_KEY is set in your environment # os.environ["OPENAI_API_KEY"] = "<YOUR_API_KEY>" # Uncomment and set if not globally configured # Enable auto-tracing for OpenAI mlflow.openai.autolog() # Set up MLflow tracking to Databricks mlflow.set_tracking_uri("databricks") mlflow.set_experiment("/Shared/docs-demo") # Create an OpenAI client connected to OpenAI SDKs client = openai.OpenAI() # Select an LLM model_name = "gpt-4o-mini"
定义并运行应用程序:
@mlflow.trace
使用修饰器可以轻松跟踪与 OpenAI 自动检测相结合的任何 Python 应用程序,以捕获对 OpenAI SDK 调用的详细信息。# Use the trace decorator to capture the application's entry point @mlflow.trace def my_app(input: str): # This call is automatically instrumented by `mlflow.openai.autolog()` response = client.chat.completions.create( model=model_name, # This example uses a Databricks hosted LLM - you can replace this with any AI Gateway or Model Serving endpoint. If you provide your own OpenAI credentials, replace with a valid OpenAI model e.g., gpt-4o, etc. messages=[ { "role": "system", "content": "You are a helpful assistant.", }, { "role": "user", "content": input, }, ], ) return response.choices[0].message.content result = my_app(input="What is MLflow?") print(result)
运行应用程序
python app.py
步骤 5:在 MLflow 中查看追踪
返回到 MLflow 实验 UI
现在,你将在 “跟踪 ”选项卡中看到生成的跟踪
单击跟踪记录以查看其详细信息
了解追踪
您刚刚创建的跟踪显示:
-
根范围:表示函数
my_app(...)
的输入- 子范围:表示 OpenAI 完成请求
- 属性:包含模型名称、令牌计数和计时信息等元数据
- 输入:发送到模型的消息
- 输出:从模型接收的响应
此简单跟踪已提供对应用程序行为的宝贵见解,例如:
- 有人问什么
- 生成了什么响应?
- 请求花费的时间
- 使用了多少个令牌(影响成本)
小窍门
对于更复杂的应用程序(如 RAG 系统或多步骤代理),MLflow 跟踪通过揭示每个组件和步骤的内部工作来提供更多价值。
后续步骤
继续您的旅程,并参考这些推荐的行动和教程。
参考指南
浏览本指南中提到的概念和功能的详细文档。