升级 ML 工作流以面向 Unity Catalog 中的模型

项目
09/09/2024

本文介绍如何将 ML 工作流迁移和升级到 Unity Catalog 中的目标模型。

要求

在开始之前，请确保满足要求中的各项要求。具体而言，请确保用于执行模型训练、部署和推理工作流的用户或主体对 Unity Catalog 中已注册的模型具有必要的权限：

训练：具有已注册的模型的所有权（这是创建新模型版本所必需的），并且对封闭目录和架构具有 USE CATALOG 和 USE SCHEMA 权限。
部署：具有已注册的模型的所有权（这是在模型上设置别名所必需的），并且对封闭目录和架构具有 USE CATALOG 和 USE SCHEMA 权限。
推理：对已注册的模型具有 EXECUTE 权限（这是使用模型版本读取和执行推理所必需的），并且对封闭目录和架构具有 USE CATALOG 和“USE SCHEMA”特权。

创建并行训练、部署和工作流

Azure Databricks 建议采用增量方法将模型训练和推理工作流升级到 Unity Catalog。借助这种方法，可以创建并行训练、部署和推理管道来利用 Unity Catalog 中的模型。如果能接受使用 Unity Catalog 得出的结果，则可以切换下游使用者来读取批量推理输出，或是增加路由到提供服务的终结点中 Unity Catalog 中模型的流量。

模型训练工作流

克隆模型训练工作流。然后，请确保：

工作流群集有权访问 Unity Catalog，并满足要求中所述的各项要求。
运行工作流的主体对 Unity Catalog 中已注册的模型具有必要的权限。

接下来，修改已克隆的工作流中的模型训练代码。可能需要克隆工作流运行的笔记本，或者在已克隆的工作流中创建并面向新的 git 分支。按照这些步骤安装必要的 MLflow 版本，将客户端配置为在训练代码中面向 Unity Catalog，并更新模型训练代码，将模型注册到 Unity Catalog。

模型部署工作流

克隆模型部署工作流，按照模型训练工作流中的类似步骤更新其计算配置，以启用对 Unity Catalog 的访问。

确保拥有已克隆工作流的主体具有必要的权限。如果部署工作流中有模型验证逻辑，请将其更新为从 UC 加载模型版本。使用别名管理生产模型推出。

模型推理工作流

批量推理工作流

按照模型训练工作流中的类似步骤克隆批量推理工作流，并更新其计算配置来启用对 Unity Catalog 的访问。确保运行已克隆的批处理推理作业的主体具有加载模型进行推理所必要的权限。

跨环境提升模型

Databricks 建议将 ML 管道部署为代码。这样就无需跨环境提升模型，因为所有生产模型都可以通过生产环境中的自动训练工作流生成。

但是，在某些情况下，跨环境重新训练模型可能过于昂贵。可以改为在 Unity Catalog 中跨已注册的模型复制模型版本，以便跨环境提升它们。

需要具有以下特权才能执行以下示例代码：

staging 和 prod 目录上的 USE CATALOG。
在 staging.ml_team 和 prod.ml_team 架构上的 USE SCHEMA。
staging.ml_team.fraud_detection 上的 EXECUTE。

此外，必须是已注册模型 prod.ml_team.fraud_detection 的所有者。

以下代码片段使用 copy_model_version MLflow 客户端 API，在 MLflow 2.8.0 及更高版本中可用。

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

client = mlflow.tracking.MlflowClient()
src_model_name = "staging.ml_team.fraud_detection"
src_model_version = "1"
src_model_uri = f"models:/{src_model_name}/{src_model_version}"
dst_model_name = "prod.ml_team.fraud_detection"
copied_model_version = client.copy_model_version(src_model_uri, dst_model_name)

模型版本进入生产环境中后，可以执行任何必要的部署前验证。然后，可以使用别名标记要部署的模型版本。

client = mlflow.tracking.MlflowClient()
client.set_registered_model_alias(name="prod.ml_team.fraud_detection", alias="Champion", version=copied_model_version.version)

在上面的示例中，只有能够从 staging.ml_team.fraud_detection 注册的模型读取和写入 prod.ml_team.fraud_detection 注册的模型的用户才能将暂存模型提升到生产环境。同一用户还可以使用别名来管理在生产环境中部署的模型版本。无需配置任何其他规则或策略即可管理模型提升和部署。

你可以自定义此流，以在与你的设置匹配的多个环境（例如 dev、qa 和 prod）中提升模型版本。访问控制是按照每个环境中的配置强制实施的。

使用作业 Webhook 手动批准模型部署

Databricks 建议在模型部署过程中使用适当的检查和测试自动执行模型部署（如果可能）。但是，如果部署生产模型确实需要执行手动批准，则可以在模型训练作业成功完成后，使用作业 Webhook 调用外部 CI/CD 系统请求手动批准，从而部署模型。提供手动批准后，CI/CD 系统随后可以部署模型版本来为流量提供服务，例如，通过在其上设置“Champion”别名。

通过