[Feat] Add OpenTelemetry tracing for RL rollout workflows

.claude/skills/xtuner-trace/SKILL.md

@@ -0,0 +1,62 @@
+---
+name: xtuner-trace
+description: Use when adding, modifying, or reviewing XTuner RL trace instrumentation for rollout, worker, session server, localhost/sandbox agent loop, judger, infer engine, Ray remote calls, HTTP proxy calls, or trace viewer fields. Helps choose the correct XTuner trace API, preserve Ray/HTTP context propagation, and add focused verification for a requested observation phase.
+---
+
+# XTuner Trace
+
+Use this skill to instrument or review a specific XTuner RL phase with trace spans, events, attributes, and context propagation.
+
+## Core Workflow
+
+1. Identify the requested observation phase and its boundary:
+   - Local code block or function only.
+   - Stable XTuner endpoint around `RolloutState`, agent item, or judger.
+   - Ray `.remote(...)` boundary.
+   - HTTP request/proxy boundary.
+   - Viewer field or payload change only.
+
+2. Read the current code before editing:
+   - Public facade: `xtuner/v1/rl/trace/__init__.py`
+   - Basic and endpoint API: `xtuner/v1/rl/trace/api.py`
+   - Propagation helpers: `xtuner/v1/rl/trace/context_propagation.py`
+   - Span names and attribute builders: `xtuner/v1/rl/trace/trace_utils.py`
+   - Existing examples near the requested phase.
+
+3. Choose the narrowest API that matches the boundary:
+   - Use `trace_span(...)` or `trace_function(...)` for local work.
+   - Use `traced_rollout_endpoint(...)` for async endpoints whose lifecycle is a `RolloutState`.
+   - Use `traced_agent_item_endpoint(...)` for agent item runner entrypoints.
+   - Use `traced_judger_endpoint(...)` for agent-loop judger entrypoints with `owner.judger`.
+   - Use `trace_remote(...)` for Ray remote submission and keep the receiver as a traced rollout endpoint or a span with `parent_carrier`.
+   - Use `traced_aiohttp_request(...)` or `instrument_aiohttp_client()` for HTTP propagation.
+
+4. Keep semantics centralized:
+   - Add new stable span names as `TRACE_SPAN_*` in `trace_utils.py`.
+   - Add reusable field mapping to `trace_utils.py` instead of scattering key construction.
+   - Use `failure_attributes(...)`, `reward_trace_attributes(...)`, `rollout_state_*_attributes(...)`, `agent_item_*_attributes(...)`, and `agent_judger_initial_attributes(...)` when they fit.
+
+5. Verify propagation, not just local span creation:
+   - Ray: the outbound call must use `trace_remote(...)`; the inbound side must restore the carrier through `traced_rollout_endpoint(...)` or `trace_span(..., parent_carrier=...)`.
+   - HTTP: inbound request context must be extracted from headers or body carrier; outbound request context must be injected into headers or request body as appropriate.
+   - Do not call OpenTelemetry SDK directly from rollout, agent, judger, or session-server business code.
+
+For concrete patterns and snippets, read [references/trace-patterns.md](references/trace-patterns.md) before making code changes.
+
+## Guardrails
+
+- Do not invent public APIs that are not exported from `xtuner.v1.rl.trace`.
+- Do not use `set_trace_attribute(...)`, `extract_trace_context(...)`, or `attach_trace_context(...)`; the current public facade does not expose them.
+- Do not make a thin endpoint helper for a phase without a stable target and final-result lifecycle. Use explicit `trace_span(...)` instead.
+- Do not manually keep `_xtuner_trace_carrier` on `RolloutState.extra_fields`; `trace_remote(...)` attaches it temporarily and restores local state.
+- Do not pass multiple `RolloutState` objects, a collection of rollout states, or no rollout state to `trace_remote(...)`; pass exactly one or use `target=...`.
+- Do not record prompts, responses, full configs, secrets, or large payloads as attributes. Prefer IDs, statuses, counts, timings, categories, and stable viewer fields.
+
+## Verification
+
+Prefer focused verification:
+
+- Run `git diff --check` after edits.
+- Run `python -m compileall -q` on touched trace and caller files.
+- Add or update behavior-oriented trace tests when changing propagation, endpoint decorators, viewer payload fields, or final-state recording.
+- For viewer changes, verify both JSONL and Jaeger-source assumptions if the change touches source or payload selection.

