开发 ZenML 自定义 Experiment Tracker Flavor从基类继承到 CLI 注册的完整实战【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenmlExperiment Tracker实验跟踪器是 ZenML 栈中负责记录、可视化和对比 ML 实验结果的组件类型本文基于 ZenML 官方文档 custom.md 并结合仓库源码系统讲解如何从零开发一个自定义 Experiment Tracker flavor包括三大基类抽象、三步式实现流程、通过 CLI 注册与验证的方法以及 flavor/config/implementation 三阶段加载机制背后的设计原理。读完本文你将能够独立实现、注册并在 ZenML 栈中使用属于自己的实验跟踪器 flavor。前置认知Experiment Tracker 在 ZenML 中的定位在动手写代码之前先明确 Experiment Tracker 在整个 ZenML 体系中的位置。根据 Experiment Trackers 总览文档Experiment Tracker 是一种可选的 Stack Component栈组件需要注册到你的 ZenML Stack 中才能生效ZenML 已经通过强制性的 Artifact Store 为流水线产物提供版本化与追踪能力但这些机制以编程方式使用为主缺乏可视化界面Experiment Tracker 的价值在于其丰富的 UI在 ZenML 中每一次流水线运行pipeline run都被视为一次实验Experiment Tracker 负责把模型、数据集、指标、参数等信息写入实验跟踪工具从而建立流水线运行与实验之间的清晰关联在 ZenML 中Experiment Tracker 与 Artifact Store 并存但职责不同前者面向可交互的浏览与对比后者面向产物的持久化存储。内置的 Experiment Tracker 均以集成integration的形式提供包括comet、mlflow、neptune、wandb、trackio等 flavor。当这些现成 flavor 无法满足你的需求时就需要按照本文的路径开发自定义 flavor。三大基类抽象Config、Implementation 与 Flavor自定义 Experiment Tracker 的核心是理解src/zenml/experiment_trackers/base_experiment_tracker.py中定义的三层抽象。与 ZenML 所有栈组件一样Experiment Tracker 遵循 配置与实现分离 的三类结构详见 custom-stack-component.md1.BaseExperimentTrackerConfig静态配置# src/zenml/experiment_trackers/base_experiment_tracker.py class BaseExperimentTrackerConfig(StackComponentConfig): Base config for experiment trackers.它继承自StackComponentConfigstack_component.py而StackComponentConfig本身是一个pydanticBaseModel。这意味着你在 Config 子类中声明的每一个字段都会自动获得类型校验与默认值能力由于 Config 不依赖具体实现ZenML 可以在不安装实验跟踪工具依赖的情况下完成栈组件的注册与校验你可以在 Config 中通过 pydantic 的model_validator编写自定义校验逻辑例如校验字段之间的组合关系是否合法参考 MLflow flavor 的校验器实现。2.BaseExperimentTracker核心实现# src/zenml/experiment_trackers/base_experiment_tracker.py class BaseExperimentTracker(StackComponent, ABC): Base class for all ZenML experiment trackers. property def config(self) - BaseExperimentTrackerConfig: return cast(BaseExperimentTrackerConfig, self._config)它继承自StackComponent抽象基类定义了实验跟踪器的公共接口。通过config属性实现类可以在运行时读取用户在注册组件时提供的全部配置值。3.BaseExperimentTrackerFlavor把两者绑定起来# src/zenml/experiment_trackers/base_experiment_tracker.py class BaseExperimentTrackerFlavor(Flavor): property def type(self) - StackComponentType: return StackComponentType.EXPERIMENT_TRACKER property def config_class(self) - Type[BaseExperimentTrackerConfig]: return BaseExperimentTrackerConfig property abstractmethod def implementation_class(self) - Type[StackComponent]: return BaseExperimentTrackerFlavor基类flavor.py要求子类必须实现name、type、implementation_class、config_class四个抽象属性。BaseExperimentTrackerFlavor已经替你固定了type StackComponentType.EXPERIMENT_TRACKER你只需提供 flavor 名称、具体 Config 类与具体实现类。三步实现自定义 Experiment Tracker Flavor根据官方文档的指引实现自定义 flavor 只需三个步骤第一步实现类继承BaseExperimentTracker创建一个类继承BaseExperimentTracker实现其中必要的抽象方法并在这里接入实验跟踪工具的真实 API。在实现中你可以通过self.config读取用户配置——这是 ZenML 约定俗成的访问方式。第二步配置类继承BaseExperimentTrackerConfigfrom typing import Optional from zenml.experiment_trackers import BaseExperimentTrackerConfig class MyExperimentTrackerConfig(BaseExperimentTrackerConfig): Configuration for the custom experiment tracker. api_endpoint: str https://api.my-tracker.example.com api_key: Optional[str] None # 敏感字段建议使用 SecretField在配置类中声明你的实验跟踪器所需的全部参数。两个实用技巧敏感字段用SecretField将api_key之类的字段声明为zenml.utils.secret_utils.SecretField类型可让用户以{{secret_name.key}}的引用形式传入密钥避免明文存储见 custom-stack-component.md 中的示例用 pydantic 校验器做组合校验由于 Config 本质是 pydantic 模型可以直接添加model_validator校验字段间的约束。内置 MLflow flavor 就是一个绝佳参照——它通过多个model_validator强制要求用户名与密码必须成对出现、Databricks 专属参数只能配合tracking_uridatabricks使用等规则源码位置。第三步Flavor 类继承BaseExperimentTrackerFlavorfrom typing import Type from zenml.experiment_trackers import ( BaseExperimentTracker, BaseExperimentTrackerConfig, BaseExperimentTrackerFlavor, ) class MyExperimentTracker(BaseExperimentTracker): ... class MyExperimentTrackerConfig(BaseExperimentTrackerConfig): ... class MyExperimentTrackerFlavor(BaseExperimentTrackerFlavor): property def name(self) - str: return my_tracker # 全局唯一避免与现有 flavor 重名 property def config_class(self) - Type[MyExperimentTrackerConfig]: return MyExperimentTrackerConfig property def implementation_class(self) - Type[MyExperimentTracker]: return MyExperimentTracker注意三点name必须全局唯一。内置 flavor 列表comet、mlflow、neptune、wandb、trackio等见 Experiment Trackers 总览自定义 flavor 不要与它们冲突推荐将实现、配置、flavor 三个类拆分为不同模块并在implementation_class属性中才导入实现类。这样即便本地没有安装实验跟踪工具的重型依赖ZenML 也能正常加载 flavor 与 Config 完成组件注册详见下文三阶段加载机制如果你需要为 flavor 提供文档链接、SDK 文档链接或仪表盘 Logo可以覆写Flavor基类中的docs_url、sdk_docs_url、logo_url属性见 flavor.py 与 MLflow flavor 的对应实现。通过 CLI 注册自定义 Flavor完成三个类的实现后通过 ZenML CLI 注册 flavor。必须使用点号dot表示法指向 flavor 类zenml experiment-tracker flavor register path.to.MyExperimentTrackerFlavor例如若你的 flavor 类MyExperimentTrackerFlavor定义在flavors/my_flavor.py中则执行zenml experiment-tracker flavor register flavors.my_flavor.MyExperimentTrackerFlavor注册成功后即可在可用 flavor 列表中看到它zenml experiment-tracker flavor list注册命令的底层实现在 stack_components.py 中CLI 会调用Client().create_flavor(sourcesource, component_typecomponent_type)完成注册并对导入失败给出明确提示。而导入与校验逻辑位于 flavor.py 的validate_flavor_source它会依次检查给定 source 能否被source_utils.load成功导入导入的类是否真的是Flavor的子类flavor 实例的type是否等于EXPERIMENT_TRACKERimplementation_class是否继承自StackComponentconfig_class是否继承自StackComponentConfig。任何一环不满足注册都会失败并抛出明确错误。关于 source 解析根目录的注意事项ZenML 以执行zenml init的目录作为 flavor 类解析的起点也就是点号路径的根。官方文档强调请遵循在仓库根目录执行zenml init的最佳实践这样团队成员才能以一致的方式解析自定义 flavor若 ZenML 在任意父目录中都没有找到已初始化的 ZenML 仓库会默认回退到当前工作目录CWD作为解析根——但最好不要依赖这种兜底机制。这一行为在 CLI 源码中有直接体现注册时若client.root为空CLI 会打印警告当前工作目录将作为 source 的相对解析根并建议执行zenml initstack_components.py。理解三阶段加载机制为什么可以不装依赖先注册原文档特别强调了自定义 flavor 中三个类在 ZenML 工作流中被加载的时机完全不同这是理解整个设计的关键类何时被导入/使用CustomExperimentTrackerFlavor通过 CLI 创建自定义 flavor 时导入并实例化CustomExperimentTrackerConfig用户注册/更新使用该 flavor 的栈组件时导入注册过程中用它对用户输入做校验因其本质是 pydantic 对象可附带自定义校验器CustomExperimentTracker只有在组件真正被使用运行流水线时才登场这种设计把 flavor 的配置与其实现解耦只要 Flavor 类和 Config 类与实际的实现类放在不同的模块/路径中即便本地没有安装实现所需的重型依赖也能正常完成 flavor 注册与组件注册/校验直到流水线真正运行时才需要加载实现类。这就是 ZenML 栈组件体系可插拔、轻量注册的核心原因之一。注册组件与接入栈把自定义 flavor 用起来自定义 flavor 注册完成后它与内置 flavor 的使用方式完全一致zenml experiment-tracker register EXPERIMENT_TRACKER_NAME \ --flavormy_tracker \ --api_endpointhttps://api.my-tracker.example.com完整 CLI 命令族还包括flavor list、flavor describe FLAVOR_NAME、list、get、describe、update、rename、delete、connect等CLI 文档。其中zenml experiment-tracker flavor describe FLAVOR_NAME会基于 flavor 的 config schema 展示该 flavor 需要哪些配置参数是排查配置问题的高效工具。随后将该组件加入栈zenml stack register STACK_NAME \ --orchestrator ORCHESTRATOR_NAME \ --artifact-store ARTIFACT_STORE_NAME \ --experiment-tracker EXPERIMENT_TRACKER_NAME在流水线代码中通过step(experiment_trackerTrue)装饰器为单个 step 显式启用实验跟踪一个栈可以挂多个 experiment tracker若只想启用某一个可在装饰器中按名称显式指定。运行时ZenML 会自动在 step 失败时将对应实验标记为失败。流水线运行后可通过以下代码获取与该 step 关联的实验跟踪器 UI 地址参考 Experiment Trackers 总览from zenml.client import Client pipeline_run Client().get_pipeline_run(PIPELINE_RUN_NAME) step pipeline_run.steps[STEP_NAME] experiment_tracker_url step.run_metadata[experiment_tracker_url].value以 MLflow flavor 为参照一个工业级实现样本若想参考一个完整、规范的 flavor 实现仓库内置的 MLflow 集成是最好的样本Flavor 定义mlflow_experiment_tracker_flavor.py —— 演示了 Config 中如何用SecretField定义tracking_username、tracking_password、tracking_token等敏感字段如何通过is_local属性判断组件是否为本地运行以及 flavor 类如何覆写display_name、docs_url、sdk_docs_url、logo_url等展示属性实现类mlflow_experiment_tracker.py —— 演示了实现类如何初始化、如何通过环境变量向底层工具传递配置、以及如何参与 step 运行的生命周期管理。其他内置参考实现还包括comet、neptune、trackio、wandb与 GCP Vertex 的 experiment tracker flavor分布在src/zenml/integrations/集成名/flavors/与experiment_trackers/目录下它们共同展示了 ZenML 官方对 flavor 代码组织、依赖隔离与校验逻辑的工程实践。注意事项与最佳实践最后汇总官方文档与源码给出的关键注意点基类抽象仍在演进中当前BaseExperimentTracker基类抽象仍在积极开发中官方目前不推荐立即扩展该基类如果确实需要实现自定义 flavor需知悉后续基类发布时可能需要重构你的实现保持配置与实现解耦将 flavor、config、implementation 三个类放在不同模块implementation_class属性内部再导入实现类是保证轻依赖注册的前提zenml init位置要一致在仓库根目录.git所在层级初始化避免依赖当前工作目录的兜底机制name 唯一且语义清晰自定义 flavor 名需全局唯一避免与内置 flavor如mlflow、wandb冲突充分测试开发完成后先在本地流水线中完整验证 flavor 的注册、组件校验、运行跟踪与失败标记行为再投入生产环境。【免费下载链接】zenmlZenML : One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考