在 Unity Catalog 中管理模型生命周期

重要

本文介绍如何使用 Unity Catalog 中的模型作为机器学习工作流的一部分来管理 ML 模型的整个生命周期。 Databricks 提供 Unity Catalog 中的 MLflow 模型注册表的托管版本。 Unity Catalog 中的模型将 Unity Catalog 的优势扩展到 ML 模型,包括跨工作区的集中访问控制、审核、世系和模型发现。 Unity Catalog 中的模型与开源 MLflow Python 客户端兼容。

Unity Catalog 中的模型的主要功能包括:

  • 模型的命名空间和治理,因此可以在环境、项目或团队级别对模型进行分组和治理(“授予数据科学家对生产模型的只读访问权限”)。
  • 按时间顺序排列的模型世系(MLflow 试验和运行在给定的时间生成模型)。
  • 模型服务
  • 模型版本控制。
  • 通过别名进行模型部署。 例如,在 prod 目录中标记模型的“冠军”版本。

如果你的工作区的默认目录配置为 Unity Catalog 中的目录,则使用 MLflow API(例如 mlflow.<model-type>.log_model(..., registered_model_name)mlflow.register_model(model_uri, name))注册的模型会默认注册到 Unity Catalog。

本文提供 Unity Catalog UI 和 API 中的模型的说明。

有关模型注册表概念的概述性介绍,请参阅使用 MLflow 进行 ML 生命周期管理

要求

  1. 必须在工作区中启用 Unity Catalog。 若要创建 Unity Catalog 元存储、在工作区中启用它和创建目录,请参阅开始使用 Unity Catalog。 如果未启用 Unity Catalog,你仍可使用经典工作区模型注册表

  2. 工作区必须附加到支持权限继承的 Unity Catalog 元存储。 这适用于在 2022 年 8 月 25 日之后创建的所有元存储。 如果在更旧的元存储上运行,请按照文档进行升级。

  3. 你必须有权在有权访问 Unity Catalog 的群集上运行命令。

  4. 若要创建新的已注册模型,除了对架构及其附带目录的拥有 USE SCHEMAUSE CATALOG 权限外,还需要对架构具有 CREATE_MODEL 权限。 CREATE_MODEL 是一个新的架构级权限,你可使用目录资源管理器 UI 或 SQL GRANT 命令来授予此权限,如下所示。

    GRANT CREATE_MODEL ON SCHEMA <schema-name> TO <principal>
    

将训练工作负载升级到 Unity Catalog

本部分提供将现有训练工作负载升级到 Unity Catalog 的说明。

安装 MLflow Python 客户端

还可以在 Databricks Runtime 11.3 LTS 及更高版本中使用 Unity Catalog 中的模型,方法是在笔记本中使用以下代码安装最新版本的 MLflow Python 客户端。

%pip install --upgrade "mlflow-skinny[databricks]"
dbutils.library.restartPython()

配置 MLflow 客户端以访问 Unity Catalog 中的模型

默认情况下,MLflow Python 客户端会在 Databricks 工作区模型注册表中创建模型。 若要升级到 Unity Catalog 中的模型,请配置 MLflow 客户端:

import mlflow
mlflow.set_registry_uri("databricks-uc")

注意

如果工作区的默认目录在 Unity Catalog(而不是 hive_metastore) 中,并且你正在使用 Databricks Runtime 13.3 LTS 或更高版本(Azure 中国区域 Databricks Runtime 15.0 ML 及更高版本)运行群集,则模型会在默认目录自动创建并从中加载,无需配置。 其他 Databricks Runtime 版本的行为没有变化。 在少数工作区中,默认目录在 2024 年 1 月之前配置为 Unity Catalog 中的目录,并在 2024 年 1 月之前使用了工作区模型注册表,这些工作区不受此行为限制。

训练和注册与 Unity Catalog 兼容的模型

所需的权限:若要创建新的已注册模型,需要对包含的架构具有 CREATE_MODELUSE SCHEMA 权限,对包含的目录具有 USE CATALOG 权限。 若要在已注册的模型下创建新的模型版本,你必须是已注册模型的所有者,并且对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限。