.claude/skills/xtuner-trace/references/trace-patterns.md

@@ -0,0 +1,175 @@
+# XTuner Trace Patterns
+
+Use this reference after the `xtuner-trace` skill triggers and before editing trace instrumentation.
+
+## Public API Surface
+
+Import these from `xtuner.v1.rl.trace` unless noted:
+
+| API | Use |
+| --- | --- |
+| `trace_span(name, attributes=None, parent_carrier=None)` | Create a span around a local phase. Pass `parent_carrier` only when restoring an inbound context. |
+| `trace_function(name=None, attributes=None)` | Wrap an entire sync or async function with one span. |
+| `trace_event(name, attributes=None)` | Add an event to the current span. |
+| `set_trace_attributes(attributes)` | Set one or more attributes on the current span. For one key, pass `{key: value}`. |
+| `inject_trace_context(carrier=None)` | Inject current W3C context into a dict carrier. |
+| `traced_rollout_endpoint(span_name, target_arg="rollout_state", initial_attributes=None)` | Wrap an async endpoint whose target and final result are a `RolloutState`. |
+| `traced_agent_item_endpoint(span_name, item_arg="item", initial_attributes=None)` | Wrap an async agent item runner endpoint. |
+| `traced_judger_endpoint(span_name="judger.run", target_arg="rollout_state")` | Wrap an async judger endpoint; owner must expose `self.judger`. |
+| `trace_remote(remote_method, *args, target=None, **kwargs)` | Call Ray `.remote(...)` while temporarily attaching the current trace carrier to exactly one `RolloutState`. |
+
+Import these from `xtuner.v1.rl.trace.context_propagation` only when working on framework or transport code:
+
+| API | Use |
+| --- | --- |
+| `traced_aiohttp_request(...)` | Create a send span, inject headers, perform an aiohttp request, and record HTTP status/send latency. |
+| `instrument_aiohttp_client()` | Patch generic aiohttp clients, such as lagent's OpenAI-compatible client, so they propagate context. |
+| `extract_trace_carrier_from_mapping(...)` | Restore inbound body carrier at HTTP proxy boundaries. |
+| `extract_trace_attributes_from_mapping(...)` | Recover upstream semantic attributes passed through request body. |
+| `remove_trace_carrier_from_mapping(...)` | Remove internal trace carrier fields before forwarding user/backend payloads. |
+
+## Decision Matrix
+
+| Requested observation | Preferred pattern |
+| --- | --- |
+| Add timing around a small local block | `with trace_span(TRACE_SPAN_..., attributes=...)` |
+| Trace a whole helper function with no custom final result | `@trace_function(...)` |
+| Trace agent loop, rollout controller, or rollout worker entrypoint | `@traced_rollout_endpoint(TRACE_SPAN_...)` |
+| Submit a traced Ray call | `await trace_remote(actor.method, rollout_state=state)` or store the returned object ref |
+| Trace localhost/sandbox agent item runner entrypoint | `@traced_agent_item_endpoint(TRACE_SPAN_...)` |
+| Trace agent infer/validate/invoke substage | Explicit `trace_span(...)` with `agent_item_initial_attributes(item)` |
+| Trace judger information and final reward | `@traced_judger_endpoint(...)` for agent-loop judger entrypoints; explicit `trace_span(...)` inside agent runner validation stages |
+| Trace HTTP proxy request phases | One parent request `trace_span(..., parent_carrier=...)`, then child spans for prepare/send/read/record |
+| Propagate context through outbound HTTP | `traced_aiohttp_request(...)`; for third-party aiohttp clients call `instrument_aiohttp_client()` before the client issues requests |
+| Add viewer-visible fields | Stable attribute builders or result recorders in `trace_utils.py`, then payload/render tests if viewer grouping changes |
+
+## Current Code Examples
+
+- Ray rollout chain:
+  - `xtuner/v1/rl/agent_loop/single_turn_agent_loop.py`
+  - `xtuner/v1/rl/rollout/controller.py`
+  - `xtuner/v1/rl/rollout/worker.py`
+  - Pattern: caller uses `trace_remote(...)`; receiver uses `@traced_rollout_endpoint(...)`.
+
+- Session server HTTP proxy:
+  - `xtuner/v1/rl/rollout/session_server.py`
+  - Pattern: `_handle_request()` restores parent context from headers/body, `_prepare_request()` records request metadata, `_forward_request()` uses `traced_aiohttp_request(...)`, `_read_upstream_response()` records stream/non-stream latency, `_record_response()` records tokenization/session trace result.
+
+- Localhost agent loop:
+  - `xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py`
+  - `xtuner/v1/rl/agent_loop/localhost_agent_loop/runner.py`
+  - `xtuner/v1/rl/agent_loop/localhost_agent_loop/stage.py`
+  - `xtuner/v1/rl/agent_loop/localhost_agent_loop/judger.py`
+  - Pattern: rollout entrypoint uses `@traced_rollout_endpoint(...)`; item runner uses `@traced_agent_item_endpoint(...)`; infer, agent invoke, validate, judger, and materialization use explicit `trace_span(...)`.
+
+## Templates
+
+### Local Phase
+
+```python
+from xtuner.v1.rl.trace import set_trace_attributes, trace_span
+from xtuner.v1.rl.trace.trace_utils import TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN, agent_item_initial_attributes
+
+attributes = agent_item_initial_attributes(item)
+attributes["stage.name"] = "infer"
+with trace_span(TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN, attributes=attributes):
+    result = await self.infer.run(item, item.infer)
+    set_trace_attributes({"agent.message_count": message_count})
+```
+
+### Failure Attributes
+
+```python
+from xtuner.v1.rl.trace import set_trace_attributes
+from xtuner.v1.rl.trace.trace_utils import failure_attributes
+
+set_trace_attributes(
+    failure_attributes(
+        error.category,
+        message=error.message,
+        type=error.type,
+    )
+)
+```
+
+### Ray Boundary
+
+```python
+from xtuner.v1.rl.trace import trace_remote, traced_rollout_endpoint
+from xtuner.v1.rl.trace.trace_utils import TRACE_SPAN_ROLLOUT_CONTROLLER_GENERATE
+
+@traced_rollout_endpoint(TRACE_SPAN_ROLLOUT_CONTROLLER_GENERATE)
+async def generate(self, rollout_state: RolloutState) -> RolloutState:
+    response_ref = trace_remote(worker.generate, rollout_state=rollout_state)
+    return await response_ref
+```
+
+Rules:
+
+- `trace_remote(...)` only propagates context; it does not create the receiver's business span.
+- The receiver must create the span, usually through `@traced_rollout_endpoint(...)`.
+- If the `RolloutState` is not directly in args/kwargs, pass `target=rollout_state`.
+
+### HTTP Proxy Boundary
+
+```python
+parent_carrier, source = _request_trace_context(headers, request_data)
+trace_attributes = extract_trace_attributes_from_mapping(request_data)
+
+with trace_span("session_server.request", attributes=trace_attributes, parent_carrier=parent_carrier):
+    prepared = await self._prepare_request(...)
+    forwarded = await self._forward_request(request, prepared)
+```
+
+Outbound aiohttp:
+
+```python
+async with traced_aiohttp_request(
+    client,
+    span_name="session_server.send_request",
+    method=request.method,
+    url=target_url,
+    headers=forward_headers,
+    data=prepared.body,
+    attributes=trace_attributes,
+) as resp:
+    return await self._read_upstream_response(request, resp, prepared)
+```
+
+Rules:
+
+- Prefer `traceparent` headers when present; otherwise use body carrier.
+- Remove internal carrier fields from request JSON before forwarding to model backends unless the downstream service is expected to consume them.
+- Keep request, prepare, send, read, and record as separate spans in proxy code; this makes latency attribution useful.
+
+## Attribute Guidance
+
+Prefer stable scalar attributes:
+
+- Identity: `xtuner.rollout_id`, `xtuner.group_id`, `xtuner.session_id`, `xtuner.task_name`
+- Status: `xtuner.status`, `agent.status`, `session.stream`
+- Timing/counts: `session.first_token_ms`, `session.response_total_ms`, `completion.tokens`, `agent.message_count`
+- Reward: `reward.score`, `reward.pass`
+- Error: `error.category`, `error.message`, `error.type`
+
+Avoid:
+
+- Full prompts, full responses, raw chat histories, stack traces, large tool payloads.
+- Secrets, bearer tokens, API keys, raw HTTP headers.
+- Per-call random field names that the viewer cannot aggregate.
+
+## Verification Targets
+
+Use the smallest meaningful checks:
+
+```bash
+python -m compileall -q xtuner/v1/rl/trace xtuner/v1/rl/rollout/session_server.py
+git diff --check
+```
+
+For propagation changes, add focused tests that assert:
+
+- Child spans share the parent `trace_id`.
+- `trace_remote(...)` cleans the temporary carrier from local `RolloutState.extra_fields`.
+- HTTP body carriers are removed before forwarding when required.
+- Final status/reward/error attributes are written after the operation completes.

examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py

@@ -15,6 +15,7 @@
 from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig
 from xtuner.v1.rl.rollout.worker import RolloutConfig
 from xtuner.v1.rl.trainer import RolloutImportanceSampling, WorkerConfig
+from xtuner.v1.rl.trace import TraceConfig
 from xtuner.v1.train.trainer import LoadCheckpointConfig
 from xtuner.v1.rl.utils import AcceleratorResourcesConfig
 from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig
@@ -298,6 +299,16 @@ def _build_dataset_cfg(specs, default_recipe):
         ),
     )
 
+trace_config = TraceConfig(
+    enabled=os.environ.get("XTUNER_TRACE_ENABLED") == "1",
+    output_dir=Path(work_dir) / "otel",
+    service_name="xtuner-agent-rollout",
+    viewer_enabled=True,
+    viewer_host="0.0.0.0",
+    viewer_port=18080,
+    viewer_jaeger_query_url="http://127.0.0.1:16686",
+)
+
 trainer = RLColocateTrainerConfig(
     resources=resources,
     train_worker_cfg=train_worker_cfg,
@@ -326,4 +337,5 @@ def _build_dataset_cfg(specs, default_recipe):
     debug_rollout_dir=debug_rollout_dir,
     debug_train=debug_train,
     debug_rollout=debug_rollout,
+    trace_config=trace_config,
 )

