Skip to content

[Feat] Add OpenTelemetry tracing for RL rollout workflows#1946

Open
YanhuiDua wants to merge 7 commits into
InternLM:mainfrom
YanhuiDua:feat/otel-trace
Open

[Feat] Add OpenTelemetry tracing for RL rollout workflows#1946
YanhuiDua wants to merge 7 commits into
InternLM:mainfrom
YanhuiDua:feat/otel-trace

Conversation

@YanhuiDua

@YanhuiDua YanhuiDua commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

概述

这个 PR 为 XTuner RL rollout 增加 tracing 能力,并通过 xtuner.v1.rl.trace 提供 XTuner 级别的 trace API

用户不需要直接依赖 OpenTelemetry SDK,也不需要关心 collector、exporter、Ray runtime env 等实现细节;只需要在 RL config 中开启 trace,并在自定义业务路径中使用 XTuner 提供的 trace API,即可把 controller、worker、
agent loop、judger、infer backend、session server 等路径上的执行信息串到同一条 trace 中。

这套能力主要用于排查 rollout 慢、失败、断链、异常传播和跨进程调用耗时问题。

使用方式

开启 trace

在 RL config 中配置 TraceConfig

from xtuner.v1.rl.trace import TraceConfig

trace_config = TraceConfig(
    enabled=True,
    output_dir="work_dirs/trace",
    service_name="xtuner-rollout",
    viewer_enabled=True,
)

Trace API

公共接口统一从 xtuner.v1.rl.trace 暴露。API 分为三层:basic API、XTuner endpoint API、context propagation API。

Basic API

Basic API 是 XTuner 封装的 trace primitive,适合普通本地阶段、临时埋点、自定义业务代码和手动 context 管理。

API 作用
trace_span(name, attributes=None, parent_carrier=None) 创建当前 span
trace_function(name=None, attributes=None) 给同步或异步函数包一层 span
trace_event(name, attributes=None) 在当前 span 上记录 event
set_trace_attribute(key, value) 写单个 attribute
set_trace_attributes(attributes) 批量写 attributes
inject_trace_context(carrier=None) 把当前 context 写入 carrier
extract_trace_context(carrier) 从 carrier 取出 TraceContext
attach_trace_context(context) 在当前 scope 激活 TraceContext

示例:

from xtuner.v1.rl.trace import trace_span, trace_event, set_trace_attributes

with trace_span("agent.localhost.infer.run", attributes={"xtuner.task_name": task_name}):
    result = await infer_engine.generate(request)
    set_trace_attributes({"xtuner.status": "ok"})
    trace_event("infer.finished", {"num_tokens": result.num_tokens})

如果一段代码只是局部阶段,或者没有统一的 XTuner 业务生命周期,就应该直接使用 basic API,而不是新增的endpoint helper。

XTuner Endpoint API

Endpoint API 面向 XTuner 已知的稳定业务入口。它不只是创建 span,还会统一处理 target 绑定、initial/final attributes、status event、failure 记录和上游 context 接入等生命周期逻辑。

API 适用场景
traced_rollout_endpoint(...) rollout controller / rollout worker 这类以 RolloutState 为核心 target 的入口
traced_agent_item_endpoint(...) agent loop 中以 agent item 为核心 target 的入口
traced_judger_endpoint(...) judger / reward 计算入口

这类 API 适合框架内部的标准入口。普通业务代码通常不需要直接处理 rollout final attributes、agent item result attributes 或 judger result attributes。

Trace SKILL

Viewer使用方式

XTuner 提供面向 rollout 的 trace viewer。它不是新的 trace 后端,而是把 Jaeger 或本地 trace JSONL 中的原始 spans 按 rollout 语义聚合成样本视图。

sample状态总览
提供每个阶段的平均耗时、有哪些sample卡在了什么阶段等信息
image

每条sample调用链
image

@YanhuiDua

Copy link
Copy Markdown
Collaborator Author

@claude review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant