MLflow Agent 接入 Databricks 工作区：Tracing 追踪与 Unity Catalog 存储配置指南

MLflow Agent 接入 Databricks 工作区：Tracing 追踪与 Unity Catalog 存储配置指南

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow

本指南基于 MLflow 仓库中mlflow agent setup的 Databricks 场景提示词模板（mlflow/agent/setup/templates/databricks.md），完整讲解如何在编码 Agent 自动接入 Databricks 工作区时完成认证校验、Experiment 绑定与 Unity Catalog（UC）Trace 存储配置。读完本文，你将掌握 Databricks SDK 认证验证的正确姿势、mlflow.set_experiment(experiment_id=...)与trace_location=UnityCatalog(...)的实战用法，以及 Trace 数据落盘 UC Delta 表所需的全部前提条件，并理解这些指令在mlflow agent setup提示词组装链路中的真实位置。

背景：这份文档出现在 Agent 工作流的哪个环节

本文所解析的databricks.md并非独立使用的手册，而是 mlflow/agent/setup/prompt.py 中按条件动态注入的提示词片段：当用户在mlflow agent setup交互流程中选择MLFLOW_TRACKING_URI=databricks（或databricks://<profile>）作为追踪后端时，CLI 会读取该模板、渲染占位符，并把渲染结果嵌入到最终交给编码 Agent 的首条用户消息中（{{ server_setup }}插槽，最终进入 python.md 的步骤 2 之前）。

整条调用链如下（均可在仓库中验证）：

1. mlflow/agent/setup/cli.py：_run_setup()检测到 tracking URI 指向 Databricks 后，调用_prompt_experiment_id()要求用户提供实验 ID 或工作区路径（以/开头时自动经get_experiment_by_name查找、不存在则create_experiment创建，见 cli.py）；
2. mlflow/agent/setup/prompt.py：将tracking_uri、experiment_id、workspace_client_args三个值渲染进databricks.md模板；
3. 渲染结果作为{{ server_setup }}拼入 python.md 的“步骤 2：配置 tracking URI”之前，随后 Agent 按 instrument.md 的硬性规则（不重复安装、不覆盖已有配置、按清单逐步执行）执行。

其中模板占位符的解析逻辑可见 prompt.py 的_PLACEHOLDER正则与_render()函数：所有{{ key }}必须在传入值字典中命中，否则抛出KeyError。这保证了注入 Agent 的提示词中{{ tracking_uri }}、{{ experiment_id }}、{{ workspace_client_args }}永远被真实值替换。

前置确认：追踪目标与实验 ID 的来历

模板开篇即申明：本次任务的全部 Trace 将发送到MLFLOW_TRACKING_URI={{ tracking_uri }}（Databricks 场景下其值为databricks或databricks://<profile>）。该值不是 Agent 自己猜测的，而是由 CLI 在启动时确定并写入提示词的：

• 若环境变量MLFLOW_TRACKING_URI已存在且指向 Databricks，CLI 直接采用（cli.py）；
• 否则在交互式后端选择菜单中，用户选择 "Connect to a Databricks workspace" 后，CLI 会提示输入 Databricks 配置 profile（留空则用默认 profile），并构造出databricks://<profile>或databricks（cli.py）。

{{ experiment_id }}同理：CLI 要求用户输入实验 ID，或输入以/开头的实验路径（如/Users/<username>/my-experiment），后者若不存在会自动创建并回显新 ID（cli.py）。模板明确要求 Agent 不要再重复配置 tracking URI（“tracking URI itself is wired in step 2 below”），因为该步骤由 python.md 的步骤 2 统一处理——可见模板之间有着严格的分工。

第一步：验证 Databricks 认证（不硬编码任何环境变量）

为什么要先验证认证

在插桩任何代码之前，必须确认当前环境能向 Databricks 工作区发出经过认证的请求。这一步做在插桩之前，是为了把“认证配置缺失”与“插桩代码 bug”两类问题隔离开，避免 Agent 在错误配置上反复调试。