examples/v1/config/rl_grpo_gsm8k_judge.py

@@ -22,6 +22,7 @@
 from xtuner.v1.rl.agent_loop_manager import AgentLoopManagerConfig, SamplerConfig, SyncProduceStrategyConfig, TaskSpecConfig
 from xtuner.v1.rl.evaluator import EvaluatorConfig
 from xtuner.v1.rl.loss import GRPOLossConfig
+from xtuner.v1.rl.trace import TraceConfig
 from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig
 
 # env
@@ -183,7 +184,18 @@
 # 7. evaluator
 evaluator_config = EvaluatorConfig(compute_metric_func=None)
 
-# 8. RL Colocate Trainer Config（CLI 通过 config["trainer"].build() 得到 Trainer）
+# 8. Trace Config
+trace_config = TraceConfig(
+    enabled=os.environ.get("XTUNER_TRACE_ENABLED") == "1",
+    output_dir=Path(work_dir) / "otel",
+    service_name="xtuner-rollout",
+    viewer_enabled=True,
+    viewer_host="0.0.0.0",
+    viewer_port=18080,
+    viewer_jaeger_query_url="http://127.0.0.1:16686",
+)
+
+# 9. RL Colocate Trainer Config（CLI 通过 config["trainer"].build() 得到 Trainer）
 trainer = RLColocateTrainerConfig(
     resources=resources,
     train_worker_cfg=train_worker_cfg,  # TODO: uniform naming of cfg and config
@@ -203,4 +215,5 @@
     work_dir=work_dir,
     seed=123,
     debug_rollout=False,
+    trace_config=trace_config,
 )

