尧图网页设计 
ZenML 自定义模型注册表Custom Model Registry开发指南从基类抽象到自定义 Flavor 落地【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml模型注册表Model Registry是 ZenML 中负责集中管理、追踪与版本化机器学习模型的 Stack 组件。当内置的 MLflow 等集成无法满足团队自定义的模型版本管理需求时你可以基于 ZenML 的BaseModelRegistry抽象基类用实现类 配置类 Flavor 类三件套开发属于自己的模型注册表 Flavor并通过 CLI 一键注册到 Stack 中。本文将结合 ZenML 开源仓库源码从核心概念、基类抽象、四步实现流程到 MLflow 参考实现与测试验证完整讲解如何开发一个自定义模型注册表让你既能照搬步骤落地也能理解底层抽象的设计意图。模型注册表在 ZenML 中的定位在深入自定义开发之前先明确模型注册表在 ZenML 生态中的角色。根据 模型注册表组件指南模型注册表是集中式存储解决方案用于跨开发与部署阶段管理与追踪 ML 模型通过记录版本、配置、指标等元数据帮助团队实现可复现性在 ZenML 中它是 Stack 组件支持对训练完成的模型进行检索、加载与部署并记录模型是在哪个 Pipeline 中训练、如何复现。ZenML 为模型注册表提供了一套统一抽象无论底层注册表工具是什么MLflow、SageMaker、自研系统等都能以一致的方式处理以下三个核心概念RegisteredModel模型的逻辑分组逻辑上的模型组作为某一模型所有版本的容器记录名称、描述、标签等信息。可以由用户创建也可以在注册新模型版本时由注册表自动创建。RegistryModelVersion模型的某个具体版本以唯一版本号标识包含名称、描述、标签、指标以及指向已记录模型产物的引用。在 ZenML 中它还携带 Pipeline 名称、Pipeline Run ID、Step 名称等溯源信息且每个版本都归属于某个 RegisteredModel。ModelVersionStage模型版本所处的生命周期状态取值为None、Staging、Production、Archived。例如一个版本先在Staging阶段测试通过后再迁移到Production阶段用于部署。这套概念正是自定义模型注册表 Flavor 要实现的契约后续所有抽象方法都围绕它们展开。基类抽象BaseModelRegistry 详解自定义模型注册表的起点是BaseModelRegistry抽象基类它位于 src/zenml/model_registries/base_model_registry.py。因为现实世界中的模型注册表形态各异有的基于数据库、有的基于对象存储、有的基于第三方平台基类刻意暴露了一个基础且通用的接口让子类按需实现。基类继承自StackComponent与ABC并通过config属性将配置暴露为BaseModelRegistryConfig类型class BaseModelRegistryConfig(StackComponentConfig): Base config for model registries. class BaseModelRegistry(StackComponent, ABC): Base class for all ZenML model registries. property def config(self) - BaseModelRegistryConfig: Returns the config of the model registry. return cast(BaseModelRegistryConfig, self._config)模型注册Model Registration方法这一组抽象方法管理注册模型这一顶层实体即 RegisteredModel 的生命周期抽象方法签名要点语义register_modelname, descriptionNone, metadataNone→RegisteredModel在注册表中注册一个模型若同名模型已存在应抛出EntityExistsErrordelete_modelname→None删除已注册模型模型不存在应抛出KeyErrorupdate_modelname, descriptionNone, metadataNone, remove_metadataNone→RegisteredModel更新模型描述与元数据支持按 key 移除元数据get_modelname→RegisteredModel获取单个注册模型list_modelsnameNone, metadataNone→List[RegisteredModel]按名称/元数据过滤列出全部注册模型模型版本Model Version方法这一组抽象方法管理某个注册模型下的具体版本即 RegistryModelVersion 的生命周期抽象方法签名要点语义register_model_versionname, versionNone, model_source_uriNone, descriptionNone, metadataNone, **kwargs→RegistryModelVersion为指定模型注册一个新版本model_source_uri指向模型产物来源delete_model_versionname, version→None删除指定版本update_model_versionname, version, descriptionNone, metadataNone, remove_metadataNone, stageNone→RegistryModelVersion更新版本描述、元数据并可迁移版本阶段stagelist_model_versionsnameNone, model_source_uriNone, metadataNone, stageNone, countNone, created_afterNone, created_beforeNone, order_by_dateNone, **kwargs→List[RegistryModelVersion]按多条件过滤列出版本支持时间范围、数量、排序get_model_versionname, version→RegistryModelVersion获取单个版本load_model_versionname, version, **kwargs→Any加载版本对应的模型对象用于推理或进一步实验get_model_uri_artifact_storemodel_version→str返回版本模型产物在 ZenML 工件存储Artifact Store中的正确 URI 格式值得注意的是get_latest_model_version并非抽象方法而是基类提供的通用实现它调用list_model_versions(namename, stagestage, order_by_datedesc, count1)并返回第一条结果未指定 stage 时返回全阶段最新版本指定 stage 时返回该阶段最新版本。这为所有子类免费提供了一个取最新版本的常用能力也示范了基类在抽象之外提供公共逻辑的模式。支撑数据模型为了统一跨注册表的返回结构基类文件还定义了三个 Pydantic 数据模型它们是自定义实现必须兼容的返回类型契约RegisteredModel字段为name必填、description、metadata是模型注册表顶层实体。RegistryModelVersion字段为version、model_source_uri模型产物 URI创建后不可更改改动需新建版本、model_format由注册表集成自动设置创建后不可更改、model_library如 TensorFlow、PyTorch 等、registered_model、description、created_at、last_updated_at、stage默认ModelVersionStage.NONE、metadata。文档注释明确指出所有模型注册表都必须用各自特定的字段扩展此类。ModelRegistryModelMetadata记录模型版本与 ZenML 的溯源关联字段包括zenml_version、zenml_run_name、zenml_pipeline_name、zenml_pipeline_uuid、zenml_pipeline_run_uuid、zenml_step_name、zenml_project通过extraallow的 Pydantic 配置额外传入的自定义键值对可通过custom_attributes属性获取成为模型历史与研发过程的记录载体。在自定义实现中你可以也应该定义继承自RegistryModelVersion、RegisteredModel的子类用于承载底层注册表的特有字段例如 MLflow 的 run_id、run_link 等再把它们作为方法的返回类型。四步构建自定义模型注册表 Flavor原文档给出了实现自定义 Flavor 的四个步骤下面结合源码逐一展开。第一步掌握核心概念自定义模型注册表建立在上述RegisteredModel、RegistryModelVersion、ModelVersionStage概念之上建议先通读 模型注册表组件指南 理解抽象语义再动手编码。第二步实现类继承 BaseModelRegistry创建一个类继承BaseModelRegistry并实现全部抽象方法。仓库测试文件 tests/integration/functional/cli/test_model_registry.py 提供了一个最小骨架示范from typing import Dict, List, Optional, Type from zenml.model_registries.base_model_registry import ( BaseModelRegistry, BaseModelRegistryFlavor, RegisteredModel, RegistryModelVersion, ) class MyModelRegistry(BaseModelRegistry): 自定义模型注册表实现。 def register_model( self, name: str, description: Optional[str] None, metadata: Optional[Dict[str, str]] None, ) - RegisteredModel: ... # 对接底层注册表 API def delete_model(self, name: str) - None: ... def update_model( self, name: str, description: Optional[str] None, metadata: Optional[Dict[str, str]] None, remove_metadata: Optional[List[str]] None, ) - RegisteredModel: ... def get_model(self, name: str) - RegisteredModel: ... def list_models( self, name: Optional[str] None, metadata: Optional[Dict[str, str]] None, ) - List[RegisteredModel]: ... def register_model_version( self, name: str, version: Optional[str] None, model_source_uri: Optional[str] None, description: Optional[str] None, metadata: Optional[ModelRegistryModelMetadata] None, **kwargs, ) - RegistryModelVersion: ... def delete_model_version(self, name: str, version: str) - None: ... def update_model_version( self, name: str, version: str, description: Optional[str] None, metadata: Optional[ModelRegistryModelMetadata] None, remove_metadata: Optional[List[str]] None, stage: Optional[ModelVersionStage] None, ) - RegistryModelVersion: ... def list_model_versions( self, name: Optional[str] None, model_source_uri: Optional[str] None, metadata: Optional[ModelRegistryModelMetadata] None, stage: Optional[ModelVersionStage] None, count: Optional[int] None, created_afterNone, created_beforeNone, order_by_date: Optional[str] None, **kwargs, ) - List[RegistryModelVersion]: ... def get_model_version(self, name: str, version: str) - RegistryModelVersion: ... def load_model_version(self, name: str, version: str, **kwargs): ... def get_model_uri_artifact_store(self, model_version) - str: ...实现建议可结合源码推导register_model_version需要校验model_source_uri是否提供MLflow 实现中缺失时直接抛ValueErrorget_model_uri_artifact_store的职责是确保返回的 URI 与当前使用的 Artifact Store 格式一致如果底层注册表无法提供该 URI应把模型保存到 ZenML 支持的其它 Artifact Store 并返回其 URI这是模型部署Serving组件依赖的关键路径。第三步配置类继承 BaseModelRegistryConfig创建一个ModelRegistryConfig类继承BaseModelRegistryConfig追加你需要的任意配置参数。由于配置对象底层是 Pydantic 对象你还可以在这里添加自定义校验器validator在用户注册/更新 Stack 组件时立即校验配置值的合法性。参考 MLflow 的实现 mlflow_model_registry_flavor.py一个空配置的基类继承同样合法from zenml.model_registries.base_model_registry import BaseModelRegistryConfig class MyModelRegistryConfig(BaseModelRegistryConfig): 我的模型注册表配置。 在此追加自定义字段例如 endpoint: Optional[str] None api_key: Optional[str] None 第四步Flavor 类继承 BaseModelRegistryFlavor将实现与配置粘合起来创建一个继承BaseModelRegistryFlavor的 Flavor 类并通过抽象属性name为 Flavor 命名。BaseModelRegistryFlavor同样定义在 base_model_registry.py已经固定了type为StackComponentType.MODEL_REGISTRY因此你只需提供name、config_class与implementation_classfrom typing import Type from zenml.model_registries.base_model_registry import BaseModelRegistryFlavor class MyModelRegistryFlavor(BaseModelRegistryFlavor): property def name(self) - str: return my_registry # 该名称将用于 --flavormy_registry property def config_class(self) - Type[MyModelRegistryConfig]: return MyModelRegistryConfig property def implementation_class(self) - Type[MyModelRegistry]: from my_package.model_registry import MyModelRegistry # 延迟导入 return MyModelRegistryMLflow 的官方实现中implementation_class使用延迟导入在属性方法内部from zenml.integrations.mlflow.model_registries import MLFlowModelRegistry这正是为了配合下面的设计原则——让 Flavor/Config 可以在不安装底层依赖的情况下被加载。通过 CLI 注册自定义 Flavor实现完成后使用以下 CLI 命令将 Flavor 注册到 ZenML注意原文档中的命令参数IMAGE-BUILDER-FLAVOR-SOURCE-PATH实为MODEL-REGISTRY-FLAVOR-SOURCE-PATH的笔误命令本身由 src/zenml/cli/model_registry.py 中的zenml model-registry子命令组提供zenml model-registry flavor register MODEL-REGISTRY-FLAVOR-SOURCE-PATH注册成功后即可像内置 Flavor 一样注册 Stack 组件并纳入 Stackzenml model-registry register my_registry --flavormy_registry zenml stack register my_stack -r my_registry ... --set设计原则Flavor / Config / Implementation 三态分离理解三个类在 ZenML 工作流中的上场时机是写出健壮自定义注册表的关键这也是原文档重点强调的注意事项MyModelRegistryFlavor在通过 CLI 创建自定义 Flavor 时被导入和使用。它是轻量的门面只携带元信息名称、类型、配置类、实现类。MyModelRegistryConfig在用户使用该 Flavor 注册/更新 Stack 组件时被导入。注册过程中该配置会被用来校验用户提供的值由于配置是 Pydantic 对象可内置自定义校验器。MyModelRegistry实现类只有在组件最终被真正使用时才被加载Step 中实际调用注册表 API 时。这种设计将 Flavor 的配置与其实现解耦只要把MyModelRegistryFlavor与MyModelRegistryConfig放在与MyModelRegistry不同的模块/路径中即使本地没有安装实现背后的重型依赖如某云厂商 SDK也依然可以完成 Flavor 与组件的注册——这解释了为什么官方 MLflow 实现的implementation_class使用函数内延迟导入见 mlflow_model_registry_flavor.py。参考实现MLFlowModelRegistry 源码解析官方以 MLflow 集成作为完整参考实现位于 mlflow_model_registry.py其中值得借鉴的工程细节包括懒加载客户端mlflow_client属性在首次访问时通过configure_mlflow()复用当前活跃 Stack 中 MLflow Experiment Tracker 的配置Client().active_stack.experiment_tracker来初始化MlflowClient实现注册表与实验追踪器共享同一配置并在类级别缓存客户端实例。Stack 校验器通过validator属性返回StackValidator要求 Stack 中必须存在EXPERIMENT_TRACKER且其 flavor 必须是mlflow否则给出明确的修复命令提示zenml stack update model_registry -e mlflow_tracker——自定义实现可以照此模式声明自己的 Stack 前置条件。阶段映射_stage_to_alias方法将 ZenML 的ModelVersionStage映射为 MLflow 别名Staging→staging、Production→champion、Archived→archived在update_model_version中根据is_mlflow_3x()分支选择 MLflow 2.x 的传统 stage 迁移 API 或 3.x 的 alias API。类型转换_cast_mlflow_version_to_model_version把 MLflow 原生ModelVersion转成 ZenML 的RegistryModelVersion同时把run_id、run_link、别名等附加信息注入ModelRegistryModelMetadata并在读取阶段对损坏的版本做容错记录 warning 后跳过。URI 处理get_model_uri_artifact_store基于当前 Artifact Store 路径拼接出.../mlflow/下的模型 URI_remove_file_scheme则处理file://scheme 与 Windows/网络路径的差异保证 URI 格式与 Artifact Store 匹配。此外MLflow 注册表与实验追踪器的联动约束、mlflow_register_model_step内置 Step 的用法以及zenml model-registry models系列 CLI 交互命令list、list-versions、get-version、register-version、delete、delete-version详见 MLflow 模型注册表文档 与 CLI 实现。验证你的自定义注册表开发完成后可从两个层面验证单元/集成测试层面仓库的 test_model_registry.py 展示了测试模式——定义继承BaseModelRegistryFlavor的ConcreteModelDeployerFlavor与继承BaseModelRegistry的ConcreteModelRegistrymock 掉各抽象方法再用 CLI 测试框架驱动zenml model-registry命令组验证组件注册、命令分发与返回结构。你也可以仿照该文件为自己的 Flavor 编写等效测试。命令行交互层面将组件接入 Stack 后通过zenml model-registry models子命令验证注册、列出与查询流程。例如# 列出注册表中全部模型 zenml model-registry models list # 列出某模型的所有版本 zenml model-registry models list-versions my-model # 查看某版本详情 zenml model-registry models get-version my-model -v 1 # 手动注册一个模型版本附带 ZenML 溯源元数据 zenml model-registry models register-version my-model \ --descriptionA new version with accuracy 98.88% \ -v 1 \ --model-urifile:///.../artifacts/model \ -m key1 value1 -m key2 value2 \ --zenml-pipeline-nametraining_pipeline \ --zenml-step-nametrainerregister-version命令见 model_registry.py内部会将--zenml-*选项与-m元数据一并封装进ModelRegistryModelMetadata再调用你实现类的register_model_version——这也是验证你的实现是否正确消费元数据契约的最直接途径。注意事项与已知限制最后有几条原文档强调的边界需要开发者留意基础抽象仍在演进模型注册表是 ZenML 中相对较新的 Stack 组件虽然功能完整但 ML 系统的模型版本管理方式千差万别API 未来可能发生变化。文档页面会随最新改动持续更新。依赖关系模型注册表通常与实验追踪器Experiment Tracker绑定——注册表本身是可选的 Stack 组件但若 Stack 中没有实验追踪器模型版本无法自动携带训练运行的溯源元数据。抽象不足时请反馈如果你发现BaseModelRegistry抽象缺乏灵活性、无法覆盖你的场景官方欢迎通过反馈渠道提交 Issue帮助推动基类演进。动手之前建议先阅读 如何实现自定义 Stack 组件 的通用指南它阐述了 ZenML 组件 Flavor 概念的基础与本文的模型注册表专项内容互为补充。【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 ZenML 自定义模型注册表（Custom Model Registry）开发指南：从基类抽象到自定义 Flavor 落地
ZenML 自定义模型注册表Custom Model Registry开发指南从基类抽象到自定义 Flavor 落地【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml模型注册表Model Registry是 ZenML 中负责集中管理、追踪与版本化机器学习模型的 Stack 组件。当内置的 MLflow 等集成无法满足团队自定义的模型版本管理需求时你可以基于 ZenML 的BaseModelRegistry抽象基类用实现类 配置类 Flavor 类三件套开发属于自己的模型注册表 Flavor并通过 CLI 一键注册到 Stack 中。本文将结合 ZenML 开源仓库源码从核心概念、基类抽象、四步实现流程到 MLflow 参考实现与测试验证完整讲解如何开发一个自定义模型注册表让你既能照搬步骤落地也能理解底层抽象的设计意图。模型注册表在 ZenML 中的定位在深入自定义开发之前先明确模型注册表在 ZenML 生态中的角色。根据 模型注册表组件指南模型注册表是集中式存储解决方案用于跨开发与部署阶段管理与追踪 ML 模型通过记录版本、配置、指标等元数据帮助团队实现可复现性在 ZenML 中它是 Stack 组件支持对训练完成的模型进行检索、加载与部署并记录模型是在哪个 Pipeline 中训练、如何复现。ZenML 为模型注册表提供了一套统一抽象无论底层注册表工具是什么MLflow、SageMaker、自研系统等都能以一致的方式处理以下三个核心概念RegisteredModel模型的逻辑分组逻辑上的模型组作为某一模型所有版本的容器记录名称、描述、标签等信息。可以由用户创建也可以在注册新模型版本时由注册表自动创建。RegistryModelVersion模型的某个具体版本以唯一版本号标识包含名称、描述、标签、指标以及指向已记录模型产物的引用。在 ZenML 中它还携带 Pipeline 名称、Pipeline Run ID、Step 名称等溯源信息且每个版本都归属于某个 RegisteredModel。ModelVersionStage模型版本所处的生命周期状态取值为None、Staging、Production、Archived。例如一个版本先在Staging阶段测试通过后再迁移到Production阶段用于部署。这套概念正是自定义模型注册表 Flavor 要实现的契约后续所有抽象方法都围绕它们展开。基类抽象BaseModelRegistry 详解自定义模型注册表的起点是BaseModelRegistry抽象基类它位于 src/zenml/model_registries/base_model_registry.py。因为现实世界中的模型注册表形态各异有的基于数据库、有的基于对象存储、有的基于第三方平台基类刻意暴露了一个基础且通用的接口让子类按需实现。基类继承自StackComponent与ABC并通过config属性将配置暴露为BaseModelRegistryConfig类型class BaseModelRegistryConfig(StackComponentConfig): Base config for model registries. class BaseModelRegistry(StackComponent, ABC): Base class for all ZenML model registries. property def config(self) - BaseModelRegistryConfig: Returns the config of the model registry. return cast(BaseModelRegistryConfig, self._config)模型注册Model Registration方法这一组抽象方法管理注册模型这一顶层实体即 RegisteredModel 的生命周期抽象方法签名要点语义register_modelname, descriptionNone, metadataNone→RegisteredModel在注册表中注册一个模型若同名模型已存在应抛出EntityExistsErrordelete_modelname→None删除已注册模型模型不存在应抛出KeyErrorupdate_modelname, descriptionNone, metadataNone, remove_metadataNone→RegisteredModel更新模型描述与元数据支持按 key 移除元数据get_modelname→RegisteredModel获取单个注册模型list_modelsnameNone, metadataNone→List[RegisteredModel]按名称/元数据过滤列出全部注册模型模型版本Model Version方法这一组抽象方法管理某个注册模型下的具体版本即 RegistryModelVersion 的生命周期抽象方法签名要点语义register_model_versionname, versionNone, model_source_uriNone, descriptionNone, metadataNone, **kwargs→RegistryModelVersion为指定模型注册一个新版本model_source_uri指向模型产物来源delete_model_versionname, version→None删除指定版本update_model_versionname, version, descriptionNone, metadataNone, remove_metadataNone, stageNone→RegistryModelVersion更新版本描述、元数据并可迁移版本阶段stagelist_model_versionsnameNone, model_source_uriNone, metadataNone, stageNone, countNone, created_afterNone, created_beforeNone, order_by_dateNone, **kwargs→List[RegistryModelVersion]按多条件过滤列出版本支持时间范围、数量、排序get_model_versionname, version→RegistryModelVersion获取单个版本load_model_versionname, version, **kwargs→Any加载版本对应的模型对象用于推理或进一步实验get_model_uri_artifact_storemodel_version→str返回版本模型产物在 ZenML 工件存储Artifact Store中的正确 URI 格式值得注意的是get_latest_model_version并非抽象方法而是基类提供的通用实现它调用list_model_versions(namename, stagestage, order_by_datedesc, count1)并返回第一条结果未指定 stage 时返回全阶段最新版本指定 stage 时返回该阶段最新版本。这为所有子类免费提供了一个取最新版本的常用能力也示范了基类在抽象之外提供公共逻辑的模式。支撑数据模型为了统一跨注册表的返回结构基类文件还定义了三个 Pydantic 数据模型它们是自定义实现必须兼容的返回类型契约RegisteredModel字段为name必填、description、metadata是模型注册表顶层实体。RegistryModelVersion字段为version、model_source_uri模型产物 URI创建后不可更改改动需新建版本、model_format由注册表集成自动设置创建后不可更改、model_library如 TensorFlow、PyTorch 等、registered_model、description、created_at、last_updated_at、stage默认ModelVersionStage.NONE、metadata。文档注释明确指出所有模型注册表都必须用各自特定的字段扩展此类。ModelRegistryModelMetadata记录模型版本与 ZenML 的溯源关联字段包括zenml_version、zenml_run_name、zenml_pipeline_name、zenml_pipeline_uuid、zenml_pipeline_run_uuid、zenml_step_name、zenml_project通过extraallow的 Pydantic 配置额外传入的自定义键值对可通过custom_attributes属性获取成为模型历史与研发过程的记录载体。在自定义实现中你可以也应该定义继承自RegistryModelVersion、RegisteredModel的子类用于承载底层注册表的特有字段例如 MLflow 的 run_id、run_link 等再把它们作为方法的返回类型。四步构建自定义模型注册表 Flavor原文档给出了实现自定义 Flavor 的四个步骤下面结合源码逐一展开。第一步掌握核心概念自定义模型注册表建立在上述RegisteredModel、RegistryModelVersion、ModelVersionStage概念之上建议先通读 模型注册表组件指南 理解抽象语义再动手编码。第二步实现类继承 BaseModelRegistry创建一个类继承BaseModelRegistry并实现全部抽象方法。仓库测试文件 tests/integration/functional/cli/test_model_registry.py 提供了一个最小骨架示范from typing import Dict, List, Optional, Type from zenml.model_registries.base_model_registry import ( BaseModelRegistry, BaseModelRegistryFlavor, RegisteredModel, RegistryModelVersion, ) class MyModelRegistry(BaseModelRegistry): 自定义模型注册表实现。 def register_model( self, name: str, description: Optional[str] None, metadata: Optional[Dict[str, str]] None, ) - RegisteredModel: ... # 对接底层注册表 API def delete_model(self, name: str) - None: ... def update_model( self, name: str, description: Optional[str] None, metadata: Optional[Dict[str, str]] None, remove_metadata: Optional[List[str]] None, ) - RegisteredModel: ... def get_model(self, name: str) - RegisteredModel: ... def list_models( self, name: Optional[str] None, metadata: Optional[Dict[str, str]] None, ) - List[RegisteredModel]: ... def register_model_version( self, name: str, version: Optional[str] None, model_source_uri: Optional[str] None, description: Optional[str] None, metadata: Optional[ModelRegistryModelMetadata] None, **kwargs, ) - RegistryModelVersion: ... def delete_model_version(self, name: str, version: str) - None: ... def update_model_version( self, name: str, version: str, description: Optional[str] None, metadata: Optional[ModelRegistryModelMetadata] None, remove_metadata: Optional[List[str]] None, stage: Optional[ModelVersionStage] None, ) - RegistryModelVersion: ... def list_model_versions( self, name: Optional[str] None, model_source_uri: Optional[str] None, metadata: Optional[ModelRegistryModelMetadata] None, stage: Optional[ModelVersionStage] None, count: Optional[int] None, created_afterNone, created_beforeNone, order_by_date: Optional[str] None, **kwargs, ) - List[RegistryModelVersion]: ... def get_model_version(self, name: str, version: str) - RegistryModelVersion: ... def load_model_version(self, name: str, version: str, **kwargs): ... def get_model_uri_artifact_store(self, model_version) - str: ...实现建议可结合源码推导register_model_version需要校验model_source_uri是否提供MLflow 实现中缺失时直接抛ValueErrorget_model_uri_artifact_store的职责是确保返回的 URI 与当前使用的 Artifact Store 格式一致如果底层注册表无法提供该 URI应把模型保存到 ZenML 支持的其它 Artifact Store 并返回其 URI这是模型部署Serving组件依赖的关键路径。第三步配置类继承 BaseModelRegistryConfig创建一个ModelRegistryConfig类继承BaseModelRegistryConfig追加你需要的任意配置参数。由于配置对象底层是 Pydantic 对象你还可以在这里添加自定义校验器validator在用户注册/更新 Stack 组件时立即校验配置值的合法性。参考 MLflow 的实现 mlflow_model_registry_flavor.py一个空配置的基类继承同样合法from zenml.model_registries.base_model_registry import BaseModelRegistryConfig class MyModelRegistryConfig(BaseModelRegistryConfig): 我的模型注册表配置。 在此追加自定义字段例如 endpoint: Optional[str] None api_key: Optional[str] None 第四步Flavor 类继承 BaseModelRegistryFlavor将实现与配置粘合起来创建一个继承BaseModelRegistryFlavor的 Flavor 类并通过抽象属性name为 Flavor 命名。BaseModelRegistryFlavor同样定义在 base_model_registry.py已经固定了type为StackComponentType.MODEL_REGISTRY因此你只需提供name、config_class与implementation_classfrom typing import Type from zenml.model_registries.base_model_registry import BaseModelRegistryFlavor class MyModelRegistryFlavor(BaseModelRegistryFlavor): property def name(self) - str: return my_registry # 该名称将用于 --flavormy_registry property def config_class(self) - Type[MyModelRegistryConfig]: return MyModelRegistryConfig property def implementation_class(self) - Type[MyModelRegistry]: from my_package.model_registry import MyModelRegistry # 延迟导入 return MyModelRegistryMLflow 的官方实现中implementation_class使用延迟导入在属性方法内部from zenml.integrations.mlflow.model_registries import MLFlowModelRegistry这正是为了配合下面的设计原则——让 Flavor/Config 可以在不安装底层依赖的情况下被加载。通过 CLI 注册自定义 Flavor实现完成后使用以下 CLI 命令将 Flavor 注册到 ZenML注意原文档中的命令参数IMAGE-BUILDER-FLAVOR-SOURCE-PATH实为MODEL-REGISTRY-FLAVOR-SOURCE-PATH的笔误命令本身由 src/zenml/cli/model_registry.py 中的zenml model-registry子命令组提供zenml model-registry flavor register MODEL-REGISTRY-FLAVOR-SOURCE-PATH注册成功后即可像内置 Flavor 一样注册 Stack 组件并纳入 Stackzenml model-registry register my_registry --flavormy_registry zenml stack register my_stack -r my_registry ... --set设计原则Flavor / Config / Implementation 三态分离理解三个类在 ZenML 工作流中的上场时机是写出健壮自定义注册表的关键这也是原文档重点强调的注意事项MyModelRegistryFlavor在通过 CLI 创建自定义 Flavor 时被导入和使用。它是轻量的门面只携带元信息名称、类型、配置类、实现类。MyModelRegistryConfig在用户使用该 Flavor 注册/更新 Stack 组件时被导入。注册过程中该配置会被用来校验用户提供的值由于配置是 Pydantic 对象可内置自定义校验器。MyModelRegistry实现类只有在组件最终被真正使用时才被加载Step 中实际调用注册表 API 时。这种设计将 Flavor 的配置与其实现解耦只要把MyModelRegistryFlavor与MyModelRegistryConfig放在与MyModelRegistry不同的模块/路径中即使本地没有安装实现背后的重型依赖如某云厂商 SDK也依然可以完成 Flavor 与组件的注册——这解释了为什么官方 MLflow 实现的implementation_class使用函数内延迟导入见 mlflow_model_registry_flavor.py。参考实现MLFlowModelRegistry 源码解析官方以 MLflow 集成作为完整参考实现位于 mlflow_model_registry.py其中值得借鉴的工程细节包括懒加载客户端mlflow_client属性在首次访问时通过configure_mlflow()复用当前活跃 Stack 中 MLflow Experiment Tracker 的配置Client().active_stack.experiment_tracker来初始化MlflowClient实现注册表与实验追踪器共享同一配置并在类级别缓存客户端实例。Stack 校验器通过validator属性返回StackValidator要求 Stack 中必须存在EXPERIMENT_TRACKER且其 flavor 必须是mlflow否则给出明确的修复命令提示zenml stack update model_registry -e mlflow_tracker——自定义实现可以照此模式声明自己的 Stack 前置条件。阶段映射_stage_to_alias方法将 ZenML 的ModelVersionStage映射为 MLflow 别名Staging→staging、Production→champion、Archived→archived在update_model_version中根据is_mlflow_3x()分支选择 MLflow 2.x 的传统 stage 迁移 API 或 3.x 的 alias API。类型转换_cast_mlflow_version_to_model_version把 MLflow 原生ModelVersion转成 ZenML 的RegistryModelVersion同时把run_id、run_link、别名等附加信息注入ModelRegistryModelMetadata并在读取阶段对损坏的版本做容错记录 warning 后跳过。URI 处理get_model_uri_artifact_store基于当前 Artifact Store 路径拼接出.../mlflow/下的模型 URI_remove_file_scheme则处理file://scheme 与 Windows/网络路径的差异保证 URI 格式与 Artifact Store 匹配。此外MLflow 注册表与实验追踪器的联动约束、mlflow_register_model_step内置 Step 的用法以及zenml model-registry models系列 CLI 交互命令list、list-versions、get-version、register-version、delete、delete-version详见 MLflow 模型注册表文档 与 CLI 实现。验证你的自定义注册表开发完成后可从两个层面验证单元/集成测试层面仓库的 test_model_registry.py 展示了测试模式——定义继承BaseModelRegistryFlavor的ConcreteModelDeployerFlavor与继承BaseModelRegistry的ConcreteModelRegistrymock 掉各抽象方法再用 CLI 测试框架驱动zenml model-registry命令组验证组件注册、命令分发与返回结构。你也可以仿照该文件为自己的 Flavor 编写等效测试。命令行交互层面将组件接入 Stack 后通过zenml model-registry models子命令验证注册、列出与查询流程。例如# 列出注册表中全部模型 zenml model-registry models list # 列出某模型的所有版本 zenml model-registry models list-versions my-model # 查看某版本详情 zenml model-registry models get-version my-model -v 1 # 手动注册一个模型版本附带 ZenML 溯源元数据 zenml model-registry models register-version my-model \ --descriptionA new version with accuracy 98.88% \ -v 1 \ --model-urifile:///.../artifacts/model \ -m key1 value1 -m key2 value2 \ --zenml-pipeline-nametraining_pipeline \ --zenml-step-nametrainerregister-version命令见 model_registry.py内部会将--zenml-*选项与-m元数据一并封装进ModelRegistryModelMetadata再调用你实现类的register_model_version——这也是验证你的实现是否正确消费元数据契约的最直接途径。注意事项与已知限制最后有几条原文档强调的边界需要开发者留意基础抽象仍在演进模型注册表是 ZenML 中相对较新的 Stack 组件虽然功能完整但 ML 系统的模型版本管理方式千差万别API 未来可能发生变化。文档页面会随最新改动持续更新。依赖关系模型注册表通常与实验追踪器Experiment Tracker绑定——注册表本身是可选的 Stack 组件但若 Stack 中没有实验追踪器模型版本无法自动携带训练运行的溯源元数据。抽象不足时请反馈如果你发现BaseModelRegistry抽象缺乏灵活性、无法覆盖你的场景官方欢迎通过反馈渠道提交 Issue帮助推动基类演进。动手之前建议先阅读 如何实现自定义 Stack 组件 的通用指南它阐述了 ZenML 组件 Flavor 概念的基础与本文的模型注册表专项内容互为补充。【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考