如何从Databricks环境外部向MLflow Registry注册模型

前提确认 #

你已完成以下准备操作：

• 安装依赖包：

  pip3 install mlflow
  pip3 install sklearn
  

• 配置环境变量：

  export MLFLOW_TRACKING_URI=databricks
  export DATABRICKS_HOST="https://mydatabricks-host"
  export DATABRICKS_TOKEN="mytoken" # 从Databricks UI的「用户设置 > 开发者 > 访问令牌」生成
  

步骤1：准备训练完成的模型（以Scikit-learn为例） #

先训练一个示例模型，或直接使用你已有的业务模型：

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# 加载并拆分数据集
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.2)

# 训练模型
model = RandomForestClassifier(n_estimators=100)
model.fit(X_train, y_train)

# 评估模型性能
y_pred = model.predict(X_test)
acc = accuracy_score(y_test, y_pred)

步骤2：启动MLflow运行并记录模型 #

在代码中初始化MLflow运行，同时记录模型参数、指标和模型本身：

# 启动MLflow运行会话
with mlflow.start_run() as run:
    # 记录模型超参数
    mlflow.log_param("n_estimators", 100)
    # 记录模型评估指标
    mlflow.log_metric("accuracy", acc)
    # 将模型上传至Databricks MLflow跟踪服务器
    mlflow.sklearn.log_model(model, "iris-rf-model")
    
    # 保存当前运行ID，用于后续模型注册
    run_id = run.info.run_id

步骤3：注册模型到Databricks MLflow模型注册表 #

提供两种注册方式，按需选择：

方式1：代码内直接注册 #

调用mlflow.register_model方法完成注册：

# 指定模型注册名称
model_name = "Iris-RandomForest-Model"
# 执行注册
result = mlflow.register_model(
    model_uri=f"runs:/{run_id}/iris-rf-model",
    name=model_name
)

# 打印注册结果
print(f"模型注册完成，版本：{result.version}，名称：{result.name}")

方式2：通过MLflow CLI注册 #

在终端执行以下命令（替换<你的运行ID>为步骤2中获取的run_id）：

mlflow models register -m "runs:/<你的运行ID>/iris-rf-model" -n "Iris-RandomForest-Model"

步骤4：验证注册结果 #

1. 登录你的Databricks工作区
2. 导航至「机器学习」>「模型注册表」
3. 在模型列表中查找你注册的模型名称，即可查看模型版本、关联的运行记录等信息

额外注意事项 #

• 确保你的Databricks令牌权限充足：需要具备MLflow跟踪、模型注册表相关操作权限（如mlflow.tracking.create、mlflow.registerModel）
• 其他框架模型适配：若使用TensorFlow、PyTorch等框架，只需替换对应的mlflow.<framework>.log_model方法，注册流程完全一致
• 环境变量兜底配置：若终端环境变量未生效，可在代码内直接设置：

  mlflow.set_tracking_uri("databricks")
  mlflow.set_registry_uri("databricks")
  import os
  os.environ["DATABRICKS_HOST"] = "https://mydatabricks-host"
  os.environ["DATABRICKS_TOKEN"] = "mytoken"
  

内容的提问来源于stack exchange，提问作者Chris Snow