如何使用Python在Databricks外部访问Model Registry的MLflow工件

前置依赖安装 #

先安装需要的第三方包：

pip install mlflow pyyaml

提前准备的参数 #

仅需提前准备好以下4个信息即可：

• Databricks工作区访问地址，格式为https://adb-<你的工作区ID>.azuredatabricks.net
• Databricks个人访问令牌（PAT，在Databricks控制台用户设置页面生成，需要授予模型注册表读取权限）
• 目标注册模型的名称
• 目标模型的版本号

实现代码 #

Databricks Model Registry原生兼容MLflow API，不需要依赖Databricks运行时或者databricks-cli，直接用MLflow客户端即可完成鉴权和工件读取：

import os
import yaml
import mlflow
from mlflow.tracking import MlflowClient

# -------------------------- 替换为你的配置 --------------------------
DATABRICKS_WORKSPACE = "https://adb-xxxx.azuredatabricks.net"
DATABRICKS_TOKEN = "dapixxxxxxxxxxxxxxxxxxxxxxxxxxxx"
TARGET_MODEL_NAME = "你的模型名称"
TARGET_MODEL_VERSION = 1  # 替换为目标版本号
# 目标文件在工件中的相对路径，根目录直接填文件名，子目录补全路径即可，比如"data/feature_spec.yml"
TARGET_FILE_PATH = "feature_spec.yml"
# -------------------------------------------------------------------

# 配置MLflow跟踪地址指向Databricks托管服务
mlflow.set_tracking_uri(DATABRICKS_WORKSPACE)
# 配置鉴权令牌
os.environ["DATABRICKS_TOKEN"] = DATABRICKS_TOKEN

# 初始化MLflow客户端
client = MlflowClient()

# 获取指定版本模型的元数据
model_meta = client.get_model_version(
    name=TARGET_MODEL_NAME,
    version=TARGET_MODEL_VERSION
)

# 下载目标工件到本地临时目录
local_file_path = client.download_artifacts(
    run_id=model_meta.run_id,
    path=TARGET_FILE_PATH,
    dst_path="./temp_model_artifacts"
)

# 读取yaml文件内容
with open(local_file_path, "r", encoding="utf-8") as f:
    feature_spec_content = yaml.safe_load(f)

# 打印验证结果
print("读取到的feature_spec内容：")
print(feature_spec_content)

常见问题说明 #

• 如果运行时报权限错误，检查个人访问令牌是否过期，以及对应令牌是否有目标模型的读取权限。
• 如果提示找不到文件，先确认TARGET_FILE_PATH的相对路径是否和模型工件里的路径一致，可以先调用client.list_artifacts(run_id=model_meta.run_id, path="")打印根目录下所有文件和文件夹，逐层确认文件位置。
• 不要把访问令牌硬编码在代码中，生产环境建议通过系统环境变量、云厂商密钥管理服务注入凭证，避免凭证泄露。

内容的提问来源于stack exchange，提问作者Minura Punchihewa