执行摘要
- 一句话:添加 Trackio 作为 rollout trace 后端
- 推荐动作:值得精读。该 PR 展示了如何以一致的设计模式扩展 verl 的追踪后端,对希望集成其他实验追踪工具的开发者有参考价值。核心实现(RolloutTraceConfig.init、rollout_trace_attr、_TrackioLoggingAdapter)逻辑清晰,适合作为模板。
功能与动机
PR body 说明:'add trace logging via Trackio, the free, local-first experiment tracking library from Hugging Face'。目的是为用户提供另一种本地优先的实验追踪选择,丰富 verl 的追踪后端生态。
实现拆解
- RolloutTraceConfig 后端注册:在 init 方法中添加 elif backend == "trackio" 分支,导入 trackio,如果当前 run 未初始化则调用 trackio.init,并将 client 设为 trackio 模块。
- rollout_trace_attr 上下文管理器的 Trackio 支持:在上下文管理器内部,针对 trackio 后端,使用 trackio.Trace 记录操作,并将 _trace_attributes 中的元数据(step、sample_index 等)传递给 metadata。
- tracking.py 的 _TrackioLoggingAdapter:新增包装类,实现 log 和 finish 方法,确保 Trackio 日志器可以按需初始化和关闭。
- ValidationGenerationsLogger 的 trackio 集成:新增 log_generations_to_trackio 方法,将验证生成样本构建为 trackio.Trace 对象并批量记录。
- 配置与文档更新:在 rollout.yaml 注释中添加 trackio,在 rollout_trace.rst 文档中增加 Trackio 使用小节和截图。
- 测试覆盖:通过 mock trackio 模块,测试 rollout trace 的 trackio 后端正常工作、使用聊天消息模式,以及 ValidationGenerationsLogger 正确 log_generations_to_trackio。
关键文件:
verl/utils/rollout_trace.py(模块 核心追踪;类别 source;类型 dependency-wiring;符号 get_backend, get_client, enable_token2text, _json_trace_content): 核心变更文件,注册 trackio 后端并在 rollout_trace_attr 中实现 Trackio Trace 记录逻辑
verl/utils/tracking.py(模块 日志适配器;类别 source;类型 core-logic;符号 _TrackioLoggingAdapter, init, log, finish): 新增 _TrackioLoggingAdapter 包装类和 log_generations_to_trackio 方法,使验证 generations 能够通过 Trackio Trace 记录
tests/utils/test_rollout_trace_on_cpu.py(模块 追踪测试;类别 test;类型 test-coverage;符号 mock_trackio_client, ChatTraceClass, run, test_rollout_trace_with_trackio_backend): 新增 mock_trackio_client fixture 以及两个端到端测试,验证 trackio backend 的正确行为和聊天消息模式
tests/utils/test_tracking_on_cpu.py(模块 日志测试;类别 test;类型 test-coverage;符号 test_validation_generations_logger_logs_trackio_traces): 新增测试文件,验证 ValidationGenerationsLogger 在 trackio 启用时将 generations 正确记录为 Trace 对象
verl/trainer/config/rollout/rollout.yaml(模块 配置文件;类别 config;类型 configuration): 在配置注释中添加 trackio 作为支持的 backend 选项,对齐文档
docs/advance/rollout_trace.rst(模块 文档;类别 docs;类型 documentation): 添加 Trackio 使用文档,包括配置参数和截图,方便用户快速上手
关键符号:RolloutTraceConfig.init, rollout_trace_attr, _TrackioLoggingAdapter.init, _TrackioLoggingAdapter.log, _TrackioLoggingAdapter.finish, ValidationGenerationsLogger.log_generations_to_trackio
关键源码片段
verl/utils/rollout_trace.py
核心变更文件,注册 trackio 后端并在 rollout_trace_attr 中实现 Trackio Trace 记录逻辑
# 文件 : verl/utils/rollout_trace.py
# RolloutTraceConfig.init 类方法,新增 trackio 后端支持
@classmethod
def init(
cls,
project_name: str,
experiment_name: str,
backend: str,
token2text: bool = False,
max_samples_per_step_per_worker: int | None = None,
):
config = cls.get_instance()
if config._initialized:
return
config.backend = backend
config.token2text = token2text
config.project_name = project_name
config.experiment_name = experiment_name
config.max_samples_per_step_per_worker = max_samples_per_step_per_worker
if backend == "weave":
import weave
config.client = weave.init(project_name)
elif backend == "mlflow":
import mlflow
mlflow.config.enable_async_logging()
config.client = mlflow
MLFLOW_TRACKING_URI = os.environ.get("MLFLOW_TRACKING_URI", "sqlite:////tmp/mlruns.db")
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
mlflow.set_experiment(project_name)
elif backend == "trackio":
import trackio
from trackio import context_vars
# 如果当前 run 未初始化(例如在验证 generations 场景可能已由 tracking.py 初始化),
# 则调用 trackio.init ,传入框架信息
if context_vars.current_run.get() is None:
trackio.init(project=project_name, name=experiment_name, config={"framework": "verl"})
config.client = trackio
else:
config.client = None
config._initialized = True
verl/utils/tracking.py
新增 _TrackioLoggingAdapter 包装类和 log_generations_to_trackio 方法,使验证 generations 能够通过 Trackio Trace 记录
# 文件 : verl/utils/tracking.py
# _TrackioLoggingAdapter 包装类,用于在 trainer 中使用 trackio 时统一接口
class _TrackioLoggingAdapter:
def __init__(self, trackio):
self.trackio = trackio
def log(self, data, step):
# 调用 trackio 的 log 方法,传入 step
self.trackio.log(data, step=step)
def finish(self):
from trackio import context_vars
# 只有在有活跃 run 时才调用 finish ,避免重复关闭
if context_vars.current_run.get() is not None:
self.trackio.finish()
# ValidationGenerationsLogger 中新增的方法,将验证生成记录为 Trackio Trace
def log_generations_to_trackio(self, samples, step):
import trackio
traces = []
for sample_index, sample in enumerate(samples):
if len(sample) >= 3:
input_text, output_text, score = sample[0], sample[1], sample[2]
else:
input_text, output_text, score = sample, "", None
# 构造符合 Trackio Trace 格式的对话消息
traces.append(
trackio.Trace(
messages=[
{"role": "user", "content": str(input_text)},
{"role": "assistant", "content": str(output_text)},
],
metadata={
"source": "validation_generations",
"sample_index": sample_index,
"score": score,
},
)
)
if traces:
trackio.log({"val/generations": traces}, step=step)
评论区精华
在 Review 中,维护者 wuxibin89 要求提供使用 Trackio 运行 e2e 训练的截图,并按照仓库的代码格式化要求格式化代码。提交者 abidlabs 照做了(添加了文档截图,并多次提交确保代码格式正确)。双方沟通后 PR 获得批准。无其他争议。
- 要求 e2e 演示截图和代码格式化 (question): 提交者 abidlabs 添加了文档截图并多次提交确保代码格式正确,之后 PR 获得批准。
风险与影响
- 风险:风险较低。引入新依赖 trackio(本地优先、开源),不存在外部 API 风险;初始化和日志逻辑使用了 try/finally 正确恢复上下文状态,不会影响其他后端;测试使用 mock 覆盖,不依赖真实 trackio 模块。但需要注意:trackio 的 API 版本变更可能导致集成代码失效;rollout_trace_attr 中新增的 _trace_attributes ContextVar 与现有 weave/mlflow 分支的嵌套使用需注意重置逻辑正确。
- 影响:用户:可通过设置 actor_rollout_ref.rollout.trace.backend=trackio 启用新的追踪后端,同时也可以在 trainer.logger 中添加 trackio 以记录验证生成。系统:代码量增加 ~369 行,但主要集中在新功能分支和测试。团队:维护者需要了解 trackio 的 API 变化,但集成模式与 mlflow/weave 一致,学习成本低。
- 风险标记:新增后端依赖, 上下文管理器变更, 测试仅 mock 覆盖
关联脉络
参与讨论