examples/v1/scripts/run_rl.sh

@@ -1,4 +1,8 @@
 set -ex
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then
+  source "${SCRIPT_DIR}/setup_trace.sh"
+fi
 ray stop --force
 # examples of usage:
 # qwen3_8B_grpo_gsm8k training: 

examples/v1/scripts/run_rl_submit.sh

@@ -1,4 +1,8 @@
 set -ex
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then
+  source "${SCRIPT_DIR}/setup_trace.sh"
+fi
 ray stop --force
 # examples of usage:
 # qwen3_8B_grpo_gsm8k training: bash examples/v1/scripts/run_rl.sh examples/v1/config/rl_qwen3_8B_grpo.py "sglang" $MODEL_PATH $DATA_PATH $EVAL_DATA_PATH

examples/v1/scripts/setup_trace.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+
+SETUP_TRACE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SETUP_TRACE_DIR}/../../.." && pwd)"
+
+export XTUNER_OTEL_ROOT="${XTUNER_OTEL_ROOT:-/tmp/xtuner_otel}"
+export PATH="${XTUNER_OTEL_ROOT}/bin:${PATH}"
+
+OTEL_INSTALL_SCRIPT="${XTUNER_OTEL_INSTALL_SCRIPT:-${REPO_ROOT}/recipe/otle/install_otel_tools.sh}"
+JAEGER_RESTART_SCRIPT="${XTUNER_JAEGER_RESTART_SCRIPT:-${REPO_ROOT}/recipe/otle/restart_jaeger_memory.sh}"
+JAEGER_CONFIG="${XTUNER_JAEGER_CONFIG:-${REPO_ROOT}/recipe/otle/jaeger/jaeger-memory.yaml}"
+TRACE_VIEWER_PORT="${XTUNER_TRACE_VIEWER_PORT:-18080}"
+
+if pgrep -f "xtuner.tools.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" >/dev/null 2>&1 ||
+   pgrep -f "xtuner.tools.trace_viewer.server.*--port=${TRACE_VIEWER_PORT}" >/dev/null 2>&1; then
+  echo "Stopping previous XTuner trace viewer on port ${TRACE_VIEWER_PORT}"
+  pkill -f "xtuner.tools.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" 2>/dev/null || true
+  pkill -f "xtuner.tools.trace_viewer.server.*--port=${TRACE_VIEWER_PORT}" 2>/dev/null || true
+  sleep 1
+fi
+
+if [ ! -d "$XTUNER_OTEL_ROOT" ]; then
+  echo "Installing XTuner OTel tools to ${XTUNER_OTEL_ROOT}"
+  bash "$OTEL_INSTALL_SCRIPT" "$XTUNER_OTEL_ROOT" || return 1 2>/dev/null || exit 1
+fi
+
+otel_collector="$(command -v otelcol-contrib || command -v otelcol || true)"
+if [ -z "$otel_collector" ]; then
+  echo "Error: XTuner trace collector not found after checking ${XTUNER_OTEL_ROOT}." >&2
+  echo "Expected otelcol-contrib or otelcol under ${XTUNER_OTEL_ROOT}/bin." >&2
+  return 1 2>/dev/null || exit 1
+fi
+
+echo "XTuner trace collector: ${otel_collector}"
+bash "$JAEGER_RESTART_SCRIPT" "$JAEGER_CONFIG" || return 1 2>/dev/null || exit 1