ZenML 与 Weights & Biases 集成实战:实验追踪、运行元数据与 Weave LLM 可观测性全指南
【免费下载链接】zenmlZenML 🙏: One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml
ZenML 官方提供wandb实验追踪器 flavor,将 Weights & Biases(W&B)平台无缝接入你的 ZenML 流水线,用于在迭代式 ML 实验和自动化流水线运行阶段统一记录、可视化模型、指标与数据集。本文以 Weights & Biases 实验追踪器文档 为核心骨架,结合仓库源码(flavor 定义、tracker 实现、run 初始化逻辑)与 单元测试 深入展开,带你完整掌握:安装与认证、stack 注册、@step接入、WandbExperimentTrackerSettings全量配置、W&B 与 Weave 双平台追踪,以及可复制的端到端代码示例。
一、什么是 Weights & Biases Experiment Tracker
Weights & Biases Experiment Tracker 是 ZenML Experiment Tracker 组件体系中的一个 flavor,由 ZenML 的 W&B 集成(WANDBintegration)提供,借助 Weights & Biases 实验追踪平台来记录并可视化流水线步骤中产生的信息(模型、参数、指标等)。
W&B 是迭代式 ML 实验阶段非常流行的平台,用于追踪和可视化实验结果。这并不意味着它不能复用于自动化的流水线运行——恰恰相反,在从实验迈向生产化工作流的过渡阶段,你同样可以用它来追踪和展示流水线产出。
适用场景:什么时候该用它
官方文档明确了以下三类使用 W&B 实验追踪器的典型场景:
- 延续既有习惯:你已经在用 W&B 追踪项目实验,希望在引入 ZenML MLOps 工作流和最佳实践后继续沿用;
- 更直观的结果浏览:你希望以更强的交互式可视化方式浏览 ZenML 流水线运行的结果(模型、指标、数据集);
- 团队协作与对外共享:你希望将 ZenML 流水线记录的 artifacts 和指标连接到 W&B,与团队、组织或外部利益相关方共享。
反之,如果你从未接触过 W&B,且更偏好自己熟悉的实验追踪工具,则可以考虑使用 Experiment Tracker 的其他 flavor。
集成组件与依赖
在源码层面,W&B 集成由 src/zenml/integrations/wandb/init.py 定义:其WandbIntegration声明了集成名称WANDB、flavor 名称wandb,并声明运行依赖为wandb>=0.12.12,<1.0.0、Pillow>=9.1.0与weave>=0.51.33,<1.0.0。也就是说,安装该集成后,W&B SDK 与 Weave SDK 会一起就绪,为后文的 Weave 追踪能力打下基础。
二、部署与认证:安装集成并注册组件
安装 W&B 集成
W&B 实验追踪器 flavor 由 W&B ZenML 集成提供,你需要先在本地安装它,才能注册 W&B 实验追踪器并将其加入 stack:
zenml integration install wandb -y必需的认证凭据
连接 W&B 平台需要配置以下凭据(对应源码中WandbExperimentTrackerConfig的字段,见 flavor 文件):
| 配置项 | 类型 | 说明 |
|---|---|---|
api_key | 必填 | 你 W&B 账户的 API key 令牌,用于向配置的 entity 与 project 写日志,认证所必需 |
project_name | 可选 | 发送新 run 的项目名称;若不指定,run 会落入 "Uncategorized"(未分类)项目 |
entity | 可选 | 接收 run 的 entity(用户名或团队名)。该 entity 必须已存在,请在开始记录 run 前先在 UI 中创建账号或团队;不指定时,run 会发送到你的默认 entity(通常是你的用户名) |
其中api_key在源码中被声明为SecretField(SecretField),这是 ZenML 对敏感字段的标记,意味着它天然支持以密文形式存储和引用。
认证方式一:Basic Authentication(基础认证,不推荐用于生产)
该方式把凭据直接作为 stack 组件属性配置:
# 注册 Weights & Biases 实验追踪器 zenml experiment-tracker register wandb_experiment_tracker --flavor=wandb \ --entity=<entity> --project_name=<project_name> --api_key=<key> # 注册 stack 并加入新的实验追踪器 zenml stack register custom_stack -e wandb_experiment_tracker ... --set⚠️警告:这种方式不推荐用于生产环境,因为凭据不会安全存储,且会在 stack 配置中明文可见。
认证方式二:ZenML Secret(推荐)
该方法要求你先配置一个 ZenML secret,安全地存储 W&B 追踪服务凭据(如何配置 secret 可参考 ZenML 的 Secrets 交互文档,仓库中对应的功能源码位于 src/zenml/secret)。
使用zenml secret create命令创建 secret:
zenml secret create wandb_secret \ --entity=<ENTITY> \ --project_name=<PROJECT_NAME> --api_key=<API_KEY>创建完成后,用 secret 引用方式配置 wandb 实验追踪器:
# 在实验追踪器组件中引用 secret 里的 entity、project 和 api-key zenml experiment-tracker register wandb_tracker \ --flavor=wandb \ --entity={{wandb_secret.entity}} \ --project_name={{wandb_secret.project_name}} \ --api_key={{wandb_secret.api_key}} ...💡 关于 ZenML Secrets 的更多信息,可阅读 ZenML 文档中「与 Secrets 交互」的相关章节,以及仓库中 secret 模块 的源码实现。
三、在流水线步骤中使用 W&B:@step 接入
要在 step 中使用 active stack 中的 W&B 实验追踪器,需要用@step装饰器显式启用 experiment tracker,然后在 step 内部像平时一样使用 W&B 的日志或自动日志(auto-logging)能力。例如下面的 TensorFlow / Keras 训练 step:
import wandb from wandb.integration.keras import WandbCallback @step(experiment_tracker="<WANDB_TRACKER_STACK_COMPONENT_NAME>") def tf_trainer( config: TrainerConfig, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray, ) -> tf.keras.Model: ... model.fit( x_train, y_train, epochs=config.epochs, validation_data=(x_val, y_val), callbacks=[ WandbCallback( log_evaluation=True, validation_steps=16, validation_data=(x_val, y_val), ) ], ) metric = ... wandb.log({"<METRIC_NAME>": metric})💡动态获取追踪器:与其硬编码实验追踪器名称,不如使用 Client 动态获取 active stack 的实验追踪器:
from zenml.client import Client experiment_tracker = Client().active_stack.experiment_tracker @step(experiment_tracker=experiment_tracker.name) def tf_trainer(...): ...
底层执行链路:ZenML 如何替你管理 wandb.init
从源码实现看,每个使用 W&B 的 step 在运行前都会经历一次完整的生命周期管理(见 wandb_experiment_tracker.py):
prepare_step_run:step 运行前,ZenML 将api_key写入环境变量WANDB_API_KEY,读取 step 的 settings,调用build_wandb_initialization构造wandb.init(...)的全部参数,然后执行wandb.init(**init_kwargs)(若启用了 Weave 且配置了project_name,还会在此处调用weave.init(project_name=...));- step 执行:step 内部可以自由使用
wandb.log、W&B callback 或 Weave 装饰器; get_step_run_metadata:step 结束后,从当前wandb.run对象读取 run 的 id、path、url、name、group、project、entity、job_type 等信息,连同实验追踪器 URL 一起写入 step 元数据(元数据键experiment_tracker_url定义于 constants.py);cleanup_step_run:step 收尾时调用wandb.finish()(step 失败时传入exit_code=1),并清除WANDB_API_KEY环境变量。
这套生命周期保证每个 ZenML step 都能获得一个干净、隔离、可追踪的 W&B run。
四、W&B UI 与运行元数据
在 W&B UI 中查看实验
W&B 提供基于 Web 的 UI,可用于查看追踪实验的详细信息。每个使用 W&B 的 ZenML step 都会创建独立的实验 run,可在 W&B UI 中逐一检查。
通过 ZenML Client 获取 run 链接
你可以通过使用实验追踪器的那个 step 的元数据,找到与特定 ZenML run 关联的 W&B 实验 URL:
from zenml.client import Client last_run = client.get_pipeline("<PIPELINE_NAME>").last_run trainer_step = last_run.steps["<STEP_NAME>"] tracking_url = trainer_step.run_metadata["experiment_tracker_url"].value print(tracking_url)也可以在 ZenML dashboard 中,将该 URL 作为使用追踪器的 step 的元数据查看。此外,你还可以直接访问全部实验 run 的概览页:https://wandb.ai/{ENTITY_NAME}/{PROJECT_NAME}/runs/。
run 的命名与标记约定
💡 每个 W&B 实验 run 的命名约定为
{pipeline_run_name}_{step_name}(例如wandb_example_pipeline-25_Apr_22-20_06_33_535737_tf_evaluator)。默认情况下,ZenML 会按 ZenML 流水线 run 名称对 W&B run 分组,附加人类可读的标签(zenml、流水线名称、流水线 run 名称),并将持久化的 ZenML 标识符存入 W&B config。
从源码 run_initialization.py 可以看到这一约定的具体实现:
- 默认显示名:
f"{info.run_name}_{info.pipeline_step_name}"(_get_run_name); - 默认分组:
info.run_name(_get_group),即所有 step 共享同一流水线 run 名称作为 W&B group; - 默认标签:
["zenml", pipeline.name, run_name]追加到用户自定义标签之后(_get_tags)。
默认写入 W&B config 的 ZenML 元数据
当 ZenML 元数据启用时,W&B config 会收到扁平化键,例如:zenml_pipeline_name、zenml_pipeline_run_name、zenml_pipeline_run_id、zenml_step_name、zenml_latest_step_run_id、zenml_latest_step_run_version。当能从 active server 推导出 ZenML dashboard 链接时,还会额外加入zenml_pipeline_run_url。这些值写入的是 W&Bconfig而非summary,从而让指标摘要保持专注于实验结果本身。
这些键与测试用例 test_default_initialization_adds_zenml_metadata_without_run_id 中断言的 config 内容完全一致,是 ZenML 与 W&B 之间可验证的契约。
五、WandbExperimentTrackerSettings 进阶配置
对于更精细的控制,可以传入WandbExperimentTrackerSettings来覆盖wandb.Settings或为 run 附加额外标签。该 settings 类的全部字段定义于 flavor 文件:
| 字段 | 类型 | 默认值 | 说明 |
|---|---|---|---|
run_name | Optional[str] | None | 用于实验追踪的 W&B run 显示名 |
run_id | Optional[str] | None | 显式 W&B run ID,用于高级 resume 流程 |
run_id_strategy | 字面量 | "wandb_generated" | 派生 W&B run ID 的策略 |
resume | Optional[Literal["allow","must","never","auto"]] | None | W&B resume 策略;ZenML 设置 run ID 时默认allow |
group | Optional[str] | None | run 所属的 W&B group |
job_type | Optional[str] | None | run 的 W&B job type |
tags | List[str] | [] | 附加到 run 的标签 |
run_config | Dict[str, Any] | {} | 用户提供的 W&B config 值 |
init_kwargs | Dict[str, Any] | {} | 非 ZenML 管理的wandb.init关键字参数 |
settings | Dict[str, Any] | {} | 额外 W&B run 配置 |
enable_zenml_metadata | bool | True | 是否在 W&B tags/config 中加入 ZenML 上下文 |
enable_zenml_dashboard_links | bool | True | 可用时是否在 W&B config 中加入 ZenML dashboard 链接 |
enable_weave | bool | False | 是否启用 Weave 集成 |
基本用法:覆盖 settings 与添加标签
import wandb from zenml.integrations.wandb.flavors.wandb_experiment_tracker_flavor import WandbExperimentTrackerSettings wandb_settings = WandbExperimentTrackerSettings( settings=wandb.Settings(...), tags=["some_tag"], enable_weave=True, # 启用 Weave 集成 ) @step( experiment_tracker="<WANDB_TRACKER_STACK_COMPONENT_NAME>", settings={ "experiment_tracker": wandb_settings } ) def my_step( x_test: np.ndarray, y_test: np.ndarray, model: tf.keras.Model, ) -> float: """Everything in this step is auto-logged""" ...控制 W&B SDK 控制台输出
ZenML 默认关闭 W&B SDK 控制台输出,等价于wandb.Settings(console="off", silent=True)。原因是:W&B 会把正常的状态与进度消息写入 stderr,并用回车符重绘进度输出;而 ZenML 日志存储会将 stderr 记录为 error 级别日志,从而产生令人困惑的日志条目,甚至触发基于错误的日志告警——即便 W&B 工作正常。wandb.log(...)记录的指标仍然会正常发送到 W&B。
如需重新显示 W&B SDK 控制台输出,覆盖这些设置即可:
wandb_settings = WandbExperimentTrackerSettings( settings=wandb.Settings( console="auto", silent=False, ), )注意:重新启用控制台输出后,W&B 的登录、同步和进度消息会出现在 ZenML step 日志中;其中部分消息可能被记录为 error 级别,因为它们由 W&B SDK 通过 stderr 输出。
从源码看,默认值定义于 run_initialization.py 的DEFAULT_WANDB_SETTINGS = {"console": "off", "silent": True},而合并逻辑是{**DEFAULT_WANDB_SETTINGS, **settings.settings},因此你的settings可以逐键覆盖默认值。测试 test_wandb_console_output_can_be_enabled_with_settings 验证了console="auto"、silent=False等覆盖行为。
管理常见 wandb.init 字段
ZenML 会替你管理常见的wandb.init(...)字段。除 stack 组件上配置的 W&B entity 和 project 之外,WandbExperimentTrackerSettings还支持:
wandb_settings = WandbExperimentTrackerSettings( run_name="optional-display-name", group="optional-group", tags=["team-a"], job_type="training", run_config={"dataset": "customers-v2"}, )标签截断与去重
ZenML 会将 W&B 标签截断到 64 个字符,并在截断后去除重复,因此前 64 个字符相同的长标签会折叠为同一个 W&B 标签。该逻辑在源码sanitize_tag与_deduplicate_tags(run_initialization.py)中实现。
最小化模式:禁用 ZenML 元数据
若希望在无 ZenML 附加分组、标签和 config 的最小模式下使用 W&B,可禁用元数据:
wandb_settings = WandbExperimentTrackerSettings( enable_zenml_metadata=False, tags=["plain-wandb"], )测试 test_minimal_mode_omits_zenml_group_tags_and_config 验证了禁用后 group 不再注入、config 只保留用户自定义键。
run ID 策略与重试语义
默认情况下,ZenML 不设置 W&B run ID(run_id_strategy="wandb_generated"),保留 W&B 生成的 ID。支持的run_id_strategy取值如下:
| 策略 | 行为 |
|---|---|
wandb_generated | 让 W&B 生成 run ID。 |
reuse_on_retry | 每个 ZenML pipeline-step 调用使用一个确定性的 W&B run ID,因此同一调用的重试会续接同一个 W&B run。 |
new_on_retry | 每次 ZenML step-run 尝试(包括重试)都使用一个新的确定性 W&B run ID。 |
场景一:重试续接同一个 W&B run。以下配置让同一 ZenML step 调用的重试都挂在同一个 W&B run 上:
retry_collapsed_settings = WandbExperimentTrackerSettings( run_id_strategy="reuse_on_retry", resume="allow", )场景二:每次尝试独立 W&B run。如果你希望每次 step-run 尝试(包括重试)都对应一个 W&B run,使用:
attempt_settings = WandbExperimentTrackerSettings( run_id_strategy="new_on_retry", )当使用run_id_strategy="new_on_retry"且未显式指定run_name时,ZenML 会把 ZenML step-run version 追加到 W&B 显示名后面,便于在 W&B UI 中区分各次重试尝试。从源码看,默认显示名会变为f"{run_name}_v{info.step_run.version}";而 run ID 分别由uuid5基于zenml:wandb:step-run:{step_run_id}或zenml:wandb:pipeline-step:{run_id}:{pipeline_step_name}确定性派生(run_initialization.py)。相关测试 test_new_on_retry_strategy_sets_deterministic_run_id_and_resume 与 test_reuse_on_retry_strategy_collapses_retries 精确验证了这两种语义。
另外,只要 ZenML 传入确定性或显式的 W&B run ID,resume就默认设为allow(源码settings.resume or ("allow" if run_id else None))。
高级用法:挂接已有 run 与透传 init_kwargs
高级用户可以将日志挂接到已存在的 W&B run,或通过init_kwargs透传不受 ZenML 管理的 W&B 初始化字段:
wandb_settings = WandbExperimentTrackerSettings( run_id="existing-wandb-run-id", init_kwargs={"notes": "Backfill evaluation metrics."}, )测试 test_unmanaged_init_kwargs_pass_through 验证了mode="offline"这类字段会被原样透传给wandb.init。
冲突校验:ZenML 尽早报错
ZenML 会在 settings 校验阶段(model_validator)就拦截非法组合,而不是等到运行时失败:
run_id不能与非默认的run_id_strategy组合;resume="never"不能与run_id_strategy="reuse_on_retry"组合,因为重试本就要复用同一个 W&B run ID;resume="must"要求显式run_id或确定性的run_id_strategy;run_config不能包含zenml_*开头的键,这些键为 ZenML 元数据保留;init_kwargs不能覆盖 ZenML 管理的wandb.init(...)键,包括entity、project、name、id、resume、group、job_type、tags、config、settings(源码_MANAGED_INIT_KWARGS集合,见 flavor 文件)。
上述规则在 test_invalid_settings_fail_early 中以参数化用例逐一覆盖。
六、集成 Weights & Biases Weave:LLM 与模型可视化
Weights & Biases Weave 是一个可定制的 dashboard 界面,用于可视化并交互式操作机器学习模型、数据和结果。ZenML 通过WandbExperimentTrackerSettings提供对 Weave 的内置支持。
启用与禁用 Weave
你可以在流水线的特定 step 上通过配置WandbExperimentTrackerSettings中的enable_weave参数(或在注册实验追踪器组件时设置)来启用或禁用 Weave:
import weave from openai import OpenAI from zenml import pipeline, step from zenml.integrations.wandb.flavors.wandb_experiment_tracker_flavor import ( WandbExperimentTrackerSettings, ) # 启用 Weave 的 settings wandb_with_weave_settings = WandbExperimentTrackerSettings( tags=["weave_enabled"], enable_weave=True, # 启用 Weave 集成 ) # 禁用 Weave 的 settings wandb_without_weave_settings = WandbExperimentTrackerSettings( tags=["weave_disabled"], enable_weave=False, # 显式禁用 Weave 集成 )在 ZenML step 中使用 Weave
使用 Weave 需要三步:
- 配置
WandbExperimentTrackerSettings并设置enable_weave=True; - 给 step 函数加上
@weave.op()装饰器; - 配置 step 使用带有 Weave settings 的 W&B 实验追踪器。
示例:
@step( experiment_tracker="wandb_weave", # 你的 W&B 实验追踪器组件名 settings={"experiment_tracker": wandb_with_weave_settings}, ) @weave.op() # Weave 装饰器 def my_step_with_weave() -> str: """This step will use Weave for enhanced visualization""" # 你的 step 实现 return "Step with Weave enabled"⚠️重要:装饰器顺序至关重要。
@weave.op()必须应用在@step装饰器之后(即更靠近函数定义)。如果顺序颠倒,step 将无法正常工作。# 正确顺序 @step(experiment_tracker="wandb_weave") @weave.op() def correct_order_step(): ... # 错误顺序 - 会引发问题 @weave.op() @step(experiment_tracker="wandb_weave") def incorrect_order_step(): ...
若要针对特定 step 显式禁用 Weave,同时保留使用@weave.op()装饰器的能力:
@step( experiment_tracker="wandb_weave", settings={"experiment_tracker": wandb_without_weave_settings}, ) @weave.op() def my_step_without_weave() -> str: """This step will not use Weave even with the @weave.op() decorator""" # 你的 step 实现 return "Step with Weave disabled"Weave 初始化行为
结合源码(wandb_experiment_tracker.py 中的_initialize_wandb)与官方文档,可以梳理出 Weave 的完整初始化规则:
- 若
enable_weave=True且 W&B 实验追踪器中指定了project_name,Weave 将使用该 project 名初始化; - 若
enable_weave=True但未指定project_name,跳过 Weave 初始化; - 若
enable_weave=False且指定了project_name(显式禁用),Weave 将以settings={"disabled": True}被禁用; - 若
enable_weave=False且未指定project_name,跳过 Weave 禁用操作。
💡 关于 Weights & Biases Weave 的能力详情,可参考 Weave 官方文档(
docs.wandb.ai/weave)。
七、端到端完整代码示例
以下两个示例演示了 ZenML W&B 集成的完整运行流程。
示例一:Transformers 微调(不使用 Weave)
该示例在 IMDB 数据集上微调 DistilBERT 做情感分类,通过 Hugging FaceTrainer的report_to=["wandb"]自动记录训练指标,并在 step 结束前用wandb.log记录最终评估结果:
from typing import Tuple from zenml import pipeline, step from zenml.client import Client from zenml.integrations.wandb.flavors.wandb_experiment_tracker_flavor import ( WandbExperimentTrackerSettings, ) from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, DistilBertForSequenceClassification, ) from datasets import load_dataset, Dataset import numpy as np from sklearn.metrics import accuracy_score, precision_recall_fscore_support import wandb # 从 active stack 获取实验追踪器 experiment_tracker = Client().active_stack.experiment_tracker @step def prepare_data() -> Tuple[Dataset, Dataset]: dataset = load_dataset("imdb") tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased") def tokenize_function(examples): return tokenizer(examples["text"], padding="max_length", truncation=True) tokenized_datasets = dataset.map(tokenize_function, batched=True) return ( tokenized_datasets["train"].shuffle(seed=42).select(range(1000)), tokenized_datasets["test"].shuffle(seed=42).select(range(100)), ) # 训练模型 @step(experiment_tracker=experiment_tracker.name) def train_model( train_dataset: Dataset, eval_dataset: Dataset ) -> DistilBertForSequenceClassification: model = AutoModelForSequenceClassification.from_pretrained( "distilbert-base-uncased", num_labels=2 ) training_args = TrainingArguments( output_dir="./results", num_train_epochs=3, per_device_train_batch_size=16, per_device_eval_batch_size=16, warmup_steps=500, weight_decay=0.01, logging_dir="./logs", evaluation_strategy="epoch", logging_steps=100, report_to=["wandb"], ) def compute_metrics(eval_pred): logits, labels = eval_pred predictions = np.argmax(logits, axis=-1) precision, recall, f1, _ = precision_recall_fscore_support( labels, predictions, average="binary" ) acc = accuracy_score(labels, predictions) return {"accuracy": acc, "f1": f1, "precision": precision, "recall": recall} trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, compute_metrics=compute_metrics, ) trainer.train() # 评估模型 eval_results = trainer.evaluate() print(f"Evaluation results: {eval_results}") # 记录最终评估结果 wandb.log({"final_evaluation": eval_results}) return model @pipeline(enable_cache=False) def fine_tuning_pipeline(): train_dataset, eval_dataset = prepare_data() model = train_model(train_dataset, eval_dataset) if __name__ == "__main__": # 运行流水线 wandb_settings = WandbExperimentTrackerSettings( tags=["distilbert", "imdb", "sentiment-analysis"], ) fine_tuning_pipeline.with_options(settings={"experiment_tracker": wandb_settings})()示例二:LLM 情感分析 + Weave 追踪
该示例展示如何用@weave.op()追踪 LLM 调用链路,同时用wandb.log记录结果统计与样本表格:
import weave from openai import OpenAI import numpy as np from sklearn.metrics import accuracy_score import pandas as pd from zenml import pipeline, step from zenml.client import Client from zenml.integrations.wandb.flavors.wandb_experiment_tracker_flavor import ( WandbExperimentTrackerSettings, ) # 从 active stack 获取实验追踪器 experiment_tracker = Client().active_stack.experiment_tracker # 创建启用 Weave 追踪的 settings weave_settings = WandbExperimentTrackerSettings( tags=["weave_example", "llm_pipeline"], enable_weave=True, ) # LLM 调用使用的 OpenAI client openai_client = OpenAI() @step def prepare_data() -> pd.DataFrame: """为 LLM 处理准备示例数据""" data = { "id": range(10), "text": [ "I love this product, it's amazing!", "This was a waste of money, terrible.", "Pretty good, but could be improved.", "Not worth the price, disappointed.", "Absolutely fantastic experience!", "It's okay, nothing special though.", "Would definitely recommend to others.", "Had some issues, but support was helpful.", "Don't buy this, it doesn't work properly.", "Perfect for my needs, very satisfied." ] } return pd.DataFrame(data) @step( experiment_tracker=experiment_tracker.name, settings={"experiment_tracker": weave_settings}, ) @weave.op() # Weave 装饰器置于 step 装饰器之后 def classify_sentiment(data: pd.DataFrame) -> pd.DataFrame: """使用 LLM 对每段文本进行情感分类""" results = [] for _, row in data.iterrows(): prompt = f"Classify the sentiment of this text as POSITIVE, NEGATIVE, or NEUTRAL: '{row['text']}'" response = openai_client.chat.completions.create( model="gpt-4o-mini", messages=[{"role": "user", "content": prompt}], temperature=0.3, ) sentiment = response.choices[0].message.content.strip() results.append({ "id": row["id"], "text": row["text"], "sentiment": sentiment, }) # 将结果转为 DataFrame result_df = pd.DataFrame(results) # 记录一些指标到 Wandb sentiments = result_df["sentiment"].value_counts() import wandb wandb.log({ "positive_count": sentiments.get("POSITIVE", 0), "negative_count": sentiments.get("NEGATIVE", 0), "neutral_count": sentiments.get("NEUTRAL", 0), "sample_data": wandb.Table(dataframe=result_df), }) return result_df @pipeline(enable_cache=False) def sentiment_analysis_pipeline(): """带 Weave 追踪的情感分析流水线""" data = prepare_data() results = classify_sentiment(data) if __name__ == "__main__": # 设置 pipeline 级别 settings pipeline_settings = { "experiment_tracker": WandbExperimentTrackerSettings( tags=["sentiment_analysis_pipeline"], enable_weave=True, ) } # 使用 settings 运行流水线 sentiment_analysis_pipeline.with_options(settings=pipeline_settings)()八、深入源码:W&B run 初始化是如何构造的
如果想彻底理解 ZenML 与 W&B 的协作机制,build_wandb_initialization是核心入口(run_initialization.py)。它按固定顺序组装wandb.init的参数:
- 以
settings.init_kwargs为基础字典; - 统一注入
entity(来自组件 config)、project(来自组件 config 的project_name)、name(_get_run_name派生)、tags(_get_tags派生)、settings(DEFAULT_WANDB_SETTINGS与用户 settings 合并); - 根据
run_id_strategy派生确定性 run ID(_get_run_id),有 ID 时注入id与resume; - 注入
group(默认取流水线 run 名称)、job_type(仅当配置时); - 注入
config:用户run_config与 ZenML 元数据合并(_get_run_config),可选的 dashboard 链接(_get_dashboard_links)仅在能从 active server 推导时加入。
对应地,wandb_experiment_tracker.py 中的get_step_run_metadata负责把 W&B run 的持久化标识(run id、run path、run url、entity、project、group、job_type 等)回写到 ZenML step 元数据,使「ZenML run → W&B run」的双向追溯成为可能——这正是第四节中通过run_metadata["experiment_tracker_url"]取回追踪 URL 的底层来源。
仓库中的集成测试 test_run_initialization.py 覆盖了默认初始化、控制台输出覆盖、标签合并、最小化模式、两种重试策略、显式 run ID、job type 透传、dashboard 链接、step 元数据以及冲突校验等十余个场景,是理解全部配置行为的权威参考。
九、总结
至此,你已经掌握了在 ZenML 中使用 Weights & Biases 的完整链路:
- 接入:
zenml integration install wandb -y安装集成,通过 Basic 认证或(推荐)ZenML Secret 认证注册wandb实验追踪器并加入 stack; - 使用:在
@step(experiment_tracker=...)中启用追踪器,step 内自由使用wandb.log、KerasWandbCallback或 Hugging Facereport_to=["wandb"]等能力; - 追踪与回溯:通过 ZenML Client 从 step 元数据取回
experiment_tracker_url,在 W&B UI 中按{pipeline_run_name}_{step_name}命名约定定位每个 run; - 精细控制:通过
WandbExperimentTrackerSettings管理 run 名称、分组、标签、job type、config、控制台输出、ZenML 元数据、run ID 策略(wandb_generated/reuse_on_retry/new_on_retry)与重试语义,并获得提前冲突校验; - LLM 可观测性:通过
enable_weave=True与@weave.op()装饰器(注意必须置于@step之下),为 LLM 流水线开启 Weave 追踪。
相关源码与测试:集成定义 src/zenml/integrations/wandb/init.py、settings 与校验 flavors/wandb_experiment_tracker_flavor.py、run 生命周期 experiment_trackers/wandb_experiment_tracker.py、run 初始化构造 experiment_trackers/run_initialization.py、行为验证 tests/integration/integrations/wandb_tests/test_run_initialization.py。
【免费下载链接】zenmlZenML 🙏: One AI Platform from Pipelines to Agents. https://zenml.io.项目地址: https://gitcode.com/GitHub_Trending/ze/zenml
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考