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 CatalogUCTrace 存储配置。读完本文你将掌握 Databricks SDK 认证验证的正确姿势、mlflow.set_experiment(experiment_id...)与trace_locationUnityCatalog(...)的实战用法以及 Trace 数据落盘 UC Delta 表所需的全部前提条件并理解这些指令在mlflow agent setup提示词组装链路中的真实位置。背景这份文档出现在 Agent 工作流的哪个环节本文所解析的databricks.md并非独立使用的手册而是 mlflow/agent/setup/prompt.py 中按条件动态注入的提示词片段当用户在mlflow agent setup交互流程中选择MLFLOW_TRACKING_URIdatabricks或databricks://profile作为追踪后端时CLI 会读取该模板、渲染占位符并把渲染结果嵌入到最终交给编码 Agent 的首条用户消息中{{ server_setup }}插槽最终进入 python.md 的步骤 2 之前。整条调用链如下均可在仓库中验证mlflow/agent/setup/cli.py_run_setup()检测到 tracking URI 指向 Databricks 后调用_prompt_experiment_id()要求用户提供实验 ID 或工作区路径以/开头时自动经get_experiment_by_name查找、不存在则create_experiment创建见 cli.pymlflow/agent/setup/prompt.py将tracking_uri、experiment_id、workspace_client_args三个值渲染进databricks.md模板渲染结果作为{{ 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已存在且指向 DatabricksCLI 直接采用cli.py否则在交互式后端选择菜单中用户选择 Connect to a Databricks workspace 后CLI 会提示输入 Databricks 配置 profile留空则用默认 profile并构造出databricks://profile或databrickscli.py。{{ experiment_id }}同理CLI 要求用户输入实验 ID或输入以/开头的实验路径如/Users/username/my-experiment后者若不存在会自动创建并回显新 IDcli.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时渲染为profileprofile默认 profile 时为空字符串prompt.py。因此实际生成的调用形如# 默认 profile from databricks.sdk import WorkspaceClient WorkspaceClient().current_user.me() # 指定 profile 为 prod from databricks.sdk import WorkspaceClient WorkspaceClient(profileprod).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 选择优先解析并传入 IDID 是工作区内最稳定、无歧义的实验标识。import mlflow mlflow.set_experiment(experiment_id1234567890123456)这一步之后后续所有mlflow.autolog()产生的 Trace 都会落到该实验下最终在 MLflow UI 中可通过“tracking URI experiment ID”构造出可打开的 Trace 链接该 URL 会在 instrument.md 的步骤 5 由 Agent 汇总报告给用户。可选进阶将 Trace 存储到 Unity Catalog Delta 表如果用户希望 Trace 由 UC Delta 表承载而非默认的跟踪服务器存储模板提供了可选代码块。前置条件缺一不可MLflow 版本 ≥ 3.11UnityCatalogtrace location 类型自 3.11.0 起标记为实验性特性见 mlflow/entities/trace_location.py 的experimental(version3.11.0)装饰器一个可用的 SQL WarehouseUC 表写入需要 SQL 仓库执行 DDL/DML用户明确要求模板强调“Skip this block entirely if the user does not ask for UC-backed traces”——只有在用户主动提出 UC 存储时才启用否则跳过。需要向用户确认的四个参数参数含义取值示例catalog_nameUC Catalog 名称mainschema_nameUC Schema 名称mlflow_tracestable_prefix表名前缀my_appSQL Warehouse ID承载写入的 SQL 仓库在 Databricks 控制台获取配置代码import mlflow from mlflow.entities.trace_location import UnityCatalog mlflow.set_experiment( experiment_id{{ experiment_id }}, trace_locationUnityCatalog( catalog_namecatalog, schema_nameschema, table_prefixprefix, ), )即实际形态如import mlflow from mlflow.entities.trace_location import UnityCatalog mlflow.set_experiment( experiment_id1234567890123456, trace_locationUnityCatalog( catalog_namemain, schema_namemlflow_traces, table_prefixmy_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 执行均有明确约定完整流程如下cli.py 启动mlflow agent setup实验性命令--agent指定编码 Agent--print可将渲染后的提示词直接打印到 stdout 供自定义调用如claude --permission-mode auto $(mlflow agent setup --agent claude --print)选定后端为 Databricks 后提示词 instrument.md通用骨架硬性规则 按清单执行 验证 汇总 python.md语言步骤安装 MLflow → 配置 tracking URI →mlflow.autolog()插桩databricks.md本文所讲的 Databricks 专属步骤Agent 执行步骤 4instrument.md 的 Verify installation端到端运行应用确认至少一条 Trace 到达{{ tracking_uri }}且无运行时错误若 HTTP 调用因服务器不可达而挂起可设MLFLOW_HTTP_REQUEST_MAX_RETRIES0与MLFLOW_HTTP_REQUEST_TIMEOUT5快速失败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确认mlflow3.11、有可用 SQL Warehouse收集 catalog/schema/table_prefix并调用带trace_locationUnityCatalog(...)的set_experiment未提出要求则整体跳过该块端到端验证至少一条 Trace 落地且无运行时错误。以上所有结论均可回到仓库源码复核模板本体在 databricks.md渲染与注入逻辑在 prompt.pyCLI 交互与参数收集在 cli.pyset_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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考