Prhub

#6423 [rollout] feat: add Trackio rollout trace logging

原始 PR 作者 abidlabs 合并时间 2026-05-25 11:44 文件变更 6 提交数 7 评论 6 代码增减 +369 / -24

执行摘要

添加 Trackio 作为 rollout trace 后端

PR body 说明:'add trace logging via Trackio, the free, local-first experiment tracking library from Hugging Face'。目的是为用户提供另一种本地优先的实验追踪选择,丰富 verl 的追踪后端生态。

值得精读。该 PR 展示了如何以一致的设计模式扩展 verl 的追踪后端,对希望集成其他实验追踪工具的开发者有参考价值。核心实现(RolloutTraceConfig.init、rollout_trace_attr、_TrackioLoggingAdapter)逻辑清晰,适合作为模板。

讨论亮点

在 Review 中,维护者 wuxibin89 要求提供使用 Trackio 运行 e2e 训练的截图,并按照仓库的代码格式化要求格式化代码。提交者 abidlabs 照做了(添加了文档截图,并多次提交确保代码格式正确)。双方沟通后 PR 获得批准。无其他争议。

实现拆解

  1. RolloutTraceConfig 后端注册:在 init 方法中添加 elif backend == "trackio" 分支,导入 trackio,如果当前 run 未初始化则调用 trackio.init,并将 client 设为 trackio 模块。
  2. rollout_trace_attr 上下文管理器的 Trackio 支持:在上下文管理器内部,针对 trackio 后端,使用 trackio.Trace 记录操作,并将 _trace_attributes 中的元数据(step、sample_index 等)传递给 metadata。
  3. tracking.py 的 _TrackioLoggingAdapter:新增包装类,实现 log 和 finish 方法,确保 Trackio 日志器可以按需初始化和关闭。
  4. ValidationGenerationsLogger 的 trackio 集成:新增 log_generations_to_trackio 方法,将验证生成样本构建为 trackio.Trace 对象并批量记录。
  5. 配置与文档更新:在 rollout.yaml 注释中添加 trackio,在 rollout_trace.rst 文档中增加 Trackio 使用小节和截图。
  6. 测试覆盖:通过 mock trackio 模块,测试 rollout trace 的 trackio 后端正常工作、使用聊天消息模式,以及 ValidationGenerationsLogger 正确 log_generations_to_trackio。
文件 模块 状态 重要度
verl/utils/rollout_trace.py 核心追踪 modified 8.38
verl/utils/tracking.py 日志适配器 modified 7.55
tests/utils/test_rollout_trace_on_cpu.py 追踪测试 modified 6.42
tests/utils/test_tracking_on_cpu.py 日志测试 added 5.43
verl/trainer/config/rollout/rollout.yaml 配置文件 modified 2.5
docs/advance/rollout_trace.rst 文档 modified 2.5

关键符号

RolloutTraceConfig.init rollout_trace_attr _TrackioLoggingAdapter.__init__ _TrackioLoggingAdapter.log _TrackioLoggingAdapter.finish ValidationGenerationsLogger.log_generations_to_trackio

关键源码片段

verl/utils/rollout_trace.py dependency-wiring

核心变更文件,注册 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 core-logic

新增 _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 Tracedef 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)

评论区精华

要求 e2e 演示截图和代码格式化 question

维护者 wuxibin89 要求提交者提供使用 Trackio 运行 e2e 训练的截图,并按照 CONTRIBUTING.md 格式化代码。

结论:提交者 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 覆盖

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论