UC 中的 ML 模型版本必须具有模型签名。 如果还没有在模型训练工作负载中记录带有签名的 MLflow 模型,可以:

  • 使用 Databricks 自动记录,它自动记录具有许多热门 ML 框架签名的模型。 有关支持的框架,请参阅 MLflow 文档
  • 使用 MLflow 2.5.0 及更高版本时,可以在 mlflow.<flavor>.log_model 调用中指定一个输入示例,随即会自动推断模型签名。 有关详细信息,请参阅 MLflow 文档

然后,以 <catalog>.<schema>.<model> 的形式将模型的三级名称传递给 MLflow API。

本部分中的示例在 prod 目录下的 ml_team 架构中创建和访问模型。

本部分中的模型训练示例创建了一个新的模型版本并将其注册到了 prod 目录中。 使用 prod 目录并不一定意味着模型版本为生产流量提供服务。 模型版本的封闭目录、架构和已注册模型反映了它的环境 (prod) 和关联的治理规则(例如,可以设置权限,以便只有管理员可以从 prod 目录中删除),但未反映其部署状态。 要管理部署状态,请使用模型别名

使用自动日志记录功能将模型注册到 Unity Catalog

from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier

# Train a sklearn model on the iris dataset
X, y = datasets.load_iris(return_X_y=True, as_frame=True)
clf = RandomForestClassifier(max_depth=7)
clf.fit(X, y)

# Note that the UC model name follows the pattern
# <catalog_name>.<schema_name>.<model_name>, corresponding to
# the catalog, schema, and registered model name
# in Unity Catalog under which to create the version
# The registered model will be created if it doesn't already exist
autolog_run = mlflow.last_active_run()
model_uri = "runs:/{}/model".format(autolog_run.info.run_id)
mlflow.register_model(model_uri, "prod.ml_team.iris_model")

使用自动推断的签名将模型注册到 Unity Catalog

MLflow 2.5.0 及更高版本支持自动推断的签名,Databricks Runtime 11.3 LTS ML 及更高版本也支持。 若要使用自动推断的签名,请使用以下代码在笔记本中安装最新的 MLflow Python 客户端:

%pip install --upgrade "mlflow-skinny[databricks]"
dbutils.library.restartPython()

以下代码演示了自动推断的签名示例。

from sklearn import datasets
from sklearn.ensemble import RandomForestClassifier

with mlflow.start_run():
    # Train a sklearn model on the iris dataset
    X, y = datasets.load_iris(return_X_y=True, as_frame=True)
    clf = RandomForestClassifier(max_depth=7)
    clf.fit(X, y)
    # Take the first row of the training dataset as the model input example.
    input_example = X.iloc[[0]]
    # Log the model and register it as a new version in UC.
    mlflow.sklearn.log_model(
        sk_model=clf,
        artifact_path="model",
        # The signature is automatically inferred from the input example and its predicted output.
        input_example=input_example,
        registered_model_name="prod.ml_team.iris_model",
    )

跟踪 Unity Catalog 中的模型的数据世系

注意

MLflow 2.11.0 及更高版本中提供了对 Unity Catalog 中的表到模型世系的支持。

在 Unity Catalog 中的表中训练模型时,可以将模型的世系跟踪到已对其进行训练和评估的上游数据集。 为此,请使用 mlflow.log_input。 这会使用生成模型的 MLflow 运行保存输入表信息。 对于使用功能存储 API 记录的模型,也会自动捕获数据世系。 请参阅查看特征存储世系

将模型注册到 Unity Catalog 时,世系信息会自动保存,并显示在目录资源管理器中模型版本 UI 的“世系”选项卡中。

以下代码展示了一个示例。

import mlflow
import pandas as pd
import pyspark.pandas as ps
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestRegressor

# Write a table to Unity Catalog
iris = load_iris()
iris_df = pd.DataFrame(iris.data, columns=iris.feature_names)
iris_df.rename(
  columns = {
    'sepal length (cm)':'sepal_length',
    'sepal width (cm)':'sepal_width',
    'petal length (cm)':'petal_length',
    'petal width (cm)':'petal_width'},
  inplace = True
)
iris_df['species'] = iris.target
ps.from_pandas(iris_df).to_table("prod.ml_team.iris", mode="overwrite")