推荐的验证方式：调用current_user.me()

模板给出的验证代码是 Databricks SDK 的标准“最小认证探测”：

from databricks.sdk import WorkspaceClient WorkspaceClient({{workspace_client_args}}).current_user.me()

这里{{ workspace_client_args }}由 CLI 渲染：选择的是databricks://<profile>时渲染为profile="<profile>"，默认 profile 时为空字符串（prompt.py）。因此实际生成的调用形如：

# 默认 profile from databricks.sdk import WorkspaceClient WorkspaceClient().current_user.me() # 指定 profile 为 prod from databricks.sdk import WorkspaceClient WorkspaceClient(profile="prod").current_user.me()

模板特别强调：不要硬性要求某个特定环境变量。Databricks SDK 的凭据解析链覆盖环境变量、~/.databrickscfg配置文件的 profile、OAuth 以及其他来源，因此只要 SDK 能成功认证即可，具体用哪种方式由用户环境决定。

认证失败时的处理流程

如果上述调用抛异常，模板要求 Agent 停下来，请求用户配置认证，并给出了三种典型途径：

• 运行databricks auth login完成交互式登录；
• 在~/.databrickscfg中配置 profile；
• 导出环境变量DATABRICKS_HOST与DATABRICKS_TOKEN。

同时有一条红线规则：绝不把密钥写进仓库文件。这与 instrument.md 中“Do not create setup-only files in the repo”的硬性规则一致，也是 Agent 提示词设计中对供应链安全的基本约束。

第二步：按 ID 绑定当前 Experiment

认证通过后，需要把当前活动实验固定下来。模板给出的做法是按实验 ID 绑定，而不是按名称：

import mlflow mlflow.set_experiment(experiment_id="{{ experiment_id }}")

set_experiment是 MLflow 的 fluent API，其签名与行为在 mlflow/tracking/fluent.py 中有完整定义：experiment_name与experiment_id二选一传入；按 ID 激活时，若 ID 不存在会抛出异常；而按名称激活时若实验不存在则会自动创建（在 Databricks 上实验名称必须是绝对路径，如/Users/<username>/my-experiment）。这也解释了为什么 CLI 选择优先解析并传入 ID：ID 是工作区内最稳定、无歧义的实验标识。

import mlflow mlflow.set_experiment(experiment_id="1234567890123456")

这一步之后，后续所有mlflow.autolog()产生的 Trace 都会落到该实验下，最终在 MLflow UI 中可通过“tracking URI + experiment ID”构造出可打开的 Trace 链接（该 URL 会在 instrument.md 的步骤 5 由 Agent 汇总报告给用户）。

可选进阶：将 Trace 存储到 Unity Catalog Delta 表

如果用户希望 Trace 由 UC Delta 表承载（而非默认的跟踪服务器存储），模板提供了可选代码块。

前置条件（缺一不可）

• MLflow 版本 ≥ 3.11：UnityCatalogtrace location 类型自 3.11.0 起标记为实验性特性（见 mlflow/entities/trace_location.py 的@experimental(version="3.11.0")装饰器）；
• 一个可用的 SQL Warehouse：UC 表写入需要 SQL 仓库执行 DDL/DML；
• 用户明确要求：模板强调“Skip this block entirely if the user does not ask for UC-backed traces”——只有在用户主动提出 UC 存储时才启用，否则跳过。

需要向用户确认的四个参数

配置代码

import mlflow from mlflow.entities.trace_location import UnityCatalog mlflow.set_experiment( experiment_id="{{ experiment_id }}", trace_location=UnityCatalog( catalog_name="<catalog>", schema_name="<schema>", table_prefix="<prefix>", ), )

即实际形态如：

import mlflow from mlflow.entities.trace_location import UnityCatalog mlflow.set_experiment( experiment_id="1234567890123456", trace_location=UnityCatalog( catalog_name="main", schema_name="mlflow_traces", table_prefix="my_app", ), )

