执行摘要
- 一句话:修复 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 修复上传逻辑。
参与讨论