# Load a Unity Catalog table, train a model, and log the input table
dataset = mlflow.data.load_delta(table_name="prod.ml_team.iris", version="0")
pd_df = dataset.df.toPandas()
X = pd_df.drop("species", axis=1)
y = pd_df["species"]
with mlflow.start_run():
    clf = RandomForestRegressor(n_estimators=100)
    clf.fit(X, y)
    mlflow.log_input(dataset, "training")

在 UI 中查看模型

所需的权限:若要在 UI 中查看已注册的模型及其模型版本,你需要对已注册的模型具有 EXECUTE 权限,并对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限

可使用目录资源管理器在 Unity Catalog 中查看和管理已注册的模型和模型版本。

控制对模型的访问

若要了解如何控制对 Unity Catalog 中注册的模型的访问,请参阅 Unity Catalog 权限和安全对象。 有关跨目录和架构组织模型的最佳做法,请参阅整理数据

可以使用授予 REST API 以编程方式配置模型权限。 配置模型权限时,请在 REST API 请求中将 securable_type 设置为 "FUNCTION"。 例如,使用 PATCH /api/2.1/unity-catalog/permissions/function/{full_name} 更新已注册的模型权限。

使用别名与标记部署和组织模型

模型别名和标记有助于在 Unity Catalog 中组织和管理模型。

模型别名让你能够将可变的命名引用分配给已注册的模型的特定版本。 可以使用别名来指示模型版本的部署状态。 例如,可以将“Champion”别名分配给当前在生产中的模型版本,并在使用该生产模型的工作负载中将此别名作为目标。 然后,可以通过将“Champion”别名重新分配给其他模型版本来更新生产模型。

标记是与已注册的模型和模型版本关联的键值对,它使你可按功能或状态标记和分类它们。 例如,可以将带有键 "task" 和值 "question-answering" 的标记(在 UI 中显示为 task:question-answering)应用于用于问题解答任务的已注册模型。 在模型版本级别,可以使用 validation_status:pending 标记正在进行部署前验证的版本,并使用 validation_status:approved 标记已经可以部署的版本。

请参阅以下部分,了解如何使用别名和标记。

在模型上设置和删除别名

所需的权限:是已注册的模型的所有者,并且对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限。

可以使用目录资源管理器 在 Unity Catalog 中设置、更新和删除模型的别名。 可以在模型详细信息页中跨已注册的模型管理别名,并在模型版本详细信息页中为特定模型版本配置别名。

要使用 MLflow 客户端 API 设置、更新和删除别名,请参阅以下示例:

from mlflow import MlflowClient
client = MlflowClient()

# create "Champion" alias for version 1 of model "prod.ml_team.iris_model"
client.set_registered_model_alias("prod.ml_team.iris_model", "Champion", 1)

# reassign the "Champion" alias to version 2
client.set_registered_model_alias("prod.ml_team.iris_model", "Champion", 2)

# get a model version by alias
client.get_model_version_by_alias("prod.ml_team.iris_model", "Champion")

# delete the alias
client.delete_registered_model_alias("prod.ml_team.iris_model", "Champion")

在模型上设置和删除标记

所需的权限:是已注册的模型的所有者,或对其具有 APPLY_TAG 权限,并且对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限。

请参阅在目录资源管理器中管理标记,以了解如何使用 UI 设置和删除标记。

要使用 MLflow 客户端 API 设置和删除标记,请参阅以下示例:

from mlflow import MlflowClient
client = MlflowClient()

# Set registered model tag
client.set_registered_model_tag("prod.ml_team.iris_model", "task", "classification")

# Delete registered model tag
client.delete_registered_model_tag("prod.ml_team.iris_model", "task")

# Set model version tag
client.set_model_version_tag("prod.ml_team.iris_model", "1", "validation_status", "approved")

# Delete model version tag
client.delete_model_version_tag("prod.ml_team.iris_model", "1", "validation_status")

已注册的模型和模型版本标记都必须满足平台范围的约束

有关别名和标记客户端 API 的更多详细信息,请参阅 MLflow API 文档

加载用于推理的模型

在推理工作负载中按别名使用模型版本

所需的权限:在已注册的模型上拥有 EXECUTE 权限,并且对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限。