UnityCatalog类在源码中的行为

trace_location参数正是set_experiment的第三个形参（fluent.py），类型限定为UnityCatalog实例。该 dataclass 定义于 mlflow/entities/trace_location.py，核心字段为catalog_name、schema_name、table_prefix（后者可空），并提供以下关键派生属性：

• schema_location：拼接为f"{catalog_name}.{schema_name}"；
• full_table_prefix：拼接为f"{catalog_name}.{schema_name}.{table_prefix}"，当table_prefix未设置时抛出MlflowException.invalid_parameter_value（错误信息为 “table_prefix is required but was not set.”）；
• full_otel_spans_table_name/full_otel_logs_table_name/full_annotations_table_name：由后端（Databricks 服务端）回填的完全限定表名（catalog.schema.table），类型标注与序列化逻辑（to_dict/from_dict/from_proto）均在同一文件中可查。

从 trace_location.py 的from_proto实现可以看到，这些 UC 位置对象通过 mlflow/utils/databricks_tracing_utils.py 的uc_table_prefix_location_from_proto与 protobuf 消息（pb.UcTablePrefixLocation）互转，即该配置最终以 protobuf 形式随实验请求发送给 Databricks 后端，由后端据此创建/定位 OTel Span 与 Logs 表。另注意类注释明确：Arclight catalog 不受支持，选用 UC Catalog 时应避开 Arclight。

与其他提示词模板的协同：一次完整的 Agent 接入流程

databricks.md只是整套提示词体系中的“Databricks 专属插槽”，其上游（CLI 组装）与下游（Agent 执行）均有明确约定，完整流程如下：

1. cli.py 启动mlflow agent setup（实验性命令，--agent指定编码 Agent，--print可将渲染后的提示词直接打印到 stdout 供自定义调用，如claude --permission-mode auto "$(mlflow agent setup --agent claude --print)"）；
2. 选定后端为 Databricks 后，提示词 = instrument.md（通用骨架：硬性规则 + 按清单执行 + 验证 + 汇总）+ python.md（语言步骤：安装 MLflow → 配置 tracking URI →mlflow.autolog()插桩）+databricks.md（本文所讲的 Databricks 专属步骤）；
3. Agent 执行步骤 4（instrument.md 的 Verify installation）：端到端运行应用，确认至少一条 Trace 到达{{ tracking_uri }}且无运行时错误；若 HTTP 调用因服务器不可达而挂起，可设MLFLOW_HTTP_REQUEST_MAX_RETRIES=0与MLFLOW_HTTP_REQUEST_TIMEOUT=5快速失败；
4. Agent 在最终汇总中报告：MLflow 版本、修改过的文件、可打开的 Trace URL（以及本地服务器场景下的 PID 与日志路径）。

实战检查清单

将本文内容归纳为可照做的核对清单：

• 确认MLFLOW_TRACKING_URI为databricks或databricks://<profile>，并已取得目标实验 ID（或可自动创建的/开头工作区路径）；
• 运行WorkspaceClient(profile=...).current_user.me()，确认 SDK 能完成认证，不抛异常；
• 认证失败时引导用户databricks auth login/ 配置~/.databrickscfg/ 导出DATABRICKS_HOST与DATABRICKS_TOKEN，且绝不将密钥写入仓库文件；
• 以 ID 绑定实验：mlflow.set_experiment(experiment_id=...)，不重复设置 tracking URI；
• 若用户要求 UC-backed Trace：确认mlflow>=3.11、有可用 SQL Warehouse，收集 catalog/schema/table_prefix，并调用带trace_location=UnityCatalog(...)的set_experiment；未提出要求则整体跳过该块；
• 端到端验证：至少一条 Trace 落地，且无运行时错误。

以上所有结论均可回到仓库源码复核：模板本体在 databricks.md，渲染与注入逻辑在 prompt.py，CLI 交互与参数收集在 cli.py，set_experiment与UnityCatalog的底层实现在 fluent.py 与 trace_location.py。

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow