Prhub

#1768 Fix uploading sglang metrics to wandb

THUDM/slime · 作者 zhuzilin · 合并时间 2026-03-26 13:23

分析状态 已生成
文件变更 5提交数 1 · 评论 0
代码增减 +70 / -18
dependency metrics wandb

执行摘要

修复 wandb 初始化时机问题,确保 SGLang 指标正确上传。

根据 slime/utils/wandb_utils.py 中的注释,primary wandb init 发生在 rollout 服务器启动之前以获取 wandb_run_id,但此时路由器地址不可用,导致无法上传 SGLang Prometheus 指标。因此,需要在服务器启动后重新初始化 wandb 以添加指标端点。

对于涉及 wandb 集成和分布式训练的开发者,推荐精读此 PR,以学习如何处理指标上传的时机问题。特别关注 reinit_wandb_primary_with_open_metrics 函数的设计。

讨论亮点

本次 PR 没有 review 评论,因此无讨论内容。

实现拆解

关键改动包括:在 wandb_utils.py 新增 reinit_wandb_primary_with_open_metrics 函数,使用 wandb.finish() 和 wandb.init() 重新初始化 primary run;在 logging_utils.py 新增 update_tracking_open_metrics 接口函数;在 rollout.py 新增 get_metrics_router_addr 公共方法,并移除 init_tracking 中的 router_addr 参数;在 train.py 和 train_async.py 中调用 update_tracking_open_metrics。

文件 模块 状态 重要度
slime/utils/wandb_utils.py utils/wandb modified 7.0
slime/ray/rollout.py ray/rollout modified 6.0
train.py training modified 5.0
train_async.py training modified 5.0
slime/utils/logging_utils.py utils/logging modified 5.0

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

reinit_wandb_primary_with_open_metrics update_tracking_open_metrics get_metrics_router_addr init_wandb_secondary

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

风险包括:重新初始化 wandb 可能增加开销或导致数据不一致(如 wandb.finish() 调用);依赖自定义 sglang_router 版本,如果使用标准版本,指标上传功能会降级(有警告日志);在 wandb_utils.py 中的新函数缺乏充分测试覆盖,可能引入回归错误。

对用户:修复了 wandb 指标上传问题,用户能正确监控训练过程。对系统:wandb 跟踪更准确,可能轻微增加初始化时间。对团队:需要确保使用自定义 sglang-router 以支持此功能,影响训练监控流程。

wandb 重新初始化开销 依赖自定义 sglang-router 版本 核心 wandb 路径变更

关联 Issue

未识别关联 Issue

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

完整报告

执行摘要

  • 一句话:修复 wandb 初始化时机问题,确保 SGLang 指标正确上传。
  • 推荐动作:对于涉及 wandb 集成和分布式训练的开发者,推荐精读此 PR,以学习如何处理指标上传的时机问题。特别关注 reinit_wandb_primary_with_open_metrics 函数的设计。

功能与动机

根据 slime/utils/wandb_utils.py 中的注释,primary wandb init 发生在 rollout 服务器启动之前以获取 wandb_run_id,但此时路由器地址不可用,导致无法上传 SGLang Prometheus 指标。因此,需要在服务器启动后重新初始化 wandb 以添加指标端点。

实现拆解

关键改动包括:在 wandb_utils.py 新增 reinit_wandb_primary_with_open_metrics 函数,使用 wandb.finish() 和 wandb.init() 重新初始化 primary run;在 logging_utils.py 新增 update_tracking_open_metrics 接口函数;在 rollout.py 新增 get_metrics_router_addr 公共方法,并移除 init_tracking 中的 router_addr 参数;在 train.py 和 train_async.py 中调用 update_tracking_open_metrics。

关键文件:

  • slime/utils/wandb_utils.py(模块 utils/wandb): 核心 wandb 初始化逻辑变更,新增重新初始化函数并移除 router_addr 参数。
  • slime/ray/rollout.py(模块 ray/rollout): 新增公共方法获取路由器地址,修改初始化顺序以支持后置 wandb 更新。
  • train.py(模块 training): 训练脚本中调用更新函数,直接影响训练流程中的 wandb 集成。
  • train_async.py(模块 training): 异步训练脚本类似修改,确保一致的行为。
  • slime/utils/logging_utils.py(模块 utils/logging): 新增接口函数 update_tracking_open_metrics,连接 wandb 更新逻辑和外部调用。

关键符号:reinit_wandb_primary_with_open_metrics, update_tracking_open_metrics, get_metrics_router_addr, init_wandb_secondary

评论区精华

本次 PR 没有 review 评论,因此无讨论内容。

  • 暂无高价值评论线程

风险与影响

  • 风险:风险包括:重新初始化 wandb 可能增加开销或导致数据不一致(如 wandb.finish() 调用);依赖自定义 sglang_router 版本,如果使用标准版本,指标上传功能会降级(有警告日志);在 wandb_utils.py 中的新函数缺乏充分测试覆盖,可能引入回归错误。
  • 影响:对用户:修复了 wandb 指标上传问题,用户能正确监控训练过程。对系统:wandb 跟踪更准确,可能轻微增加初始化时间。对团队:需要确保使用自定义 sglang-router 以支持此功能,影响训练监控流程。
  • 风险标记:wandb重新初始化开销, 依赖自定义sglang-router版本, 核心wandb路径变更

关联脉络

  • PR #1770 use zhuzilin/sgl-router for sglang-router: 两者都涉及自定义 sglang-router 版本以支持 wandb 指标上传,PR 1770 提供了依赖基础,而本 PR 修复上传逻辑。

参与讨论