可以写入通过别名引用模型版本的批推理工作负载。 例如,下面的代码段加载并应用“冠军”模型版本进行批推理。 如果“冠军”版本更新为引用新的模型版本,则批推理工作负载会在下一次执行时自动挑拣它。 这允许将模型部署与批推理工作负载分离。

import mlflow.pyfunc
model_version_uri = "models:/prod.ml_team.iris_model@Champion"
champion_version = mlflow.pyfunc.load_model(model_version_uri)
champion_version.predict(test_x)

还可以写入部署工作流以通过别名获取模型版本,并使用模型服务 REST API 更新模型服务端点以服务该版本:

import mlflow
import requests
client = mlflow.tracking.MlflowClient()
champion_version = client.get_model_version_by_alias("prod.ml_team.iris_model", "Champion")
# Invoke the model serving REST API to update endpoint to serve the current "Champion" version
model_name = champion_version.name
model_version = champion_version.version
requests.request(...)

在推理工作负载中按版本号使用模型版本

还可以按版本号加载模型版本:

import mlflow.pyfunc
# Load version 1 of the model "prod.ml_team.iris_model"
model_version_uri = "models:/prod.ml_team.iris_model/1"
first_version = mlflow.pyfunc.load_model(model_version_uri)
first_version.predict(test_x)

跨工作区共享模型

与同一区域中的用户共享模型

只要拥有适当的权限,就可以从任何工作区访问 Unity Catalog 中的模型,前提是该工作区附加到了包含该模型的元存储。 例如,可以从开发工作区的 prod 目录中访问模型,以便于将新开发的模型与产品基线进行比较。

若要在创建的已注册模型上与其他用户协作(共享写入权限),必须将模型的所有权授予包含你自己和要协作的用户的组。 协作者还必须对包含模型的目录和架构具有 USE CATALOGUSE SCHEMA 权限。 有关详细信息,请参阅 Unity Catalog 权限和安全对象

与其他区域或帐户中的用户共享模型

若要与其他区域或帐户中的用户共享模型,请使用增量共享 Databricks 到 Databricks 共享流。 请参阅将模型添加到共享(面向提供方)和在 Databricks 到 Databricks 模型中获取访问权限(面向接收方)。 作为接收方,基于共享创建目录后,可以像访问 Unity Catalog 中的任何其他模型一样访问该共享目录中的模型。

对模型或模型版本进行批注

所需的权限:是已注册的模型的所有者,并且对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限。

可以通过对模型或模型版本进行批注来提供与其相关的信息。 例如,你可能想提供有关所用方法和算法的相关问题或信息的概述性介绍。

使用 UI 对模型或模型版本进行批注

请参阅使用 Markdown 注释在 Catalog Explorer 中记录数据

使用 API 对模型或模型版本进行批注

若要更新已模型模型的说明,请使用 MLflow 客户端 API update_registered_model() 方法:

client = MlflowClient()
client.update_registered_model(
  name="<model-name>",
  description="<description>"
)

若要更新模型版本说明,请使用 MLflow 客户端 API update_model_version() 方法:

client = MlflowClient()
client.update_model_version(
  name="<model-name>",
  version=<model-version>,
  description="<description>"
)

重命名模型(仅 API)

所需的权限:是已注册的模型的所有者、包含已注册模型的架构具有 CREATE_MODEL 权限,并且对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限。

若要重命名已注册模型,请使用 MLflow 客户端 API rename_registered_model() 方法:

client=MlflowClient()
client.rename_registered_model("<model-name>", "<new-model-name>")

删除模型或模型版本

所需的权限:是已注册的模型的所有者,并且对包含该模型的架构和目录具有 USE SCHEMAUSE CATALOG 权限。

可以使用目录资源管理器 UI 或 API 删除已注册模型或已注册模型内的模型版本。

使用 API 删除模型版本或模型

警告

不能撤消此操作。 删除模型时,会删除 Unity Catalog 存储的所有模型项目以及与该已注册模型关联的所有元数据。

删除模型版本

若要删除模型版本,请使用 MLflow 客户端 API delete_model_version() 方法:

# Delete versions 1,2, and 3 of the model
client = MlflowClient()
versions=[1, 2, 3]
for version in versions:
  client.delete_model_version(name="<model-name>", version=version)

