跳到主要内容

MLflow 模型注册表

MLflow 模型注册中心 (Model Registry) 是一个集中式模型存储库，提供了一套 API 和用户界面，旨在协作管理机器学习模型的全生命周期。它提供血缘关系（即由哪个 MLflow 实验和运行产生该模型）、版本控制、别名、元数据标记和注释支持，确保您在从开发到生产部署的每个阶段都能获得全面的信息。

为什么需要模型注册中心？​

随着机器学习项目的复杂性和规模不断增加，跨不同环境、团队和迭代手动管理模型变得越来越容易出错且效率低下。MLflow 模型注册中心通过提供一个集中、结构化的系统来组织和管理机器学习模型在整个生命周期中的状态，从而解决了这一挑战。

使用模型注册中心可获得以下好处

  • 🗂️ 版本控制：注册中心会自动跟踪每个模型的版本，允许团队比较不同迭代、回滚到之前的状态，并并行管理多个版本（例如：预发布版 vs. 生产版）。
  • 🧬 模型血缘与可追溯性：每个注册的模型版本都与产生它的 MLflow 运行、已记录的模型或 Notebook 相关联，实现了完全的可重现性。您可以准确追溯模型的训练方式、所使用的数据及参数。
  • 🚀 生产就绪工作流：模型别名（如 @champion）和标签等功能简化了部署工作流的管理，能够以可控且可审计的方式将模型提升至实验、预发布或生产环境。
  • 🛡️ 治理与合规：通过结构化元数据、标记以及基于角色的访问控制（与 Databricks 等后端或托管的 MLflow 服务结合使用时），模型注册中心支持企业级机器学习运维所需的治理要求。

无论您是独立数据科学家还是大型 ML 平台团队的一员，模型注册中心都是扩展可靠且可维护的机器学习系统的基础组件。

概念​

模型注册中心引入了一些概念来描述并促进 MLflow 模型的全生命周期。

概念描述
模型

MLflow 模型是通过某种模型风格的 mlflow.<model_flavor>.log_model() 方法或 MLflow 3.0 引入的 mlflow.create_external_model() API 创建的。一旦记录完成，该模型即可注册到模型注册中心。

注册模型 (Registered Model)

MLflow 模型可以注册到模型注册中心。一个注册模型拥有唯一的名称，包含版本、别名、标签及其他元数据。

模型版本 (Model Version)

每个注册模型可以有一个或多个版本。当新模型添加到模型注册中心时，它被添加为版本 1。每个以相同模型名称注册的新模型都会递增版本号。模型版本具有标签，可用于跟踪模型版本的属性（例如 pre_deploy_checks: "PASSED"）。

模型 URI

您可以使用以下格式的 URI 来引用已注册模型：models:/<model-name>/<model-version>。例如，如果您有一个名为 "MyModel" 且版本为 1 的注册模型，则引用该模型的 URI 为：models:/MyModel/1。

模型别名 (Model Alias)

模型别名允许您为已注册模型的特定版本分配一个可变的、有名称的引用。通过为特定模型版本分配别名，您可以使用该别名通过模型 URI 或模型注册中心 API 引用该版本。例如，您可以创建一个名为 champion 的别名，指向名为 MyModel 的模型的版本 1。之后，您就可以使用 URI models:/MyModel@champion 来引用 MyModel 的版本 1。

别名对于模型部署特别有用。例如，您可以为用于生产流量的模型版本分配 champion 别名，并在生产工作负载中引用此别名。随后，您只需将 champion 别名重新分配给另一个模型版本，即可更新生产流量所使用的模型。

标签

标签是您与已注册模型及模型版本关联的键值对，允许您按功能或状态对其进行标记和分类。例如，您可以为旨在进行问答任务的已注册模型添加键为 "task"、值为 "question-answering" 的标签（在 UI 中显示为 task:question-answering）。在模型版本级别，您可以为正在进行预部署验证的版本标记 validation_status:pending，为已批准部署的版本标记 validation_status:approved。

注释与描述

您可以使用 Markdown 对顶级模型及其每个版本进行单独注释，包括描述以及对团队有用的任何相关信息，例如算法描述、所用数据集或特定版本建模方法中涉及的整体方法论。

模型注册中心的实践​

MLflow 模型注册中心既可用于开源 (OSS) MLflow，也可用于 Databricks 等托管平台。根据环境的不同，注册中心提供的集成、治理和协作功能级别有所差异。

OSS MLflow 中的模型注册中心​

在开源版 MLflow 中，模型注册中心为版本控制和管理 ML 模型提供了 UI 和 API。您可以注册模型、跟踪版本、添加标签和描述，并使模型在预发布 (Staging) 和生产 (Production) 等阶段之间转换。

在 MLflow 中注册模型

使用 MLflow Python API 注册模型​

MLflow 提供了多种注册模型版本的方法

text
# Option 1: specify `registered_model_name` parameter when logging a model
mlflow.<flavor>.log_model(..., registered_model_name="<YOUR_MODEL_NAME>")

# Option 2: register a logged model
mlflow.register_model(model_uri="<YOUR_MODEL_URI>", name="<YOUR_MODEL_NAME>")

注册模型后，您可以使用模型名称和版本将其加载回来

text
mlflow.<flavor>.load_model("models:/<YOUR_MODEL_NAME>/<YOUR_MODEL_VERSION>")

要了解有关 OSS 模型注册中心的更多信息，请参阅模型注册中心教程。

Databricks 中的模型注册中心​

Databricks 通过将模型注册中心与 Unity Catalog (UC) 集成来扩展 MLflow 的功能，从而实现集中式治理、细粒度的访问控制以及跨工作区协作。

Unity Catalog 集成的关键好处包括

  • 🛡️ 增强治理：对模型资产应用访问策略和权限控制。
  • 🌐 跨工作区访问：注册一次模型，即可在多个 Databricks 工作区中访问。
  • 🔗 模型血缘：跟踪用于创建每个模型的 Notebook、数据集和实验。
  • 🔍 发现与复用：从共享目录中浏览和复用生产级模型。
在 Databricks UC 中注册模型

使用 MLflow Python API 注册模型到 Databricks UC​

前提条件：将 tracking uri 设置为 Databricks

python
import mlflow

mlflow.set_registry_uri("databricks-uc")

使用 MLflow API 注册模型

text
# Option 1: specify `registered_model_name` parameter when logging a model
mlflow.<flavor>.log_model(..., registered_model_name="<YOUR_MODEL_NAME>")

# Option 2: register a logged model
mlflow.register_model(model_uri="<YOUR_MODEL_URI>", name="<YOUR_MODEL_NAME>")
警告

UC 中的 ML 模型版本必须具有模型签名 (model signature)。如果您想为已经记录或保存的模型设置签名，可以使用 mlflow.models.set_signature() API。

注册模型后，您可以使用模型名称和版本将其加载回来

text
mlflow.<flavor>.load_model("models:/<YOUR_MODEL_NAME>/<YOUR_MODEL_VERSION>")

有关详细信息，请参阅有关管理模型生命周期的 Databricks 文档。