删除模型

若要删除模型,请使用 MLflow 客户端 API delete_registered_model() 方法:

client = MlflowClient()
client.delete_registered_model(name="<model-name>")

列出和搜索模型

可以使用 MLflow 的 search_registered_models() Python API 在 Unity Catalog 中列出已注册模型:

client=MlflowClient()
client.search_registered_models()

还可以使用 search_model_versions() 方法搜索特定的型号名称并列出其版本详细信息:

from pprint import pprint

client=MlflowClient()
[pprint(mv) for mv in client.search_model_versions("name='<model-name>'")]

示例

此示例演示如何使用 Unity Catalog 中的模型来生成机器学习应用程序。

Unity Catalog 中的模型示例

将工作流和模型迁移到 Unity Catalog

Azure Databricks 建议使用 Unity Catalog 中的模型加强治理、在工作区和环境之间实现轻松共享,并享受到更加灵活的 MLOps 工作流。 该表比较了工作区模型注册表和 Unity Catalog 的功能。

功能 工作区模型注册表(旧版) Unity Catalog 中的模型(建议)
通过命名别名引用模型版本 模型注册阶段:将模型版本移动到四个固定阶段之一,以便在该阶段引用它们。 无法重命名或添加阶段。 模型注册表别名:为每个注册模型创建最多 10 个自定义且可重新分配的模型版本命名引用。
为模型创建访问控制的环境 模型注册表阶段:使用一个注册模型中的阶段来表示其模型版本的环境,仅对四个固定阶段中的两个(StagingProduction)进行访问控制。 已注册模型:为 MLOps 工作流中的每个环境创建一个已注册模型,利用三级命名空间和 Unity Catalog 的权限来表达治理。
跨环境提升模型(部署模型) 使用 transition_model_version_stage() MLflow 客户端 API 将模型版本移动到不同阶段,这可能会破坏引用前一阶段的工作流。 使用 copy_model_version() MLflow 客户端 API 将模型版本从一个注册模型复制到另一个注册模型。
跨工作区访问和共享模型 跨工作区手动导出和导入模型,或使用个人访问令牌和工作区机密范围配置与远程模型注册表的连接。 现成地访问同一帐户中跨工作区的模型。 无需任何配置。
配置权限 在工作区级别设置权限。 在帐户级别设置权限,从而在工作区之间应用一致的治理。
访问 Databricks 市场中的模型 不可用。 将模型从 Databricks 市场加载到 Unity Catalog 元存储中,并跨工作区访问它们。

下面链接的文章介绍如何将工作区模型注册表中的工作流(模型训练和批量推理作业)和模型迁移到 Unity Catalog。

Unity 目录支持的限制

  • Unity Catalog 中的模型不支持阶段。 Databricks 建议使用 Unity Catalog 中的三级命名空间来表达模型所处的环境,并使用别名来促进模型的部署。 有关详细信息,请参阅升级指南
  • Unity Catalog 中的模型不支持 Webhook。 请参阅升级指南中建议的替代方案。
  • Unity Catalog 中的模型不支持某些搜索 API 字段和运算符。 通过使用支持的筛选器调用搜索 API 并扫描结果,可以缓解这种情况。 下面是一些示例:
    • search_model_versionssearch_registered_models 客户端 API 中不支持 order_by 参数。
    • search_model_versionssearch_registered_models 不支持基于标记的筛选器 (tags.mykey = 'myvalue')。
    • search_model_versionssearch_registered_models 不支持除完全相等以外的运算符(例如 LIKEILIKE!=)。
    • 不支持按名称(例如 MlflowClient().search_registered_models(filter_string="name='main.default.mymodel'"))搜索已注册的模型。 若要按名称提取特定已注册模型,请使用 get_registered_model
  • Unity Catalog 不支持有关已注册模型和模型版本的电子邮件通知和注释讨论线程。
  • Unity Catalog 中的模型不支持活动日志。 但是,你可以使用审核日志跟踪 Unity Catalog 中模型的活动。
  • 对于通过“增量共享”共享的模型,search_registered_models 可能会返回过时的结果。 若要确保返回最新结果,请使用 Databricks CLI 或 SDK 列出架构中的模型。