From ab4356e39e95dcce1857ce5e0efef681c81df21a Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Thu, 9 Jul 2026 12:48:59 +0000 Subject: [PATCH 1/6] Add trace runtime, API, viewer, and skills --- .claude/skills/xtuner-trace/SKILL.md | 83 ++ .../xtuner-trace/references/trace-patterns.md | 414 ++++++ examples/v1/scripts/setup_trace.sh | 35 + recipe/otle/README.md | 54 + recipe/otle/install_otel_tools.sh | 44 + recipe/otle/jaeger/jaeger-memory.yaml | 40 + recipe/otle/restart_jaeger_memory.sh | 87 ++ tests/rl/test_trace.py | 163 +++ tests/rl/trace_utils.py | 183 +++ xtuner/tools/trace_viewer/__init__.py | 18 + xtuner/tools/trace_viewer/payload.py | 1258 +++++++++++++++++ xtuner/tools/trace_viewer/render.py | 551 ++++++++ xtuner/tools/trace_viewer/server.py | 474 +++++++ xtuner/tools/trace_viewer/source.py | 90 ++ xtuner/v1/rl/trace/__init__.py | 50 + xtuner/v1/rl/trace/api.py | 334 +++++ xtuner/v1/rl/trace/otel_utils.py | 139 ++ xtuner/v1/rl/trace/runtime.py | 769 ++++++++++ 18 files changed, 4786 insertions(+) create mode 100644 .claude/skills/xtuner-trace/SKILL.md create mode 100644 .claude/skills/xtuner-trace/references/trace-patterns.md create mode 100644 examples/v1/scripts/setup_trace.sh create mode 100644 recipe/otle/README.md create mode 100755 recipe/otle/install_otel_tools.sh create mode 100644 recipe/otle/jaeger/jaeger-memory.yaml create mode 100755 recipe/otle/restart_jaeger_memory.sh create mode 100644 tests/rl/test_trace.py create mode 100644 tests/rl/trace_utils.py create mode 100644 xtuner/tools/trace_viewer/__init__.py create mode 100644 xtuner/tools/trace_viewer/payload.py create mode 100644 xtuner/tools/trace_viewer/render.py create mode 100644 xtuner/tools/trace_viewer/server.py create mode 100644 xtuner/tools/trace_viewer/source.py create mode 100644 xtuner/v1/rl/trace/__init__.py create mode 100644 xtuner/v1/rl/trace/api.py create mode 100644 xtuner/v1/rl/trace/otel_utils.py create mode 100644 xtuner/v1/rl/trace/runtime.py diff --git a/.claude/skills/xtuner-trace/SKILL.md b/.claude/skills/xtuner-trace/SKILL.md new file mode 100644 index 0000000000..ecd6055cf2 --- /dev/null +++ b/.claude/skills/xtuner-trace/SKILL.md @@ -0,0 +1,83 @@ +--- +name: xtuner-trace +description: Use when the user wants to instrument, inspect, or debug XTuner traces to reconstruct a sample/request execution path, follow cross-process call chains, or identify latency hotspots, bottleneck stages, and abnormal paths from span data. +--- + +# XTuner Trace + +Use this skill for the current trace layer in this repository: OpenTelemetry runtime setup, the basic public API, local setup scripts, and the trace viewer. + +## Current Boundary + +Keep this package infrastructure-only: + +- Runtime/configuration: `xtuner/v1/rl/trace/runtime.py` +- OTel SDK adapter: `xtuner/v1/rl/trace/otel_utils.py` +- Public facade: `xtuner/v1/rl/trace/__init__.py` +- Basic API: `xtuner/v1/rl/trace/api.py` +- Viewer: `xtuner/tools/trace_viewer/` +- Local tooling: `recipe/otle/` and `examples/v1/scripts/setup_trace.sh` + +Do not add rollout, agent, judger, Ray remote, HTTP proxy, reward, status, or session-server business semantics to the basic trace package. + +## Basic API + +These are the interfaces defined in `xtuner/v1/rl/trace/api.py` and re-exported +from `xtuner.v1.rl.trace`: + +- `trace_span(name, attributes=None, parent_carrier=None)` +- `trace_function(name=None, attributes=None)` +- `trace_event(name, attributes=None)` +- `set_trace_attributes(attributes)` +- `inject_trace_context(carrier=None)` + +## Runtime API + +Use these `xtuner/v1/rl/trace/runtime.py` interfaces for explicit trace +runtime setup: + +- `TraceConfig` +- `configure_trace(...)` +- `close_trace()` + +## Add Trace Workflow + +When adding trace instrumentation to an XTuner run: + +1. Ask for or locate the launch script and training config before editing. +2. In the launch script, source `examples/v1/scripts/setup_trace.sh` when + `XTUNER_TRACE_ENABLED=1`. +3. In the training config, add `TraceConfig` and set `enabled=True`; set + `viewer_enabled=True` when the user needs interactive inspection. +4. Before adding any `trace_span(...)` instrumentation, you mask ask the user which + stages they want to observe and which metrics each stage should expose. + Do not infer default stages unless the user explicitly asks you to choose. +5. Add `trace_span(...)` only around the user-confirmed observed stages. Put fields known at span + start in initial attributes, and update runtime or final fields with + `set_trace_attributes(...)`. +6. For cross-process or request boundaries, inject a carrier with + `inject_trace_context(...)` and pass it to downstream + `trace_span(..., parent_carrier=carrier)`. +7. Keep transport-specific propagation at the caller boundary; do not move + rollout, agent, judger, Ray, or HTTP semantics into the basic trace package. +8. Ensure the main training log includes the trace output path, viewer URL, and + a restart command for the viewer. + +## Guardrails + +- Do not reintroduce `trace_remote`, `traced_rollout_endpoint`, `traced_agent_item_endpoint`, or `traced_judger_endpoint`. +- Do not recreate `trace_utils.py`, `context_propagation.py`, span-name registries, or business attribute builders under `xtuner/v1/rl/trace`. +- Do not import `RolloutState`, agent item classes, judgers, rollout workers, Ray actors, aiohttp clients, or trainer configs from the basic trace package. +- Do not call OpenTelemetry SDK directly from business code; use the basic API only when trace instrumentation is explicitly requested. +- Do not record prompts, responses, full configs, secrets, raw headers, stack traces, or large payloads as attributes. +- Viewer stage grouping must come from span attributes such as `xtuner.stage`, `stage`, or `stage.name`; otherwise fall back to the raw span name. + +For concrete patterns, read [references/trace-patterns.md](references/trace-patterns.md) before editing trace code. + +## Verification + +Use focused checks: + +- `PYTHONPATH=. python -m compileall -q xtuner/v1/rl/trace xtuner/tools/trace_viewer` +- `PYTHONPATH=. python -m unittest discover -s tests/rl -p 'test_trace*.py' -v` when the trace tests exist in the worktree. +- `git diff --check` diff --git a/.claude/skills/xtuner-trace/references/trace-patterns.md b/.claude/skills/xtuner-trace/references/trace-patterns.md new file mode 100644 index 0000000000..f4f0d87d58 --- /dev/null +++ b/.claude/skills/xtuner-trace/references/trace-patterns.md @@ -0,0 +1,414 @@ +# XTuner Trace Patterns + +These patterns describe the retained trace surface in this branch: runtime, basic API, viewer, and local setup tooling. + +## Basic API + +Use the public facade: + +```python +from xtuner.v1.rl.trace import ( + inject_trace_context, + set_trace_attributes, + trace_event, + trace_function, + trace_span, +) +``` + +### Local Span + +```python +with trace_span("phase.name", attributes={"xtuner.stage": "phase"}): + ... + set_trace_attributes({"phase.count": 3}) +``` + +### Decorated Function + +```python +@trace_function("phase.load") +def load_item(path: str) -> object: + ... +``` + +### Parent Carrier + +```python +carrier = inject_trace_context() + +with trace_span("child.phase", parent_carrier=carrier): + ... +``` + +`parent_carrier` is a basic W3C context carrier. It is not tied to `RolloutState`, Ray, aiohttp, or any XTuner business object. + +## End-to-End Trace Wiring + +Use this pattern when adding trace support to an XTuner training run. + +### Launch Script + +Enable trace bootstrap from the launcher, guarded by `XTUNER_TRACE_ENABLED`: + +```bash +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then + source "${SCRIPT_DIR}/setup_trace.sh" +fi +``` + +Keep training stdout/stderr in the main training log so runtime trace messages +are visible: + +```bash +python xtuner/v1/train/cli/rl.py \ + --config "$CONFIG_PATH" \ + --num-workers "$XTUNER_RL_NUM_WORKERS" \ + 2>&1 | tee -a "${WORK_DIR}/training_log_${current_time}.txt" +``` + +### Training Config + +Add `TraceConfig` to the training config and pass it to the trainer config: + +```python +import os +from pathlib import Path + +from xtuner.v1.rl.trace import TraceConfig + + +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( + ..., + trace_config=trace_config, +) +``` + +### Stage Spans + +Before adding spans, ask the user which stages they want to observe and which +metrics each stage should expose. + +Put fields known when entering the stage, especially fields used for grouping +or filtering, in `trace_span(..., attributes=...)`. Put results known only +after or during execution, such as status, duration, token counts, errors, and +reward values, in `set_trace_attributes(...)`. This lets the viewer reconstruct +the call chain and identify latency hotspots or abnormal stages. + +Infer stage: + +```python +import time + + +started_at = time.monotonic() +attributes = { + "xtuner.stage": "infer", + "xtuner.task_name": task_name, + "xtuner.sample_id": sample_id, + "rollout.backend": backend, +} + +with trace_span("agent.infer", attributes=attributes): + result = await run_infer(...) + set_trace_attributes( + { + "xtuner.status": "ok" if result.ok else "error", + "prompt.tokens": result.prompt_tokens, + "completion.tokens": result.completion_tokens, + "stage.duration_ms": int((time.monotonic() - started_at) * 1000), + } + ) +``` + +Resource acquire stage: + +```python +with trace_span( + "sandbox.acquire", + attributes={ + "xtuner.stage": "acquire", + "xtuner.sample_id": sample_id, + "sandbox.name": sandbox_name, + }, +): + client = await pool.get(sandbox_name) + set_trace_attributes( + { + "sandbox.env_id": env_id, + "sandbox.image": image, + "sandbox.reused": reused, + } + ) +``` + +Validation or reward stage: + +```python +with trace_span( + "sample.validate", + attributes={ + "xtuner.stage": "validate", + "xtuner.sample_id": sample_id, + "validator.name": validator_name, + }, +): + record = await validate(...) + set_trace_attributes( + { + "xtuner.status": "ok" if record.ok else "error", + "reward.score": record.score, + "reward.passed": record.passed, + } + ) + + if record.error is not None: + set_trace_attributes( + { + "error": True, + "error.type": record.error.type, + "error.message": record.error.message, + } + ) +``` + +Do not record prompts, responses, full configs, secrets, raw headers, stack +traces, or large payloads. + +### Cross-Boundary Propagation + +When a traced operation crosses a process, actor, or HTTP boundary, keep the +propagation helper in the business or transport module that owns that boundary. +Do not put Ray, HTTP, rollout, agent, or judger helpers back into +`xtuner/v1/rl/trace`. + +#### Ray Remote Calls + +For Ray calls, inject the current context into a plain carrier and temporarily +attach that carrier to a serializable domain object or explicit call argument. + +```python +from collections.abc import Mapping +from contextlib import contextmanager +from typing import Any + +from xtuner.v1.rl.trace import inject_trace_context, trace_span + + +_TRACE_CARRIER_FIELD = "_xtuner_trace_carrier" + + +@contextmanager +def attach_trace_carrier_temporarily(target: Any, carrier: Mapping[str, str]): + if not carrier: + yield + return + + extra_fields = getattr(target, "extra_fields", None) + if extra_fields is None: + extra_fields = {} + target.extra_fields = extra_fields + had_previous = _TRACE_CARRIER_FIELD in extra_fields + previous = extra_fields.get(_TRACE_CARRIER_FIELD) + extra_fields[_TRACE_CARRIER_FIELD] = dict(carrier) + try: + yield + finally: + if had_previous: + extra_fields[_TRACE_CARRIER_FIELD] = previous + else: + extra_fields.pop(_TRACE_CARRIER_FIELD, None) + + +def call_ray_remote_with_trace(remote_method, *args, trace_target: Any, **kwargs): + carrier: dict[str, str] = {} + inject_trace_context(carrier) + with attach_trace_carrier_temporarily(trace_target, carrier): + return remote_method.remote(*args, **kwargs) + + +def pop_trace_parent_carrier(target: Any) -> dict[str, str] | None: + extra_fields = getattr(target, "extra_fields", None) + if not isinstance(extra_fields, dict): + return None + carrier = extra_fields.pop(_TRACE_CARRIER_FIELD, None) + if not isinstance(carrier, Mapping): + return None + return {str(key): str(value) for key, value in carrier.items()} +``` + +On the receiver side, attach the carrier to the new span: + +```python +parent_carrier = pop_trace_parent_carrier(trace_target) + +with trace_span("worker.stage", attributes=attributes, parent_carrier=parent_carrier): + ... +``` + +Keep the helper domain-specific. Validate that exactly one trace target is used +when the transport needs one target object; reject collections if restoring the +carrier to the right child span would be ambiguous. + +#### HTTP Calls + +For HTTP calls, prefer W3C headers. If a third-party client or proxy may drop +custom headers, also put a copied carrier into the JSON body under an internal +field and remove it before forwarding the payload downstream. + +```python +from collections.abc import Mapping +from typing import Any + +from xtuner.v1.rl.trace import inject_trace_context, trace_span + + +_TRACE_CARRIER_FIELD = "_xtuner_trace_carrier" + + +def inject_trace_carrier_into_json_body(kwargs: dict[str, Any], carrier: Mapping[str, str]) -> None: + if not carrier: + return + payload = kwargs.get("json") + if not isinstance(payload, dict): + return + payload = dict(payload) + payload.setdefault(_TRACE_CARRIER_FIELD, dict(carrier)) + kwargs["json"] = payload + + +def extract_trace_carrier_from_mapping(payload: Mapping[str, Any] | None) -> dict[str, str] | None: + if not isinstance(payload, Mapping): + return None + carrier = payload.get(_TRACE_CARRIER_FIELD) + if not isinstance(carrier, Mapping): + return None + return {str(key): str(value) for key, value in carrier.items()} + + +def remove_trace_carrier_from_mapping(payload: dict[str, Any]) -> None: + payload.pop(_TRACE_CARRIER_FIELD, None) + + +def extract_http_parent_carrier(headers: Mapping[str, Any], payload: Mapping[str, Any] | None) -> dict[str, str] | None: + header_carrier = { + str(key): str(value) + for key, value in headers.items() + if str(key).lower() in {"traceparent", "tracestate", "baggage"} + } + return header_carrier or extract_trace_carrier_from_mapping(payload) +``` + +Caller: + +```python +headers = dict(headers or {}) +request_kwargs = {"json": payload} +carrier: dict[str, str] = {} +inject_trace_context(carrier) +headers.update(carrier) +inject_trace_carrier_into_json_body(request_kwargs, carrier) + +with trace_span("http.client.request", attributes={"http.method": "POST", "http.url": url}): + response = await client.post(url, headers=headers, **request_kwargs) +``` + +Receiver: + +```python +headers = dict(request.headers) +payload = await request.json() +parent_carrier = extract_http_parent_carrier(headers, payload) + +with trace_span("http.server.request", attributes=attributes, parent_carrier=parent_carrier): + remove_trace_carrier_from_mapping(payload) + ... +``` + +Do not record raw headers or full payloads as span attributes. Only record +derived fields such as method, route, status code, stage, IDs, and timing. + +### Viewer Output + +When `viewer_enabled=True`, the trace runtime logs the trace JSONL path, viewer +URL, and restart command. Keep those runtime log lines visible in the main +training log and report them to the user after the run. + +Useful reference points from the full trace implementation: + +- `examples/v1/scripts/run_rl_run.sh` +- `examples/v1/scripts/setup_trace.sh` +- `examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py` + +## Runtime + +Configure tracing explicitly: + +```python +from xtuner.v1.rl.trace import TraceConfig, close_trace, configure_trace + +runtime = configure_trace( + TraceConfig( + enabled=True, + output_dir="work_dirs/example/otel", + service_name="xtuner", + viewer_enabled=True, + ) +) + +try: + ... +finally: + close_trace() +``` + +The runtime owns OTel collector setup, trace JSONL output, live JSONL output, and optional viewer process startup. + +## Viewer + +The viewer reads Jaeger-style traces or OTel JSONL converted into Jaeger-style payloads. It should not depend on hard-coded rollout/agent span registries. + +Stage display rules: + +1. Use `xtuner.stage` when present. +2. Else use `stage` when present. +3. Else use `stage.name` when present. +4. Else fall back to the raw span name. + +Useful stable attributes: + +- IDs: `xtuner.rollout_id`, `xtuner.group_id`, `xtuner.session_id`, `xtuner.task_name` +- Status: `xtuner.status` +- Stage: `xtuner.stage`, `stage`, `stage.name` +- Counts/timing: `prompt.tokens`, `completion.tokens`, `http.status_code` +- Errors: `error`, `error.message`, `exception.type` + +Avoid recording prompts, responses, full configs, secrets, raw headers, or large payloads. + +## Local Setup + +`examples/v1/scripts/setup_trace.sh` and `recipe/otle/` are local helper tooling for installing OTel collector binaries and starting local Jaeger/viewer dependencies. Keep these as setup assets; do not use them to add automatic trace behavior to training configs or launch scripts unless that integration is explicitly requested. + +## Removed Semantics + +Do not use or recreate these removed surfaces in this branch: + +- `trace_remote` +- `traced_rollout_endpoint` +- `traced_agent_item_endpoint` +- `traced_judger_endpoint` +- `xtuner/v1/rl/trace/trace_utils.py` +- `xtuner/v1/rl/trace/context_propagation.py` +- fixed `TRACE_SPAN_*` registries +- rollout/agent/judger attribute builders diff --git a/examples/v1/scripts/setup_trace.sh b/examples/v1/scripts/setup_trace.sh new file mode 100644 index 0000000000..560f01eaa8 --- /dev/null +++ b/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 diff --git a/recipe/otle/README.md b/recipe/otle/README.md new file mode 100644 index 0000000000..b3d432e87e --- /dev/null +++ b/recipe/otle/README.md @@ -0,0 +1,54 @@ +# XTuner OTel Trace + +XTuner exports rollout traces through OpenTelemetry. For local inspection, start +Jaeger with `jaeger/jaeger-memory.yaml`, enable trace in the training config, and +open the XTuner rollout viewer. + +The reference Jaeger config exposes: + +- Jaeger UI and Query API: `http://127.0.0.1:16686` +- OTLP gRPC receiver: `http://127.0.0.1:14317` +- OTLP HTTP receiver: `http://127.0.0.1:14318/v1/traces` + +Install local binaries: + +```bash +bash recipe/otle/install_otel_tools.sh +export PATH=/tmp/xtuner_otel/bin:$PATH +``` + +Start Jaeger: + +```bash +jaeger --config recipe/otle/jaeger/jaeger-memory.yaml +``` + +For local smoke tests, restart the in-memory Jaeger before each experiment so +old services, operations, and trace ids cannot be mixed with the new run: + +```bash +bash recipe/otle/restart_jaeger_memory.sh +``` + +Run XTuner with trace enabled: + +```bash +export XTUNER_TRACE_ENABLED=1 +export XTUNER_TRACE_SERVICE_NAME=xtuner-rollout +``` + +By default, XTuner starts a local collector that writes +`/traces/traces.jsonl` and forwards spans to the reference Jaeger +OTLP gRPC endpoint `http://127.0.0.1:14317`. + +Open the rollout viewer: + +```bash +python -m xtuner.tools.trace_viewer.server \ + --jaeger-query-url http://127.0.0.1:16686 \ + --service xtuner-rollout +``` + +The viewer is a thin Jaeger Query API adapter. Jaeger remains the trace backend; +XTuner only groups spans by rollout metadata such as `xtuner.rollout_id`, +`xtuner.group_id`, `xtuner.task_name`, and `xtuner.status`. diff --git a/recipe/otle/install_otel_tools.sh b/recipe/otle/install_otel_tools.sh new file mode 100755 index 0000000000..b8f3c0c86e --- /dev/null +++ b/recipe/otle/install_otel_tools.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="${1:-/tmp/xtuner_otel}" +BIN_DIR="${ROOT}/bin" +DOWNLOAD_DIR="${ROOT}/downloads" +OTEL_VERSION="${OTEL_VERSION:-0.128.0}" +JAEGER_VERSION="${JAEGER_VERSION:-2.19.0}" + +mkdir -p "${BIN_DIR}" "${DOWNLOAD_DIR}" + +download_and_extract() { + local url="$1" + local archive="$2" + + if [ ! -f "${DOWNLOAD_DIR}/${archive}" ]; then + curl -L --fail --retry 3 --output "${DOWNLOAD_DIR}/${archive}" "${url}" + fi + tar -xzf "${DOWNLOAD_DIR}/${archive}" -C "${BIN_DIR}" +} + +download_and_extract \ + "https://github.com/open-telemetry/opentelemetry-collector-releases/releases/download/v${OTEL_VERSION}/otelcol_${OTEL_VERSION}_linux_amd64.tar.gz" \ + "otelcol_${OTEL_VERSION}_linux_amd64.tar.gz" + +download_and_extract \ + "https://github.com/open-telemetry/opentelemetry-collector-releases/releases/download/v${OTEL_VERSION}/otelcol-contrib_${OTEL_VERSION}_linux_amd64.tar.gz" \ + "otelcol-contrib_${OTEL_VERSION}_linux_amd64.tar.gz" + +download_and_extract \ + "https://github.com/jaegertracing/jaeger/releases/download/v${JAEGER_VERSION}/jaeger-${JAEGER_VERSION}-linux-amd64.tar.gz" \ + "jaeger-${JAEGER_VERSION}-linux-amd64.tar.gz" + +if [ ! -f "${BIN_DIR}/jaeger" ] && [ -f "${BIN_DIR}/jaeger-${JAEGER_VERSION}-linux-amd64/jaeger" ]; then + ln -sfn "${BIN_DIR}/jaeger-${JAEGER_VERSION}-linux-amd64/jaeger" "${BIN_DIR}/jaeger" +fi + +chmod +x "${BIN_DIR}/otelcol" "${BIN_DIR}/otelcol-contrib" "${BIN_DIR}/jaeger" + +echo "Installed:" +"${BIN_DIR}/otelcol" --version +"${BIN_DIR}/otelcol-contrib" --version +"${BIN_DIR}/jaeger" version +echo "Add to PATH: export PATH=${BIN_DIR}:\$PATH" diff --git a/recipe/otle/jaeger/jaeger-memory.yaml b/recipe/otle/jaeger/jaeger-memory.yaml new file mode 100644 index 0000000000..533ad65925 --- /dev/null +++ b/recipe/otle/jaeger/jaeger-memory.yaml @@ -0,0 +1,40 @@ +receivers: + otlp: + protocols: + grpc: + endpoint: 0.0.0.0:14317 + http: + endpoint: 0.0.0.0:14318 + +processors: + batch: + +exporters: + jaeger_storage_exporter: + trace_storage: memstore + +extensions: + jaeger_storage: + backends: + memstore: + memory: + max_traces: 100000 + jaeger_query: + storage: + traces: memstore + base_path: / + http: + endpoint: 0.0.0.0:16686 + grpc: + endpoint: 0.0.0.0:16685 + +service: + telemetry: + metrics: + level: none + extensions: [jaeger_storage, jaeger_query] + pipelines: + traces: + receivers: [otlp] + processors: [batch] + exporters: [jaeger_storage_exporter] diff --git a/recipe/otle/restart_jaeger_memory.sh b/recipe/otle/restart_jaeger_memory.sh new file mode 100755 index 0000000000..b35eb8930f --- /dev/null +++ b/recipe/otle/restart_jaeger_memory.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Explicit local reset tool: restarting this process clears Jaeger in-memory traces. +# Do not use this against a shared Jaeger deployment. + +ROOT="${XTUNER_OTEL_ROOT:-/tmp/xtuner_otel}" +JAEGER_BIN="${JAEGER_BIN:-${ROOT}/bin/jaeger}" +CONFIG="${1:-recipe/otle/jaeger/jaeger-memory.yaml}" +PID_FILE="${XTUNER_JAEGER_PID_FILE:-/tmp/xtuner_jaeger_memory.pid}" +LOG_FILE="${XTUNER_JAEGER_LOG_FILE:-/tmp/xtuner_jaeger_memory.log}" +QUERY_URL="${XTUNER_JAEGER_QUERY_URL:-http://127.0.0.1:16686}" +WAIT_TIMEOUT_S="${XTUNER_JAEGER_WAIT_TIMEOUT_S:-30}" + +if ! command -v "${JAEGER_BIN}" >/dev/null 2>&1; then + echo "Jaeger binary not found: ${JAEGER_BIN}" >&2 + echo "Install it first: bash recipe/otle/install_otel_tools.sh" >&2 + exit 1 +fi + +if [ ! -f "${CONFIG}" ]; then + echo "Jaeger config not found: ${CONFIG}" >&2 + exit 1 +fi + +if ! command -v setsid >/dev/null 2>&1; then + echo "setsid is required to detach Jaeger from the launcher process group." >&2 + exit 1 +fi + +stop_pid() { + local pid="$1" + if ! kill -0 "${pid}" >/dev/null 2>&1; then + return + fi + kill "${pid}" >/dev/null 2>&1 || true + for _ in $(seq 1 50); do + if ! kill -0 "${pid}" >/dev/null 2>&1; then + return + fi + sleep 0.1 + done + kill -9 "${pid}" >/dev/null 2>&1 || true +} + +if [ -f "${PID_FILE}" ]; then + old_pid="$(cat "${PID_FILE}")" + if [ -n "${old_pid}" ]; then + stop_pid "${old_pid}" + fi + rm -f "${PID_FILE}" +fi + +if command -v pgrep >/dev/null 2>&1; then + while IFS= read -r pid; do + [ -n "${pid}" ] || continue + if [ "${pid}" = "$$" ]; then + continue + fi + stop_pid "${pid}" + done < <(pgrep -f "jaeger.*jaeger-memory.yaml" || true) +fi + +mkdir -p "$(dirname "${PID_FILE}")" "$(dirname "${LOG_FILE}")" +setsid "${JAEGER_BIN}" --config "${CONFIG}" >"${LOG_FILE}" 2>&1 "${PID_FILE}" + +deadline=$((SECONDS + WAIT_TIMEOUT_S)) +until curl -fsS "${QUERY_URL}/api/services" >/dev/null 2>&1; do + if ! kill -0 "${new_pid}" >/dev/null 2>&1; then + echo "Jaeger exited before becoming ready. Log: ${LOG_FILE}" >&2 + exit 1 + fi + if [ "${SECONDS}" -ge "${deadline}" ]; then + echo "Timed out waiting for Jaeger Query API at ${QUERY_URL}. Log: ${LOG_FILE}" >&2 + exit 1 + fi + sleep 0.5 +done + +echo "Jaeger in-memory storage restarted." +echo "PID file: ${PID_FILE}" +echo "Log file: ${LOG_FILE}" +echo "Jaeger UI: ${QUERY_URL}" +echo "OTLP gRPC: http://127.0.0.1:14317" +echo "OTLP HTTP: http://127.0.0.1:14318/v1/traces" diff --git a/tests/rl/test_trace.py b/tests/rl/test_trace.py new file mode 100644 index 0000000000..afe6ab6ac3 --- /dev/null +++ b/tests/rl/test_trace.py @@ -0,0 +1,163 @@ +import json +import os +import subprocess +import sys +import unittest +from pathlib import Path + + +def _run_trace_utils(repo_root: Path, command: str) -> dict: + env = os.environ.copy() + env["PYTHONPATH"] = os.fspath(repo_root) + os.pathsep + env.get("PYTHONPATH", "") + result = subprocess.run( + [sys.executable, os.fspath(Path(__file__).with_name("trace_utils.py")), command], + cwd=repo_root, + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True, + ) + return json.loads(result.stdout.strip().splitlines()[-1]) + + +class TestTrace(unittest.TestCase): + def test_trace_span_records_attributes_events_and_errors(self): + repo_root = Path(__file__).resolve().parents[2] + output = _run_trace_utils(repo_root, "record-span") + + self.assertEqual(output["success_attributes"]["xtuner.stage"], "unit") + self.assertEqual(output["success_attributes"]["unit.count"], 1) + self.assertEqual(output["success_events"], ["unit.event"]) + self.assertEqual(output["failure_status"], "ERROR") + self.assertEqual(output["failure_attributes"]["error"], True) + self.assertEqual(output["failure_attributes"]["error.type"], "RuntimeError") + self.assertEqual(output["failure_attributes"]["error.message"], "boom") + + def test_injected_parent_carrier_links_child_span_in_another_process(self): + repo_root = Path(__file__).resolve().parents[2] + output = _run_trace_utils(repo_root, "parent-child") + + self.assertEqual(output["child"]["trace_id"], output["parent_trace_id"]) + self.assertEqual(output["child"]["parent_span_id"], output["parent_span_id"]) + + def test_nested_trace_span_preserves_parent_to_child_order(self): + repo_root = Path(__file__).resolve().parents[2] + output = _run_trace_utils(repo_root, "nested-span-order") + + self.assertEqual(output["child_parent_span_id"], output["parent_span_id"]) + self.assertEqual(output["span_name_paths"]["order.parent"], ["order.parent"]) + self.assertEqual(output["span_name_paths"]["order.child"], ["order.parent", "order.child"]) + self.assertEqual( + output["live_sequence"], + [ + { + "event": "start", + "span_name": "order.parent", + "span_name_path": ["order.parent"], + }, + { + "event": "start", + "span_name": "order.child", + "span_name_path": ["order.parent", "order.child"], + }, + { + "event": "end", + "span_name": "order.child", + "span_name_path": ["order.parent", "order.child"], + }, + { + "event": "end", + "span_name": "order.parent", + "span_name_path": ["order.parent"], + }, + ], + ) + + def test_live_viewer_uses_span_name_path_for_display_chain(self): + from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces + + payload = build_rollout_view_payload_from_jaeger_traces( + [], + live_records=[ + { + "event": "start", + "time_s": 1.0, + "trace_id": "trace-1", + "span_id": "span-1", + "span_name": "child.phase", + "span_name_path": ["parent.phase", "child.phase"], + "attributes": { + "xtuner.rollout_id": "rollout-1", + "xtuner.status": "running", + }, + } + ], + train_step="all", + ) + + self.assertEqual( + [node["name"] for node in payload["samples"][0]["display_path"]], + ["parent.phase", "child.phase"], + ) + self.assertEqual(payload["samples"][0]["chain"], "parent.phase -> child.phase") + + def test_viewer_filters_latest_train_step_and_renders_payload(self): + from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces + from xtuner.tools.trace_viewer.render import render_rollout_trace_html + + traces = [ + { + "traceID": "trace-1", + "processes": {"p1": {"serviceName": "xtuner-test", "tags": []}}, + "spans": [ + { + "traceID": "trace-1", + "spanID": "span-1", + "operationName": "old.operation", + "processID": "p1", + "startTime": 1_000, + "duration": 1_000, + "tags": [ + {"key": "xtuner.rollout_id", "value": "rollout-1"}, + {"key": "xtuner.producer_future_step", "value": 1}, + {"key": "xtuner.stage", "value": "stage_one"}, + ], + } + ], + }, + { + "traceID": "trace-2", + "processes": {"p1": {"serviceName": "xtuner-test", "tags": []}}, + "spans": [ + { + "traceID": "trace-2", + "spanID": "span-2", + "operationName": "new.operation", + "processID": "p1", + "startTime": 2_000, + "duration": 1_000, + "tags": [ + {"key": "xtuner.rollout_id", "value": "rollout-2"}, + {"key": "xtuner.producer_future_step", "value": 2}, + {"key": "xtuner.stage", "value": "stage_two"}, + ], + } + ], + }, + ] + + payload = build_rollout_view_payload_from_jaeger_traces(traces) + html = render_rollout_trace_html(payload) + + self.assertEqual(payload["selected_train_step"], 2) + self.assertEqual(payload["available_train_steps"], [1, 2]) + self.assertEqual(payload["sample_count"], 1) + self.assertEqual(payload["samples"][0]["rollout_id"], "rollout-2") + self.assertEqual(payload["samples"][0]["stage"], "stage_two") + self.assertIn("XTuner Rollout Trace Viewer", html) + self.assertIn("stage_two", html) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/rl/trace_utils.py b/tests/rl/trace_utils.py new file mode 100644 index 0000000000..14a50df5a6 --- /dev/null +++ b/tests/rl/trace_utils.py @@ -0,0 +1,183 @@ +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path +from unittest import mock + +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor + +try: + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +except ImportError: + from opentelemetry.sdk.trace.export import InMemorySpanExporter + +import xtuner.v1.rl.trace.api as trace_api + + +def _json_safe(value): + if isinstance(value, (str, bool, int, float)) or value is None: + return value + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + return str(value) + + +def _install_in_memory_exporter(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(provider) + return exporter + + +def _emit(payload: dict) -> None: + print(json.dumps(payload, sort_keys=True)) + + +def record_span() -> None: + exporter = _install_in_memory_exporter() + + with ( + mock.patch.object(trace_api, "_ensure_trace_runtime_from_env"), + mock.patch.object(trace_api, "is_trace_enabled", return_value=True), + ): + with trace_api.trace_span("unit.parent", attributes={"xtuner.stage": "unit"}): + trace_api.set_trace_attributes({"unit.count": 1}) + trace_api.trace_event("unit.event", {"ok": True}) + + try: + with trace_api.trace_span("unit.failure"): + raise RuntimeError("boom") + except RuntimeError: + pass + + spans = {span.name: span for span in exporter.get_finished_spans()} + success = spans["unit.parent"] + failure = spans["unit.failure"] + _emit( + { + "success_attributes": { + key: _json_safe(value) for key, value in success.attributes.items() + }, + "success_events": [event.name for event in success.events], + "failure_status": failure.status.status_code.name, + "failure_attributes": { + key: _json_safe(value) for key, value in failure.attributes.items() + }, + } + ) + + +def child_span() -> None: + carrier = json.loads(os.environ["XTUNER_TEST_TRACE_CARRIER"]) + exporter = _install_in_memory_exporter() + + with ( + mock.patch.object(trace_api, "_ensure_trace_runtime_from_env"), + mock.patch.object(trace_api, "is_trace_enabled", return_value=True), + ): + with trace_api.trace_span("child.phase", parent_carrier=carrier): + pass + + (span,) = exporter.get_finished_spans() + _emit( + { + "trace_id": f"{span.context.trace_id:032x}", + "span_id": f"{span.context.span_id:016x}", + "parent_span_id": f"{span.parent.span_id:016x}" if span.parent else None, + "span_name_path": list(span.attributes.get("xtuner.span_name_path") or []), + } + ) + + +def parent_child() -> None: + exporter = _install_in_memory_exporter() + + with ( + mock.patch.object(trace_api, "_ensure_trace_runtime_from_env"), + mock.patch.object(trace_api, "is_trace_enabled", return_value=True), + ): + with trace_api.trace_span("parent.phase"): + carrier = trace_api.inject_trace_context({}) + env = os.environ.copy() + env["XTUNER_TEST_TRACE_CARRIER"] = json.dumps(carrier) + child_result = subprocess.run( + [sys.executable, os.fspath(Path(__file__).resolve()), "child-span"], + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True, + ) + + (parent_span,) = exporter.get_finished_spans() + child = json.loads(child_result.stdout.strip().splitlines()[-1]) + _emit( + { + "parent_trace_id": f"{parent_span.context.trace_id:032x}", + "parent_span_id": f"{parent_span.context.span_id:016x}", + "carrier": carrier, + "child": child, + } + ) + + +def nested_span_order() -> None: + exporter = _install_in_memory_exporter() + + with tempfile.TemporaryDirectory() as tmpdir: + live_path = Path(tmpdir) / "trace-live.jsonl" + with mock.patch.dict(os.environ, {"XTUNER_OTEL_LIVE_JSONL_PATH": os.fspath(live_path)}): + with ( + mock.patch.object(trace_api, "_ensure_trace_runtime_from_env"), + mock.patch.object(trace_api, "is_trace_enabled", return_value=True), + ): + with trace_api.trace_span("order.parent"): + with trace_api.trace_span("order.child"): + pass + + live_records = [json.loads(line) for line in live_path.read_text(encoding="utf-8").splitlines()] + + spans = {span.name: span for span in exporter.get_finished_spans()} + parent = spans["order.parent"] + child = spans["order.child"] + _emit( + { + "parent_span_id": f"{parent.context.span_id:016x}", + "child_parent_span_id": f"{child.parent.span_id:016x}" if child.parent else None, + "span_name_paths": { + name: list(span.attributes.get("xtuner.span_name_path") or []) + for name, span in spans.items() + }, + "live_sequence": [ + { + "event": record.get("event"), + "span_name": record.get("span_name"), + "span_name_path": record.get("span_name_path"), + } + for record in live_records + ], + } + ) + + +def main() -> None: + command = sys.argv[1] if len(sys.argv) > 1 else "" + if command == "record-span": + record_span() + elif command == "child-span": + child_span() + elif command == "parent-child": + parent_child() + elif command == "nested-span-order": + nested_span_order() + else: + raise SystemExit(f"unknown trace utils command: {command}") + + +if __name__ == "__main__": + main() diff --git a/xtuner/tools/trace_viewer/__init__.py b/xtuner/tools/trace_viewer/__init__.py new file mode 100644 index 0000000000..a2ef41bd45 --- /dev/null +++ b/xtuner/tools/trace_viewer/__init__.py @@ -0,0 +1,18 @@ +from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces +from xtuner.tools.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html +from xtuner.tools.trace_viewer.source import ( + JaegerQuerySource, + JsonlTraceSource, + fetch_jaeger_traces, + normalize_jaeger_query_url, +) + +__all__ = [ + "JaegerQuerySource", + "JsonlTraceSource", + "build_rollout_view_payload_from_jaeger_traces", + "fetch_jaeger_traces", + "normalize_jaeger_query_url", + "render_rollout_trace_html", + "write_rollout_trace_html", +] diff --git a/xtuner/tools/trace_viewer/payload.py b/xtuner/tools/trace_viewer/payload.py new file mode 100644 index 0000000000..3cdb5daa43 --- /dev/null +++ b/xtuner/tools/trace_viewer/payload.py @@ -0,0 +1,1258 @@ +from __future__ import annotations + +import json +import math +import time +from collections import Counter, defaultdict +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +_SPAN_NAME_PATH_ATTRIBUTE = "xtuner.span_name_path" +_LEGACY_LOGICAL_PATH_ATTRIBUTE = "xtuner.logical_path" + +_INITIAL_SAMPLE_STATUSES = {"init"} +_NON_TERMINAL_SAMPLE_STATUSES = {"", "init", "pending", "queued", "running", "scheduled", "started", "unknown"} +_ERROR_SAMPLE_STATUSES = {"aborted", "error", "exception", "failed", "timeout", "timed_out"} +_TERMINAL_STAGE_STATUSES = { + "completed", + "failed", + "aborted", + "timeout", + "timed_out", + "expired", + "stale", + "filtered", +} + + +def build_rollout_view_payload_from_jaeger_traces( + traces: Iterable[dict[str, Any]], + *, + jaeger_query_url: str | None = None, + jaeger_link_url: str | None = None, + live_records: Iterable[dict[str, Any]] | None = None, + service_name: str | None = None, + run_id: str | None = None, + train_step: Any = "latest", +) -> dict[str, Any]: + samples_by_key: dict[tuple[str, Any], dict[str, Any]] = {} + jaeger_trace_link_base_url = _jaeger_trace_link_base_url(jaeger_query_url, jaeger_link_url) + + for trace_data in traces: + trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") + if not trace_id: + continue + process_metadata = _process_metadata(trace_data) + span_entries = [] + entries_by_span_id: dict[str, dict[str, Any]] = {} + for span in trace_data.get("spans") or []: + process = process_metadata.get(str(span.get("processID") or ""), {}) + span_service_name = process.get("service_name") + if service_name is not None and span_service_name != service_name: + continue + tags = _tags_to_dict(span.get("tags") or []) + span_run_id = tags.get("run.id") or process.get("run_id") + if run_id is not None and span_run_id != run_id: + continue + entry = { + "span": span, + "tags": tags, + "service_name": span_service_name, + "run_id": span_run_id, + "span_id": _span_id(span), + "trace_id": trace_id, + } + span_entries.append(entry) + if entry["span_id"]: + entries_by_span_id[entry["span_id"]] = entry + + for entry in span_entries: + span = entry["span"] + tags = entry["tags"] + rollout_id, sample_tags = _resolve_rollout_sample(entry, entries_by_span_id) + if rollout_id is None: + continue + span_service_name = entry["service_name"] + span_run_id = entry["run_id"] + sample_key = (trace_id, rollout_id) + sample = samples_by_key.setdefault( + sample_key, + { + "trace_id": trace_id, + "rollout_id": rollout_id, + "group_id": sample_tags.get("xtuner.group_id"), + "producer_future_step": _producer_future_step(sample_tags), + "task_name": sample_tags.get("xtuner.task_name"), + "status": sample_tags.get("xtuner.status"), + "service_name": span_service_name, + "run_id": span_run_id, + "jaeger_url": _jaeger_trace_url(jaeger_trace_link_base_url, trace_id), + "spans": [], + }, + ) + _merge_sample_fields(sample, sample_tags) + if span_run_id is not None: + sample["run_id"] = span_run_id + sample["spans"].append(_span_payload(span, tags, service_name=span_service_name, run_id=span_run_id)) + + _merge_live_records(samples_by_key, live_records or (), jaeger_trace_link_base_url) + + generated_at_s = time.time() + samples = [] + for sample in samples_by_key.values(): + sample["spans"].sort(key=lambda item: (item["start_time_us"], item["span_id"])) + sample["span_count"] = len(sample["spans"]) + _apply_live_state(sample, generated_at_s) + _apply_sample_display_status(sample) + _apply_sample_reward_filter(sample) + sample["stage"] = _sample_stage(sample) + samples.append(sample) + + samples.sort(key=lambda item: (str(item.get("group_id")), str(item.get("rollout_id")), item["trace_id"])) + base_payload = { + "title": "XTuner Rollout Trace Viewer", + "generated_at_s": generated_at_s, + "source": "jaeger", + "jaeger_query_url": _normalize_jaeger_query_url(jaeger_query_url), + "jaeger_link_url": jaeger_trace_link_base_url, + "service_name": service_name, + "run_id": run_id, + "available_train_steps": _available_train_steps(samples), + "samples": samples, + } + return filter_rollout_view_payload_by_train_step(base_payload, train_step) + + +def load_jaeger_traces_from_otel_jsonl(trace_jsonl_path: Path | str) -> list[dict[str, Any]]: + traces_by_id: dict[str, dict[str, Any]] = {} + process_ids: dict[tuple[str, str, tuple[tuple[str, str], ...]], str] = {} + path = Path(trace_jsonl_path).expanduser() + if not path.is_file(): + return [] + + for record in _iter_jsonl_records(path): + if not isinstance(record, dict): + continue + jaeger_traces = _jaeger_traces_from_json_record(record) + if jaeger_traces is not None: + for trace_data in jaeger_traces: + _merge_jaeger_trace(traces_by_id, trace_data) + continue + for resource_span in record.get("resourceSpans") or []: + if not isinstance(resource_span, dict): + continue + resource_attrs = _otel_attributes_to_dict( + (resource_span.get("resource") or {}).get("attributes") or [] + ) + service_name = str(resource_attrs.get("service.name") or "unknown") + process_tags = _dict_to_jaeger_tags(resource_attrs) + for scope_span in _otel_scope_spans(resource_span): + for otel_span in scope_span.get("spans") or []: + if not isinstance(otel_span, dict): + continue + trace_id = str( + otel_span.get("traceId") + or otel_span.get("traceID") + or otel_span.get("trace_id") + or "" + ) + span_id = str( + otel_span.get("spanId") + or otel_span.get("spanID") + or otel_span.get("span_id") + or "" + ) + if not trace_id or not span_id: + continue + trace_data = traces_by_id.setdefault( + trace_id, + { + "traceID": trace_id, + "processes": {}, + "spans": [], + }, + ) + process_key = ( + trace_id, + service_name, + tuple(sorted((str(key), str(value)) for key, value in resource_attrs.items())), + ) + process_id = process_ids.get(process_key) + if process_id is None: + process_id = f"p{len(trace_data['processes']) + 1}" + process_ids[process_key] = process_id + trace_data["processes"][process_id] = { + "serviceName": service_name, + "tags": process_tags, + } + trace_data["spans"].append( + _otel_span_to_jaeger_span(otel_span, trace_id, span_id, process_id) + ) + return list(traces_by_id.values()) + + +def load_live_trace_records(live_jsonl_path: Path | str | None) -> list[dict[str, Any]]: + if live_jsonl_path is None: + return [] + path = Path(live_jsonl_path).expanduser() + if not path.is_file(): + return [] + return list(_iter_jsonl_records(path)) + + +def _iter_jsonl_records(path: Path) -> Iterable[dict[str, Any]]: + with path.open("r", encoding="utf-8") as handle: + for line in handle: + stripped = line.strip() + if not stripped: + continue + try: + payload = json.loads(stripped) + except json.JSONDecodeError: + continue + if isinstance(payload, dict): + yield payload + + +def _jaeger_traces_from_json_record(record: dict[str, Any]) -> list[dict[str, Any]] | None: + data = record.get("data") + if isinstance(data, list): + return [trace_data for trace_data in data if isinstance(trace_data, dict)] + if record.get("traceID") is not None and isinstance(record.get("spans"), list): + return [record] + return None + + +def _merge_jaeger_trace(traces_by_id: dict[str, dict[str, Any]], trace_data: dict[str, Any]) -> None: + trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") + if not trace_id: + return + target = traces_by_id.setdefault(trace_id, {"traceID": trace_id, "processes": {}, "spans": []}) + if isinstance(trace_data.get("processes"), dict): + target["processes"].update(trace_data["processes"]) + target["spans"].extend([span for span in trace_data.get("spans") or [] if isinstance(span, dict)]) + + +def _otel_scope_spans(resource_span: dict[str, Any]) -> list[dict[str, Any]]: + scope_spans = resource_span.get("scopeSpans") + if isinstance(scope_spans, list): + return [scope_span for scope_span in scope_spans if isinstance(scope_span, dict)] + legacy_scope_spans = resource_span.get("instrumentationLibrarySpans") + if isinstance(legacy_scope_spans, list): + return [scope_span for scope_span in legacy_scope_spans if isinstance(scope_span, dict)] + return [] + + +def _otel_span_to_jaeger_span( + otel_span: dict[str, Any], + trace_id: str, + span_id: str, + process_id: str, +) -> dict[str, Any]: + attributes = _otel_attributes_to_dict(otel_span.get("attributes") or []) + tags = _dict_to_jaeger_tags(attributes) + tags.extend(_otel_status_tags(otel_span.get("status") or {})) + start_ns = _int_from_otel_time( + otel_span.get("startTimeUnixNano") + or otel_span.get("start_time_unix_nano") + or otel_span.get("startTime") + or 0 + ) + end_ns = _int_from_otel_time( + otel_span.get("endTimeUnixNano") + or otel_span.get("end_time_unix_nano") + or otel_span.get("endTime") + or start_ns + ) + parent_span_id = otel_span.get("parentSpanId") or otel_span.get("parent_span_id") + references = [] + if parent_span_id is not None and str(parent_span_id): + references.append({"refType": "CHILD_OF", "traceID": trace_id, "spanID": str(parent_span_id)}) + return { + "traceID": trace_id, + "spanID": span_id, + "operationName": str(otel_span.get("name") or otel_span.get("operationName") or "unknown"), + "processID": process_id, + "startTime": start_ns // 1_000, + "duration": max(0, end_ns - start_ns) // 1_000, + "references": references, + "tags": tags, + } + + +def _otel_status_tags(status: dict[str, Any]) -> list[dict[str, Any]]: + code = str(status.get("code") or status.get("statusCode") or "STATUS_CODE_UNSET") + if code in {"STATUS_CODE_ERROR", "ERROR", "2"}: + normalized = "ERROR" + elif code in {"STATUS_CODE_OK", "OK", "1"}: + normalized = "OK" + else: + normalized = "UNSET" + tags = [{"key": "otel.status_code", "type": "string", "value": normalized}] + message = status.get("message") or status.get("description") + if message: + tags.append({"key": "otel.status_description", "type": "string", "value": str(message)}) + if normalized == "ERROR": + tags.append({"key": "error.message", "type": "string", "value": str(message)}) + return tags + + +def _otel_attributes_to_dict(attributes: Any) -> dict[str, Any]: + if isinstance(attributes, dict): + return dict(attributes) + result: dict[str, Any] = {} + for attribute in attributes or []: + if not isinstance(attribute, dict): + continue + key = attribute.get("key") + if key is None: + continue + result[str(key)] = _otel_any_value_to_python(attribute.get("value")) + return result + + +def _otel_any_value_to_python(value: Any) -> Any: + if not isinstance(value, dict): + return value + if "stringValue" in value: + return value["stringValue"] + if "boolValue" in value: + return bool(value["boolValue"]) + if "intValue" in value: + return _int_or_original(value["intValue"]) + if "doubleValue" in value: + return float(value["doubleValue"]) + if "bytesValue" in value: + return value["bytesValue"] + if "arrayValue" in value: + return [_otel_any_value_to_python(item) for item in (value["arrayValue"].get("values") or [])] + if "kvlistValue" in value: + return { + str(item.get("key")): _otel_any_value_to_python(item.get("value")) + for item in (value["kvlistValue"].get("values") or []) + if isinstance(item, dict) and item.get("key") is not None + } + return value + + +def _dict_to_jaeger_tags(attributes: dict[str, Any]) -> list[dict[str, Any]]: + return [ + {"key": str(key), "type": _jaeger_tag_type(value), "value": value} + for key, value in attributes.items() + if value is not None + ] + + +def _jaeger_tag_type(value: Any) -> str: + if isinstance(value, bool): + return "bool" + if isinstance(value, int): + return "int64" + if isinstance(value, float): + return "float64" + return "string" + + +def _int_from_otel_time(value: Any) -> int: + parsed = _int_or_original(value) + return parsed if isinstance(parsed, int) and not isinstance(parsed, bool) else 0 + + +def _int_or_original(value: Any) -> int | Any: + if isinstance(value, bool): + return value + if isinstance(value, int): + return value + try: + return int(str(value)) + except (TypeError, ValueError): + return value + + +def _merge_sample_fields(sample: dict[str, Any], tags: dict[str, Any]) -> None: + for sample_key, tag_key in ( + ("rollout_id", "xtuner.rollout_id"), + ("group_id", "xtuner.group_id"), + ("task_name", "xtuner.task_name"), + ): + if tags.get(tag_key) is not None: + sample[sample_key] = tags[tag_key] + _merge_sample_status(sample, tags.get("xtuner.status")) + producer_future_step = _producer_future_step(tags) + if producer_future_step is not None: + sample["producer_future_step"] = producer_future_step + + +def _merge_sample_status(sample: dict[str, Any], status: Any) -> None: + if status is None: + return + current_status = sample.get("status") + if current_status is None or _sample_status_priority(status) >= _sample_status_priority(current_status): + sample["status"] = status + + +def _apply_sample_display_status(sample: dict[str, Any]) -> None: + status = str(sample.get("status") or "").strip().lower() + if status not in _INITIAL_SAMPLE_STATUSES: + return + if _sample_has_observed_stage(sample): + sample["status"] = "running" + + +def _sample_has_observed_stage(sample: dict[str, Any]) -> bool: + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict) and current_stage.get("name"): + return True + return bool(sample.get("spans")) + + +def _sample_status_priority(status: Any) -> int: + normalized = str(status or "").strip().lower() + if normalized in _ERROR_SAMPLE_STATUSES: + return 3 + if normalized in _TERMINAL_STAGE_STATUSES: + return 2 + if normalized in _NON_TERMINAL_SAMPLE_STATUSES: + return 0 + return 1 + + +def _merge_live_records( + samples_by_key: dict[tuple[str, Any], dict[str, Any]], + live_records: Iterable[dict[str, Any]], + jaeger_trace_link_base_url: str | None, +) -> None: + for record in live_records: + if not isinstance(record, dict): + continue + trace_id = str(record.get("trace_id") or "") + attributes = record.get("attributes") + if not trace_id or not isinstance(attributes, dict): + continue + rollout_id = attributes.get("xtuner.rollout_id") + if rollout_id is None: + continue + sample_key = (trace_id, rollout_id) + sample = samples_by_key.setdefault( + sample_key, + { + "trace_id": trace_id, + "rollout_id": rollout_id, + "group_id": attributes.get("xtuner.group_id"), + "producer_future_step": _producer_future_step(attributes), + "task_name": attributes.get("xtuner.task_name"), + "status": attributes.get("xtuner.status"), + "service_name": None, + "run_id": None, + "jaeger_url": _jaeger_trace_url(jaeger_trace_link_base_url, trace_id), + "spans": [], + }, + ) + _merge_sample_fields(sample, attributes) + sample.setdefault("_live_records", []).append(record) + + +def _apply_live_state(sample: dict[str, Any], generated_at_s: float) -> None: + live_states = _live_span_states(sample.pop("_live_records", [])) + active_states = [state for state in live_states if state.get("status") == "running"] + active_states.sort(key=lambda state: (len(state.get("span_name_path") or []), float(state.get("started_at_s") or 0.0))) + current_state = active_states[-1] if active_states else None + if current_state is not None: + started_at_s = float(current_state.get("started_at_s") or generated_at_s) + span_name = str(current_state.get("span_name") or "") + sample["current_stage"] = { + "name": span_name, + "stage": str(current_state.get("stage") or span_name or "unknown"), + "status": "running", + "elapsed_ms": round(max(0.0, generated_at_s - started_at_s) * 1000.0, 3), + "started_at_s": started_at_s, + } + else: + sample["current_stage"] = None + sample["live_spans"] = live_states + sample["display_path"] = _build_display_path(sample, live_states, current_state, generated_at_s) + sample["chain"] = " -> ".join(node["name"] for node in sample["display_path"]) + + +def _live_span_states(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + states: dict[str, dict[str, Any]] = {} + for index, record in enumerate(records): + span_name = str(record.get("span_name") or "") + attributes = record.get("attributes") if isinstance(record.get("attributes"), dict) else {} + span_id = str(record.get("span_id") or "") + if not span_id: + span_id = f"live:{span_name}:{index}" + state = states.setdefault( + span_id, + { + "span_id": span_id, + "span_name": span_name, + "stage": _stage_from_span_name_and_attributes(span_name, attributes), + "span_name_path": _span_name_path_from_value( + record.get("span_name_path") or record.get("logical_path") + ), + "attributes": attributes, + "status": "running", + }, + ) + event = str(record.get("event") or "") + if event == "start": + state["started_at_s"] = _float_or_none(record.get("time_s")) + state["status"] = "running" + elif event == "end": + state["ended_at_s"] = _float_or_none(record.get("time_s")) + state["duration_ms"] = _float_or_none(record.get("duration_ms")) + state["status"] = str(record.get("status") or "completed") + if record.get("error_message"): + state["error_message"] = record["error_message"] + if record.get("trace_id") is not None: + state["trace_id"] = record["trace_id"] + return sorted(states.values(), key=lambda state: (float(state.get("started_at_s") or 0.0), str(state.get("span_id")))) + + +def _build_display_path( + sample: dict[str, Any], + live_states: list[dict[str, Any]], + current_state: dict[str, Any] | None, + generated_at_s: float, +) -> list[dict[str, Any]]: + spans = sample.get("spans") or [] + spans_by_name = {str(span.get("name") or ""): span for span in spans} + live_by_name = {str(state.get("span_name") or ""): state for state in live_states} + path = _display_path_names(spans, live_states, current_state) + nodes = [] + for name in path: + span = spans_by_name.get(name) + live_state = live_by_name.get(name) + if current_state is not None and name == current_state.get("span_name"): + started_at_s = float(current_state.get("started_at_s") or generated_at_s) + nodes.append( + { + "name": name, + "stage": current_state.get("stage") or name, + "source": "live", + "status": "running", + "elapsed_ms": round(max(0.0, generated_at_s - started_at_s) * 1000.0, 3), + } + ) + elif span is not None: + nodes.append( + { + "name": name, + "stage": _span_semantic_stage(span), + "source": "span", + "status": "done" if str(span.get("status") or "").upper() != "ERROR" else "error", + "duration_ms": span.get("duration_ms"), + } + ) + elif live_state is not None and live_state.get("status") == "running": + nodes.append({"name": name, "stage": live_state.get("stage") or name, "source": "live", "status": "active"}) + else: + nodes.append({"name": name, "source": "logical", "status": "inferred"}) + return nodes + + +def _display_path_names( + spans: list[dict[str, Any]], + live_states: list[dict[str, Any]], + current_state: dict[str, Any] | None, +) -> list[str]: + path = _span_name_path_from_value(current_state.get("span_name_path")) if current_state is not None else [] + if not path: + path = [str(span.get("name") or "") for span in spans if span.get("name")] + if not path: + path = [str(state.get("span_name") or "") for state in live_states if state.get("span_name")] + if not path: + span_paths = [ + _span_name_path_from_span_attributes(span.get("attributes") or {}) + for span in spans + ] + path = max(span_paths, key=len, default=[]) + if not path: + live_paths = [_span_name_path_from_value(state.get("span_name_path")) for state in live_states] + path = max(live_paths, key=len, default=[]) + return _unique_path(path) + + +def _span_name_path_from_span_attributes(attributes: dict[str, Any]) -> list[str]: + return _span_name_path_from_value( + attributes.get(_SPAN_NAME_PATH_ATTRIBUTE) or attributes.get(_LEGACY_LOGICAL_PATH_ATTRIBUTE) + ) + + +def _span_name_path_from_value(value: Any) -> list[str]: + if isinstance(value, str): + try: + decoded = json.loads(value) + except json.JSONDecodeError: + decoded = [part.strip() for part in value.split("->")] + value = decoded + if isinstance(value, (list, tuple)): + return [str(item).strip() for item in value if isinstance(item, str) and item.strip()] + return [] + + +def _unique_path(path: list[str]) -> list[str]: + result = [] + for name in path: + if not result or result[-1] != name: + result.append(name) + return result + + +def _producer_future_step(tags: dict[str, Any]) -> Any: + value = tags.get("xtuner.producer_future_step") + return value if value is not None else tags.get("xtuner.train_step") + + +def _build_step_group_summaries(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: + steps: dict[str, dict[str, Any]] = {} + for sample in samples: + producer_future_step = sample.get("producer_future_step") + group_id = sample.get("group_id") + step_key = _summary_key(producer_future_step) + group_key = _summary_key(group_id) + + step_summary = steps.setdefault( + step_key, + { + "producer_future_step": producer_future_step, + "sample_count": 0, + "groups": {}, + }, + ) + step_summary["sample_count"] += 1 + groups = step_summary["groups"] + group_summary = groups.setdefault( + group_key, + { + "group_id": group_id, + "sample_count": 0, + "statuses": Counter(), + "stages": Counter(), + "rollout_ids": [], + }, + ) + group_summary["sample_count"] += 1 + group_summary["statuses"][str(sample.get("status") or "unknown")] += 1 + if _should_show_group_stage(sample): + group_summary["stages"][str(sample.get("stage") or "unknown")] += 1 + group_summary["rollout_ids"].append(sample.get("rollout_id")) + + summaries = [] + for step_summary in steps.values(): + groups = [] + for group_summary in step_summary["groups"].values(): + group_summary["statuses"] = dict(sorted(group_summary["statuses"].items())) + group_summary["stages"] = dict(sorted(group_summary["stages"].items())) + group_summary["rollout_ids"].sort(key=lambda value: str(value)) + groups.append(group_summary) + groups.sort(key=lambda item: _sortable_summary_value(item["group_id"])) + summaries.append( + { + "producer_future_step": step_summary["producer_future_step"], + "group_count": sum(1 for group in groups if group["group_id"] is not None), + "sample_count": step_summary["sample_count"], + "groups": groups, + } + ) + summaries.sort(key=lambda item: _sortable_summary_value(item["producer_future_step"])) + return summaries + + +def _available_train_steps(samples: list[dict[str, Any]]) -> list[Any]: + values = {sample.get("producer_future_step") for sample in samples if sample.get("producer_future_step") is not None} + return sorted(values, key=_sortable_summary_value) + + +def _select_train_step(requested: Any, available_steps: list[Any]) -> Any: + if not available_steps: + return "all" + if requested is None or requested == "": + requested = "latest" + requested_text = str(requested).strip() + if requested_text.lower() == "all": + return "all" + if requested_text.lower() == "latest": + return max(available_steps, key=_sortable_summary_value) + for step in available_steps: + if str(step) == requested_text: + return step + return max(available_steps, key=_sortable_summary_value) + + +def _filter_samples_by_train_step(samples: list[dict[str, Any]], selected_train_step: Any) -> list[dict[str, Any]]: + if selected_train_step == "all": + return samples + return [sample for sample in samples if str(sample.get("producer_future_step")) == str(selected_train_step)] + + +def filter_rollout_view_payload_by_train_step(payload: dict[str, Any], train_step: Any = "latest") -> dict[str, Any]: + samples = list(payload.get("samples") or []) + generated_at_s = float(payload.get("generated_at_s") or time.time()) + available_train_steps = list(payload.get("available_train_steps") or _available_train_steps(samples)) + selected_train_step = _select_train_step(train_step, available_train_steps) + visible_samples = _filter_samples_by_train_step(samples, selected_train_step) + + status_counts: Counter[str] = Counter() + stage_counts: Counter[str] = Counter() + group_ids: set[Any] = set() + for sample in visible_samples: + status = str(sample.get("status") or "unknown") + status_counts[status] += 1 + stage_counts[str(sample["stage"])] += 1 + if sample.get("group_id") is not None: + group_ids.add(sample["group_id"]) + + step_group_summaries = _build_step_group_summaries(visible_samples) + filtered_payload = dict(payload) + filtered_payload.update( + { + "generated_at_s": generated_at_s, + "selected_train_step": selected_train_step, + "available_train_steps": available_train_steps, + "total_sample_count": len(samples), + "sample_count": len(visible_samples), + "group_count": len(group_ids), + "step_count": len(step_group_summaries), + "step_group_summaries": step_group_summaries, + "stage_occupancy": _build_stage_occupancy(visible_samples, generated_at_s), + "stage_duration_summaries": _build_stage_duration_summaries(visible_samples), + "status_counts": dict(sorted(status_counts.items())), + "stage_counts": dict(sorted(stage_counts.items())), + "samples": visible_samples, + } + ) + return filtered_payload + + +def _build_stage_duration_summaries(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: + durations_by_stage: dict[str, list[float]] = defaultdict(list) + raw_durations_by_stage: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list)) + errors_by_stage: dict[str, list[dict[str, Any]]] = defaultdict(list) + for sample in samples: + for span in sample.get("spans") or []: + try: + duration_s = float(span.get("duration_ms") or 0.0) / 1000.0 + except (TypeError, ValueError): + continue + stage = _span_semantic_stage(span) + raw_span_name = _span_raw_name(span) + durations_by_stage[stage].append(duration_s) + raw_durations_by_stage[stage][raw_span_name].append(duration_s) + error = _extract_sample_error(sample) + if error is not None: + errors_by_stage[str(error.get("stage") or sample.get("stage") or "unknown")].append( + { + **error, + "rollout_id": sample.get("rollout_id"), + "group_id": sample.get("group_id"), + "producer_future_step": sample.get("producer_future_step"), + "jaeger_url": sample.get("jaeger_url"), + } + ) + + summaries = [] + for stage, durations in durations_by_stage.items(): + durations.sort() + top_errors = _summarize_stage_errors(errors_by_stage.get(stage, [])) + summaries.append( + { + "stage": stage, + "span_count": len(durations), + "avg_duration_s": _round_duration(sum(durations) / len(durations)), + "p50_duration_s": _round_duration(_nearest_rank_percentile(durations, 0.50)), + "p95_duration_s": _round_duration(_nearest_rank_percentile(durations, 0.95)), + "max_duration_s": _round_duration(durations[-1]), + "error_count": sum(error["sample_count"] for error in top_errors), + "top_errors": top_errors, + "raw_spans": _summarize_raw_span_durations(raw_durations_by_stage.get(stage, {})), + } + ) + for stage, errors in errors_by_stage.items(): + if stage in durations_by_stage: + continue + top_errors = _summarize_stage_errors(errors) + summaries.append( + { + "stage": stage, + "span_count": 0, + "avg_duration_s": 0.0, + "p50_duration_s": 0.0, + "p95_duration_s": 0.0, + "max_duration_s": 0.0, + "error_count": sum(error["sample_count"] for error in top_errors), + "top_errors": top_errors, + "raw_spans": [], + } + ) + return summaries + + +def _summarize_raw_span_durations(raw_durations: dict[str, list[float]]) -> list[dict[str, Any]]: + rows = [] + for span_name, durations in raw_durations.items(): + if not durations: + continue + durations = sorted(durations) + rows.append( + { + "span": span_name, + "span_count": len(durations), + "avg_duration_s": _round_duration(sum(durations) / len(durations)), + "max_duration_s": _round_duration(durations[-1]), + } + ) + rows.sort(key=lambda item: (-item["span_count"], str(item["span"]))) + return rows + + +def _summarize_stage_errors(errors: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[tuple[str, str, Any], dict[str, Any]] = {} + for error in errors: + key = ( + str(error.get("error_type") or "error"), + str(error.get("message") or ""), + error.get("http_status_code"), + ) + summary = grouped.setdefault( + key, + { + "error_type": key[0], + "message": key[1], + "http_status_code": key[2], + "sample_count": 0, + "rollout_ids": [], + "groups": [], + "steps": [], + "jaeger_urls": [], + }, + ) + summary["sample_count"] += 1 + _append_unique(summary["rollout_ids"], error.get("rollout_id")) + _append_unique(summary["groups"], error.get("group_id")) + _append_unique(summary["steps"], error.get("producer_future_step")) + _append_unique(summary["jaeger_urls"], error.get("jaeger_url")) + result = list(grouped.values()) + for summary in result: + summary["rollout_ids"].sort(key=lambda value: str(value)) + summary["groups"].sort(key=lambda value: str(value)) + summary["steps"].sort(key=lambda value: str(value)) + result.sort(key=lambda item: (-item["sample_count"], str(item["error_type"]), str(item["message"]))) + return result + + +def _build_stage_occupancy(samples: list[dict[str, Any]], generated_at_s: float) -> list[dict[str, Any]]: + buckets: dict[str, dict[str, Any]] = {} + for sample in samples: + stage, raw_span_name = _stage_bucket_for_sample(sample) + latest_start_s = _latest_span_start_s(sample) + age_s = 0.0 if stage in _TERMINAL_STAGE_STATUSES else max(0.0, generated_at_s - latest_start_s) + bucket = buckets.setdefault( + stage, + { + "stage": stage, + "sample_count": 0, + "_group_ids": set(), + "_raw_spans": {}, + "oldest_age_s": 0.0, + "oldest_rollout_id": None, + "oldest_group_id": None, + "oldest_producer_future_step": None, + }, + ) + bucket["sample_count"] += 1 + if sample.get("group_id") is not None: + bucket["_group_ids"].add(sample["group_id"]) + if raw_span_name: + raw_spans = bucket["_raw_spans"] + raw_bucket = raw_spans.setdefault(raw_span_name, {"span": raw_span_name, "sample_count": 0, "_group_ids": set()}) + raw_bucket["sample_count"] += 1 + if sample.get("group_id") is not None: + raw_bucket["_group_ids"].add(sample["group_id"]) + if age_s >= bucket["oldest_age_s"]: + bucket["oldest_age_s"] = _round_duration(age_s) + bucket["oldest_rollout_id"] = sample.get("rollout_id") + bucket["oldest_group_id"] = sample.get("group_id") + bucket["oldest_producer_future_step"] = sample.get("producer_future_step") + rows = list(buckets.values()) + for row in rows: + row["group_count"] = len(row.pop("_group_ids")) + row["raw_spans"] = _summarize_raw_span_occupancy(row.pop("_raw_spans")) + return rows + + +def _summarize_raw_span_occupancy(raw_spans: dict[str, dict[str, Any]]) -> list[dict[str, Any]]: + rows = [] + for raw_bucket in raw_spans.values(): + rows.append( + { + "span": raw_bucket["span"], + "sample_count": raw_bucket["sample_count"], + "group_count": len(raw_bucket["_group_ids"]), + } + ) + rows.sort(key=lambda item: (-item["sample_count"], str(item["span"]))) + return rows + + +def _stage_bucket_for_sample(sample: dict[str, Any]) -> tuple[str, str | None]: + status = str(sample.get("status") or "").strip().lower() + if status in _TERMINAL_STAGE_STATUSES: + return status, None + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict): + raw_stage_name = str(current_stage.get("name") or "").strip() + semantic_stage = str(current_stage.get("stage") or "").strip() + if semantic_stage: + return semantic_stage, raw_stage_name or None + if raw_stage_name: + return raw_stage_name, raw_stage_name + spans = sample.get("spans") or [] + if not spans: + return status or "unknown", None + latest_span = spans[-1] + if latest_span.get("rollout_backend"): + return str(latest_span.get("stage") or "llm_generate"), _span_raw_name(latest_span) + if status in {"pending", "queued", "scheduled"}: + return "scheduled", None + return _span_semantic_stage(latest_span), _span_raw_name(latest_span) + + +def _latest_span_start_s(sample: dict[str, Any]) -> float: + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict) and current_stage.get("started_at_s") is not None: + try: + return float(current_stage["started_at_s"]) + except (TypeError, ValueError): + pass + spans = sample.get("spans") or [] + if not spans: + return time.time() + return max(float(span.get("start_time_us") or 0) / 1_000_000.0 for span in spans) + + +def _apply_sample_reward_filter(sample: dict[str, Any]) -> None: + values: dict[str, Any] = {} + for span in sample.get("spans") or []: + attrs = span.get("attributes") or {} + for target, keys in ( + ("reward_score", ("reward.score", "reward_score")), + ("reward_pass", ("reward.pass",)), + ("filter_decision", ("filter.decision",)), + ("filter_reason", ("filter.reason",)), + ("train_included", ("train.included",)), + ("oversample_source", ("oversample.source",)), + ("drop_reason", ("drop.reason",)), + ): + for key in keys: + if key in attrs: + values[target] = attrs[key] + if "reward_score" in values: + values["reward_score"] = _to_float(values["reward_score"]) + if "reward_pass" in values: + values["reward_pass"] = _to_bool(values["reward_pass"]) + if "train_included" in values: + values["train_included"] = _to_bool(values["train_included"]) + sample.update(values) + + +def _extract_sample_error(sample: dict[str, Any]) -> dict[str, Any] | None: + fallback_span_name = str(sample.get("stage") or "unknown") + for span in sample.get("spans") or []: + attrs = span.get("attributes") or {} + http_status = attrs.get("http.status_code") + is_http_error = False + try: + is_http_error = http_status is not None and int(http_status) >= 400 + except (TypeError, ValueError): + is_http_error = False + has_error = ( + str(span.get("status") or "").upper() == "ERROR" + or _to_bool(attrs.get("error")) is True + or any(key.startswith("error.") for key in attrs) + or any(key.startswith("exception.") for key in attrs) + or is_http_error + ) + if not has_error: + continue + return { + "error_type": attrs.get("exception.type") + or attrs.get("error.type") + or attrs.get("xtuner.error_type") + or ("HTTPError" if is_http_error else str(sample.get("status") or "error")), + "message": attrs.get("xtuner.error_msg") + or attrs.get("error.message") + or attrs.get("exception.message") + or (f"http_status={http_status}" if is_http_error else ""), + "http_status_code": http_status, + "span_name": span.get("name") or fallback_span_name, + "stage": _span_semantic_stage(span), + } + status = str(sample.get("status") or "").lower() + if status in _ERROR_SAMPLE_STATUSES: + return { + "error_type": status, + "message": "", + "http_status_code": None, + "span_name": fallback_span_name, + "stage": fallback_span_name, + } + return None + + +def _nearest_rank_percentile(sorted_values: list[float], percentile: float) -> float: + if not sorted_values: + return 0.0 + index = max(0, min(len(sorted_values) - 1, math.ceil(percentile * len(sorted_values)) - 1)) + return sorted_values[index] + + +def _round_duration(value: float) -> float: + return round(value, 3) + + +def _span_raw_name(span: dict[str, Any]) -> str: + return str(span.get("name") or "unknown") + + +def _span_semantic_stage(span: dict[str, Any]) -> str: + name = _span_raw_name(span) + attributes = span.get("attributes") + return _stage_from_span_name_and_attributes(name, attributes if isinstance(attributes, dict) else None) + + +def _stage_from_span_name_and_attributes(span_name: str, attributes: dict[str, Any] | None = None) -> str: + attrs = attributes or {} + for key in ("xtuner.stage", "stage", "stage.name"): + stage = str(attrs.get(key) or "").strip() + if stage: + return stage + return span_name or "unknown" + + +def _sample_stage(sample: dict[str, Any]) -> str: + status = str(sample.get("status") or "").strip().lower() + if status and status not in _NON_TERMINAL_SAMPLE_STATUSES: + return status + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict): + semantic_stage = str(current_stage.get("stage") or "").strip() + if semantic_stage: + return semantic_stage + if current_stage.get("name"): + return str(current_stage["name"]) + + spans = sample.get("spans") or [] + for span in spans: + attributes = span.get("attributes") or {} + error_value = str(attributes.get("error") or "").strip().lower() + if str(span.get("status") or "").upper() == "ERROR" or error_value == "true": + return "error" + if spans: + return _span_semantic_stage(spans[-1]) + return status or "unknown" + + +def _should_show_group_stage(sample: dict[str, Any]) -> bool: + status = str(sample.get("status") or "unknown").strip().lower() + stage = str(sample.get("stage") or "unknown").strip().lower() + return stage != status + + +def _summary_key(value: Any) -> str: + return "" if value is None else str(value) + + +def _sortable_summary_value(value: Any) -> tuple[int, float | str]: + if value is None: + return (1, "") + if isinstance(value, bool): + return (0, str(value)) + if isinstance(value, (int, float)): + return (0, float(value)) + text = str(value) + try: + return (0, float(text)) + except ValueError: + return (0, text) + + +def _span_payload( + span: dict[str, Any], + tags: dict[str, Any], + *, + service_name: str | None, + run_id: str | None, +) -> dict[str, Any]: + name = str(span.get("operationName") or span.get("name") or "unknown") + attributes = { + key: value + for key, value in tags.items() + if key.startswith("xtuner.") + or key.startswith("agent.") + or key.startswith("session.") + or key.startswith("judger.") + or key.startswith("http.") + or key.startswith("error.") + or key.startswith("exception.") + or key.startswith("filter.") + or key.startswith("reward.") + or key.startswith("oversample.") + or key.startswith("drop.") + or key.startswith("train.") + or key.startswith("stage.") + or key in {"error", "rollout.backend", "prompt.tokens", "completion.tokens", "reward_score"} + } + return { + "name": name, + "stage": _stage_from_span_name_and_attributes(name, attributes), + "span_id": str(span.get("spanID") or span.get("span_id") or ""), + "parent_span_id": _parent_span_id(span), + "start_time_us": int(span.get("startTime") or 0), + "duration_ms": float(span.get("duration") or 0) / 1000.0, + "status": tags.get("otel.status_code") or tags.get("status.code") or "UNSET", + "service_name": service_name, + "run_id": run_id, + "rollout_backend": tags.get("rollout.backend"), + "attributes": attributes, + } + + +def _span_id(span: dict[str, Any]) -> str: + return str(span.get("spanID") or span.get("span_id") or "") + + +def _resolve_rollout_sample( + entry: dict[str, Any], + entries_by_span_id: dict[str, dict[str, Any]], +) -> tuple[Any | None, dict[str, Any]]: + tags = entry["tags"] + + ancestor_sample: tuple[Any, dict[str, Any]] | None = None + visited: set[str] = set() + parent_span_id = _parent_span_id(entry["span"]) + while parent_span_id and parent_span_id not in visited: + visited.add(parent_span_id) + parent = entries_by_span_id.get(parent_span_id) + if parent is None: + break + parent_tags = parent["tags"] + parent_rollout_id = parent_tags.get("xtuner.rollout_id") + if parent_rollout_id is not None: + ancestor_sample = (parent_rollout_id, parent_tags) + parent_span_id = _parent_span_id(parent["span"]) + + if ancestor_sample is not None: + return ancestor_sample + + rollout_id = tags.get("xtuner.rollout_id") + if rollout_id is not None: + return rollout_id, tags + return None, {} + + +def _process_metadata(trace_data: dict[str, Any]) -> dict[str, dict[str, Any]]: + metadata: dict[str, dict[str, Any]] = {} + for process_id, process in (trace_data.get("processes") or {}).items(): + if not isinstance(process, dict): + continue + tags = _tags_to_dict(process.get("tags") or []) + metadata[str(process_id)] = { + "service_name": str(process["serviceName"]) if process.get("serviceName") is not None else None, + "run_id": tags.get("run.id"), + } + return metadata + + +def _tags_to_dict(tags: list[dict[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for tag in tags: + key = tag.get("key") + if key is None: + continue + result[str(key)] = tag.get("value") + return result + + +def _parent_span_id(span: dict[str, Any]) -> str | None: + for reference in span.get("references") or []: + if reference.get("refType") == "CHILD_OF" and reference.get("spanID") is not None: + return str(reference["spanID"]) + return None + + +def _normalize_jaeger_query_url(jaeger_query_url: str | None) -> str | None: + if jaeger_query_url is None: + return None + stripped = jaeger_query_url.strip() + return stripped.rstrip("/") if stripped else None + + +def _jaeger_trace_url(jaeger_query_url: str | None, trace_id: str) -> str | None: + base = _normalize_jaeger_query_url(jaeger_query_url) + if base is None: + return None + return f"{base}/trace/{trace_id}" + + +def _jaeger_trace_link_base_url(jaeger_query_url: str | None, jaeger_link_url: str | None) -> str | None: + return _normalize_jaeger_query_url(jaeger_link_url) or _normalize_jaeger_query_url(jaeger_query_url) + + +def _append_unique(values: list[Any], value: Any) -> None: + if value is not None and value not in values: + values.append(value) + + +def _to_float(value: Any) -> float | Any: + try: + return float(value) + except (TypeError, ValueError): + return value + + +def _float_or_none(value: Any) -> float | None: + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _to_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if value is None: + return None + text = str(value).strip().lower() + if text in {"1", "true", "yes", "y", "on"}: + return True + if text in {"0", "false", "no", "n", "off"}: + return False + return None + + +def _int_value(value: Any) -> int | None: + if isinstance(value, bool) or value is None: + return None + if isinstance(value, int): + return value + text = str(value).strip() + if not text: + return None + try: + parsed = float(text) + except ValueError: + return None + if parsed.is_integer(): + return int(parsed) + return None + + +__all__ = [ + "build_rollout_view_payload_from_jaeger_traces", + "filter_rollout_view_payload_by_train_step", + "load_jaeger_traces_from_otel_jsonl", + "load_live_trace_records", +] diff --git a/xtuner/tools/trace_viewer/render.py b/xtuner/tools/trace_viewer/render.py new file mode 100644 index 0000000000..6aeb9f7ff3 --- /dev/null +++ b/xtuner/tools/trace_viewer/render.py @@ -0,0 +1,551 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + + +def render_rollout_trace_html( + payload: dict[str, Any], + *, + live: bool = False, + api_url: str = "/api/trace", + refresh_interval_s: float = 2.0, +) -> str: + data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + return ( + _HTML_TEMPLATE.replace("__TRACE_DATA__", data) + .replace("__LIVE_MODE__", json.dumps(live)) + .replace("__TRACE_API_URL__", json.dumps(api_url)) + .replace("__REFRESH_INTERVAL_MS__", str(int(refresh_interval_s * 1000))) + ) + + +def write_rollout_trace_html(payload: dict[str, Any], output_path: Path) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(render_rollout_trace_html(payload), encoding="utf-8") + + +_HTML_TEMPLATE = r""" + + + + + XTuner Rollout Trace Viewer + + + +
+
+

XTuner Rollout Trace Viewer

+
+
+
+
+
+
+ +
+ +
+ +
+

Stage Occupancy

+
+ + + +
StageSamplesGroups
+
+
+ +
+

Stage Durations

+
+ + + +
StageCountAvg sP50 sP95 sMax sErrorsTop Error
+
+
+ +
+

Samples

+
+ + +
+
+ + + + + + + + + + + + + +
SampleStatusGroupStepRewardPath / CurrentJaeger
+
+
+ +
+ + + +""" + + +__all__ = ["render_rollout_trace_html", "write_rollout_trace_html"] diff --git a/xtuner/tools/trace_viewer/server.py b/xtuner/tools/trace_viewer/server.py new file mode 100644 index 0000000000..e9dda50448 --- /dev/null +++ b/xtuner/tools/trace_viewer/server.py @@ -0,0 +1,474 @@ +from __future__ import annotations + +import argparse +import http.server +import json +import os +import threading +import time +from collections.abc import Callable +from pathlib import Path +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.parse import parse_qs, urlsplit +from urllib.request import Request, urlopen + +from xtuner.tools.trace_viewer.payload import ( + build_rollout_view_payload_from_jaeger_traces, + filter_rollout_view_payload_by_train_step, + load_live_trace_records, +) +from xtuner.tools.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html +from xtuner.tools.trace_viewer.source import ( + JAEGER_DEFAULT_LIMIT, + JAEGER_DEFAULT_LOOKBACK_S, + JAEGER_DEFAULT_QUERY_URL, + JaegerQuerySource, + JsonlTraceSource, + normalize_jaeger_query_url, + require_jaeger_query_url, +) + + +_JAEGER_PROXY_PREFIX = "/jaeger" +_PROXY_TIMEOUT_S = 10.0 +_HOP_BY_HOP_HEADERS = { + "connection", + "content-length", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +} + + +class TraceViewerHandle: + def __init__( + self, + server: http.server.ThreadingHTTPServer, + thread: threading.Thread, + *, + host: str, + port: int, + url: str, + ) -> None: + self.server = server + self.thread = thread + self.host = host + self.port = port + self.url = url + + def close(self) -> None: + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + + +class _TraceViewerPayloadCache: + def __init__( + self, + load_base_payload: Callable[[], dict[str, Any]], + *, + source_signature: Callable[[], Any] | None = None, + max_age_s: float | None = None, + ) -> None: + self._load_base_payload = load_base_payload + self._source_signature = source_signature or (lambda: None) + self._max_age_s = max_age_s + self._lock = threading.Lock() + self._signature: Any = object() + self._loaded_at_s = 0.0 + self._base_payload: dict[str, Any] | None = None + self._payloads_by_step: dict[str, dict[str, Any]] = {} + + def get(self, train_step: str | int | None = "latest") -> dict[str, Any]: + with self._lock: + signature = self._source_signature() + now = time.monotonic() + expired = self._max_age_s is not None and now - self._loaded_at_s >= self._max_age_s + if self._base_payload is None or signature != self._signature or expired: + self._signature = signature + self._base_payload = self._load_base_payload() + self._loaded_at_s = now + self._payloads_by_step.clear() + cache_key = _train_step_cache_key(train_step) + payload = self._payloads_by_step.get(cache_key) + if payload is None: + payload = filter_rollout_view_payload_by_train_step(self._base_payload, train_step) + self._payloads_by_step[cache_key] = payload + return payload + + +def _train_step_cache_key(train_step: str | int | None) -> str: + if train_step is None: + return "latest" + text = str(train_step).strip() + return text or "latest" + + +def fetch_jaeger_traces( + jaeger_query_url: str, + *, + service_name: str, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + timeout_s: float = 5.0, +) -> list[dict[str, Any]]: + return JaegerQuerySource( + query_url=jaeger_query_url, + service_name=service_name, + lookback_s=lookback_s, + limit=limit, + timeout_s=timeout_s, + ).load() + + +def fetch_rollout_view_payload( + jaeger_query_url: str, + *, + jaeger_link_url: str | None = None, + live_jsonl_path: Path | str | None = None, + service_name: str, + run_id: str | None = None, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + train_step: str | int | None = "latest", +) -> dict[str, Any]: + traces = JaegerQuerySource( + query_url=jaeger_query_url, + service_name=service_name, + lookback_s=lookback_s, + limit=limit, + ).load() + payload = build_rollout_view_payload_from_jaeger_traces( + traces, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=jaeger_link_url, + live_records=load_live_trace_records(live_jsonl_path), + service_name=service_name, + run_id=run_id, + train_step=train_step, + ) + payload["service_name"] = service_name + payload["run_id"] = run_id + payload["lookback_s"] = lookback_s + payload["limit"] = limit + return payload + + +def fetch_rollout_view_payload_from_trace_jsonl( + trace_jsonl_path: Path | str, + *, + jaeger_query_url: str | None = None, + jaeger_link_url: str | None = None, + live_jsonl_path: Path | str | None = None, + service_name: str | None = None, + run_id: str | None = None, + train_step: str | int | None = "latest", +) -> dict[str, Any]: + traces = JsonlTraceSource(trace_jsonl_path).load() + payload = build_rollout_view_payload_from_jaeger_traces( + traces, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=jaeger_link_url, + live_records=load_live_trace_records(live_jsonl_path), + service_name=service_name, + run_id=run_id, + train_step=train_step, + ) + payload["source"] = "trace_jsonl" + payload["trace_jsonl_path"] = os.fspath(Path(trace_jsonl_path).expanduser()) + payload["service_name"] = service_name + payload["run_id"] = run_id + return payload + + +def start_rollout_trace_viewer( + jaeger_query_url: str | None = JAEGER_DEFAULT_QUERY_URL, + *, + jaeger_link_url: str | None = None, + service_name: str, + run_id: str | None = None, + trace_jsonl_path: Path | str | None = None, + live_jsonl_path: Path | str | None = None, + payload_output_path: Path | str | None = None, + host: str = "127.0.0.1", + port: int = 0, + refresh_interval_s: float = 2.0, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + train_step: str | int | None = "latest", +) -> TraceViewerHandle: + jaeger_query_url = normalize_jaeger_query_url(jaeger_query_url) + jaeger_link_url = normalize_jaeger_query_url(jaeger_link_url) + viewer_jaeger_link_url = jaeger_link_url or (_JAEGER_PROXY_PREFIX if jaeger_query_url is not None else None) + if trace_jsonl_path is None: + jaeger_query_url = require_jaeger_query_url(jaeger_query_url) + + def load_base_payload() -> dict[str, Any]: + if trace_jsonl_path is not None: + payload = fetch_rollout_view_payload_from_trace_jsonl( + trace_jsonl_path, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=viewer_jaeger_link_url, + live_jsonl_path=live_jsonl_path, + service_name=service_name, + run_id=run_id, + train_step="all", + ) + else: + payload = fetch_rollout_view_payload( + jaeger_query_url, + jaeger_link_url=viewer_jaeger_link_url, + live_jsonl_path=live_jsonl_path, + service_name=service_name, + run_id=run_id, + lookback_s=lookback_s, + limit=limit, + train_step="all", + ) + return payload + + def current_source_signature() -> Any: + if trace_jsonl_path is None: + return ("jaeger", _source_signature(live_jsonl_path)) + return _source_signature(trace_jsonl_path, live_jsonl_path) + + payload_cache = _TraceViewerPayloadCache( + load_base_payload, + source_signature=current_source_signature, + max_age_s=max(refresh_interval_s, 5.0) if trace_jsonl_path is None else None, + ) + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self) -> None: + parsed = urlsplit(self.path) + path = parsed.path + if path == _JAEGER_PROXY_PREFIX or path.startswith(f"{_JAEGER_PROXY_PREFIX}/"): + self._proxy_jaeger() + return + if path in {"/", "/index.html"}: + html_body = render_rollout_trace_html( + self._payload(self._query_train_step(parsed.query)), + live=True, + api_url="/api/trace", + refresh_interval_s=refresh_interval_s, + ) + self._send_bytes(html_body.encode("utf-8"), "text/html; charset=utf-8") + return + if path == "/api/trace": + self._send_json(self._payload(self._query_train_step(parsed.query))) + return + self.send_error(404) + + def _query_train_step(self, query: str) -> str | int | None: + values = parse_qs(query).get("train_step") + if not values: + return train_step + return values[-1] + + def _payload(self, selected_train_step: str | int | None) -> dict[str, Any]: + payload = payload_cache.get(selected_train_step) + if payload_output_path is not None: + _write_payload_json(payload, payload_output_path) + return payload + + def _send_json(self, payload: dict[str, Any]) -> None: + self._send_bytes(json.dumps(payload, ensure_ascii=False).encode("utf-8"), "application/json") + + def _proxy_jaeger(self) -> None: + if jaeger_query_url is None: + self.send_error(502, "Jaeger query URL is not configured") + return + try: + target_url = _jaeger_proxy_target_url(jaeger_query_url, self.path) + request = Request( + target_url, + headers={ + "Accept": self.headers.get("Accept", "*/*"), + "User-Agent": self.headers.get("User-Agent", "XTunerTraceViewer"), + }, + ) + with urlopen(request, timeout=_PROXY_TIMEOUT_S) as response: + self._send_proxy_response(response.status, response.headers.items(), response.read()) + except HTTPError as exc: + self._send_proxy_response(exc.code, exc.headers.items(), exc.read()) + except (OSError, URLError) as exc: + self.send_error(502, f"Failed to proxy Jaeger request: {exc}") + + def _send_proxy_response(self, status: int, headers: Any, body: bytes) -> None: + self.send_response(status) + for key, value in headers: + if key.lower() in _HOP_BY_HOP_HEADERS: + continue + self.send_header(key, value) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _send_bytes(self, body: bytes, content_type: str) -> None: + self.send_response(200) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: Any) -> None: + return + + server = http.server.ThreadingHTTPServer((host, port), Handler) + server_host, server_port = server.server_address + display_host = server_host or host + thread = threading.Thread(target=server.serve_forever, name="XTunerRolloutTraceViewer", daemon=True) + thread.start() + return TraceViewerHandle( + server=server, + thread=thread, + host=display_host, + port=server_port, + url=f"http://{display_host}:{server_port}", + ) + + +def _write_payload_json(payload: dict[str, Any], output_path: Path | str) -> None: + path = Path(output_path).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + +def _source_signature(*paths: Path | str | None) -> tuple[tuple[str, int | None, int | None], ...]: + signatures = [] + for value in paths: + if value is None: + continue + path = Path(value).expanduser() + try: + stat = path.stat() + except OSError: + signatures.append((os.fspath(path), None, None)) + continue + signatures.append((os.fspath(path), stat.st_mtime_ns, stat.st_size)) + return tuple(signatures) + + +def _jaeger_proxy_target_url(jaeger_query_url: str, request_path: str) -> str: + parsed = urlsplit(request_path) + path = parsed.path + if path == _JAEGER_PROXY_PREFIX: + jaeger_path = "/" + else: + jaeger_path = path.removeprefix(_JAEGER_PROXY_PREFIX) or "/" + target = f"{require_jaeger_query_url(jaeger_query_url)}{jaeger_path}" + if parsed.query: + target = f"{target}?{parsed.query}" + return target + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Serve or render an XTuner rollout trace viewer backed by Jaeger or JSONL." + ) + parser.add_argument("--jaeger-query-url", default=JAEGER_DEFAULT_QUERY_URL) + parser.add_argument("--jaeger-link-url", default=None) + parser.add_argument("--trace-jsonl", type=Path, default=None) + parser.add_argument("--live-jsonl", type=Path, default=None) + parser.add_argument("--service", "--service-name", dest="service", default="xtuner-rollout") + parser.add_argument("--run-id", default=None) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=0) + parser.add_argument("--lookback", type=int, default=JAEGER_DEFAULT_LOOKBACK_S) + parser.add_argument("--limit", type=int, default=JAEGER_DEFAULT_LIMIT) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument("--payload-output", type=Path, default=None) + parser.add_argument("--train-step", default="latest", help="Initial train step to render: latest, all, or a step value.") + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + + if args.output is not None: + if args.trace_jsonl is not None: + payload = fetch_rollout_view_payload_from_trace_jsonl( + args.trace_jsonl, + jaeger_query_url=args.jaeger_query_url, + jaeger_link_url=args.jaeger_link_url, + live_jsonl_path=args.live_jsonl, + service_name=args.service, + run_id=args.run_id, + train_step=args.train_step, + ) + else: + payload = fetch_rollout_view_payload( + args.jaeger_query_url, + jaeger_link_url=args.jaeger_link_url, + live_jsonl_path=args.live_jsonl, + service_name=args.service, + run_id=args.run_id, + lookback_s=args.lookback, + limit=args.limit, + train_step=args.train_step, + ) + if args.payload_output is not None: + _write_payload_json(payload, args.payload_output) + write_rollout_trace_html(payload, args.output) + print(args.output) + if args.payload_output is not None: + print(args.payload_output) + return + + handle = start_rollout_trace_viewer( + args.jaeger_query_url, + jaeger_link_url=args.jaeger_link_url, + service_name=args.service, + run_id=args.run_id, + trace_jsonl_path=args.trace_jsonl, + live_jsonl_path=args.live_jsonl, + payload_output_path=args.payload_output, + host=args.host, + port=args.port, + lookback_s=args.lookback, + limit=args.limit, + train_step=args.train_step, + ) + print(f"XTuner Rollout Trace Viewer: {handle.url}", flush=True) + if args.trace_jsonl is not None: + print(f"Trace JSONL: {args.trace_jsonl}", flush=True) + if args.live_jsonl is not None: + print(f"Live JSONL: {args.live_jsonl}", flush=True) + jaeger_query_url = normalize_jaeger_query_url(args.jaeger_query_url) + if jaeger_query_url is not None: + print(f"Jaeger Trace Viewer: {jaeger_query_url}", flush=True) + print(f"Jaeger Same-Origin Proxy: {handle.url}{_JAEGER_PROXY_PREFIX}/", flush=True) + jaeger_link_url = normalize_jaeger_query_url(args.jaeger_link_url) + if jaeger_link_url is not None: + print(f"Jaeger Open Links: {jaeger_link_url}", flush=True) + if args.payload_output is not None: + print(f"Viewer Payload JSON: {args.payload_output}", flush=True) + try: + handle.thread.join() + except KeyboardInterrupt: + pass + finally: + handle.close() + + +if __name__ == "__main__": + main() + + +__all__ = [ + "TraceViewerHandle", + "_TraceViewerPayloadCache", + "build_rollout_view_payload_from_jaeger_traces", + "fetch_jaeger_traces", + "fetch_rollout_view_payload", + "fetch_rollout_view_payload_from_trace_jsonl", + "render_rollout_trace_html", + "start_rollout_trace_viewer", + "write_rollout_trace_html", +] diff --git a/xtuner/tools/trace_viewer/source.py b/xtuner/tools/trace_viewer/source.py new file mode 100644 index 0000000000..3c82ef641c --- /dev/null +++ b/xtuner/tools/trace_viewer/source.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.parse import urlencode +from urllib.request import Request, urlopen + +from xtuner.tools.trace_viewer.payload import load_jaeger_traces_from_otel_jsonl + + +JAEGER_DEFAULT_QUERY_URL = "http://127.0.0.1:16686" +JAEGER_DEFAULT_LOOKBACK_S = 60 * 60 +JAEGER_DEFAULT_LIMIT = 500 + + +@dataclass(frozen=True) +class JaegerQuerySource: + query_url: str + service_name: str + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S + limit: int = JAEGER_DEFAULT_LIMIT + timeout_s: float = 5.0 + + def load(self) -> list[dict[str, Any]]: + return fetch_jaeger_traces( + self.query_url, + service_name=self.service_name, + lookback_s=self.lookback_s, + limit=self.limit, + timeout_s=self.timeout_s, + ) + + +@dataclass(frozen=True) +class JsonlTraceSource: + trace_jsonl_path: Path | str + + def load(self) -> list[dict[str, Any]]: + return load_jaeger_traces_from_otel_jsonl(self.trace_jsonl_path) + + +def fetch_jaeger_traces( + jaeger_query_url: str, + *, + service_name: str, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + timeout_s: float = 5.0, +) -> list[dict[str, Any]]: + base_url = require_jaeger_query_url(jaeger_query_url) + query = urlencode( + { + "service": service_name, + "lookback": f"{max(1, lookback_s)}s", + "limit": max(1, limit), + } + ) + request = Request(f"{base_url}/api/traces?{query}", headers={"Accept": "application/json"}) + with urlopen(request, timeout=timeout_s) as response: + payload = json.loads(response.read().decode("utf-8")) + traces = payload.get("data") if isinstance(payload, dict) else None + return traces if isinstance(traces, list) else [] + + +def normalize_jaeger_query_url(jaeger_query_url: str | None) -> str | None: + if jaeger_query_url is None: + return None + stripped = jaeger_query_url.strip() + return stripped.rstrip("/") if stripped else None + + +def require_jaeger_query_url(jaeger_query_url: str | None) -> str: + normalized = normalize_jaeger_query_url(jaeger_query_url) + if normalized is None: + raise ValueError("jaeger_query_url is required") + return normalized + + +__all__ = [ + "JAEGER_DEFAULT_LIMIT", + "JAEGER_DEFAULT_LOOKBACK_S", + "JAEGER_DEFAULT_QUERY_URL", + "JaegerQuerySource", + "JsonlTraceSource", + "fetch_jaeger_traces", + "normalize_jaeger_query_url", + "require_jaeger_query_url", +] diff --git a/xtuner/v1/rl/trace/__init__.py b/xtuner/v1/rl/trace/__init__.py new file mode 100644 index 0000000000..41b3562632 --- /dev/null +++ b/xtuner/v1/rl/trace/__init__.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from importlib import import_module +from typing import Any + + +_LAZY_EXPORTS = { + "TraceConfig": (".runtime", "TraceConfig"), + "TraceRuntime": (".runtime", "TraceRuntime"), + "close_trace": (".runtime", "close_trace"), + "configure_trace": (".runtime", "configure_trace"), + "configure_trace_runtime": (".runtime", "configure_trace_runtime"), + "current_trace_runtime": (".runtime", "current_trace_runtime"), + "get_trace_env_vars": (".runtime", "get_trace_env_vars"), + "inject_trace_context": (".api", "inject_trace_context"), + "set_trace_attributes": (".api", "set_trace_attributes"), + "trace_event": (".api", "trace_event"), + "trace_function": (".api", "trace_function"), + "trace_span": (".api", "trace_span"), +} + + +__all__ = [ + "TraceConfig", + "TraceRuntime", + "close_trace", + "configure_trace", + "configure_trace_runtime", + "current_trace_runtime", + "get_trace_env_vars", + "inject_trace_context", + "set_trace_attributes", + "trace_event", + "trace_function", + "trace_span", +] + + +def __getattr__(name: str) -> Any: + try: + module_name, attr_name = _LAZY_EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + value = getattr(import_module(module_name, __name__), attr_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(__all__) diff --git a/xtuner/v1/rl/trace/api.py b/xtuner/v1/rl/trace/api.py new file mode 100644 index 0000000000..a56f9e9db8 --- /dev/null +++ b/xtuner/v1/rl/trace/api.py @@ -0,0 +1,334 @@ +from __future__ import annotations + +import contextlib +import contextvars +import inspect +import json +import os +import time +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from functools import wraps +from pathlib import Path +from typing import Any, TypeVar + +from . import otel_utils +from .runtime import ( + ensure_trace_runtime_from_env as _ensure_trace_runtime_from_env, +) +from .runtime import is_trace_enabled + + +F = TypeVar("F", bound=Callable[..., Any]) + + +_SPAN_NAME_PATH_BAGGAGE_KEY = "xtuner.span_name_path" +_SPAN_NAME_PATH_ATTRIBUTE = "xtuner.span_name_path" +_LIVE_TRACE_PATH_ENV = "XTUNER_OTEL_LIVE_JSONL_PATH" +# OTel context propagation carries trace/span IDs, not the current span names. +# XTuner keeps this local path so child spans and remote carriers can expose a +# readable execution chain for the viewer. +_CURRENT_SPAN_NAME_PATH: contextvars.ContextVar[tuple[str, ...]] = contextvars.ContextVar( + "xtuner_trace_span_name_path", + default=(), +) + + +@contextmanager +def trace_span( + name: str, + attributes: Mapping[str, Any] | None = None, + *, + parent_carrier: Mapping[str, str] | None = None, +): + """Create a current trace span. + + XTuner wraps OTel spans with runtime auto-initialization, trace-enabled + gating, attribute normalization, failure recording, live JSONL progress, + and a ``xtuner.span_name_path`` attribute for viewer-friendly call chains. + ``parent_carrier`` accepts the W3C carrier produced by + ``inject_trace_context`` across a process or request boundary. + """ + + span_name, normalized_attributes = _validate_trace_inputs( + name=name, + attributes=attributes, + ) + assert span_name is not None + _ensure_trace_runtime_from_env() + if not is_trace_enabled(): + yield + return + + with _attach_parent_carrier(parent_carrier): + span_name_path = (*_CURRENT_SPAN_NAME_PATH.get(), span_name) + normalized_attributes = dict(normalized_attributes) + normalized_attributes.setdefault(_SPAN_NAME_PATH_ATTRIBUTE, span_name_path) + path_token = _CURRENT_SPAN_NAME_PATH.set(span_name_path) + started_at_s = time.time() + status = "completed" + error_message: str | None = None + try: + with otel_utils.start_span(span_name, attributes=normalized_attributes): + span_ids = otel_utils.current_span_ids() + _record_live_span( + "start", + span_name=span_name, + span_name_path=span_name_path, + attributes=normalized_attributes, + span_ids=span_ids, + started_at_s=started_at_s, + ) + try: + yield + except Exception as exc: + status = "error" + error_message = str(exc) + otel_utils.record_failure(exc) + raise + finally: + _record_live_span( + "end", + span_name=span_name, + span_name_path=span_name_path, + attributes=normalized_attributes, + span_ids=span_ids, + started_at_s=started_at_s, + ended_at_s=time.time(), + status=status, + error_message=error_message, + ) + finally: + _CURRENT_SPAN_NAME_PATH.reset(path_token) + + +def trace_function( + name: str | None = None, + attributes: Mapping[str, Any] | None = None, +) -> Callable[[F], F]: + """Decorate a sync or async function with ``trace_span``. + + This provides the same XTuner additions as ``trace_span`` while keeping + call sites free of direct OTel SDK usage. + """ + + def decorator(func: F) -> F: + span_name, _ = _validate_trace_inputs( + name=name or f"{func.__module__}.{func.__qualname__}", + attributes=None, + ) + assert span_name is not None + + if inspect.iscoroutinefunction(func): + + @wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + with trace_span(span_name, attributes=attributes): + return await func(*args, **kwargs) + + return async_wrapper # type: ignore[return-value] + + @wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + with trace_span(span_name, attributes=attributes): + return func(*args, **kwargs) + + return wrapper # type: ignore[return-value] + + return decorator + + +def trace_event(name: str, attributes: Mapping[str, Any] | None = None) -> None: + """Add an event to the current span through XTuner's normalized API.""" + + event_name, normalized_attributes = _validate_trace_inputs( + name=name, + attributes=attributes, + ) + assert event_name is not None + if not is_trace_enabled(): + return + otel_utils.add_event(event_name, attributes=normalized_attributes) + + +def set_trace_attributes(attributes: Mapping[str, Any] | None) -> None: + """Set current-span attributes after applying XTuner value + normalization.""" + + _, normalized_attributes = _validate_trace_inputs(attributes=attributes) + if not is_trace_enabled(): + return + otel_utils.set_attributes(normalized_attributes) + + +def inject_trace_context(carrier: dict[str, str] | None = None) -> dict[str, str]: + """Inject OTel context plus XTuner's span-name path into a W3C carrier. + + OTel's trace context links parent/child spans. XTuner adds W3C Baggage for + ``xtuner.span_name_path`` so downstream spans can keep a readable chain. + """ + + target = carrier if carrier is not None else {} + _ensure_trace_runtime_from_env() + if not is_trace_enabled(): + return target + span_name_path = _CURRENT_SPAN_NAME_PATH.get() + context = None + if span_name_path: + context = otel_utils.context_with_baggage( + _SPAN_NAME_PATH_BAGGAGE_KEY, + json.dumps(span_name_path, separators=(",", ":")), + ) + otel_utils.inject_otel_context(context=context, carrier=target) + return target + + +@contextmanager +def _attach_parent_carrier(parent_carrier: Mapping[str, str] | None): + if parent_carrier is None: + yield + return + + parent_context = otel_utils.extract_otel_context(parent_carrier) + token = otel_utils.attach_otel_context(parent_context) + span_name_path = _extract_span_name_path(parent_context, otel_utils=otel_utils) + path_token = _CURRENT_SPAN_NAME_PATH.set(span_name_path) if span_name_path else None + try: + yield + finally: + if path_token is not None: + _CURRENT_SPAN_NAME_PATH.reset(path_token) + otel_utils.detach_otel_context(token) + + +def _extract_span_name_path(parent_context: Any, *, otel_utils: Any) -> tuple[str, ...]: + payload = otel_utils.get_baggage(_SPAN_NAME_PATH_BAGGAGE_KEY, context=parent_context) + if not payload: + return () + return _parse_span_name_path_payload(payload) + + +def _parse_span_name_path_payload(payload: object) -> tuple[str, ...]: + try: + value = json.loads(str(payload)) + except json.JSONDecodeError: + return () + if not isinstance(value, list): + return () + return tuple(str(item).strip() for item in value if isinstance(item, str) and item.strip()) + + +def _validate_trace_inputs( + *, + name: str | None = None, + attributes: Mapping[str, Any] | None = None, +) -> tuple[str | None, dict[str, Any]]: + normalized_name: str | None = None + if name is not None: + if not isinstance(name, str): + raise TypeError("trace name must be a string") + normalized_name = name.strip() + if not normalized_name: + raise ValueError("trace name cannot be empty") + + if attributes is None: + return normalized_name, {} + if not isinstance(attributes, Mapping): + raise TypeError("attributes must be a mapping") + + normalized = {} + for key, value in attributes.items(): + if not isinstance(key, str): + raise TypeError("attribute key must be a string") + normalized_key = key.strip() + if not normalized_key: + raise ValueError("attribute key cannot be empty") + normalized[normalized_key] = _normalize_trace_attribute_value(value) + return normalized_name, normalized + + +def _normalize_trace_attribute_value(value: Any) -> Any: + def normalize_scalar(item: Any) -> str | bool | int | float: + if isinstance(item, bool): + return item + if isinstance(item, int): + return item if -(2**63) <= item <= 2**63 - 1 else str(item) + if isinstance(item, (str, float)): + return item + return f"<{type(item).__name__}>" + + if isinstance(value, bool): + return value + if isinstance(value, int): + return value if -(2**63) <= value <= 2**63 - 1 else str(value) + if isinstance(value, (str, float)): + return value + if isinstance(value, (list, tuple)): + items = tuple(normalize_scalar(item) for item in value) + if not items: + return items + if all(isinstance(item, str) for item in items): + return items + if all(isinstance(item, bool) for item in items): + return items + if all(isinstance(item, int) and not isinstance(item, bool) for item in items): + return items + if all(isinstance(item, float) for item in items): + return items + return tuple(str(item) for item in items) + return f"<{type(value).__name__}>" + + +def _record_live_span( + event: str, + *, + span_name: str, + span_name_path: tuple[str, ...], + attributes: Mapping[str, Any], + span_ids: Mapping[str, str] | None, + started_at_s: float, + ended_at_s: float | None = None, + status: str | None = None, + error_message: str | None = None, +) -> None: + live_path = os.environ.get(_LIVE_TRACE_PATH_ENV) + if not live_path: + return + record: dict[str, Any] = { + "event": event, + "time_s": ended_at_s if ended_at_s is not None else started_at_s, + "span_name": span_name, + "span_name_path": list(span_name_path), + "attributes": {key: _json_safe_trace_value(value) for key, value in attributes.items()}, + } + if span_ids is not None: + record.update(span_ids) + if status is not None: + record["status"] = status + if ended_at_s is not None: + record["duration_ms"] = round(max(0.0, ended_at_s - started_at_s) * 1000.0, 3) + if error_message: + record["error_message"] = error_message + + with contextlib.suppress(Exception): + path = Path(live_path).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def _json_safe_trace_value(value: Any) -> Any: + if isinstance(value, (str, bool, int, float)) or value is None: + return value + if isinstance(value, (list, tuple)): + return [_json_safe_trace_value(item) for item in value] + return str(value) + + +__all__ = [ + "inject_trace_context", + "set_trace_attributes", + "trace_event", + "trace_function", + "trace_span", +] diff --git a/xtuner/v1/rl/trace/otel_utils.py b/xtuner/v1/rl/trace/otel_utils.py new file mode 100644 index 0000000000..c109f8b484 --- /dev/null +++ b/xtuner/v1/rl/trace/otel_utils.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +from typing import Any, Mapping, MutableMapping + +from opentelemetry import baggage, propagate, trace +from opentelemetry import context as otel_context +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.trace import Status, StatusCode + + +def configure_tracer_provider( + *, + service_name: str, + run_id: str, + endpoint: str, + protocol: str = "grpc", +) -> TracerProvider: + if protocol != "grpc": + raise ValueError(f"Unsupported OTel trace export protocol: {protocol!r}") + try: + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter + except ImportError as exc: + raise RuntimeError( + "XTuner OTel tracing requires the official OpenTelemetry OTLP gRPC trace exporter. " + "Install `opentelemetry-exporter-otlp-proto-grpc` before enabling trace." + ) from exc + resource = Resource.create({"service.name": service_name, "run.id": run_id}) + provider = TracerProvider(resource=resource) + provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter(endpoint=endpoint, insecure=True))) + trace.set_tracer_provider(provider) + return provider + + +def inject_otel_context( + context: Any | None = None, + carrier: MutableMapping[str, str] | None = None, +) -> MutableMapping[str, str]: + """Inject the current or provided OTel context into a W3C carrier.""" + + carrier = carrier if carrier is not None else {} + propagate.inject(carrier, context=context) + return carrier + + +def extract_otel_context( + carrier: Mapping[str, str], + context: Any | None = None, +) -> Any: + """Extract W3C TraceContext from a carrier into an OTel context.""" + + return propagate.extract(carrier, context=context) + + +def context_with_baggage(name: str, value: object, context: Any | None = None) -> Any: + """Return a context with one W3C Baggage item attached.""" + + return baggage.set_baggage(name, value, context=context) + + +def get_baggage(name: str, context: Any | None = None) -> object | None: + """Read one W3C Baggage item from a context.""" + + return baggage.get_baggage(name, context=context) + + +def attach_otel_context(context: Any) -> object: + """Attach extracted context to the current execution scope.""" + + return otel_context.attach(context) + + +def detach_otel_context(token: object) -> None: + """Detach a previously attached OTel context token.""" + + otel_context.detach(token) + + +def start_span(name: str, *, attributes: Mapping[str, Any] | None = None): + """Start a current OTel span with XTuner-managed exception handling.""" + + return trace.get_tracer("xtuner").start_as_current_span( + name, + attributes=attributes, + record_exception=False, + set_status_on_exception=False, + ) + + +def current_span_ids() -> dict[str, str] | None: + span = trace.get_current_span() + span_context = span.get_span_context() + if not span_context.is_valid: + return None + return { + "trace_id": f"{span_context.trace_id:032x}", + "span_id": f"{span_context.span_id:016x}", + } + + +def add_event(name: str, *, attributes: Mapping[str, Any] | None = None) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + span.add_event(name, attributes=attributes) + + +def set_attributes(attributes: Mapping[str, Any]) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + for key, value in attributes.items(): + span.set_attribute(key, value) + + +def set_error_status(message: str | None = None) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + span.set_attribute("error", True) + description = message or "error" + span.set_attribute("error.message", description) + span.set_status(Status(StatusCode.ERROR, description)) + + +def record_failure(exc: BaseException) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + error_attributes = { + "error.type": type(exc).__name__, + "error.message": str(exc), + "error": True, + } + for key, value in error_attributes.items(): + span.set_attribute(key, value) + span.record_exception(exc, attributes=error_attributes) + span.set_status(Status(StatusCode.ERROR, str(exc))) diff --git a/xtuner/v1/rl/trace/runtime.py b/xtuner/v1/rl/trace/runtime.py new file mode 100644 index 0000000000..77e48f9285 --- /dev/null +++ b/xtuner/v1/rl/trace/runtime.py @@ -0,0 +1,769 @@ +from __future__ import annotations + +import atexit +import contextlib +import errno +import os +import shlex +import shutil +import socket +import subprocess +import sys +import time +import uuid +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Literal, Mapping + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from xtuner.v1.rl.utils.misc import find_free_ports +from xtuner.v1.utils import get_logger + + +logger = get_logger() + +TraceRuntimeMode = Literal["disabled", "driver", "inherited"] + + +TRACE_ENV_KEYS = ( + "XTUNER_OTEL_ENABLED", + "XTUNER_OTEL_OUTPUT_DIR", + "XTUNER_OTEL_RUN_ID", + "XTUNER_OTEL_RUN_DIR", + "XTUNER_OTEL_JSONL_PATH", + "XTUNER_OTEL_LIVE_JSONL_PATH", + "OTEL_TRACES_EXPORTER", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_EXPORTER_OTLP_PROTOCOL", + "OTEL_SERVICE_NAME", +) + + +DEFAULT_JAEGER_OTLP_GRPC_ENDPOINT = "127.0.0.1:14317" +_TRACE_VIEWER_SERVER_MODULE = "xtuner.tools.trace_viewer.server" +_READY_TIMEOUT_S = 3.0 +_READY_POLL_INTERVAL_S = 0.05 +_LOG_TAIL_CHARS = 4000 + +_OTELCOL_CONFIG_YAML_TEMPLATE = """ +receivers: + otlp: + protocols: + grpc: + endpoint: 0.0.0.0:{port} + +exporters: + file: + path: {output_path} + rotation: + max_megabytes: 64 + max_days: 0 + max_backups: 0 + +service: + telemetry: + metrics: + level: none + pipelines: + traces: + receivers: [otlp] + exporters: [{exporters}] +""".lstrip() + +_OTELCOL_OTLP_GRPC_EXPORTER_YAML_TEMPLATE = """ + otlp/jaeger: + endpoint: {endpoint} + tls: + insecure: true +""".rstrip() + + +def _configure_tracer_provider( + *, + service_name: str, + run_id: str, + endpoint: str, + protocol: str, +) -> Any: + try: + from xtuner.v1.rl.trace.otel_utils import configure_tracer_provider + except ModuleNotFoundError as exc: + if exc.name == "opentelemetry" or (exc.name or "").startswith("opentelemetry."): + raise RuntimeError( + "XTuner OTel tracing requires OpenTelemetry packages. " + "Install `opentelemetry-sdk` and `opentelemetry-exporter-otlp-proto-grpc` " + "before enabling trace." + ) from exc + raise + return configure_tracer_provider( + service_name=service_name, + run_id=run_id, + endpoint=endpoint, + protocol=protocol, + ) + + +class TraceConfig(BaseModel): + """Public rollout tracing configuration. + + The interface is intentionally XTuner-level. OTel endpoint, exporter, and collector choices are runtime + implementation details. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + + enabled: bool = False + output_dir: Path | str | None = Field(default=None) + service_name: str = "xtuner-rollout" + viewer_enabled: bool = False + viewer_host: str = "127.0.0.1" + viewer_port: int = Field(default=0, ge=0, le=65535) + viewer_jaeger_query_url: str | None = None + viewer_jaeger_link_url: str | None = None + + @field_validator("output_dir") + @classmethod + def _expand_output_dir(cls, value: Path | str | None) -> Path | None: + if value is None: + return None + return Path(value).expanduser() + + +@dataclass(frozen=True) +class TraceRuntime: + enabled: bool + mode: TraceRuntimeMode + run_id: str + run_dir: Path + trace_jsonl_path: Path + live_jsonl_path: Path + service_name: str + trace_viewer_url: str | None = None + trace_viewer_port: int | None = None + + +@dataclass +class _OTelCollector: + port: int + output_path: Path + _binary_path: str = field(repr=False) + _config_path: Path = field(repr=False) + _stdout_path: Path = field(repr=False) + _stderr_path: Path = field(repr=False) + _process: subprocess.Popen | None = field(repr=False) + + @classmethod + def start( + cls, + *, + port: int, + output_path: Path, + ) -> _OTelCollector: + otelcol = shutil.which("otelcol-contrib") or shutil.which("otelcol") + if otelcol is None: + raise RuntimeError( + "XTuner OTel tracing requires `otelcol-contrib` or `otelcol` on PATH. " + "Install an official OpenTelemetry Collector binary before enabling trace." + ) + + jaeger_exporter = "\n" + _OTELCOL_OTLP_GRPC_EXPORTER_YAML_TEMPLATE.format( + endpoint=DEFAULT_JAEGER_OTLP_GRPC_ENDPOINT + ) + + config_path = output_path.parent / "otelcol.yaml" + config_yaml = _OTELCOL_CONFIG_YAML_TEMPLATE.format( + port=port, + output_path=output_path, + exporters="file, otlp/jaeger", + ).replace("\nservice:", f"{jaeger_exporter}\n\nservice:") + config_path.write_text(config_yaml, encoding="utf-8") + stdout_path = output_path.parent / "otelcol.stdout.log" + stderr_path = output_path.parent / "otelcol.stderr.log" + with stdout_path.open("wb") as stdout_file, stderr_path.open("wb") as stderr_file: + process = subprocess.Popen( + [otelcol, "--config", os.fspath(config_path)], + stdout=stdout_file, + stderr=stderr_file, + ) + collector = cls( + port=port, + output_path=output_path, + _binary_path=otelcol, + _config_path=config_path, + _stdout_path=stdout_path, + _stderr_path=stderr_path, + _process=process, + ) + try: + collector._wait_until_ready() + except Exception: + collector.close() + raise + return collector + + def _wait_until_ready(self) -> None: + deadline = time.monotonic() + _READY_TIMEOUT_S + last_error: OSError | None = None + while time.monotonic() < deadline: + process = self._process + if process is None: + stderr_tail = "" + with contextlib.suppress(OSError): + stderr_tail = self._stderr_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + "OpenTelemetry collector failed to start: collector process is not available. " + f"binary={self._binary_path}, config={self._config_path}, port={self.port}, " + f"stderr_tail={stderr_tail!r}" + ) + exit_code = process.poll() + if exit_code is not None: + stderr_tail = "" + with contextlib.suppress(OSError): + stderr_tail = self._stderr_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + f"OpenTelemetry collector failed to start: collector exited with code {exit_code}. " + f"binary={self._binary_path}, config={self._config_path}, port={self.port}, " + f"stderr_tail={stderr_tail!r}" + ) + try: + with socket.create_connection(("127.0.0.1", self.port), timeout=0.1): + return + except OSError as exc: + last_error = exc + time.sleep(_READY_POLL_INTERVAL_S) + + detail = f"collector did not become ready within {_READY_TIMEOUT_S:.1f}s" + if last_error is not None: + detail += f"; last connection error: {last_error}" + stderr_tail = "" + with contextlib.suppress(OSError): + stderr_tail = self._stderr_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + f"OpenTelemetry collector failed to start: {detail}. " + f"binary={self._binary_path}, config={self._config_path}, port={self.port}, " + f"stderr_tail={stderr_tail!r}" + ) + + def close(self) -> None: + process = self._process + self._process = None + if process is None: + return + if process.poll() is None: + process.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + if process.poll() is None: + process.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + + +@dataclass +class _TraceViewerProcess: + host: str + port: int + url: str + log_path: Path + _process: subprocess.Popen | None = field(repr=False) + command: list[str] = field(default_factory=list, repr=False) + + @classmethod + def start( + cls, + *, + trace_jsonl_path: Path, + live_jsonl_path: Path, + jaeger_query_url: str | None, + jaeger_link_url: str | None, + service_name: str, + run_id: str, + host: str, + port: int, + ) -> _TraceViewerProcess: + if port == 0: + port = find_free_ports(nums=1, host=host)[0] + _ensure_trace_viewer_port_available(host=host, port=port, run_id=run_id) + command = _build_trace_viewer_command( + trace_jsonl_path=trace_jsonl_path, + live_jsonl_path=live_jsonl_path, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=jaeger_link_url, + service_name=service_name, + run_id=run_id, + host=host, + port=port, + ) + log_path = trace_jsonl_path.parent / "viewer.log" + with log_path.open("ab") as log_file: + process = subprocess.Popen( + command, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + viewer = cls( + host=host, + port=port, + url=f"http://{host}:{port}", + log_path=log_path, + command=command, + _process=process, + ) + try: + viewer._wait_until_ready() + except Exception: + viewer.close() + raise + return viewer + + def _wait_until_ready(self) -> None: + deadline = time.monotonic() + _READY_TIMEOUT_S + last_error: OSError | None = None + while time.monotonic() < deadline: + self._raise_if_process_exited() + connect_host = "127.0.0.1" if self.host in {"", "0.0.0.0"} else self.host + try: + with socket.create_connection((connect_host, self.port), timeout=0.1): + time.sleep(_READY_POLL_INTERVAL_S) + self._raise_if_process_exited() + return + except OSError as exc: + last_error = exc + time.sleep(_READY_POLL_INTERVAL_S) + + detail = f"viewer did not become ready within {_READY_TIMEOUT_S:.1f}s" + if last_error is not None: + detail += f"; last connection error: {last_error}" + log_tail = "" + with contextlib.suppress(OSError): + log_tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError(f"XTuner trace viewer failed to start: {detail}, url={self.url}, log_tail={log_tail!r}") + + def _raise_if_process_exited(self) -> None: + process = self._process + if process is None: + raise RuntimeError(f"XTuner trace viewer failed to start: process is not available, log={self.log_path}") + exit_code = process.poll() + if exit_code is None: + return + + log_tail = "" + with contextlib.suppress(OSError): + log_tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + f"XTuner trace viewer failed to start: exited with code {exit_code}, url={self.url}, log_tail={log_tail!r}" + ) + + def restart_command(self) -> str: + return shlex.join(self.command) + + def close(self) -> None: + process = self._process + self._process = None + if process is None: + return + if process.poll() is None: + process.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + if process.poll() is None: + process.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + + +def _ensure_trace_viewer_port_available(*, host: str, port: int, run_id: str) -> None: + existing_viewer = _find_existing_trace_viewer_process_on_port(port) + if existing_viewer is not None: + pid = existing_viewer.get("pid", "unknown") + cmdline = existing_viewer.get("cmdline", "") + raise RuntimeError( + f"XTuner trace viewer port {port} already has an existing XTuner trace viewer process: " + f"pid={pid}, cmdline={cmdline!r}. " + f"Stop the old viewer/training process or set XTUNER_TRACE_VIEWER_PORT to another port " + f"before starting run_id={run_id}." + ) + + bind_host = host or "0.0.0.0" + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((bind_host, port)) + except OSError as exc: + if exc.errno == errno.EADDRINUSE: + raise RuntimeError( + f"XTuner trace viewer port {port} is already in use on host {host!r}. " + "Stop the process using that port or set XTUNER_TRACE_VIEWER_PORT to another port." + ) from exc + raise + + +def _find_existing_trace_viewer_process_on_port(port: int) -> Mapping[str, Any] | None: + listening_inodes = _listening_socket_inodes_on_port(port) + if not listening_inodes: + return None + + proc_root = Path("/proc") + with contextlib.suppress(OSError): + proc_dirs = list(proc_root.iterdir()) + for proc_dir in proc_dirs: + if not proc_dir.name.isdigit(): + continue + cmdline = _read_process_cmdline(proc_dir) + if _TRACE_VIEWER_SERVER_MODULE not in cmdline: + continue + if _process_has_socket_inode(proc_dir, listening_inodes): + return {"pid": int(proc_dir.name), "cmdline": cmdline} + return None + + +def _listening_socket_inodes_on_port(port: int) -> set[str]: + inodes: set[str] = set() + for proc_net_path in (Path("/proc/net/tcp"), Path("/proc/net/tcp6")): + with contextlib.suppress(OSError, ValueError): + for line in proc_net_path.read_text(encoding="utf-8", errors="replace").splitlines()[1:]: + parts = line.split() + if len(parts) < 10: + continue + local_address, state, inode = parts[1], parts[3], parts[9] + if state != "0A": + continue + _, port_hex = local_address.rsplit(":", 1) + if int(port_hex, 16) == port: + inodes.add(inode) + return inodes + + +def _read_process_cmdline(proc_dir: Path) -> str: + with contextlib.suppress(OSError): + raw_cmdline = (proc_dir / "cmdline").read_bytes() + return " ".join(part.decode("utf-8", errors="replace") for part in raw_cmdline.split(b"\0") if part) + return "" + + +def _process_has_socket_inode(proc_dir: Path, socket_inodes: set[str]) -> bool: + fd_dir = proc_dir / "fd" + with contextlib.suppress(OSError): + for fd_path in fd_dir.iterdir(): + with contextlib.suppress(OSError): + target = os.readlink(fd_path) + if target.startswith("socket:[") and target.endswith("]") and target[8:-1] in socket_inodes: + return True + return False + + +def _build_trace_viewer_command( + *, + trace_jsonl_path: Path, + live_jsonl_path: Path, + jaeger_query_url: str | None, + jaeger_link_url: str | None, + service_name: str, + run_id: str, + host: str, + port: int, +) -> list[str]: + command = [ + sys.executable, + "-m", + "xtuner.tools.trace_viewer.server", + "--trace-jsonl", + os.fspath(trace_jsonl_path), + "--live-jsonl", + os.fspath(live_jsonl_path), + "--service", + service_name, + "--run-id", + run_id, + "--host", + host, + "--port", + str(port), + ] + if jaeger_query_url is not None: + command.extend(["--jaeger-query-url", jaeger_query_url]) + if jaeger_link_url is not None: + command.extend(["--jaeger-link-url", jaeger_link_url]) + return command + + +@dataclass +class _TraceRuntimeHandle: + runtime: TraceRuntime + endpoint: str + env_vars: dict[str, str] + collector_port: int | None = None + collector: _OTelCollector | None = None + provider: Any | None = None + viewer_host: str | None = None + viewer_port: int = 0 + viewer_jaeger_query_url: str | None = None + viewer_jaeger_link_url: str | None = None + viewer: _TraceViewerProcess | None = None + + def start(self) -> None: + apply_trace_env(self.env_vars) + if not self.runtime.enabled: + logger.info("XTuner OTel tracing disabled.") + return + try: + if self.runtime.mode == "driver": + if self.collector_port is None: + raise RuntimeError("driver trace runtime requires a collector port") + self.collector = _OTelCollector.start( + port=self.collector_port, + output_path=self.runtime.trace_jsonl_path, + ) + self.provider = _configure_tracer_provider( + service_name=self.runtime.service_name, + run_id=self.runtime.run_id, + endpoint=self.endpoint, + protocol=self.env_vars["OTEL_EXPORTER_OTLP_PROTOCOL"], + ) + if self.runtime.mode == "driver" and self.viewer_host is not None: + self.viewer = _TraceViewerProcess.start( + trace_jsonl_path=self.runtime.trace_jsonl_path, + live_jsonl_path=self.runtime.live_jsonl_path, + jaeger_query_url=self.viewer_jaeger_query_url, + jaeger_link_url=self.viewer_jaeger_link_url, + service_name=self.runtime.service_name, + run_id=self.runtime.run_id, + host=self.viewer_host, + port=self.viewer_port, + ) + self.runtime = replace( + self.runtime, + trace_viewer_url=self.viewer.url, + trace_viewer_port=self.viewer.port, + ) + except Exception: + self.close(stop_viewer=True) + clear_trace_env() + raise + logger.info( + f"XTuner OTel tracing enabled: run_id={self.runtime.run_id}, endpoint={self.endpoint}, " + f"traces={self.runtime.trace_jsonl_path}" + ) + if self.viewer is not None: + logger.info( + f"XTuner trace viewer enabled: url={self.viewer.url}, host={self.viewer.host}, " + f"port={self.viewer.port}. Restart with: {self.viewer.restart_command()}" + ) + + def close(self, *, stop_viewer: bool = True) -> None: + viewer = self.viewer + self.viewer = None + if viewer is not None and stop_viewer: + logger.info("XTuner trace viewer stopped with training process.") + with contextlib.suppress(Exception): + viewer.close() + + provider = self.provider + self.provider = None + if provider is not None: + with contextlib.suppress(Exception): + provider.shutdown() + + collector = self.collector + self.collector = None + if collector is not None: + with contextlib.suppress(Exception): + collector.close() + + +_RUNTIME: _TraceRuntimeHandle | None = None +_ATEXIT_REGISTERED = False + + +def configure_trace(config: TraceConfig | None = None) -> TraceRuntime: + return configure_trace_runtime(config or TraceConfig()) + + +def configure_trace_runtime(config: TraceConfig) -> TraceRuntime: + """Configure Layer1 OTel runtime for the current process.""" + + global _RUNTIME + + close_trace() + runtime_handle = _build_trace_runtime_handle(config) + runtime_handle.start() + _RUNTIME = runtime_handle + if runtime_handle.runtime.enabled: + register_atexit_once(close_trace) + return runtime_handle.runtime + + +def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: + if not config.enabled: + return _TraceRuntimeHandle( + runtime=TraceRuntime( + enabled=False, + mode="disabled", + run_id="", + run_dir=Path(), + trace_jsonl_path=Path(), + live_jsonl_path=Path(), + service_name=config.service_name, + trace_viewer_url=None, + trace_viewer_port=None, + ), + endpoint="", + env_vars={}, + ) + + output_dir = Path(config.output_dir or Path.cwd() / "otel_traces").expanduser() + timestamp = time.strftime("%Y%m%d-%H%M%S", time.localtime()) + run_id = f"{timestamp}-{os.getpid()}-{uuid.uuid4().hex[:8]}" + run_dir = output_dir / run_id + traces_dir = run_dir / "traces" + traces_dir.mkdir(parents=True, exist_ok=True) + trace_jsonl_path = traces_dir / "traces.jsonl" + live_jsonl_path = traces_dir / "live.jsonl" + + try: + port = find_free_ports(nums=1, host="127.0.0.1", start_port=4317, end_port=4318)[0] + except RuntimeError: + port = find_free_ports(nums=1, host="127.0.0.1")[0] + endpoint = f"http://127.0.0.1:{port}" + protocol = "grpc" + + env_vars = { + "XTUNER_OTEL_ENABLED": "1", + "XTUNER_OTEL_OUTPUT_DIR": os.fspath(output_dir), + "XTUNER_OTEL_RUN_ID": run_id, + "XTUNER_OTEL_RUN_DIR": os.fspath(run_dir), + "XTUNER_OTEL_JSONL_PATH": os.fspath(trace_jsonl_path), + "XTUNER_OTEL_LIVE_JSONL_PATH": os.fspath(live_jsonl_path), + "OTEL_TRACES_EXPORTER": "otlp", + "OTEL_EXPORTER_OTLP_ENDPOINT": endpoint, + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": endpoint, + "OTEL_EXPORTER_OTLP_PROTOCOL": protocol, + "OTEL_SERVICE_NAME": config.service_name, + } + return _TraceRuntimeHandle( + runtime=TraceRuntime( + enabled=True, + mode="driver", + run_id=run_id, + run_dir=run_dir, + trace_jsonl_path=trace_jsonl_path, + live_jsonl_path=live_jsonl_path, + service_name=config.service_name, + trace_viewer_url=None, + trace_viewer_port=None, + ), + endpoint=endpoint, + env_vars=env_vars, + collector_port=port, + viewer_host=config.viewer_host if config.viewer_enabled else None, + viewer_port=config.viewer_port, + viewer_jaeger_query_url=config.viewer_jaeger_query_url, + viewer_jaeger_link_url=config.viewer_jaeger_link_url, + ) + + +def get_trace_env_vars() -> dict[str, str]: + """Return the env vars that should be injected into child Ray processes.""" + + if _RUNTIME is None or not _RUNTIME.runtime.enabled: + return get_trace_env_vars_from_env() + return dict(_RUNTIME.env_vars) + + +def current_trace_runtime() -> TraceRuntime | None: + """Return the active trace runtime owned by this process, if any.""" + + if _RUNTIME is None: + return None + return _RUNTIME.runtime + + +def get_trace_env_vars_from_env() -> dict[str, str]: + """Return inherited trace env vars before process-local runtime exists.""" + + if os.environ.get("XTUNER_OTEL_ENABLED") != "1": + return {} + return {key: os.environ[key] for key in TRACE_ENV_KEYS if key in os.environ} + + +def ensure_trace_runtime_from_env() -> bool: + """Lazily configure trace runtime in Ray child processes from inherited + env.""" + + global _RUNTIME + + env_vars = get_trace_env_vars_from_env() + if _RUNTIME is not None and _RUNTIME.runtime.enabled: + return True + if not env_vars: + return False + + env_vars = {key: str(env_vars[key]) for key in TRACE_ENV_KEYS if key in env_vars} + if "OTEL_EXPORTER_OTLP_ENDPOINT" not in env_vars: + return False + env_vars.setdefault("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", env_vars["OTEL_EXPORTER_OTLP_ENDPOINT"]) + env_vars.setdefault("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc") + env_vars.setdefault("OTEL_TRACES_EXPORTER", "otlp") + + run_dir = Path(env_vars.get("XTUNER_OTEL_RUN_DIR") or Path.cwd()).expanduser() + trace_jsonl_path = Path(env_vars.get("XTUNER_OTEL_JSONL_PATH") or run_dir / "traces" / "traces.jsonl").expanduser() + live_jsonl_path = Path( + env_vars.get("XTUNER_OTEL_LIVE_JSONL_PATH") or run_dir / "traces" / "live.jsonl" + ).expanduser() + runtime_handle = _TraceRuntimeHandle( + runtime=TraceRuntime( + enabled=True, + mode="inherited", + run_id=env_vars.get("XTUNER_OTEL_RUN_ID", ""), + run_dir=run_dir, + trace_jsonl_path=trace_jsonl_path, + live_jsonl_path=live_jsonl_path, + service_name=env_vars.get("OTEL_SERVICE_NAME", "xtuner-rollout"), + trace_viewer_url=None, + trace_viewer_port=None, + ), + endpoint=env_vars["OTEL_EXPORTER_OTLP_ENDPOINT"], + env_vars=env_vars, + ) + runtime_handle.start() + _RUNTIME = runtime_handle + register_atexit_once(close_trace) + return True + + +def is_trace_enabled() -> bool: + """Return whether XTuner trace runtime is enabled in this process.""" + + return _RUNTIME is not None and _RUNTIME.runtime.enabled + + +def close_trace() -> None: + """Flush provider state and stop local trace processes owned by this + process.""" + + global _RUNTIME + + runtime = _RUNTIME + _RUNTIME = None + if runtime is not None: + runtime.close() + clear_trace_env() + + +def apply_trace_env(env_vars: Mapping[str, str]) -> None: + clear_trace_env() + os.environ.update(env_vars) + + +def clear_trace_env() -> None: + for key in TRACE_ENV_KEYS: + os.environ.pop(key, None) + + +def register_atexit_once(close_fn) -> None: + global _ATEXIT_REGISTERED + if not _ATEXIT_REGISTERED: + atexit.register(close_fn) + _ATEXIT_REGISTERED = True From 1a58fc928d7e97196060c53d8c3d62b7b5337c71 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 13 Jul 2026 07:42:35 +0000 Subject: [PATCH 2/6] add trace_rollout_endpoint as starter rollout chain --- .claude/skills/xtuner-trace/SKILL.md | 23 +- .../xtuner-trace/references/trace-patterns.md | 29 ++- examples/v1/config/rl_grpo_gsm8k_judge.py | 13 ++ recipe/otle/README.md | 1 - xtuner/v1/rl/agent_loop/agent_loop.py | 4 + .../rl/agent_loop/single_turn_agent_loop.py | 7 +- xtuner/v1/rl/rollout/controller.py | 14 +- xtuner/v1/rl/rollout/worker.py | 11 +- xtuner/v1/rl/trace/rollout_api.py | 204 ++++++++++++++++++ xtuner/v1/rl/trace/runtime.py | 5 + xtuner/v1/train/rl_trainer.py | 11 + 11 files changed, 312 insertions(+), 10 deletions(-) create mode 100644 xtuner/v1/rl/trace/rollout_api.py diff --git a/.claude/skills/xtuner-trace/SKILL.md b/.claude/skills/xtuner-trace/SKILL.md index ecd6055cf2..b95ad2b214 100644 --- a/.claude/skills/xtuner-trace/SKILL.md +++ b/.claude/skills/xtuner-trace/SKILL.md @@ -15,10 +15,11 @@ Keep this package infrastructure-only: - OTel SDK adapter: `xtuner/v1/rl/trace/otel_utils.py` - Public facade: `xtuner/v1/rl/trace/__init__.py` - Basic API: `xtuner/v1/rl/trace/api.py` +- Rollout starter preset: `xtuner/v1/rl/trace/rollout_api.py` - Viewer: `xtuner/tools/trace_viewer/` - Local tooling: `recipe/otle/` and `examples/v1/scripts/setup_trace.sh` -Do not add rollout, agent, judger, Ray remote, HTTP proxy, reward, status, or session-server business semantics to the basic trace package. +Do not add rollout, agent, judger, Ray remote, HTTP proxy, reward, status, or session-server business semantics to `api.py`, `runtime.py`, or `otel_utils.py`. Keep rollout-specific starter behavior inside `rollout_api.py` and gated by `TraceConfig.enable_rollout_trace`. ## Basic API @@ -40,6 +41,17 @@ runtime setup: - `configure_trace(...)` - `close_trace()` +## Rollout Starter Preset + +`TraceConfig(enable_rollout_trace=True)` enables the built-in rollout starter +trace. Its helpers live in `xtuner/v1/rl/trace/rollout_api.py`: + +- `trace_rollout_endpoint(...)` +- `trace_rollout_remote(...)` + +Keep this layer narrow: it may depend on `RolloutState` and Ray call boundaries, +but it must not leak those dependencies into the basic API. + ## Add Trace Workflow When adding trace instrumentation to an XTuner run: @@ -48,7 +60,9 @@ When adding trace instrumentation to an XTuner run: 2. In the launch script, source `examples/v1/scripts/setup_trace.sh` when `XTUNER_TRACE_ENABLED=1`. 3. In the training config, add `TraceConfig` and set `enabled=True`; set - `viewer_enabled=True` when the user needs interactive inspection. + `viewer_enabled=True` when the user needs interactive inspection. Set + `enable_rollout_trace=True` only when the user wants the built-in rollout + starter trace. 4. Before adding any `trace_span(...)` instrumentation, you mask ask the user which stages they want to observe and which metrics each stage should expose. Do not infer default stages unless the user explicitly asks you to choose. @@ -65,7 +79,10 @@ When adding trace instrumentation to an XTuner run: ## Guardrails -- Do not reintroduce `trace_remote`, `traced_rollout_endpoint`, `traced_agent_item_endpoint`, or `traced_judger_endpoint`. +- Do not reintroduce the old generic `trace_remote`, `traced_rollout_endpoint`, + `traced_agent_item_endpoint`, or `traced_judger_endpoint` surfaces. The + supported rollout starter helpers are `trace_rollout_endpoint` and + `trace_rollout_remote` in `rollout_api.py`. - Do not recreate `trace_utils.py`, `context_propagation.py`, span-name registries, or business attribute builders under `xtuner/v1/rl/trace`. - Do not import `RolloutState`, agent item classes, judgers, rollout workers, Ray actors, aiohttp clients, or trainer configs from the basic trace package. - Do not call OpenTelemetry SDK directly from business code; use the basic API only when trace instrumentation is explicitly requested. diff --git a/.claude/skills/xtuner-trace/references/trace-patterns.md b/.claude/skills/xtuner-trace/references/trace-patterns.md index f4f0d87d58..6963590eaa 100644 --- a/.claude/skills/xtuner-trace/references/trace-patterns.md +++ b/.claude/skills/xtuner-trace/references/trace-patterns.md @@ -74,19 +74,18 @@ Add `TraceConfig` to the training config and pass it to the trainer config: ```python import os -from pathlib import Path from xtuner.v1.rl.trace import TraceConfig 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", + enable_rollout_trace=True, ) trainer = RLColocateTrainerConfig( @@ -95,6 +94,10 @@ trainer = RLColocateTrainerConfig( ) ``` +When `output_dir` is omitted in an RL trainer config, the trainer places trace +artifacts under the current experiment directory, for example +`work_dirs/name/20260713071356/otel/`. + ### Stage Spans Before adding spans, ask the user which stages they want to observe and which @@ -400,7 +403,27 @@ Avoid recording prompts, responses, full configs, secrets, raw headers, or large `examples/v1/scripts/setup_trace.sh` and `recipe/otle/` are local helper tooling for installing OTel collector binaries and starting local Jaeger/viewer dependencies. Keep these as setup assets; do not use them to add automatic trace behavior to training configs or launch scripts unless that integration is explicitly requested. -## Removed Semantics +## Rollout Starter Preset + +The built-in rollout starter trace is opt-in: + +```python +trace_config = TraceConfig( + enabled=True, + viewer_enabled=True, + enable_rollout_trace=True, +) +``` + +Its rollout-specific helpers live in `xtuner/v1/rl/trace/rollout_api.py`: + +- `trace_rollout_endpoint(...)` +- `trace_rollout_remote(...)` + +Use this only for the starter rollout chain. For custom instrumentation, ask +which stages and metrics the user wants before adding spans. + +## Removed Legacy Semantics Do not use or recreate these removed surfaces in this branch: diff --git a/examples/v1/config/rl_grpo_gsm8k_judge.py b/examples/v1/config/rl_grpo_gsm8k_judge.py index c7c3255f53..409d48b951 100644 --- a/examples/v1/config/rl_grpo_gsm8k_judge.py +++ b/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 @@ -30,6 +31,7 @@ data_path = os.environ["DATA_PATH"] eval_data_path = os.environ["EVAL_DATA_PATH"] enable_return_routed_experts = os.environ.get("ENABLE_RETURN_ROUTED_EXPERTS", "0") +enable_trace = True NNODE = int(os.environ.get("WORLD_SIZE", "1")) # basic settings @@ -183,6 +185,16 @@ # 7. evaluator evaluator_config = EvaluatorConfig(compute_metric_func=None) +trace_config = TraceConfig( + enabled=enable_trace, + 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", + enable_rollout_trace=True, +) + # 8. RL Colocate Trainer Config（CLI 通过 config["trainer"].build() 得到 Trainer） trainer = RLColocateTrainerConfig( resources=resources, @@ -203,4 +215,5 @@ work_dir=work_dir, seed=123, debug_rollout=False, + trace_config=trace_config, ) diff --git a/recipe/otle/README.md b/recipe/otle/README.md index b3d432e87e..545c987e92 100644 --- a/recipe/otle/README.md +++ b/recipe/otle/README.md @@ -34,7 +34,6 @@ Run XTuner with trace enabled: ```bash export XTUNER_TRACE_ENABLED=1 -export XTUNER_TRACE_SERVICE_NAME=xtuner-rollout ``` By default, XTuner starts a local collector that writes diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index fff3658bb0..6a69d0f80c 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -13,6 +13,9 @@ from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.rollout.constants import AGENT_LOOP_RAY_GENERATE_MAX_CONCURRENCY +from xtuner.v1.rl.trace.rollout_api import ( + trace_rollout_endpoint, +) from xtuner.v1.rl.utils import ( JUDGER_PAUSE_JUDGE_TASK_TIMEOUT_S, CPUActorLauncher, @@ -204,6 +207,7 @@ async def run_judger(self, rollout_state: RolloutState) -> RolloutState: ... @overload async def run_judger(self, rollout_state: list[RolloutState]) -> list[RolloutState]: ... + @trace_rollout_endpoint("judger.run") async def run_judger(self, rollout_state: RolloutState | list[RolloutState]) -> RolloutState | list[RolloutState]: assert self.judger is not None if isinstance(rollout_state, list): diff --git a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py index 86e1539da8..2368d1c3b5 100644 --- a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py +++ b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py @@ -1,6 +1,7 @@ from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController +from xtuner.v1.rl.trace.rollout_api import trace_rollout_endpoint, trace_rollout_remote from .agent_loop import AgentLoop, AgentLoopConfig @@ -64,6 +65,7 @@ def __init__( enable_batch_judge=enable_batch_judge, ) + @trace_rollout_endpoint("single_turn_agent_loop.run") async def generate_sample( self, rollout_state: RolloutState, @@ -74,7 +76,10 @@ async def generate_sample( # 推理引擎generate, 生成的结果会覆盖到 rollout_state.response_ids 上 assert self.rollout_ctl is not None - rollout_state = await self.rollout_ctl.generate.remote(rollout_state) # type: ignore[attr-defined] + rollout_state = await trace_rollout_remote( + self.rollout_ctl.generate, # type: ignore[attr-defined] + rollout_state, + ) # 非 COMPLETED 状态（如被截断、放弃等）直接早退，不触发打分 if rollout_state.status != Status.COMPLETED: return rollout_state diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index ae4572334c..028abb1acd 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -7,6 +7,7 @@ from ray.util.placement_group import PlacementGroup from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.trace.rollout_api import trace_rollout_endpoint, trace_rollout_remote from xtuner.v1.rl.utils import AutoAcceleratorWorkers from xtuner.v1.utils import XTUNER_DETERMINISTIC, get_logger @@ -89,6 +90,7 @@ def validate_registered_workers_to_proxy(self) -> None: self.proxy_manager.validate_registered_session_urls() @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) + @trace_rollout_endpoint("rollout.controller.generate") async def generate(self, rollout_state: RolloutState) -> RolloutState: if XTUNER_DETERMINISTIC: sample_params = rollout_state.sample_params.model_copy(deep=True) @@ -104,7 +106,10 @@ async def generate(self, rollout_state: RolloutState) -> RolloutState: rollout_state.error_msg = "No active rollout worker available." return rollout_state - response_ref = worker.generate.remote(rollout_state=rollout_state) # type: ignore[attr-defined] + response_ref = trace_rollout_remote( + worker.generate, # type: ignore[attr-defined] + rollout_state=rollout_state, + ) try: response_rollout_state = await asyncio.wait_for( response_ref, @@ -197,10 +202,17 @@ def _build_remote_worker_cls(self, worker_base_cls): assert self.config.rollout_max_batch_size_per_instance is not None, ( "rollout_max_batch_size_per_instance must be set before building RolloutWorker." ) + from xtuner.v1.rl.trace import get_trace_env_vars + + trace_env_vars = get_trace_env_vars() + ray_kwargs = {} + if trace_env_vars: + ray_kwargs["runtime_env"] = {"env_vars": trace_env_vars} return ray.remote( concurrency_groups={ ROLLOUT_CONCURRENCY_GROUP_GENERATE: ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY, }, + **ray_kwargs, )(worker_base_cls) def _init_workers(self, placement_group: PlacementGroup) -> RolloutWorkerRegistry: diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index c50ffa6bfb..3a79c9309a 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -28,6 +28,9 @@ reset_rollout_response, update_status_from_finish_reason, ) +from xtuner.v1.rl.trace.rollout_api import ( + trace_rollout_endpoint, +) from xtuner.v1.rl.utils import ( AutoAcceleratorWorkers, CPUResourcesConfig, @@ -525,19 +528,24 @@ def build(self, placement_group: "PlacementGroup"): import ray from xtuner.v1.rl.rollout.controller import RolloutController + from xtuner.v1.rl.trace import get_trace_env_vars num_workers = 1 register_cpu_resources( name="rollout_controller", cpu_resources=CPUResourcesConfig(num_workers=num_workers), ) + trace_env_vars = get_trace_env_vars() + actor_options: dict[str, Any] = {"num_cpus": num_workers} + if trace_env_vars: + actor_options["runtime_env"] = {"env_vars": trace_env_vars} return ( ray.remote( concurrency_groups={ ROLLOUT_CONCURRENCY_GROUP_GENERATE: ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY, }, )(RolloutController) - .options(num_cpus=num_workers) + .options(**actor_options) .remote(self, placement_group) ) @@ -968,6 +976,7 @@ async def _decode_routed_experts(self, routed_experts: Any) -> Any: return routed_experts @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) + @trace_rollout_endpoint("rollout.worker.generate") async def generate(self, rollout_state: RolloutState) -> RolloutState: request_max_tokens = rollout_state.sample_params.max_tokens try: diff --git a/xtuner/v1/rl/trace/rollout_api.py b/xtuner/v1/rl/trace/rollout_api.py new file mode 100644 index 0000000000..2476177577 --- /dev/null +++ b/xtuner/v1/rl/trace/rollout_api.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +import inspect +import os +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from functools import wraps +from typing import Any, Protocol, TypeVar + +from xtuner.v1.data_proto.rl_data import RolloutState +from xtuner.v1.utils import get_logger + +from . import api as trace_api +from .runtime import is_trace_enabled + + +F = TypeVar("F", bound=Callable[..., Any]) + +TRACE_ROLLOUT_ENABLED_ENV = "XTUNER_TRACE_ENABLE_ROLLOUT" +TRACE_CARRIER_EXTRA_FIELD = "_xtuner_trace_carrier" +TRACE_METADATA_EXTRA_FIELD = "_xtuner_trace_metadata" + + +class _RayRemoteMethod(Protocol): + def remote(self, *args: Any, **kwargs: Any) -> Any: ... + + +def is_rollout_trace_enabled() -> bool: + return os.environ.get(TRACE_ROLLOUT_ENABLED_ENV) == "1" and is_trace_enabled() + + +def trace_rollout_endpoint( + span_name: str, + *, + target_arg: str = "rollout_state", + initial_attributes: Callable[[Any, Any], Mapping[str, Any]] | None = None, +) -> Callable[[F], F]: + def decorator(func: F) -> F: + if not inspect.iscoroutinefunction(func): + raise TypeError("trace_rollout_endpoint() only supports async functions") + + signature = inspect.signature(func) + + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + if not is_rollout_trace_enabled(): + return await func(*args, **kwargs) + + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + if target_arg not in bound.arguments: + raise TypeError(f"trace_rollout_endpoint target argument {target_arg!r} was not bound") + target_value = bound.arguments[target_arg] + if isinstance(target_value, (list, tuple)) or not isinstance(target_value, RolloutState): + get_logger().warning( + f"XTuner rollout trace disabled for this endpoint: span={span_name!r}, " + f"target_arg={target_arg!r}, target_type={type(target_value).__name__}." + ) + return await func(*args, **kwargs) + rollout_state = target_value + + owner = bound.arguments.get("self", args[0] if args else None) + attributes = ( + dict(initial_attributes(owner, target_value)) + if initial_attributes is not None + else rollout_state_initial_attributes(rollout_state) + ) + attributes.setdefault("xtuner.stage", span_name) + + parent_carrier = extract_rollout_trace_parent_carrier(rollout_state) + with trace_api.trace_span(span_name, attributes=attributes, parent_carrier=parent_carrier): + result = await func(*args, **kwargs) + result_rollout_state = result if isinstance(result, RolloutState) else rollout_state + trace_api.set_trace_attributes(rollout_state_final_attributes(result_rollout_state)) + return result + + return wrapper # type: ignore[return-value] + + return decorator + + +def trace_rollout_remote( + remote_method: _RayRemoteMethod, + *args: Any, + target: RolloutState | None = None, + **kwargs: Any, +) -> Any: + if not is_rollout_trace_enabled(): + return remote_method.remote(*args, **kwargs) + + rollout_state = _resolve_rollout_state_target(args, kwargs, target=target, owner="trace_rollout_remote") + carrier = trace_api.inject_trace_context({}) + with attach_rollout_trace_carrier(rollout_state, carrier): + return remote_method.remote(*args, **kwargs) + + +@contextmanager +def attach_rollout_trace_carrier(rollout_state: RolloutState, carrier: Mapping[str, str]): + if not carrier: + yield + return + + extra_fields = rollout_state.extra_fields + had_previous_carrier = TRACE_CARRIER_EXTRA_FIELD in extra_fields + previous_carrier = extra_fields.get(TRACE_CARRIER_EXTRA_FIELD) + extra_fields[TRACE_CARRIER_EXTRA_FIELD] = dict(carrier) + try: + yield + finally: + if had_previous_carrier: + extra_fields[TRACE_CARRIER_EXTRA_FIELD] = previous_carrier + else: + extra_fields.pop(TRACE_CARRIER_EXTRA_FIELD, None) + + +def extract_rollout_trace_parent_carrier(rollout_state: RolloutState) -> dict[str, str] | None: + extra_fields = rollout_state.extra_fields + carrier = extra_fields.pop(TRACE_CARRIER_EXTRA_FIELD, None) + if not isinstance(carrier, Mapping): + return None + return {str(key): str(value) for key, value in carrier.items()} + + +def rollout_state_initial_attributes(rollout_state: RolloutState) -> dict[str, Any]: + metadata = _trace_metadata(rollout_state) + attributes: dict[str, Any] = { + "xtuner.status": rollout_state.status.value, + "xtuner.rollout_id": rollout_state.rollout_id, + "xtuner.group_id": rollout_state.group_id, + "xtuner.session_id": rollout_state.session_id, + "xtuner.task_name": rollout_state.task_name or metadata.get("task_name"), + "xtuner.producer_future_step": metadata.get("producer_future_step"), + "xtuner.model_step": metadata.get("model_step"), + } + if rollout_state.prompt_ids is not None: + attributes["prompt.tokens"] = len(rollout_state.prompt_ids) + return {key: value for key, value in attributes.items() if value is not None} + + +def rollout_state_final_attributes(rollout_state: RolloutState) -> dict[str, Any]: + attributes = rollout_state_initial_attributes(rollout_state) + attributes.update( + { + "finish_reason": rollout_state.finish_reason, + "error.message": rollout_state.error_msg, + } + ) + if rollout_state.response_ids is not None: + attributes["completion.tokens"] = len(rollout_state.response_ids) + reward = rollout_state.reward + if isinstance(reward, Mapping): + attributes.update( + { + "reward.score": reward.get("score"), + "reward.pass": reward.get("pass"), + } + ) + return {key: value for key, value in attributes.items() if value is not None} + + +def _trace_metadata(rollout_state: RolloutState) -> dict[str, Any]: + metadata = rollout_state.extra_fields.get(TRACE_METADATA_EXTRA_FIELD) + if not isinstance(metadata, Mapping): + metadata = rollout_state.extra_fields + return {str(key): value for key, value in metadata.items()} + + +def _resolve_rollout_state_target( + args: tuple[Any, ...], + kwargs: Mapping[str, Any], + *, + target: RolloutState | None = None, + owner: str, +) -> RolloutState: + if target is not None: + if not isinstance(target, RolloutState): + raise TypeError(f"{owner} target must be a RolloutState") + return target + + rollout_states: list[RolloutState] = [] + for value in (*args, *kwargs.values()): + if isinstance(value, RolloutState): + rollout_states.append(value) + elif isinstance(value, (list, tuple, set, frozenset)) and any( + isinstance(item, RolloutState) for item in value + ): + raise TypeError(f"{owner} supports a single RolloutState, not a RolloutState collection") + + if len(rollout_states) != 1: + raise ValueError(f"{owner} requires exactly one RolloutState argument") + return rollout_states[0] + + +__all__ = [ + "TRACE_CARRIER_EXTRA_FIELD", + "TRACE_METADATA_EXTRA_FIELD", + "TRACE_ROLLOUT_ENABLED_ENV", + "extract_rollout_trace_parent_carrier", + "is_rollout_trace_enabled", + "rollout_state_final_attributes", + "rollout_state_initial_attributes", + "trace_rollout_endpoint", + "trace_rollout_remote", +] diff --git a/xtuner/v1/rl/trace/runtime.py b/xtuner/v1/rl/trace/runtime.py index 77e48f9285..753ff208d9 100644 --- a/xtuner/v1/rl/trace/runtime.py +++ b/xtuner/v1/rl/trace/runtime.py @@ -33,6 +33,7 @@ "XTUNER_OTEL_RUN_DIR", "XTUNER_OTEL_JSONL_PATH", "XTUNER_OTEL_LIVE_JSONL_PATH", + "XTUNER_TRACE_ENABLE_ROLLOUT", "OTEL_TRACES_EXPORTER", "OTEL_EXPORTER_OTLP_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", @@ -122,6 +123,7 @@ class TraceConfig(BaseModel): viewer_port: int = Field(default=0, ge=0, le=65535) viewer_jaeger_query_url: str | None = None viewer_jaeger_link_url: str | None = None + enable_rollout_trace: bool = False @field_validator("output_dir") @classmethod @@ -636,6 +638,7 @@ def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: "XTUNER_OTEL_RUN_DIR": os.fspath(run_dir), "XTUNER_OTEL_JSONL_PATH": os.fspath(trace_jsonl_path), "XTUNER_OTEL_LIVE_JSONL_PATH": os.fspath(live_jsonl_path), + "XTUNER_TRACE_ENABLE_ROLLOUT": "1" if config.enable_rollout_trace else "0", "OTEL_TRACES_EXPORTER": "otlp", "OTEL_EXPORTER_OTLP_ENDPOINT": endpoint, "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": endpoint, @@ -736,6 +739,8 @@ def ensure_trace_runtime_from_env() -> bool: def is_trace_enabled() -> bool: """Return whether XTuner trace runtime is enabled in this process.""" + if _RUNTIME is None: + ensure_trace_runtime_from_env() return _RUNTIME is not None and _RUNTIME.runtime.enabled diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index eecb796ddd..d3a989a988 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -41,6 +41,7 @@ ) from xtuner.v1.rl.rollout.controller import RolloutControllerProxy from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.trace import TraceConfig, close_trace, configure_trace from xtuner.v1.rl.trainer.controller import TrainingController from xtuner.v1.rl.trainer.worker import WorkerConfig, WorkerLogItem from xtuner.v1.rl.utils import ( @@ -360,6 +361,7 @@ class BaseRLTrainerConfig(BaseModel): debug_train: bool = False skip_checkpoint_validation: bool = False exp_tracker: Literal["tensorboard", "jsonl"] = "tensorboard" + trace_config: TraceConfig = Field(default_factory=TraceConfig) @model_validator(mode="after") def _validate_sync_intervals(self): @@ -583,6 +585,7 @@ def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: self._init_load_source(cfg) self._init_save_config(cfg) log_dir = self._init_logger(cfg, logger_tag) + self._init_trace(cfg) self._save_runtime_environment(log_dir) self._init_train_state(cfg) self._init_train_worker_config(cfg, log_dir) @@ -633,6 +636,12 @@ def _init_logger(self, cfg: BaseRLTrainerConfig, logger_tag: str) -> Path: patch_default_save_plan() return log_dir + def _init_trace(self, cfg: BaseRLTrainerConfig) -> None: + trace_config = cfg.trace_config + if trace_config.output_dir is None: + trace_config = trace_config.model_copy(update={"output_dir": self.exp_dir / "otel"}) + self._trace_runtime = configure_trace(trace_config) + def _save_runtime_environment(self, log_dir: Path) -> None: if get_rank() != 0: return @@ -1627,6 +1636,7 @@ def fit(self): self._fit() finally: self._exp_tracker.close() + close_trace() def _fit(self): self.logger.info("Start RL training") @@ -1855,6 +1865,7 @@ def fit(self): return asyncio_run(self._fit()) finally: self._exp_tracker.close() + close_trace() async def _get_batch_or_raise_producer_failure( self, From 6c054eec25df10e891ed0750a3816af689995ab9 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 13 Jul 2026 10:47:38 +0000 Subject: [PATCH 3/6] simplify trace viewer --- .claude/skills/xtuner-trace/SKILL.md | 16 +- recipe/otle/README.md | 8 +- xtuner/tools/trace_viewer/__init__.py | 18 - xtuner/tools/trace_viewer/payload.py | 787 ++++++------------ xtuner/tools/trace_viewer/render.py | 99 +-- xtuner/tools/trace_viewer/server.py | 218 ++--- xtuner/tools/trace_viewer/source.py | 90 -- .../v1/rl/agent_loop_manager/produce_utils.py | 3 + xtuner/v1/rl/trace/api.py | 81 +- xtuner/v1/rl/trace/rollout_api.py | 65 +- xtuner/v1/rl/trace/runtime.py | 17 +- 11 files changed, 423 insertions(+), 979 deletions(-) delete mode 100644 xtuner/tools/trace_viewer/__init__.py delete mode 100644 xtuner/tools/trace_viewer/source.py diff --git a/.claude/skills/xtuner-trace/SKILL.md b/.claude/skills/xtuner-trace/SKILL.md index b95ad2b214..cf559e56e9 100644 --- a/.claude/skills/xtuner-trace/SKILL.md +++ b/.claude/skills/xtuner-trace/SKILL.md @@ -52,6 +52,18 @@ trace. Its helpers live in `xtuner/v1/rl/trace/rollout_api.py`: Keep this layer narrow: it may depend on `RolloutState` and Ray call boundaries, but it must not leak those dependencies into the basic API. +Rollout endpoint `xtuner.span_name_path` is a sample call chain, not just the +current process context stack. `trace_rollout_endpoint(...)` should append the +endpoint span name to an internal trace-only chain carried in +`RolloutState.extra_fields`, pass that state through remote rollout boundaries, +and clean the internal field at the outermost endpoint. For example, after a +sample generates and then judges, `judger.run` should record a path like +`single_turn_agent_loop.run -> rollout.controller.generate -> rollout.worker.generate -> judger.run` +even though the local active context at judge time may only be +`single_turn_agent_loop.run -> judger.run`. Keep this behavior in +`rollout_api.py` and viewer tests; do not add rollout call-chain semantics to +the basic trace API. + ## Add Trace Workflow When adding trace instrumentation to an XTuner run: @@ -79,10 +91,6 @@ When adding trace instrumentation to an XTuner run: ## Guardrails -- Do not reintroduce the old generic `trace_remote`, `traced_rollout_endpoint`, - `traced_agent_item_endpoint`, or `traced_judger_endpoint` surfaces. The - supported rollout starter helpers are `trace_rollout_endpoint` and - `trace_rollout_remote` in `rollout_api.py`. - Do not recreate `trace_utils.py`, `context_propagation.py`, span-name registries, or business attribute builders under `xtuner/v1/rl/trace`. - Do not import `RolloutState`, agent item classes, judgers, rollout workers, Ray actors, aiohttp clients, or trainer configs from the basic trace package. - Do not call OpenTelemetry SDK directly from business code; use the basic API only when trace instrumentation is explicitly requested. diff --git a/recipe/otle/README.md b/recipe/otle/README.md index 545c987e92..f912ae008d 100644 --- a/recipe/otle/README.md +++ b/recipe/otle/README.md @@ -44,10 +44,12 @@ Open the rollout viewer: ```bash python -m xtuner.tools.trace_viewer.server \ + --trace-jsonl /traces/traces.jsonl \ --jaeger-query-url http://127.0.0.1:16686 \ --service xtuner-rollout ``` -The viewer is a thin Jaeger Query API adapter. Jaeger remains the trace backend; -XTuner only groups spans by rollout metadata such as `xtuner.rollout_id`, -`xtuner.group_id`, `xtuner.task_name`, and `xtuner.status`. +The viewer reads `traces.jsonl` directly and uses Jaeger only for trace links +and the optional same-origin proxy. XTuner groups spans by rollout metadata such +as `xtuner.rollout_id`, `xtuner.group_id`, `xtuner.task_name`, and +`xtuner.status`. diff --git a/xtuner/tools/trace_viewer/__init__.py b/xtuner/tools/trace_viewer/__init__.py deleted file mode 100644 index a2ef41bd45..0000000000 --- a/xtuner/tools/trace_viewer/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces -from xtuner.tools.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html -from xtuner.tools.trace_viewer.source import ( - JaegerQuerySource, - JsonlTraceSource, - fetch_jaeger_traces, - normalize_jaeger_query_url, -) - -__all__ = [ - "JaegerQuerySource", - "JsonlTraceSource", - "build_rollout_view_payload_from_jaeger_traces", - "fetch_jaeger_traces", - "normalize_jaeger_query_url", - "render_rollout_trace_html", - "write_rollout_trace_html", -] diff --git a/xtuner/tools/trace_viewer/payload.py b/xtuner/tools/trace_viewer/payload.py index 3cdb5daa43..b80a421707 100644 --- a/xtuner/tools/trace_viewer/payload.py +++ b/xtuner/tools/trace_viewer/payload.py @@ -1,3 +1,10 @@ +"""Build the XTuner trace-viewer payload from Jaeger-shaped traces. + +This module owns XTuner semantics: rollout grouping, stage names, train-step +filtering, duration summaries, error summaries, and display paths. It does not +serve HTTP or render HTML. +""" + from __future__ import annotations import json @@ -31,23 +38,44 @@ def build_rollout_view_payload_from_jaeger_traces( *, jaeger_query_url: str | None = None, jaeger_link_url: str | None = None, - live_records: Iterable[dict[str, Any]] | None = None, service_name: str | None = None, run_id: str | None = None, train_step: Any = "latest", ) -> dict[str, Any]: samples_by_key: dict[tuple[str, Any], dict[str, Any]] = {} - jaeger_trace_link_base_url = _jaeger_trace_link_base_url(jaeger_query_url, jaeger_link_url) + jaeger_trace_link_base_url = _normalize_jaeger_query_url(jaeger_link_url) or _normalize_jaeger_query_url( + jaeger_query_url + ) for trace_data in traces: + if not isinstance(trace_data, dict): + raise ValueError(f"Jaeger trace must be an object, got {type(trace_data).__name__}") trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") if not trace_id: - continue + raise ValueError("Jaeger trace is missing traceID") process_metadata = _process_metadata(trace_data) span_entries = [] entries_by_span_id: dict[str, dict[str, Any]] = {} - for span in trace_data.get("spans") or []: - process = process_metadata.get(str(span.get("processID") or ""), {}) + spans = trace_data.get("spans") or [] + if not isinstance(spans, list): + raise ValueError(f"Jaeger trace {trace_id} has non-list spans") + if spans and not process_metadata: + raise ValueError(f"Jaeger trace {trace_id} has spans but no processes") + for span in spans: + if not isinstance(span, dict): + raise ValueError(f"Jaeger trace {trace_id} span must be an object") + span_id = str(span.get("spanID") or span.get("span_id") or "") + if not span_id: + raise ValueError(f"Jaeger trace {trace_id} span is missing spanID") + process_id = str(span.get("processID") or "") + if process_metadata: + if not process_id: + raise ValueError(f"Jaeger trace {trace_id} span {span_id} is missing processID") + process = process_metadata.get(process_id) + if process is None: + raise ValueError(f"Jaeger trace {trace_id} span {span_id} references unknown processID {process_id}") + else: + process = {} span_service_name = process.get("service_name") if service_name is not None and span_service_name != service_name: continue @@ -60,7 +88,7 @@ def build_rollout_view_payload_from_jaeger_traces( "tags": tags, "service_name": span_service_name, "run_id": span_run_id, - "span_id": _span_id(span), + "span_id": span_id, "trace_id": trace_id, } span_entries.append(entry) @@ -84,27 +112,46 @@ def build_rollout_view_payload_from_jaeger_traces( "group_id": sample_tags.get("xtuner.group_id"), "producer_future_step": _producer_future_step(sample_tags), "task_name": sample_tags.get("xtuner.task_name"), - "status": sample_tags.get("xtuner.status"), + "status": _sample_status_from_span(span, tags), "service_name": span_service_name, "run_id": span_run_id, - "jaeger_url": _jaeger_trace_url(jaeger_trace_link_base_url, trace_id), + "jaeger_url": f"{jaeger_trace_link_base_url}/trace/{trace_id}" + if jaeger_trace_link_base_url is not None + else None, "spans": [], }, ) - _merge_sample_fields(sample, sample_tags) + for sample_key, tag_key in ( + ("rollout_id", "xtuner.rollout_id"), + ("group_id", "xtuner.group_id"), + ("task_name", "xtuner.task_name"), + ): + if sample_tags.get(tag_key) is not None: + sample[sample_key] = sample_tags[tag_key] + status = _sample_status_from_span(span, tags) + if status is not None: + current_status = sample.get("status") + if current_status is None or _sample_status_priority(status) >= _sample_status_priority(current_status): + sample["status"] = status + producer_future_step = _producer_future_step(sample_tags) + if producer_future_step is not None: + sample["producer_future_step"] = producer_future_step if span_run_id is not None: sample["run_id"] = span_run_id sample["spans"].append(_span_payload(span, tags, service_name=span_service_name, run_id=span_run_id)) - _merge_live_records(samples_by_key, live_records or (), jaeger_trace_link_base_url) - generated_at_s = time.time() samples = [] for sample in samples_by_key.values(): sample["spans"].sort(key=lambda item: (item["start_time_us"], item["span_id"])) sample["span_count"] = len(sample["spans"]) - _apply_live_state(sample, generated_at_s) - _apply_sample_display_status(sample) + if sample.get("status") is None: + if _has_finished_sample_root_span(sample["spans"]): + sample["status"] = "completed" + elif sample.get("spans"): + sample["status"] = "running" + sample["display_path"] = _build_display_path(sample) + sample["chain"] = " -> ".join(node["name"] for node in sample["display_path"]) _apply_sample_reward_filter(sample) sample["stage"] = _sample_stage(sample) samples.append(sample) @@ -129,28 +176,29 @@ def load_jaeger_traces_from_otel_jsonl(trace_jsonl_path: Path | str) -> list[dic process_ids: dict[tuple[str, str, tuple[tuple[str, str], ...]], str] = {} path = Path(trace_jsonl_path).expanduser() if not path.is_file(): - return [] + raise FileNotFoundError(f"trace_jsonl path does not exist: {path}") - for record in _iter_jsonl_records(path): - if not isinstance(record, dict): - continue - jaeger_traces = _jaeger_traces_from_json_record(record) + for line_no, record in _iter_jsonl_records(path): + context = f"{path}:{line_no}" + jaeger_traces = _jaeger_traces_from_json_record(record, context=context) if jaeger_traces is not None: for trace_data in jaeger_traces: - _merge_jaeger_trace(traces_by_id, trace_data) + _merge_jaeger_trace(traces_by_id, trace_data, context=context) continue - for resource_span in record.get("resourceSpans") or []: - if not isinstance(resource_span, dict): - continue - resource_attrs = _otel_attributes_to_dict( - (resource_span.get("resource") or {}).get("attributes") or [] - ) + if "resourceSpans" not in record: + raise ValueError(f"{context}: trace record must contain Jaeger data or OTLP resourceSpans") + resource_spans = record.get("resourceSpans") + if not isinstance(resource_spans, list): + raise ValueError(f"{context}: OTLP resourceSpans must be a list") + for resource_index, resource_span in enumerate(resource_spans): + resource_attrs = _otel_attributes_to_dict((resource_span.get("resource") or {}).get("attributes") or []) service_name = str(resource_attrs.get("service.name") or "unknown") process_tags = _dict_to_jaeger_tags(resource_attrs) - for scope_span in _otel_scope_spans(resource_span): - for otel_span in scope_span.get("spans") or []: - if not isinstance(otel_span, dict): - continue + scope_spans = resource_span.get("scopeSpans") + if scope_spans is None: + scope_spans = resource_span.get("instrumentationLibrarySpans") or [] + for scope_index, scope_span in enumerate(scope_spans): + for span_index, otel_span in enumerate(scope_span.get("spans") or []): trace_id = str( otel_span.get("traceId") or otel_span.get("traceID") @@ -164,7 +212,10 @@ def load_jaeger_traces_from_otel_jsonl(trace_jsonl_path: Path | str) -> list[dic or "" ) if not trace_id or not span_id: - continue + raise ValueError( + f"{context}: resourceSpans[{resource_index}].scopeSpans[{scope_index}].spans[{span_index}] " + "is missing traceId or spanId" + ) trace_data = traces_by_id.setdefault( trace_id, { @@ -192,56 +243,47 @@ def load_jaeger_traces_from_otel_jsonl(trace_jsonl_path: Path | str) -> list[dic return list(traces_by_id.values()) -def load_live_trace_records(live_jsonl_path: Path | str | None) -> list[dict[str, Any]]: - if live_jsonl_path is None: - return [] - path = Path(live_jsonl_path).expanduser() - if not path.is_file(): - return [] - return list(_iter_jsonl_records(path)) - - -def _iter_jsonl_records(path: Path) -> Iterable[dict[str, Any]]: +def _iter_jsonl_records(path: Path) -> Iterable[tuple[int, dict[str, Any]]]: with path.open("r", encoding="utf-8") as handle: - for line in handle: + for line_no, line in enumerate(handle, start=1): stripped = line.strip() if not stripped: continue try: payload = json.loads(stripped) - except json.JSONDecodeError: - continue - if isinstance(payload, dict): - yield payload + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid JSON in {path}:{line_no}: {exc.msg}") from exc + if not isinstance(payload, dict): + raise ValueError(f"{path}:{line_no}: trace JSONL record must be an object") + yield line_no, payload -def _jaeger_traces_from_json_record(record: dict[str, Any]) -> list[dict[str, Any]] | None: +def _jaeger_traces_from_json_record(record: dict[str, Any], *, context: str) -> list[dict[str, Any]] | None: data = record.get("data") - if isinstance(data, list): - return [trace_data for trace_data in data if isinstance(trace_data, dict)] + if data is not None: + if not isinstance(data, list): + raise ValueError(f"{context}: Jaeger data must be a list") + return data if record.get("traceID") is not None and isinstance(record.get("spans"), list): return [record] + if record.get("traceID") is not None: + raise ValueError(f"{context}: Jaeger trace spans must be a list") return None -def _merge_jaeger_trace(traces_by_id: dict[str, dict[str, Any]], trace_data: dict[str, Any]) -> None: +def _merge_jaeger_trace(traces_by_id: dict[str, dict[str, Any]], trace_data: dict[str, Any], *, context: str) -> None: trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") if not trace_id: - return + raise ValueError(f"{context}: Jaeger trace is missing traceID") target = traces_by_id.setdefault(trace_id, {"traceID": trace_id, "processes": {}, "spans": []}) - if isinstance(trace_data.get("processes"), dict): - target["processes"].update(trace_data["processes"]) - target["spans"].extend([span for span in trace_data.get("spans") or [] if isinstance(span, dict)]) - - -def _otel_scope_spans(resource_span: dict[str, Any]) -> list[dict[str, Any]]: - scope_spans = resource_span.get("scopeSpans") - if isinstance(scope_spans, list): - return [scope_span for scope_span in scope_spans if isinstance(scope_span, dict)] - legacy_scope_spans = resource_span.get("instrumentationLibrarySpans") - if isinstance(legacy_scope_spans, list): - return [scope_span for scope_span in legacy_scope_spans if isinstance(scope_span, dict)] - return [] + processes = trace_data.get("processes") or {} + if not isinstance(processes, dict): + raise ValueError(f"{context}: Jaeger trace {trace_id} processes must be an object") + target["processes"].update(processes) + spans = trace_data.get("spans") or [] + if not isinstance(spans, list): + raise ValueError(f"{context}: Jaeger trace {trace_id} spans must be a list") + target["spans"].extend(spans) def _otel_span_to_jaeger_span( @@ -253,18 +295,20 @@ def _otel_span_to_jaeger_span( attributes = _otel_attributes_to_dict(otel_span.get("attributes") or []) tags = _dict_to_jaeger_tags(attributes) tags.extend(_otel_status_tags(otel_span.get("status") or {})) - start_ns = _int_from_otel_time( + start_time = ( otel_span.get("startTimeUnixNano") or otel_span.get("start_time_unix_nano") or otel_span.get("startTime") or 0 ) - end_ns = _int_from_otel_time( + start_ns = int(str(start_time)) + end_time = ( otel_span.get("endTimeUnixNano") or otel_span.get("end_time_unix_nano") or otel_span.get("endTime") or start_ns ) + end_ns = int(str(end_time)) parent_span_id = otel_span.get("parentSpanId") or otel_span.get("parent_span_id") references = [] if parent_span_id is not None and str(parent_span_id): @@ -281,7 +325,7 @@ def _otel_span_to_jaeger_span( } -def _otel_status_tags(status: dict[str, Any]) -> list[dict[str, Any]]: +def _otel_status_tags(status: Any) -> list[dict[str, Any]]: code = str(status.get("code") or status.get("statusCode") or "STATUS_CODE_UNSET") if code in {"STATUS_CODE_ERROR", "ERROR", "2"}: normalized = "ERROR" @@ -303,16 +347,15 @@ def _otel_attributes_to_dict(attributes: Any) -> dict[str, Any]: return dict(attributes) result: dict[str, Any] = {} for attribute in attributes or []: - if not isinstance(attribute, dict): - continue key = attribute.get("key") - if key is None: - continue - result[str(key)] = _otel_any_value_to_python(attribute.get("value")) + if key is not None: + result[str(key)] = _otel_any_value_to_python(attribute.get("value")) return result def _otel_any_value_to_python(value: Any) -> Any: + if value is None: + return None if not isinstance(value, dict): return value if "stringValue" in value: @@ -320,7 +363,7 @@ def _otel_any_value_to_python(value: Any) -> Any: if "boolValue" in value: return bool(value["boolValue"]) if "intValue" in value: - return _int_or_original(value["intValue"]) + return int(str(value["intValue"])) if "doubleValue" in value: return float(value["doubleValue"]) if "bytesValue" in value: @@ -337,74 +380,20 @@ def _otel_any_value_to_python(value: Any) -> Any: def _dict_to_jaeger_tags(attributes: dict[str, Any]) -> list[dict[str, Any]]: - return [ - {"key": str(key), "type": _jaeger_tag_type(value), "value": value} - for key, value in attributes.items() - if value is not None - ] - - -def _jaeger_tag_type(value: Any) -> str: - if isinstance(value, bool): - return "bool" - if isinstance(value, int): - return "int64" - if isinstance(value, float): - return "float64" - return "string" - - -def _int_from_otel_time(value: Any) -> int: - parsed = _int_or_original(value) - return parsed if isinstance(parsed, int) and not isinstance(parsed, bool) else 0 - - -def _int_or_original(value: Any) -> int | Any: - if isinstance(value, bool): - return value - if isinstance(value, int): - return value - try: - return int(str(value)) - except (TypeError, ValueError): - return value - - -def _merge_sample_fields(sample: dict[str, Any], tags: dict[str, Any]) -> None: - for sample_key, tag_key in ( - ("rollout_id", "xtuner.rollout_id"), - ("group_id", "xtuner.group_id"), - ("task_name", "xtuner.task_name"), - ): - if tags.get(tag_key) is not None: - sample[sample_key] = tags[tag_key] - _merge_sample_status(sample, tags.get("xtuner.status")) - producer_future_step = _producer_future_step(tags) - if producer_future_step is not None: - sample["producer_future_step"] = producer_future_step - - -def _merge_sample_status(sample: dict[str, Any], status: Any) -> None: - if status is None: - return - current_status = sample.get("status") - if current_status is None or _sample_status_priority(status) >= _sample_status_priority(current_status): - sample["status"] = status - - -def _apply_sample_display_status(sample: dict[str, Any]) -> None: - status = str(sample.get("status") or "").strip().lower() - if status not in _INITIAL_SAMPLE_STATUSES: - return - if _sample_has_observed_stage(sample): - sample["status"] = "running" - - -def _sample_has_observed_stage(sample: dict[str, Any]) -> bool: - current_stage = sample.get("current_stage") - if isinstance(current_stage, dict) and current_stage.get("name"): - return True - return bool(sample.get("spans")) + tags = [] + for key, value in attributes.items(): + if value is None: + continue + if isinstance(value, bool): + tag_type = "bool" + elif isinstance(value, int): + tag_type = "int64" + elif isinstance(value, float): + tag_type = "float64" + else: + tag_type = "string" + tags.append({"key": str(key), "type": tag_type, "value": value}) + return tags def _sample_status_priority(status: Any) -> int: @@ -418,125 +407,39 @@ def _sample_status_priority(status: Any) -> int: return 1 -def _merge_live_records( - samples_by_key: dict[tuple[str, Any], dict[str, Any]], - live_records: Iterable[dict[str, Any]], - jaeger_trace_link_base_url: str | None, -) -> None: - for record in live_records: - if not isinstance(record, dict): - continue - trace_id = str(record.get("trace_id") or "") - attributes = record.get("attributes") - if not trace_id or not isinstance(attributes, dict): - continue - rollout_id = attributes.get("xtuner.rollout_id") - if rollout_id is None: - continue - sample_key = (trace_id, rollout_id) - sample = samples_by_key.setdefault( - sample_key, - { - "trace_id": trace_id, - "rollout_id": rollout_id, - "group_id": attributes.get("xtuner.group_id"), - "producer_future_step": _producer_future_step(attributes), - "task_name": attributes.get("xtuner.task_name"), - "status": attributes.get("xtuner.status"), - "service_name": None, - "run_id": None, - "jaeger_url": _jaeger_trace_url(jaeger_trace_link_base_url, trace_id), - "spans": [], - }, - ) - _merge_sample_fields(sample, attributes) - sample.setdefault("_live_records", []).append(record) - - -def _apply_live_state(sample: dict[str, Any], generated_at_s: float) -> None: - live_states = _live_span_states(sample.pop("_live_records", [])) - active_states = [state for state in live_states if state.get("status") == "running"] - active_states.sort(key=lambda state: (len(state.get("span_name_path") or []), float(state.get("started_at_s") or 0.0))) - current_state = active_states[-1] if active_states else None - if current_state is not None: - started_at_s = float(current_state.get("started_at_s") or generated_at_s) - span_name = str(current_state.get("span_name") or "") - sample["current_stage"] = { - "name": span_name, - "stage": str(current_state.get("stage") or span_name or "unknown"), - "status": "running", - "elapsed_ms": round(max(0.0, generated_at_s - started_at_s) * 1000.0, 3), - "started_at_s": started_at_s, - } - else: - sample["current_stage"] = None - sample["live_spans"] = live_states - sample["display_path"] = _build_display_path(sample, live_states, current_state, generated_at_s) - sample["chain"] = " -> ".join(node["name"] for node in sample["display_path"]) - - -def _live_span_states(records: list[dict[str, Any]]) -> list[dict[str, Any]]: - states: dict[str, dict[str, Any]] = {} - for index, record in enumerate(records): - span_name = str(record.get("span_name") or "") - attributes = record.get("attributes") if isinstance(record.get("attributes"), dict) else {} - span_id = str(record.get("span_id") or "") - if not span_id: - span_id = f"live:{span_name}:{index}" - state = states.setdefault( - span_id, - { - "span_id": span_id, - "span_name": span_name, - "stage": _stage_from_span_name_and_attributes(span_name, attributes), - "span_name_path": _span_name_path_from_value( - record.get("span_name_path") or record.get("logical_path") - ), - "attributes": attributes, - "status": "running", - }, - ) - event = str(record.get("event") or "") - if event == "start": - state["started_at_s"] = _float_or_none(record.get("time_s")) - state["status"] = "running" - elif event == "end": - state["ended_at_s"] = _float_or_none(record.get("time_s")) - state["duration_ms"] = _float_or_none(record.get("duration_ms")) - state["status"] = str(record.get("status") or "completed") - if record.get("error_message"): - state["error_message"] = record["error_message"] - if record.get("trace_id") is not None: - state["trace_id"] = record["trace_id"] - return sorted(states.values(), key=lambda state: (float(state.get("started_at_s") or 0.0), str(state.get("span_id")))) - - -def _build_display_path( - sample: dict[str, Any], - live_states: list[dict[str, Any]], - current_state: dict[str, Any] | None, - generated_at_s: float, -) -> list[dict[str, Any]]: +def _sample_status_from_span(span: dict[str, Any], tags: dict[str, Any]) -> Any: + status = tags.get("xtuner.status") + if status is None: + return None + normalized = str(status).strip().lower() + if normalized in _ERROR_SAMPLE_STATUSES: + return status + return None + + +def _has_finished_sample_root_span(spans: list[dict[str, Any]]) -> bool: + return any(_is_sample_root_span(span) for span in spans) + + +def _is_sample_root_span(span: dict[str, Any]) -> bool: + attributes = span.get("attributes") or {} + span_path = _span_name_path(attributes) + if span_path: + return len(span_path) == 1 + return span.get("parent_span_id") is None + + +def _build_display_path(sample: dict[str, Any]) -> list[dict[str, Any]]: spans = sample.get("spans") or [] spans_by_name = {str(span.get("name") or ""): span for span in spans} - live_by_name = {str(state.get("span_name") or ""): state for state in live_states} - path = _display_path_names(spans, live_states, current_state) + if str(sample.get("status") or "").strip().lower() == "running" and not _has_finished_sample_root_span(spans): + path = _running_display_path_names(spans) + else: + path = _display_path_names(spans) nodes = [] for name in path: span = spans_by_name.get(name) - live_state = live_by_name.get(name) - if current_state is not None and name == current_state.get("span_name"): - started_at_s = float(current_state.get("started_at_s") or generated_at_s) - nodes.append( - { - "name": name, - "stage": current_state.get("stage") or name, - "source": "live", - "status": "running", - "elapsed_ms": round(max(0.0, generated_at_s - started_at_s) * 1000.0, 3), - } - ) - elif span is not None: + if span is not None: nodes.append( { "name": name, @@ -546,121 +449,70 @@ def _build_display_path( "duration_ms": span.get("duration_ms"), } ) - elif live_state is not None and live_state.get("status") == "running": - nodes.append({"name": name, "stage": live_state.get("stage") or name, "source": "live", "status": "active"}) else: nodes.append({"name": name, "source": "logical", "status": "inferred"}) return nodes -def _display_path_names( - spans: list[dict[str, Any]], - live_states: list[dict[str, Any]], - current_state: dict[str, Any] | None, -) -> list[str]: - path = _span_name_path_from_value(current_state.get("span_name_path")) if current_state is not None else [] +def _running_display_path_names(spans: list[dict[str, Any]]) -> list[str]: + roots = [] + for span in spans: + attributes = span.get("attributes") or {} + span_path = _span_name_path(attributes) + if span_path: + roots.append(span_path[0]) + elif span.get("parent_span_id") is None and span.get("name"): + roots.append(str(span["name"])) + unique_roots = [] + for name in roots: + if name not in unique_roots: + unique_roots.append(name) + return unique_roots + + +def _display_path_names(spans: list[dict[str, Any]]) -> list[str]: + span_paths = [] + for span in spans: + attributes = span.get("attributes") or {} + span_paths.append(_span_name_path(attributes)) + path = [] + for span_path in span_paths: + if not span_path: + continue + common_prefix_len = 0 + while ( + common_prefix_len < len(path) + and common_prefix_len < len(span_path) + and path[common_prefix_len] == span_path[common_prefix_len] + ): + common_prefix_len += 1 + path.extend(span_path[common_prefix_len:]) if not path: path = [str(span.get("name") or "") for span in spans if span.get("name")] - if not path: - path = [str(state.get("span_name") or "") for state in live_states if state.get("span_name")] - if not path: - span_paths = [ - _span_name_path_from_span_attributes(span.get("attributes") or {}) - for span in spans - ] - path = max(span_paths, key=len, default=[]) - if not path: - live_paths = [_span_name_path_from_value(state.get("span_name_path")) for state in live_states] - path = max(live_paths, key=len, default=[]) - return _unique_path(path) - - -def _span_name_path_from_span_attributes(attributes: dict[str, Any]) -> list[str]: - return _span_name_path_from_value( - attributes.get(_SPAN_NAME_PATH_ATTRIBUTE) or attributes.get(_LEGACY_LOGICAL_PATH_ATTRIBUTE) - ) + unique_path = [] + for name in path: + if not unique_path or unique_path[-1] != name: + unique_path.append(name) + return unique_path -def _span_name_path_from_value(value: Any) -> list[str]: +def _span_name_path(attributes: dict[str, Any]) -> list[str]: + value = attributes.get(_SPAN_NAME_PATH_ATTRIBUTE) or attributes.get(_LEGACY_LOGICAL_PATH_ATTRIBUTE) if isinstance(value, str): try: - decoded = json.loads(value) + value = json.loads(value) except json.JSONDecodeError: - decoded = [part.strip() for part in value.split("->")] - value = decoded + value = [part.strip() for part in value.split("->")] if isinstance(value, (list, tuple)): return [str(item).strip() for item in value if isinstance(item, str) and item.strip()] return [] -def _unique_path(path: list[str]) -> list[str]: - result = [] - for name in path: - if not result or result[-1] != name: - result.append(name) - return result - - def _producer_future_step(tags: dict[str, Any]) -> Any: value = tags.get("xtuner.producer_future_step") return value if value is not None else tags.get("xtuner.train_step") -def _build_step_group_summaries(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: - steps: dict[str, dict[str, Any]] = {} - for sample in samples: - producer_future_step = sample.get("producer_future_step") - group_id = sample.get("group_id") - step_key = _summary_key(producer_future_step) - group_key = _summary_key(group_id) - - step_summary = steps.setdefault( - step_key, - { - "producer_future_step": producer_future_step, - "sample_count": 0, - "groups": {}, - }, - ) - step_summary["sample_count"] += 1 - groups = step_summary["groups"] - group_summary = groups.setdefault( - group_key, - { - "group_id": group_id, - "sample_count": 0, - "statuses": Counter(), - "stages": Counter(), - "rollout_ids": [], - }, - ) - group_summary["sample_count"] += 1 - group_summary["statuses"][str(sample.get("status") or "unknown")] += 1 - if _should_show_group_stage(sample): - group_summary["stages"][str(sample.get("stage") or "unknown")] += 1 - group_summary["rollout_ids"].append(sample.get("rollout_id")) - - summaries = [] - for step_summary in steps.values(): - groups = [] - for group_summary in step_summary["groups"].values(): - group_summary["statuses"] = dict(sorted(group_summary["statuses"].items())) - group_summary["stages"] = dict(sorted(group_summary["stages"].items())) - group_summary["rollout_ids"].sort(key=lambda value: str(value)) - groups.append(group_summary) - groups.sort(key=lambda item: _sortable_summary_value(item["group_id"])) - summaries.append( - { - "producer_future_step": step_summary["producer_future_step"], - "group_count": sum(1 for group in groups if group["group_id"] is not None), - "sample_count": step_summary["sample_count"], - "groups": groups, - } - ) - summaries.sort(key=lambda item: _sortable_summary_value(item["producer_future_step"])) - return summaries - - def _available_train_steps(samples: list[dict[str, Any]]) -> list[Any]: values = {sample.get("producer_future_step") for sample in samples if sample.get("producer_future_step") is not None} return sorted(values, key=_sortable_summary_value) @@ -682,51 +534,56 @@ def _select_train_step(requested: Any, available_steps: list[Any]) -> Any: return max(available_steps, key=_sortable_summary_value) -def _filter_samples_by_train_step(samples: list[dict[str, Any]], selected_train_step: Any) -> list[dict[str, Any]]: - if selected_train_step == "all": - return samples - return [sample for sample in samples if str(sample.get("producer_future_step")) == str(selected_train_step)] - - def filter_rollout_view_payload_by_train_step(payload: dict[str, Any], train_step: Any = "latest") -> dict[str, Any]: samples = list(payload.get("samples") or []) generated_at_s = float(payload.get("generated_at_s") or time.time()) available_train_steps = list(payload.get("available_train_steps") or _available_train_steps(samples)) + requested_train_step = _requested_train_step(train_step) selected_train_step = _select_train_step(train_step, available_train_steps) - visible_samples = _filter_samples_by_train_step(samples, selected_train_step) + if selected_train_step == "all": + visible_samples = samples + else: + visible_samples = [ + sample for sample in samples if str(sample.get("producer_future_step")) == str(selected_train_step) + ] status_counts: Counter[str] = Counter() - stage_counts: Counter[str] = Counter() group_ids: set[Any] = set() + visible_steps: set[Any] = set() for sample in visible_samples: status = str(sample.get("status") or "unknown") status_counts[status] += 1 - stage_counts[str(sample["stage"])] += 1 if sample.get("group_id") is not None: group_ids.add(sample["group_id"]) + if sample.get("producer_future_step") is not None: + visible_steps.add(sample["producer_future_step"]) - step_group_summaries = _build_step_group_summaries(visible_samples) filtered_payload = dict(payload) filtered_payload.update( { "generated_at_s": generated_at_s, + "requested_train_step": requested_train_step, "selected_train_step": selected_train_step, "available_train_steps": available_train_steps, - "total_sample_count": len(samples), "sample_count": len(visible_samples), "group_count": len(group_ids), - "step_count": len(step_group_summaries), - "step_group_summaries": step_group_summaries, - "stage_occupancy": _build_stage_occupancy(visible_samples, generated_at_s), + "step_count": len(visible_steps), "stage_duration_summaries": _build_stage_duration_summaries(visible_samples), "status_counts": dict(sorted(status_counts.items())), - "stage_counts": dict(sorted(stage_counts.items())), "samples": visible_samples, } ) return filtered_payload +def _requested_train_step(train_step: Any) -> Any: + if train_step is None: + return "latest" + if isinstance(train_step, str) and not train_step.strip(): + return "latest" + return train_step + + def _build_stage_duration_summaries(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: durations_by_stage: dict[str, list[float]] = defaultdict(list) raw_durations_by_stage: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list)) @@ -843,96 +700,6 @@ def _summarize_stage_errors(errors: list[dict[str, Any]]) -> list[dict[str, Any] return result -def _build_stage_occupancy(samples: list[dict[str, Any]], generated_at_s: float) -> list[dict[str, Any]]: - buckets: dict[str, dict[str, Any]] = {} - for sample in samples: - stage, raw_span_name = _stage_bucket_for_sample(sample) - latest_start_s = _latest_span_start_s(sample) - age_s = 0.0 if stage in _TERMINAL_STAGE_STATUSES else max(0.0, generated_at_s - latest_start_s) - bucket = buckets.setdefault( - stage, - { - "stage": stage, - "sample_count": 0, - "_group_ids": set(), - "_raw_spans": {}, - "oldest_age_s": 0.0, - "oldest_rollout_id": None, - "oldest_group_id": None, - "oldest_producer_future_step": None, - }, - ) - bucket["sample_count"] += 1 - if sample.get("group_id") is not None: - bucket["_group_ids"].add(sample["group_id"]) - if raw_span_name: - raw_spans = bucket["_raw_spans"] - raw_bucket = raw_spans.setdefault(raw_span_name, {"span": raw_span_name, "sample_count": 0, "_group_ids": set()}) - raw_bucket["sample_count"] += 1 - if sample.get("group_id") is not None: - raw_bucket["_group_ids"].add(sample["group_id"]) - if age_s >= bucket["oldest_age_s"]: - bucket["oldest_age_s"] = _round_duration(age_s) - bucket["oldest_rollout_id"] = sample.get("rollout_id") - bucket["oldest_group_id"] = sample.get("group_id") - bucket["oldest_producer_future_step"] = sample.get("producer_future_step") - rows = list(buckets.values()) - for row in rows: - row["group_count"] = len(row.pop("_group_ids")) - row["raw_spans"] = _summarize_raw_span_occupancy(row.pop("_raw_spans")) - return rows - - -def _summarize_raw_span_occupancy(raw_spans: dict[str, dict[str, Any]]) -> list[dict[str, Any]]: - rows = [] - for raw_bucket in raw_spans.values(): - rows.append( - { - "span": raw_bucket["span"], - "sample_count": raw_bucket["sample_count"], - "group_count": len(raw_bucket["_group_ids"]), - } - ) - rows.sort(key=lambda item: (-item["sample_count"], str(item["span"]))) - return rows - - -def _stage_bucket_for_sample(sample: dict[str, Any]) -> tuple[str, str | None]: - status = str(sample.get("status") or "").strip().lower() - if status in _TERMINAL_STAGE_STATUSES: - return status, None - current_stage = sample.get("current_stage") - if isinstance(current_stage, dict): - raw_stage_name = str(current_stage.get("name") or "").strip() - semantic_stage = str(current_stage.get("stage") or "").strip() - if semantic_stage: - return semantic_stage, raw_stage_name or None - if raw_stage_name: - return raw_stage_name, raw_stage_name - spans = sample.get("spans") or [] - if not spans: - return status or "unknown", None - latest_span = spans[-1] - if latest_span.get("rollout_backend"): - return str(latest_span.get("stage") or "llm_generate"), _span_raw_name(latest_span) - if status in {"pending", "queued", "scheduled"}: - return "scheduled", None - return _span_semantic_stage(latest_span), _span_raw_name(latest_span) - - -def _latest_span_start_s(sample: dict[str, Any]) -> float: - current_stage = sample.get("current_stage") - if isinstance(current_stage, dict) and current_stage.get("started_at_s") is not None: - try: - return float(current_stage["started_at_s"]) - except (TypeError, ValueError): - pass - spans = sample.get("spans") or [] - if not spans: - return time.time() - return max(float(span.get("start_time_us") or 0) / 1_000_000.0 for span in spans) - - def _apply_sample_reward_filter(sample: dict[str, Any]) -> None: values: dict[str, Any] = {} for span in sample.get("spans") or []: @@ -950,7 +717,10 @@ def _apply_sample_reward_filter(sample: dict[str, Any]) -> None: if key in attrs: values[target] = attrs[key] if "reward_score" in values: - values["reward_score"] = _to_float(values["reward_score"]) + try: + values["reward_score"] = float(values["reward_score"]) + except (TypeError, ValueError): + pass if "reward_pass" in values: values["reward_pass"] = _to_bool(values["reward_pass"]) if "train_included" in values: @@ -1034,15 +804,8 @@ def _stage_from_span_name_and_attributes(span_name: str, attributes: dict[str, A def _sample_stage(sample: dict[str, Any]) -> str: status = str(sample.get("status") or "").strip().lower() - if status and status not in _NON_TERMINAL_SAMPLE_STATUSES: + if status in _ERROR_SAMPLE_STATUSES: return status - current_stage = sample.get("current_stage") - if isinstance(current_stage, dict): - semantic_stage = str(current_stage.get("stage") or "").strip() - if semantic_stage: - return semantic_stage - if current_stage.get("name"): - return str(current_stage["name"]) spans = sample.get("spans") or [] for span in spans: @@ -1055,16 +818,6 @@ def _sample_stage(sample: dict[str, Any]) -> str: return status or "unknown" -def _should_show_group_stage(sample: dict[str, Any]) -> bool: - status = str(sample.get("status") or "unknown").strip().lower() - stage = str(sample.get("stage") or "unknown").strip().lower() - return stage != status - - -def _summary_key(value: Any) -> str: - return "" if value is None else str(value) - - def _sortable_summary_value(value: Any) -> tuple[int, float | str]: if value is None: return (1, "") @@ -1086,7 +839,15 @@ def _span_payload( service_name: str | None, run_id: str | None, ) -> dict[str, Any]: - name = str(span.get("operationName") or span.get("name") or "unknown") + span_id = str(span.get("spanID") or span.get("span_id") or "") + name_value = span.get("operationName") or span.get("name") + if not name_value: + raise ValueError(f"Jaeger span {span_id} is missing operationName") + if span.get("startTime") is None: + raise ValueError(f"Jaeger span {span_id} is missing startTime") + if span.get("duration") is None: + raise ValueError(f"Jaeger span {span_id} is missing duration") + name = str(name_value) attributes = { key: value for key, value in tags.items() @@ -1108,10 +869,10 @@ def _span_payload( return { "name": name, "stage": _stage_from_span_name_and_attributes(name, attributes), - "span_id": str(span.get("spanID") or span.get("span_id") or ""), + "span_id": span_id, "parent_span_id": _parent_span_id(span), - "start_time_us": int(span.get("startTime") or 0), - "duration_ms": float(span.get("duration") or 0) / 1000.0, + "start_time_us": int(span["startTime"]), + "duration_ms": float(span["duration"]) / 1000.0, "status": tags.get("otel.status_code") or tags.get("status.code") or "UNSET", "service_name": service_name, "run_id": run_id, @@ -1120,10 +881,6 @@ def _span_payload( } -def _span_id(span: dict[str, Any]) -> str: - return str(span.get("spanID") or span.get("span_id") or "") - - def _resolve_rollout_sample( entry: dict[str, Any], entries_by_span_id: dict[str, dict[str, Any]], @@ -1155,29 +912,44 @@ def _resolve_rollout_sample( def _process_metadata(trace_data: dict[str, Any]) -> dict[str, dict[str, Any]]: metadata: dict[str, dict[str, Any]] = {} - for process_id, process in (trace_data.get("processes") or {}).items(): + processes = trace_data.get("processes") or {} + if not isinstance(processes, dict): + trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") + raise ValueError(f"Jaeger trace {trace_id} processes must be an object") + for process_id, process in processes.items(): if not isinstance(process, dict): - continue + raise ValueError(f"Jaeger process {process_id} must be an object") tags = _tags_to_dict(process.get("tags") or []) + if process.get("serviceName") is None: + raise ValueError(f"Jaeger process {process_id} is missing serviceName") metadata[str(process_id)] = { - "service_name": str(process["serviceName"]) if process.get("serviceName") is not None else None, + "service_name": str(process["serviceName"]), "run_id": tags.get("run.id"), } return metadata -def _tags_to_dict(tags: list[dict[str, Any]]) -> dict[str, Any]: +def _tags_to_dict(tags: Any) -> dict[str, Any]: + if not isinstance(tags, list): + raise ValueError(f"Jaeger tags must be a list, got {type(tags).__name__}") result: dict[str, Any] = {} - for tag in tags: + for index, tag in enumerate(tags): + if not isinstance(tag, dict): + raise ValueError(f"Jaeger tags[{index}] must be an object") key = tag.get("key") if key is None: - continue + raise ValueError(f"Jaeger tags[{index}] is missing key") result[str(key)] = tag.get("value") return result def _parent_span_id(span: dict[str, Any]) -> str | None: - for reference in span.get("references") or []: + references = span.get("references") or [] + if not isinstance(references, list): + raise ValueError("Jaeger span references must be a list") + for index, reference in enumerate(references): + if not isinstance(reference, dict): + raise ValueError(f"Jaeger span references[{index}] must be an object") if reference.get("refType") == "CHILD_OF" and reference.get("spanID") is not None: return str(reference["spanID"]) return None @@ -1190,36 +962,11 @@ def _normalize_jaeger_query_url(jaeger_query_url: str | None) -> str | None: return stripped.rstrip("/") if stripped else None -def _jaeger_trace_url(jaeger_query_url: str | None, trace_id: str) -> str | None: - base = _normalize_jaeger_query_url(jaeger_query_url) - if base is None: - return None - return f"{base}/trace/{trace_id}" - - -def _jaeger_trace_link_base_url(jaeger_query_url: str | None, jaeger_link_url: str | None) -> str | None: - return _normalize_jaeger_query_url(jaeger_link_url) or _normalize_jaeger_query_url(jaeger_query_url) - - def _append_unique(values: list[Any], value: Any) -> None: if value is not None and value not in values: values.append(value) -def _to_float(value: Any) -> float | Any: - try: - return float(value) - except (TypeError, ValueError): - return value - - -def _float_or_none(value: Any) -> float | None: - try: - return float(value) - except (TypeError, ValueError): - return None - - def _to_bool(value: Any) -> bool | None: if isinstance(value, bool): return value @@ -1233,26 +980,8 @@ def _to_bool(value: Any) -> bool | None: return None -def _int_value(value: Any) -> int | None: - if isinstance(value, bool) or value is None: - return None - if isinstance(value, int): - return value - text = str(value).strip() - if not text: - return None - try: - parsed = float(text) - except ValueError: - return None - if parsed.is_integer(): - return int(parsed) - return None - - __all__ = [ "build_rollout_view_payload_from_jaeger_traces", "filter_rollout_view_payload_by_train_step", "load_jaeger_traces_from_otel_jsonl", - "load_live_trace_records", ] diff --git a/xtuner/tools/trace_viewer/render.py b/xtuner/tools/trace_viewer/render.py index 6aeb9f7ff3..acfa730787 100644 --- a/xtuner/tools/trace_viewer/render.py +++ b/xtuner/tools/trace_viewer/render.py @@ -1,3 +1,5 @@ +"""Render the XTuner trace viewer HTML from an already-built payload.""" + from __future__ import annotations import json @@ -8,14 +10,14 @@ def render_rollout_trace_html( payload: dict[str, Any], *, - live: bool = False, + auto_refresh: bool = False, api_url: str = "/api/trace", refresh_interval_s: float = 2.0, ) -> str: data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) return ( _HTML_TEMPLATE.replace("__TRACE_DATA__", data) - .replace("__LIVE_MODE__", json.dumps(live)) + .replace("__AUTO_REFRESH__", json.dumps(auto_refresh)) .replace("__TRACE_API_URL__", json.dumps(api_url)) .replace("__REFRESH_INTERVAL_MS__", str(int(refresh_interval_s * 1000))) ) @@ -111,7 +113,6 @@ def write_rollout_trace_html(payload: dict[str, Any], output_path: Path) -> None .raw-span-item { overflow-wrap: anywhere; } .raw-span-details { margin-top: 3px; } .path-summary { display: grid; gap: 6px; min-width: 360px; } - .path-current { display: flex; flex-wrap: wrap; gap: 6px; align-items: center; } .path-chain { display: flex; flex-wrap: wrap; gap: 4px; align-items: center; line-height: 1.8; } .path-node { display: inline-flex; @@ -160,16 +161,6 @@ def write_rollout_trace_html(payload: dict[str, Any], output_path: Path) -> None
-
-

Stage Occupancy

-
- - - -
StageSamplesGroups
-
-
-

Stage Durations

@@ -195,7 +186,7 @@ def write_rollout_trace_html(payload: dict[str, Any], output_path: Path) -> None Group Step Reward - Path / Current + Path Jaeger @@ -207,20 +198,19 @@ def write_rollout_trace_html(payload: dict[str, Any], output_path: Path) -> None diff --git a/xtuner/tools/trace_viewer/server.py b/xtuner/tools/trace_viewer/server.py index e9dda50448..94a531dacf 100644 --- a/xtuner/tools/trace_viewer/server.py +++ b/xtuner/tools/trace_viewer/server.py @@ -1,3 +1,10 @@ +"""CLI and HTTP server for the XTuner trace viewer. + +The server reads ``traces.jsonl``, asks ``payload.py`` to build the XTuner view +model, and asks ``render.py`` to render HTML. It is the only module that owns +HTTP refresh, static HTML output, and the optional same-origin Jaeger proxy. +""" + from __future__ import annotations import argparse @@ -16,21 +23,13 @@ from xtuner.tools.trace_viewer.payload import ( build_rollout_view_payload_from_jaeger_traces, filter_rollout_view_payload_by_train_step, - load_live_trace_records, + load_jaeger_traces_from_otel_jsonl, ) from xtuner.tools.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html -from xtuner.tools.trace_viewer.source import ( - JAEGER_DEFAULT_LIMIT, - JAEGER_DEFAULT_LOOKBACK_S, - JAEGER_DEFAULT_QUERY_URL, - JaegerQuerySource, - JsonlTraceSource, - normalize_jaeger_query_url, - require_jaeger_query_url, -) _JAEGER_PROXY_PREFIX = "/jaeger" +JAEGER_DEFAULT_QUERY_URL = "http://127.0.0.1:16686" _PROXY_TIMEOUT_S = 10.0 _HOP_BY_HOP_HEADERS = { "connection", @@ -45,6 +44,20 @@ } +def normalize_jaeger_query_url(jaeger_query_url: str | None) -> str | None: + if jaeger_query_url is None: + return None + stripped = jaeger_query_url.strip() + return stripped.rstrip("/") if stripped else None + + +def require_jaeger_query_url(jaeger_query_url: str | None) -> str: + normalized = normalize_jaeger_query_url(jaeger_query_url) + if normalized is None: + raise ValueError("jaeger_query_url is required") + return normalized + + class TraceViewerHandle: def __init__( self, @@ -109,72 +122,20 @@ def _train_step_cache_key(train_step: str | int | None) -> str: return text or "latest" -def fetch_jaeger_traces( - jaeger_query_url: str, - *, - service_name: str, - lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, - limit: int = JAEGER_DEFAULT_LIMIT, - timeout_s: float = 5.0, -) -> list[dict[str, Any]]: - return JaegerQuerySource( - query_url=jaeger_query_url, - service_name=service_name, - lookback_s=lookback_s, - limit=limit, - timeout_s=timeout_s, - ).load() - - -def fetch_rollout_view_payload( - jaeger_query_url: str, - *, - jaeger_link_url: str | None = None, - live_jsonl_path: Path | str | None = None, - service_name: str, - run_id: str | None = None, - lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, - limit: int = JAEGER_DEFAULT_LIMIT, - train_step: str | int | None = "latest", -) -> dict[str, Any]: - traces = JaegerQuerySource( - query_url=jaeger_query_url, - service_name=service_name, - lookback_s=lookback_s, - limit=limit, - ).load() - payload = build_rollout_view_payload_from_jaeger_traces( - traces, - jaeger_query_url=jaeger_query_url, - jaeger_link_url=jaeger_link_url, - live_records=load_live_trace_records(live_jsonl_path), - service_name=service_name, - run_id=run_id, - train_step=train_step, - ) - payload["service_name"] = service_name - payload["run_id"] = run_id - payload["lookback_s"] = lookback_s - payload["limit"] = limit - return payload - - def fetch_rollout_view_payload_from_trace_jsonl( trace_jsonl_path: Path | str, *, jaeger_query_url: str | None = None, jaeger_link_url: str | None = None, - live_jsonl_path: Path | str | None = None, service_name: str | None = None, run_id: str | None = None, train_step: str | int | None = "latest", ) -> dict[str, Any]: - traces = JsonlTraceSource(trace_jsonl_path).load() + traces = load_jaeger_traces_from_otel_jsonl(trace_jsonl_path) payload = build_rollout_view_payload_from_jaeger_traces( traces, jaeger_query_url=jaeger_query_url, jaeger_link_url=jaeger_link_url, - live_records=load_live_trace_records(live_jsonl_path), service_name=service_name, run_id=run_id, train_step=train_step, @@ -192,56 +153,35 @@ def start_rollout_trace_viewer( jaeger_link_url: str | None = None, service_name: str, run_id: str | None = None, - trace_jsonl_path: Path | str | None = None, - live_jsonl_path: Path | str | None = None, - payload_output_path: Path | str | None = None, + trace_jsonl_path: Path | str, host: str = "127.0.0.1", port: int = 0, refresh_interval_s: float = 2.0, - lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, - limit: int = JAEGER_DEFAULT_LIMIT, train_step: str | int | None = "latest", ) -> TraceViewerHandle: jaeger_query_url = normalize_jaeger_query_url(jaeger_query_url) jaeger_link_url = normalize_jaeger_query_url(jaeger_link_url) viewer_jaeger_link_url = jaeger_link_url or (_JAEGER_PROXY_PREFIX if jaeger_query_url is not None else None) - if trace_jsonl_path is None: - jaeger_query_url = require_jaeger_query_url(jaeger_query_url) def load_base_payload() -> dict[str, Any]: - if trace_jsonl_path is not None: - payload = fetch_rollout_view_payload_from_trace_jsonl( - trace_jsonl_path, - jaeger_query_url=jaeger_query_url, - jaeger_link_url=viewer_jaeger_link_url, - live_jsonl_path=live_jsonl_path, - service_name=service_name, - run_id=run_id, - train_step="all", - ) - else: - payload = fetch_rollout_view_payload( - jaeger_query_url, - jaeger_link_url=viewer_jaeger_link_url, - live_jsonl_path=live_jsonl_path, - service_name=service_name, - run_id=run_id, - lookback_s=lookback_s, - limit=limit, - train_step="all", - ) + payload = fetch_rollout_view_payload_from_trace_jsonl( + trace_jsonl_path, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=viewer_jaeger_link_url, + service_name=service_name, + run_id=run_id, + train_step="all", + ) return payload def current_source_signature() -> Any: - if trace_jsonl_path is None: - return ("jaeger", _source_signature(live_jsonl_path)) - return _source_signature(trace_jsonl_path, live_jsonl_path) + return _source_signature(trace_jsonl_path) payload_cache = _TraceViewerPayloadCache( load_base_payload, source_signature=current_source_signature, - max_age_s=max(refresh_interval_s, 5.0) if trace_jsonl_path is None else None, ) + payload_cache.get(train_step) class Handler(http.server.BaseHTTPRequestHandler): def do_GET(self) -> None: @@ -253,7 +193,7 @@ def do_GET(self) -> None: if path in {"/", "/index.html"}: html_body = render_rollout_trace_html( self._payload(self._query_train_step(parsed.query)), - live=True, + auto_refresh=True, api_url="/api/trace", refresh_interval_s=refresh_interval_s, ) @@ -271,10 +211,7 @@ def _query_train_step(self, query: str) -> str | int | None: return values[-1] def _payload(self, selected_train_step: str | int | None) -> dict[str, Any]: - payload = payload_cache.get(selected_train_step) - if payload_output_path is not None: - _write_payload_json(payload, payload_output_path) - return payload + return payload_cache.get(selected_train_step) def _send_json(self, payload: dict[str, Any]) -> None: self._send_bytes(json.dumps(payload, ensure_ascii=False).encode("utf-8"), "application/json") @@ -334,25 +271,10 @@ def log_message(self, format: str, *args: Any) -> None: ) -def _write_payload_json(payload: dict[str, Any], output_path: Path | str) -> None: - path = Path(output_path).expanduser() - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") - - -def _source_signature(*paths: Path | str | None) -> tuple[tuple[str, int | None, int | None], ...]: - signatures = [] - for value in paths: - if value is None: - continue - path = Path(value).expanduser() - try: - stat = path.stat() - except OSError: - signatures.append((os.fspath(path), None, None)) - continue - signatures.append((os.fspath(path), stat.st_mtime_ns, stat.st_size)) - return tuple(signatures) +def _source_signature(path_value: Path | str) -> tuple[str, int, int]: + path = Path(path_value).expanduser() + stat = path.stat() + return (os.fspath(path), stat.st_mtime_ns, stat.st_size) def _jaeger_proxy_target_url(jaeger_query_url: str, request_path: str) -> str: @@ -368,57 +290,36 @@ def _jaeger_proxy_target_url(jaeger_query_url: str, request_path: str) -> str: return target -def _parse_args() -> argparse.Namespace: +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Serve or render an XTuner rollout trace viewer backed by Jaeger or JSONL." + description="Serve or render an XTuner rollout trace viewer backed by traces.jsonl." ) parser.add_argument("--jaeger-query-url", default=JAEGER_DEFAULT_QUERY_URL) parser.add_argument("--jaeger-link-url", default=None) - parser.add_argument("--trace-jsonl", type=Path, default=None) - parser.add_argument("--live-jsonl", type=Path, default=None) + parser.add_argument("--trace-jsonl", type=Path, required=True) parser.add_argument("--service", "--service-name", dest="service", default="xtuner-rollout") parser.add_argument("--run-id", default=None) parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=0) - parser.add_argument("--lookback", type=int, default=JAEGER_DEFAULT_LOOKBACK_S) - parser.add_argument("--limit", type=int, default=JAEGER_DEFAULT_LIMIT) parser.add_argument("--output", type=Path, default=None) - parser.add_argument("--payload-output", type=Path, default=None) parser.add_argument("--train-step", default="latest", help="Initial train step to render: latest, all, or a step value.") - return parser.parse_args() + return parser.parse_args(argv) def main() -> None: args = _parse_args() if args.output is not None: - if args.trace_jsonl is not None: - payload = fetch_rollout_view_payload_from_trace_jsonl( - args.trace_jsonl, - jaeger_query_url=args.jaeger_query_url, - jaeger_link_url=args.jaeger_link_url, - live_jsonl_path=args.live_jsonl, - service_name=args.service, - run_id=args.run_id, - train_step=args.train_step, - ) - else: - payload = fetch_rollout_view_payload( - args.jaeger_query_url, - jaeger_link_url=args.jaeger_link_url, - live_jsonl_path=args.live_jsonl, - service_name=args.service, - run_id=args.run_id, - lookback_s=args.lookback, - limit=args.limit, - train_step=args.train_step, - ) - if args.payload_output is not None: - _write_payload_json(payload, args.payload_output) + payload = fetch_rollout_view_payload_from_trace_jsonl( + args.trace_jsonl, + jaeger_query_url=args.jaeger_query_url, + jaeger_link_url=args.jaeger_link_url, + service_name=args.service, + run_id=args.run_id, + train_step=args.train_step, + ) write_rollout_trace_html(payload, args.output) print(args.output) - if args.payload_output is not None: - print(args.payload_output) return handle = start_rollout_trace_viewer( @@ -427,19 +328,12 @@ def main() -> None: service_name=args.service, run_id=args.run_id, trace_jsonl_path=args.trace_jsonl, - live_jsonl_path=args.live_jsonl, - payload_output_path=args.payload_output, host=args.host, port=args.port, - lookback_s=args.lookback, - limit=args.limit, train_step=args.train_step, ) print(f"XTuner Rollout Trace Viewer: {handle.url}", flush=True) - if args.trace_jsonl is not None: - print(f"Trace JSONL: {args.trace_jsonl}", flush=True) - if args.live_jsonl is not None: - print(f"Live JSONL: {args.live_jsonl}", flush=True) + print(f"Trace JSONL: {args.trace_jsonl}", flush=True) jaeger_query_url = normalize_jaeger_query_url(args.jaeger_query_url) if jaeger_query_url is not None: print(f"Jaeger Trace Viewer: {jaeger_query_url}", flush=True) @@ -447,8 +341,6 @@ def main() -> None: jaeger_link_url = normalize_jaeger_query_url(args.jaeger_link_url) if jaeger_link_url is not None: print(f"Jaeger Open Links: {jaeger_link_url}", flush=True) - if args.payload_output is not None: - print(f"Viewer Payload JSON: {args.payload_output}", flush=True) try: handle.thread.join() except KeyboardInterrupt: @@ -465,8 +357,6 @@ def main() -> None: "TraceViewerHandle", "_TraceViewerPayloadCache", "build_rollout_view_payload_from_jaeger_traces", - "fetch_jaeger_traces", - "fetch_rollout_view_payload", "fetch_rollout_view_payload_from_trace_jsonl", "render_rollout_trace_html", "start_rollout_trace_viewer", diff --git a/xtuner/tools/trace_viewer/source.py b/xtuner/tools/trace_viewer/source.py deleted file mode 100644 index 3c82ef641c..0000000000 --- a/xtuner/tools/trace_viewer/source.py +++ /dev/null @@ -1,90 +0,0 @@ -from __future__ import annotations - -import json -from dataclasses import dataclass -from pathlib import Path -from typing import Any -from urllib.parse import urlencode -from urllib.request import Request, urlopen - -from xtuner.tools.trace_viewer.payload import load_jaeger_traces_from_otel_jsonl - - -JAEGER_DEFAULT_QUERY_URL = "http://127.0.0.1:16686" -JAEGER_DEFAULT_LOOKBACK_S = 60 * 60 -JAEGER_DEFAULT_LIMIT = 500 - - -@dataclass(frozen=True) -class JaegerQuerySource: - query_url: str - service_name: str - lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S - limit: int = JAEGER_DEFAULT_LIMIT - timeout_s: float = 5.0 - - def load(self) -> list[dict[str, Any]]: - return fetch_jaeger_traces( - self.query_url, - service_name=self.service_name, - lookback_s=self.lookback_s, - limit=self.limit, - timeout_s=self.timeout_s, - ) - - -@dataclass(frozen=True) -class JsonlTraceSource: - trace_jsonl_path: Path | str - - def load(self) -> list[dict[str, Any]]: - return load_jaeger_traces_from_otel_jsonl(self.trace_jsonl_path) - - -def fetch_jaeger_traces( - jaeger_query_url: str, - *, - service_name: str, - lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, - limit: int = JAEGER_DEFAULT_LIMIT, - timeout_s: float = 5.0, -) -> list[dict[str, Any]]: - base_url = require_jaeger_query_url(jaeger_query_url) - query = urlencode( - { - "service": service_name, - "lookback": f"{max(1, lookback_s)}s", - "limit": max(1, limit), - } - ) - request = Request(f"{base_url}/api/traces?{query}", headers={"Accept": "application/json"}) - with urlopen(request, timeout=timeout_s) as response: - payload = json.loads(response.read().decode("utf-8")) - traces = payload.get("data") if isinstance(payload, dict) else None - return traces if isinstance(traces, list) else [] - - -def normalize_jaeger_query_url(jaeger_query_url: str | None) -> str | None: - if jaeger_query_url is None: - return None - stripped = jaeger_query_url.strip() - return stripped.rstrip("/") if stripped else None - - -def require_jaeger_query_url(jaeger_query_url: str | None) -> str: - normalized = normalize_jaeger_query_url(jaeger_query_url) - if normalized is None: - raise ValueError("jaeger_query_url is required") - return normalized - - -__all__ = [ - "JAEGER_DEFAULT_LIMIT", - "JAEGER_DEFAULT_LOOKBACK_S", - "JAEGER_DEFAULT_QUERY_URL", - "JaegerQuerySource", - "JsonlTraceSource", - "fetch_jaeger_traces", - "normalize_jaeger_query_url", - "require_jaeger_query_url", -] diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 33f1e115b2..d3d8101d85 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -144,6 +144,9 @@ async def generate_group( enable_partial_rollout: bool = False, ) -> list[RolloutState]: # strategy 不关心 agent_loop 是 ray actor 还是本地对象。 + for item in rollout_state: + item.extra_fields["producer_future_step"] = self.train_step + start = time.perf_counter() if isinstance(self.agent_loop, ray.actor.ActorHandle): result = await self.agent_loop.generate_group.remote( diff --git a/xtuner/v1/rl/trace/api.py b/xtuner/v1/rl/trace/api.py index a56f9e9db8..c0d17c8096 100644 --- a/xtuner/v1/rl/trace/api.py +++ b/xtuner/v1/rl/trace/api.py @@ -1,15 +1,11 @@ from __future__ import annotations -import contextlib import contextvars import inspect import json -import os -import time from collections.abc import Callable, Mapping from contextlib import contextmanager from functools import wraps -from pathlib import Path from typing import Any, TypeVar from . import otel_utils @@ -24,7 +20,6 @@ _SPAN_NAME_PATH_BAGGAGE_KEY = "xtuner.span_name_path" _SPAN_NAME_PATH_ATTRIBUTE = "xtuner.span_name_path" -_LIVE_TRACE_PATH_ENV = "XTUNER_OTEL_LIVE_JSONL_PATH" # OTel context propagation carries trace/span IDs, not the current span names. # XTuner keeps this local path so child spans and remote carriers can expose a # readable execution chain for the viewer. @@ -44,8 +39,8 @@ def trace_span( """Create a current trace span. XTuner wraps OTel spans with runtime auto-initialization, trace-enabled - gating, attribute normalization, failure recording, live JSONL progress, - and a ``xtuner.span_name_path`` attribute for viewer-friendly call chains. + gating, attribute normalization, failure recording, and a + ``xtuner.span_name_path`` attribute for viewer-friendly call chains. ``parent_carrier`` accepts the W3C carrier produced by ``inject_trace_context`` across a process or request boundary. """ @@ -65,39 +60,13 @@ def trace_span( normalized_attributes = dict(normalized_attributes) normalized_attributes.setdefault(_SPAN_NAME_PATH_ATTRIBUTE, span_name_path) path_token = _CURRENT_SPAN_NAME_PATH.set(span_name_path) - started_at_s = time.time() - status = "completed" - error_message: str | None = None try: with otel_utils.start_span(span_name, attributes=normalized_attributes): - span_ids = otel_utils.current_span_ids() - _record_live_span( - "start", - span_name=span_name, - span_name_path=span_name_path, - attributes=normalized_attributes, - span_ids=span_ids, - started_at_s=started_at_s, - ) try: yield except Exception as exc: - status = "error" - error_message = str(exc) otel_utils.record_failure(exc) raise - finally: - _record_live_span( - "end", - span_name=span_name, - span_name_path=span_name_path, - attributes=normalized_attributes, - span_ids=span_ids, - started_at_s=started_at_s, - ended_at_s=time.time(), - status=status, - error_message=error_message, - ) finally: _CURRENT_SPAN_NAME_PATH.reset(path_token) @@ -279,52 +248,6 @@ def normalize_scalar(item: Any) -> str | bool | int | float: return f"<{type(value).__name__}>" -def _record_live_span( - event: str, - *, - span_name: str, - span_name_path: tuple[str, ...], - attributes: Mapping[str, Any], - span_ids: Mapping[str, str] | None, - started_at_s: float, - ended_at_s: float | None = None, - status: str | None = None, - error_message: str | None = None, -) -> None: - live_path = os.environ.get(_LIVE_TRACE_PATH_ENV) - if not live_path: - return - record: dict[str, Any] = { - "event": event, - "time_s": ended_at_s if ended_at_s is not None else started_at_s, - "span_name": span_name, - "span_name_path": list(span_name_path), - "attributes": {key: _json_safe_trace_value(value) for key, value in attributes.items()}, - } - if span_ids is not None: - record.update(span_ids) - if status is not None: - record["status"] = status - if ended_at_s is not None: - record["duration_ms"] = round(max(0.0, ended_at_s - started_at_s) * 1000.0, 3) - if error_message: - record["error_message"] = error_message - - with contextlib.suppress(Exception): - path = Path(live_path).expanduser() - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("a", encoding="utf-8") as handle: - handle.write(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n") - - -def _json_safe_trace_value(value: Any) -> Any: - if isinstance(value, (str, bool, int, float)) or value is None: - return value - if isinstance(value, (list, tuple)): - return [_json_safe_trace_value(item) for item in value] - return str(value) - - __all__ = [ "inject_trace_context", "set_trace_attributes", diff --git a/xtuner/v1/rl/trace/rollout_api.py b/xtuner/v1/rl/trace/rollout_api.py index 2476177577..e20a4cccdd 100644 --- a/xtuner/v1/rl/trace/rollout_api.py +++ b/xtuner/v1/rl/trace/rollout_api.py @@ -11,6 +11,7 @@ from xtuner.v1.utils import get_logger from . import api as trace_api +from . import otel_utils from .runtime import is_trace_enabled @@ -18,7 +19,7 @@ TRACE_ROLLOUT_ENABLED_ENV = "XTUNER_TRACE_ENABLE_ROLLOUT" TRACE_CARRIER_EXTRA_FIELD = "_xtuner_trace_carrier" -TRACE_METADATA_EXTRA_FIELD = "_xtuner_trace_metadata" +TRACE_CALL_CHAIN_EXTRA_FIELD = "_xtuner_trace_call_chain" class _RayRemoteMethod(Protocol): @@ -66,13 +67,20 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any: else rollout_state_initial_attributes(rollout_state) ) attributes.setdefault("xtuner.stage", span_name) - parent_carrier = extract_rollout_trace_parent_carrier(rollout_state) - with trace_api.trace_span(span_name, attributes=attributes, parent_carrier=parent_carrier): - result = await func(*args, **kwargs) - result_rollout_state = result if isinstance(result, RolloutState) else rollout_state - trace_api.set_trace_attributes(rollout_state_final_attributes(result_rollout_state)) - return result + + with _attach_rollout_call_chain(rollout_state, span_name, parent_carrier) as ( + call_chain, + cleanup_call_chain_on_exit, + ): + attributes["xtuner.span_name_path"] = call_chain + with trace_api.trace_span(span_name, attributes=attributes, parent_carrier=parent_carrier): + result = await func(*args, **kwargs) + result_rollout_state = result if isinstance(result, RolloutState) else rollout_state + trace_api.set_trace_attributes(rollout_state_final_attributes(result_rollout_state)) + if cleanup_call_chain_on_exit and isinstance(result, RolloutState): + result.extra_fields.pop(TRACE_CALL_CHAIN_EXTRA_FIELD, None) + return result return wrapper # type: ignore[return-value] @@ -122,15 +130,14 @@ def extract_rollout_trace_parent_carrier(rollout_state: RolloutState) -> dict[st def rollout_state_initial_attributes(rollout_state: RolloutState) -> dict[str, Any]: - metadata = _trace_metadata(rollout_state) + extra_fields = rollout_state.extra_fields attributes: dict[str, Any] = { "xtuner.status": rollout_state.status.value, "xtuner.rollout_id": rollout_state.rollout_id, "xtuner.group_id": rollout_state.group_id, "xtuner.session_id": rollout_state.session_id, - "xtuner.task_name": rollout_state.task_name or metadata.get("task_name"), - "xtuner.producer_future_step": metadata.get("producer_future_step"), - "xtuner.model_step": metadata.get("model_step"), + "xtuner.task_name": rollout_state.task_name or extra_fields.get("task_name"), + "xtuner.producer_future_step": extra_fields.get("producer_future_step"), } if rollout_state.prompt_ids is not None: attributes["prompt.tokens"] = len(rollout_state.prompt_ids) @@ -158,11 +165,36 @@ def rollout_state_final_attributes(rollout_state: RolloutState) -> dict[str, Any return {key: value for key, value in attributes.items() if value is not None} -def _trace_metadata(rollout_state: RolloutState) -> dict[str, Any]: - metadata = rollout_state.extra_fields.get(TRACE_METADATA_EXTRA_FIELD) - if not isinstance(metadata, Mapping): - metadata = rollout_state.extra_fields - return {str(key): value for key, value in metadata.items()} +@contextmanager +def _attach_rollout_call_chain( + rollout_state: RolloutState, + span_name: str, + parent_carrier: Mapping[str, str] | None, +): + extra_fields = rollout_state.extra_fields + cleanup_call_chain_on_exit = TRACE_CALL_CHAIN_EXTRA_FIELD not in extra_fields + call_chain = (*_rollout_call_chain(rollout_state, parent_carrier), span_name) + extra_fields[TRACE_CALL_CHAIN_EXTRA_FIELD] = list(call_chain) + try: + yield call_chain, cleanup_call_chain_on_exit + finally: + if cleanup_call_chain_on_exit: + rollout_state.extra_fields.pop(TRACE_CALL_CHAIN_EXTRA_FIELD, None) + + +def _rollout_call_chain( + rollout_state: RolloutState, + parent_carrier: Mapping[str, str] | None, +) -> tuple[str, ...]: + value = rollout_state.extra_fields.get(TRACE_CALL_CHAIN_EXTRA_FIELD) + if value is None and parent_carrier: + parent_context = otel_utils.extract_otel_context(parent_carrier) + value = trace_api._extract_span_name_path(parent_context, otel_utils=otel_utils) + if isinstance(value, str): + return tuple(part.strip() for part in value.split("->") if part.strip()) + if isinstance(value, (list, tuple)): + return tuple(str(item).strip() for item in value if str(item).strip()) + return () def _resolve_rollout_state_target( @@ -193,7 +225,6 @@ def _resolve_rollout_state_target( __all__ = [ "TRACE_CARRIER_EXTRA_FIELD", - "TRACE_METADATA_EXTRA_FIELD", "TRACE_ROLLOUT_ENABLED_ENV", "extract_rollout_trace_parent_carrier", "is_rollout_trace_enabled", diff --git a/xtuner/v1/rl/trace/runtime.py b/xtuner/v1/rl/trace/runtime.py index 753ff208d9..63138627ae 100644 --- a/xtuner/v1/rl/trace/runtime.py +++ b/xtuner/v1/rl/trace/runtime.py @@ -32,7 +32,6 @@ "XTUNER_OTEL_RUN_ID", "XTUNER_OTEL_RUN_DIR", "XTUNER_OTEL_JSONL_PATH", - "XTUNER_OTEL_LIVE_JSONL_PATH", "XTUNER_TRACE_ENABLE_ROLLOUT", "OTEL_TRACES_EXPORTER", "OTEL_EXPORTER_OTLP_ENDPOINT", @@ -140,7 +139,6 @@ class TraceRuntime: run_id: str run_dir: Path trace_jsonl_path: Path - live_jsonl_path: Path service_name: str trace_viewer_url: str | None = None trace_viewer_port: int | None = None @@ -277,7 +275,6 @@ def start( cls, *, trace_jsonl_path: Path, - live_jsonl_path: Path, jaeger_query_url: str | None, jaeger_link_url: str | None, service_name: str, @@ -290,7 +287,6 @@ def start( _ensure_trace_viewer_port_available(host=host, port=port, run_id=run_id) command = _build_trace_viewer_command( trace_jsonl_path=trace_jsonl_path, - live_jsonl_path=live_jsonl_path, jaeger_query_url=jaeger_query_url, jaeger_link_url=jaeger_link_url, service_name=service_name, @@ -459,7 +455,6 @@ def _process_has_socket_inode(proc_dir: Path, socket_inodes: set[str]) -> bool: def _build_trace_viewer_command( *, trace_jsonl_path: Path, - live_jsonl_path: Path, jaeger_query_url: str | None, jaeger_link_url: str | None, service_name: str, @@ -473,8 +468,6 @@ def _build_trace_viewer_command( "xtuner.tools.trace_viewer.server", "--trace-jsonl", os.fspath(trace_jsonl_path), - "--live-jsonl", - os.fspath(live_jsonl_path), "--service", service_name, "--run-id", @@ -527,7 +520,6 @@ def start(self) -> None: if self.runtime.mode == "driver" and self.viewer_host is not None: self.viewer = _TraceViewerProcess.start( trace_jsonl_path=self.runtime.trace_jsonl_path, - live_jsonl_path=self.runtime.live_jsonl_path, jaeger_query_url=self.viewer_jaeger_query_url, jaeger_link_url=self.viewer_jaeger_link_url, service_name=self.runtime.service_name, @@ -606,7 +598,6 @@ def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: run_id="", run_dir=Path(), trace_jsonl_path=Path(), - live_jsonl_path=Path(), service_name=config.service_name, trace_viewer_url=None, trace_viewer_port=None, @@ -622,7 +613,7 @@ def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: traces_dir = run_dir / "traces" traces_dir.mkdir(parents=True, exist_ok=True) trace_jsonl_path = traces_dir / "traces.jsonl" - live_jsonl_path = traces_dir / "live.jsonl" + trace_jsonl_path.touch(exist_ok=True) try: port = find_free_ports(nums=1, host="127.0.0.1", start_port=4317, end_port=4318)[0] @@ -637,7 +628,6 @@ def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: "XTUNER_OTEL_RUN_ID": run_id, "XTUNER_OTEL_RUN_DIR": os.fspath(run_dir), "XTUNER_OTEL_JSONL_PATH": os.fspath(trace_jsonl_path), - "XTUNER_OTEL_LIVE_JSONL_PATH": os.fspath(live_jsonl_path), "XTUNER_TRACE_ENABLE_ROLLOUT": "1" if config.enable_rollout_trace else "0", "OTEL_TRACES_EXPORTER": "otlp", "OTEL_EXPORTER_OTLP_ENDPOINT": endpoint, @@ -652,7 +642,6 @@ def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: run_id=run_id, run_dir=run_dir, trace_jsonl_path=trace_jsonl_path, - live_jsonl_path=live_jsonl_path, service_name=config.service_name, trace_viewer_url=None, trace_viewer_port=None, @@ -712,9 +701,6 @@ def ensure_trace_runtime_from_env() -> bool: run_dir = Path(env_vars.get("XTUNER_OTEL_RUN_DIR") or Path.cwd()).expanduser() trace_jsonl_path = Path(env_vars.get("XTUNER_OTEL_JSONL_PATH") or run_dir / "traces" / "traces.jsonl").expanduser() - live_jsonl_path = Path( - env_vars.get("XTUNER_OTEL_LIVE_JSONL_PATH") or run_dir / "traces" / "live.jsonl" - ).expanduser() runtime_handle = _TraceRuntimeHandle( runtime=TraceRuntime( enabled=True, @@ -722,7 +708,6 @@ def ensure_trace_runtime_from_env() -> bool: run_id=env_vars.get("XTUNER_OTEL_RUN_ID", ""), run_dir=run_dir, trace_jsonl_path=trace_jsonl_path, - live_jsonl_path=live_jsonl_path, service_name=env_vars.get("OTEL_SERVICE_NAME", "xtuner-rollout"), trace_viewer_url=None, trace_viewer_port=None, From 896549895fa5fcda9131ea9283cc73d3d063d51b Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Tue, 14 Jul 2026 03:52:48 +0000 Subject: [PATCH 4/6] fix ci --- tests/rl/test_trace.py | 78 ++++++++++++++------------------ tests/rl/trace_utils.py | 28 +++--------- xtuner/v1/rl/trace/otel_utils.py | 41 ++++++++++++++--- 3 files changed, 76 insertions(+), 71 deletions(-) diff --git a/tests/rl/test_trace.py b/tests/rl/test_trace.py index afe6ab6ac3..f7f59513c8 100644 --- a/tests/rl/test_trace.py +++ b/tests/rl/test_trace.py @@ -48,53 +48,45 @@ def test_nested_trace_span_preserves_parent_to_child_order(self): self.assertEqual(output["child_parent_span_id"], output["parent_span_id"]) self.assertEqual(output["span_name_paths"]["order.parent"], ["order.parent"]) self.assertEqual(output["span_name_paths"]["order.child"], ["order.parent", "order.child"]) - self.assertEqual( - output["live_sequence"], - [ - { - "event": "start", - "span_name": "order.parent", - "span_name_path": ["order.parent"], - }, - { - "event": "start", - "span_name": "order.child", - "span_name_path": ["order.parent", "order.child"], - }, - { - "event": "end", - "span_name": "order.child", - "span_name_path": ["order.parent", "order.child"], - }, - { - "event": "end", - "span_name": "order.parent", - "span_name_path": ["order.parent"], - }, - ], - ) - def test_live_viewer_uses_span_name_path_for_display_chain(self): + def test_viewer_uses_span_name_path_for_display_chain(self): from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces - payload = build_rollout_view_payload_from_jaeger_traces( - [], - live_records=[ - { - "event": "start", - "time_s": 1.0, - "trace_id": "trace-1", - "span_id": "span-1", - "span_name": "child.phase", - "span_name_path": ["parent.phase", "child.phase"], - "attributes": { - "xtuner.rollout_id": "rollout-1", - "xtuner.status": "running", + traces = [ + { + "traceID": "trace-1", + "processes": {"p1": {"serviceName": "xtuner-test", "tags": []}}, + "spans": [ + { + "traceID": "trace-1", + "spanID": "span-1", + "operationName": "parent.phase", + "processID": "p1", + "startTime": 1_000, + "duration": 2_000, + "tags": [ + {"key": "xtuner.rollout_id", "value": "rollout-1"}, + {"key": "xtuner.span_name_path", "value": ["parent.phase"]}, + ], }, - } - ], - train_step="all", - ) + { + "traceID": "trace-1", + "spanID": "span-2", + "operationName": "child.phase", + "processID": "p1", + "startTime": 2_000, + "duration": 1_000, + "references": [{"refType": "CHILD_OF", "traceID": "trace-1", "spanID": "span-1"}], + "tags": [ + {"key": "xtuner.rollout_id", "value": "rollout-1"}, + {"key": "xtuner.span_name_path", "value": ["parent.phase", "child.phase"]}, + ], + }, + ], + } + ] + + payload = build_rollout_view_payload_from_jaeger_traces(traces, train_step="all") self.assertEqual( [node["name"] for node in payload["samples"][0]["display_path"]], diff --git a/tests/rl/trace_utils.py b/tests/rl/trace_utils.py index 14a50df5a6..b92fc20a04 100644 --- a/tests/rl/trace_utils.py +++ b/tests/rl/trace_utils.py @@ -2,7 +2,6 @@ import os import subprocess import sys -import tempfile from pathlib import Path from unittest import mock @@ -129,18 +128,13 @@ def parent_child() -> None: def nested_span_order() -> None: exporter = _install_in_memory_exporter() - with tempfile.TemporaryDirectory() as tmpdir: - live_path = Path(tmpdir) / "trace-live.jsonl" - with mock.patch.dict(os.environ, {"XTUNER_OTEL_LIVE_JSONL_PATH": os.fspath(live_path)}): - with ( - mock.patch.object(trace_api, "_ensure_trace_runtime_from_env"), - mock.patch.object(trace_api, "is_trace_enabled", return_value=True), - ): - with trace_api.trace_span("order.parent"): - with trace_api.trace_span("order.child"): - pass - - live_records = [json.loads(line) for line in live_path.read_text(encoding="utf-8").splitlines()] + with ( + mock.patch.object(trace_api, "_ensure_trace_runtime_from_env"), + mock.patch.object(trace_api, "is_trace_enabled", return_value=True), + ): + with trace_api.trace_span("order.parent"): + with trace_api.trace_span("order.child"): + pass spans = {span.name: span for span in exporter.get_finished_spans()} parent = spans["order.parent"] @@ -153,14 +147,6 @@ def nested_span_order() -> None: name: list(span.attributes.get("xtuner.span_name_path") or []) for name, span in spans.items() }, - "live_sequence": [ - { - "event": record.get("event"), - "span_name": record.get("span_name"), - "span_name_path": record.get("span_name_path"), - } - for record in live_records - ], } ) diff --git a/xtuner/v1/rl/trace/otel_utils.py b/xtuner/v1/rl/trace/otel_utils.py index c109f8b484..2677cc9095 100644 --- a/xtuner/v1/rl/trace/otel_utils.py +++ b/xtuner/v1/rl/trace/otel_utils.py @@ -1,13 +1,10 @@ from __future__ import annotations -from typing import Any, Mapping, MutableMapping +from typing import TYPE_CHECKING, Any, Mapping, MutableMapping -from opentelemetry import baggage, propagate, trace -from opentelemetry import context as otel_context -from opentelemetry.sdk.resources import Resource -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor -from opentelemetry.trace import Status, StatusCode + +if TYPE_CHECKING: + from opentelemetry.sdk.trace import TracerProvider def configure_tracer_provider( @@ -20,7 +17,11 @@ def configure_tracer_provider( if protocol != "grpc": raise ValueError(f"Unsupported OTel trace export protocol: {protocol!r}") try: + from opentelemetry import trace from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter + from opentelemetry.sdk.resources import Resource + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor except ImportError as exc: raise RuntimeError( "XTuner OTel tracing requires the official OpenTelemetry OTLP gRPC trace exporter. " @@ -39,6 +40,8 @@ def inject_otel_context( ) -> MutableMapping[str, str]: """Inject the current or provided OTel context into a W3C carrier.""" + from opentelemetry import propagate + carrier = carrier if carrier is not None else {} propagate.inject(carrier, context=context) return carrier @@ -50,36 +53,48 @@ def extract_otel_context( ) -> Any: """Extract W3C TraceContext from a carrier into an OTel context.""" + from opentelemetry import propagate + return propagate.extract(carrier, context=context) def context_with_baggage(name: str, value: object, context: Any | None = None) -> Any: """Return a context with one W3C Baggage item attached.""" + from opentelemetry import baggage + return baggage.set_baggage(name, value, context=context) def get_baggage(name: str, context: Any | None = None) -> object | None: """Read one W3C Baggage item from a context.""" + from opentelemetry import baggage + return baggage.get_baggage(name, context=context) def attach_otel_context(context: Any) -> object: """Attach extracted context to the current execution scope.""" + from opentelemetry import context as otel_context + return otel_context.attach(context) def detach_otel_context(token: object) -> None: """Detach a previously attached OTel context token.""" + from opentelemetry import context as otel_context + otel_context.detach(token) def start_span(name: str, *, attributes: Mapping[str, Any] | None = None): """Start a current OTel span with XTuner-managed exception handling.""" + from opentelemetry import trace + return trace.get_tracer("xtuner").start_as_current_span( name, attributes=attributes, @@ -89,6 +104,8 @@ def start_span(name: str, *, attributes: Mapping[str, Any] | None = None): def current_span_ids() -> dict[str, str] | None: + from opentelemetry import trace + span = trace.get_current_span() span_context = span.get_span_context() if not span_context.is_valid: @@ -100,6 +117,8 @@ def current_span_ids() -> dict[str, str] | None: def add_event(name: str, *, attributes: Mapping[str, Any] | None = None) -> None: + from opentelemetry import trace + span = trace.get_current_span() if not span.is_recording(): return @@ -107,6 +126,8 @@ def add_event(name: str, *, attributes: Mapping[str, Any] | None = None) -> None def set_attributes(attributes: Mapping[str, Any]) -> None: + from opentelemetry import trace + span = trace.get_current_span() if not span.is_recording(): return @@ -115,6 +136,9 @@ def set_attributes(attributes: Mapping[str, Any]) -> None: def set_error_status(message: str | None = None) -> None: + from opentelemetry import trace + from opentelemetry.trace import Status, StatusCode + span = trace.get_current_span() if not span.is_recording(): return @@ -125,6 +149,9 @@ def set_error_status(message: str | None = None) -> None: def record_failure(exc: BaseException) -> None: + from opentelemetry import trace + from opentelemetry.trace import Status, StatusCode + span = trace.get_current_span() if not span.is_recording(): return From 0c0819b30634032a4ab7b6e7af5f20b21dc476a5 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Wed, 15 Jul 2026 07:21:38 +0000 Subject: [PATCH 5/6] mv trace_viewer to recipe --- .claude/skills/xtuner-trace/SKILL.md | 21 +- .../xtuner-trace/references/trace-patterns.md | 25 +- examples/v1/config/rl_grpo_gsm8k_judge.py | 8 +- examples/v1/scripts/setup_trace.sh | 8 +- recipe/otle/README.md | 9 +- .../tools => recipe}/trace_viewer/payload.py | 0 .../tools => recipe}/trace_viewer/render.py | 0 .../tools => recipe}/trace_viewer/server.py | 4 +- tests/rl/test_trace.py | 6 +- .../v1/rl/agent_loop_manager/produce_utils.py | 2 + xtuner/v1/rl/trace/runtime.py | 282 ++++-------------- 11 files changed, 105 insertions(+), 260 deletions(-) rename {xtuner/tools => recipe}/trace_viewer/payload.py (100%) rename {xtuner/tools => recipe}/trace_viewer/render.py (100%) rename {xtuner/tools => recipe}/trace_viewer/server.py (98%) diff --git a/.claude/skills/xtuner-trace/SKILL.md b/.claude/skills/xtuner-trace/SKILL.md index cf559e56e9..1b56d7d2a9 100644 --- a/.claude/skills/xtuner-trace/SKILL.md +++ b/.claude/skills/xtuner-trace/SKILL.md @@ -16,11 +16,20 @@ Keep this package infrastructure-only: - Public facade: `xtuner/v1/rl/trace/__init__.py` - Basic API: `xtuner/v1/rl/trace/api.py` - Rollout starter preset: `xtuner/v1/rl/trace/rollout_api.py` -- Viewer: `xtuner/tools/trace_viewer/` +- Viewer: `recipe/trace_viewer/` - Local tooling: `recipe/otle/` and `examples/v1/scripts/setup_trace.sh` Do not add rollout, agent, judger, Ray remote, HTTP proxy, reward, status, or session-server business semantics to `api.py`, `runtime.py`, or `otel_utils.py`. Keep rollout-specific starter behavior inside `rollout_api.py` and gated by `TraceConfig.enable_rollout_trace`. +## Viewer Choice + +Choose the viewer by the observation question: + +- Use `recipe/trace_viewer` when the user wants rollout-oriented inspection, such as comparing samples within a training step by stage duration, status, reward/filter metadata, or recorded sample call chain. +- Use Jaeger when the question is not sample/step-centric, such as inspecting raw spans, service boundaries, cross-process causality, backend/client calls, collector/export behavior, or custom instrumentation that does not emit the rollout/sample attributes expected by the recipe viewer. + +The recipe viewer reads `traces.jsonl` and adds XTuner rollout/sample aggregation. Jaeger remains the raw trace drill-down surface for the same traced run. + ## Basic API These are the interfaces defined in `xtuner/v1/rl/trace/api.py` and re-exported @@ -72,10 +81,10 @@ When adding trace instrumentation to an XTuner run: 2. In the launch script, source `examples/v1/scripts/setup_trace.sh` when `XTUNER_TRACE_ENABLED=1`. 3. In the training config, add `TraceConfig` and set `enabled=True`; set - `viewer_enabled=True` when the user needs interactive inspection. Set + `xtuner_viewer_enabled=True` when the user needs interactive sample/step inspection. Set `enable_rollout_trace=True` only when the user wants the built-in rollout starter trace. -4. Before adding any `trace_span(...)` instrumentation, you mask ask the user which +4. Before adding any `trace_span(...)` instrumentation, you must ask the user which stages they want to observe and which metrics each stage should expose. Do not infer default stages unless the user explicitly asks you to choose. 5. Add `trace_span(...)` only around the user-confirmed observed stages. Put fields known at span @@ -86,8 +95,8 @@ When adding trace instrumentation to an XTuner run: `trace_span(..., parent_carrier=carrier)`. 7. Keep transport-specific propagation at the caller boundary; do not move rollout, agent, judger, Ray, or HTTP semantics into the basic trace package. -8. Ensure the main training log includes the trace output path, viewer URL, and - a restart command for the viewer. +8. Ensure the main training log includes the trace output path, XTuner viewer URL, + and equivalent manual viewer command when the XTuner viewer is enabled. ## Guardrails @@ -103,6 +112,6 @@ For concrete patterns, read [references/trace-patterns.md](references/trace-patt Use focused checks: -- `PYTHONPATH=. python -m compileall -q xtuner/v1/rl/trace xtuner/tools/trace_viewer` +- `PYTHONPATH=. python -m compileall -q xtuner/v1/rl/trace recipe/trace_viewer` - `PYTHONPATH=. python -m unittest discover -s tests/rl -p 'test_trace*.py' -v` when the trace tests exist in the worktree. - `git diff --check` diff --git a/.claude/skills/xtuner-trace/references/trace-patterns.md b/.claude/skills/xtuner-trace/references/trace-patterns.md index 6963590eaa..ca62e8234b 100644 --- a/.claude/skills/xtuner-trace/references/trace-patterns.md +++ b/.claude/skills/xtuner-trace/references/trace-patterns.md @@ -81,10 +81,10 @@ from xtuner.v1.rl.trace import TraceConfig trace_config = TraceConfig( enabled=os.environ.get("XTUNER_TRACE_ENABLED") == "1", 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", + xtuner_viewer_enabled=True, + xtuner_viewer_host="0.0.0.0", + xtuner_viewer_port=18080, + xtuner_viewer_jaeger_query_url="http://127.0.0.1:16686", enable_rollout_trace=True, ) @@ -344,9 +344,12 @@ derived fields such as method, route, status code, stage, IDs, and timing. ### Viewer Output -When `viewer_enabled=True`, the trace runtime logs the trace JSONL path, viewer -URL, and restart command. Keep those runtime log lines visible in the main -training log and report them to the user after the run. +When `xtuner_viewer_enabled=True`, the trace runtime logs the trace JSONL path, +the XTuner viewer URL, and the equivalent manual viewer command. The viewer +defaults to port `18080`; `xtuner_viewer_port=0` is still available when a run +should pick a free port automatically. The viewer process inherits the training +stdout/stderr; keep those runtime log lines visible in the main training log and +report them to the user after the run. Useful reference points from the full trace implementation: @@ -366,7 +369,7 @@ runtime = configure_trace( enabled=True, output_dir="work_dirs/example/otel", service_name="xtuner", - viewer_enabled=True, + xtuner_viewer_enabled=True, ) ) @@ -376,7 +379,7 @@ finally: close_trace() ``` -The runtime owns OTel collector setup, trace JSONL output, live JSONL output, and optional viewer process startup. +The runtime owns OTel collector setup, trace JSONL output, and optional lightweight XTuner viewer startup. ## Viewer @@ -401,7 +404,7 @@ Avoid recording prompts, responses, full configs, secrets, raw headers, or large ## Local Setup -`examples/v1/scripts/setup_trace.sh` and `recipe/otle/` are local helper tooling for installing OTel collector binaries and starting local Jaeger/viewer dependencies. Keep these as setup assets; do not use them to add automatic trace behavior to training configs or launch scripts unless that integration is explicitly requested. +`examples/v1/scripts/setup_trace.sh` and `recipe/otle/` are local helper tooling for installing OTel collector binaries, starting the local Jaeger dependency, and clearing stale `recipe.trace_viewer.server` processes on `XTUNER_TRACE_VIEWER_PORT` or the default viewer port `18080`. Keep these as setup assets; do not use them to add automatic trace behavior to training configs or launch scripts unless that integration is explicitly requested. ## Rollout Starter Preset @@ -410,7 +413,7 @@ The built-in rollout starter trace is opt-in: ```python trace_config = TraceConfig( enabled=True, - viewer_enabled=True, + xtuner_viewer_enabled=True, enable_rollout_trace=True, ) ``` diff --git a/examples/v1/config/rl_grpo_gsm8k_judge.py b/examples/v1/config/rl_grpo_gsm8k_judge.py index 409d48b951..1f68946a81 100644 --- a/examples/v1/config/rl_grpo_gsm8k_judge.py +++ b/examples/v1/config/rl_grpo_gsm8k_judge.py @@ -188,10 +188,10 @@ trace_config = TraceConfig( enabled=enable_trace, 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", + xtuner_viewer_enabled=True, + xtuner_viewer_host="0.0.0.0", + xtuner_viewer_port=18080, + xtuner_viewer_jaeger_query_url="http://127.0.0.1:16686", enable_rollout_trace=True, ) diff --git a/examples/v1/scripts/setup_trace.sh b/examples/v1/scripts/setup_trace.sh index 560f01eaa8..3db641f384 100644 --- a/examples/v1/scripts/setup_trace.sh +++ b/examples/v1/scripts/setup_trace.sh @@ -11,11 +11,11 @@ JAEGER_RESTART_SCRIPT="${XTUNER_JAEGER_RESTART_SCRIPT:-${REPO_ROOT}/recipe/otle/ 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 +if pgrep -f "recipe.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" >/dev/null 2>&1 || + pgrep -f "recipe.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 + pkill -f "recipe.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" 2>/dev/null || true + pkill -f "recipe.trace_viewer.server.*--port=${TRACE_VIEWER_PORT}" 2>/dev/null || true sleep 1 fi diff --git a/recipe/otle/README.md b/recipe/otle/README.md index f912ae008d..6dda038c98 100644 --- a/recipe/otle/README.md +++ b/recipe/otle/README.md @@ -40,10 +40,17 @@ By default, XTuner starts a local collector that writes `/traces/traces.jsonl` and forwards spans to the reference Jaeger OTLP gRPC endpoint `http://127.0.0.1:14317`. +Set `TraceConfig(xtuner_viewer_enabled=True, ...)` to start the XTuner rollout +viewer with the trace runtime. The viewer output goes to the same terminal or +training log as the training process. The default viewer port is `18080`. +`examples/v1/scripts/setup_trace.sh` clears stale `recipe.trace_viewer.server` +processes on `XTUNER_TRACE_VIEWER_PORT` or `18080` before preparing the local +trace dependencies. + Open the rollout viewer: ```bash -python -m xtuner.tools.trace_viewer.server \ +python -m recipe.trace_viewer.server \ --trace-jsonl /traces/traces.jsonl \ --jaeger-query-url http://127.0.0.1:16686 \ --service xtuner-rollout diff --git a/xtuner/tools/trace_viewer/payload.py b/recipe/trace_viewer/payload.py similarity index 100% rename from xtuner/tools/trace_viewer/payload.py rename to recipe/trace_viewer/payload.py diff --git a/xtuner/tools/trace_viewer/render.py b/recipe/trace_viewer/render.py similarity index 100% rename from xtuner/tools/trace_viewer/render.py rename to recipe/trace_viewer/render.py diff --git a/xtuner/tools/trace_viewer/server.py b/recipe/trace_viewer/server.py similarity index 98% rename from xtuner/tools/trace_viewer/server.py rename to recipe/trace_viewer/server.py index 94a531dacf..2ad5f07923 100644 --- a/xtuner/tools/trace_viewer/server.py +++ b/recipe/trace_viewer/server.py @@ -20,12 +20,12 @@ from urllib.parse import parse_qs, urlsplit from urllib.request import Request, urlopen -from xtuner.tools.trace_viewer.payload import ( +from recipe.trace_viewer.payload import ( build_rollout_view_payload_from_jaeger_traces, filter_rollout_view_payload_by_train_step, load_jaeger_traces_from_otel_jsonl, ) -from xtuner.tools.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html +from recipe.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html _JAEGER_PROXY_PREFIX = "/jaeger" diff --git a/tests/rl/test_trace.py b/tests/rl/test_trace.py index f7f59513c8..c2782fa546 100644 --- a/tests/rl/test_trace.py +++ b/tests/rl/test_trace.py @@ -50,7 +50,7 @@ def test_nested_trace_span_preserves_parent_to_child_order(self): self.assertEqual(output["span_name_paths"]["order.child"], ["order.parent", "order.child"]) def test_viewer_uses_span_name_path_for_display_chain(self): - from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces + from recipe.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces traces = [ { @@ -95,8 +95,8 @@ def test_viewer_uses_span_name_path_for_display_chain(self): self.assertEqual(payload["samples"][0]["chain"], "parent.phase -> child.phase") def test_viewer_filters_latest_train_step_and_renders_payload(self): - from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces - from xtuner.tools.trace_viewer.render import render_rollout_trace_html + from recipe.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces + from recipe.trace_viewer.render import render_rollout_trace_html traces = [ { diff --git a/xtuner/v1/rl/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index d3d8101d85..b80a2f9e99 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -144,6 +144,8 @@ async def generate_group( enable_partial_rollout: bool = False, ) -> list[RolloutState]: # strategy 不关心 agent_loop 是 ray actor 还是本地对象。 + # Trace viewer uses this producer-side future step to group samples by + # the training step they are being generated for. for item in rollout_state: item.extra_fields["producer_future_step"] = self.train_step diff --git a/xtuner/v1/rl/trace/runtime.py b/xtuner/v1/rl/trace/runtime.py index 63138627ae..893bbee61a 100644 --- a/xtuner/v1/rl/trace/runtime.py +++ b/xtuner/v1/rl/trace/runtime.py @@ -2,7 +2,6 @@ import atexit import contextlib -import errno import os import shlex import shutil @@ -11,9 +10,10 @@ import sys import time import uuid +from collections.abc import Mapping from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Any, Literal, Mapping +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator @@ -42,10 +42,11 @@ DEFAULT_JAEGER_OTLP_GRPC_ENDPOINT = "127.0.0.1:14317" -_TRACE_VIEWER_SERVER_MODULE = "xtuner.tools.trace_viewer.server" +_TRACE_VIEWER_SERVER_MODULE = "recipe.trace_viewer.server" _READY_TIMEOUT_S = 3.0 _READY_POLL_INTERVAL_S = 0.05 _LOG_TAIL_CHARS = 4000 +_TRACE_VIEWER_STARTUP_CHECK_TIMEOUT_S = 0.5 _OTELCOL_CONFIG_YAML_TEMPLATE = """ receivers: @@ -117,11 +118,10 @@ class TraceConfig(BaseModel): enabled: bool = False output_dir: Path | str | None = Field(default=None) service_name: str = "xtuner-rollout" - viewer_enabled: bool = False - viewer_host: str = "127.0.0.1" - viewer_port: int = Field(default=0, ge=0, le=65535) - viewer_jaeger_query_url: str | None = None - viewer_jaeger_link_url: str | None = None + xtuner_viewer_enabled: bool = False + xtuner_viewer_host: str = "127.0.0.1" + xtuner_viewer_port: int = Field(default=18080, ge=0, le=65535) + xtuner_viewer_jaeger_query_url: str | None = None enable_rollout_trace: bool = False @field_validator("output_dir") @@ -261,202 +261,10 @@ def close(self) -> None: process.wait(timeout=5) -@dataclass -class _TraceViewerProcess: - host: str - port: int - url: str - log_path: Path - _process: subprocess.Popen | None = field(repr=False) - command: list[str] = field(default_factory=list, repr=False) - - @classmethod - def start( - cls, - *, - trace_jsonl_path: Path, - jaeger_query_url: str | None, - jaeger_link_url: str | None, - service_name: str, - run_id: str, - host: str, - port: int, - ) -> _TraceViewerProcess: - if port == 0: - port = find_free_ports(nums=1, host=host)[0] - _ensure_trace_viewer_port_available(host=host, port=port, run_id=run_id) - command = _build_trace_viewer_command( - trace_jsonl_path=trace_jsonl_path, - jaeger_query_url=jaeger_query_url, - jaeger_link_url=jaeger_link_url, - service_name=service_name, - run_id=run_id, - host=host, - port=port, - ) - log_path = trace_jsonl_path.parent / "viewer.log" - with log_path.open("ab") as log_file: - process = subprocess.Popen( - command, - stdout=log_file, - stderr=subprocess.STDOUT, - start_new_session=True, - ) - viewer = cls( - host=host, - port=port, - url=f"http://{host}:{port}", - log_path=log_path, - command=command, - _process=process, - ) - try: - viewer._wait_until_ready() - except Exception: - viewer.close() - raise - return viewer - - def _wait_until_ready(self) -> None: - deadline = time.monotonic() + _READY_TIMEOUT_S - last_error: OSError | None = None - while time.monotonic() < deadline: - self._raise_if_process_exited() - connect_host = "127.0.0.1" if self.host in {"", "0.0.0.0"} else self.host - try: - with socket.create_connection((connect_host, self.port), timeout=0.1): - time.sleep(_READY_POLL_INTERVAL_S) - self._raise_if_process_exited() - return - except OSError as exc: - last_error = exc - time.sleep(_READY_POLL_INTERVAL_S) - - detail = f"viewer did not become ready within {_READY_TIMEOUT_S:.1f}s" - if last_error is not None: - detail += f"; last connection error: {last_error}" - log_tail = "" - with contextlib.suppress(OSError): - log_tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] - raise RuntimeError(f"XTuner trace viewer failed to start: {detail}, url={self.url}, log_tail={log_tail!r}") - - def _raise_if_process_exited(self) -> None: - process = self._process - if process is None: - raise RuntimeError(f"XTuner trace viewer failed to start: process is not available, log={self.log_path}") - exit_code = process.poll() - if exit_code is None: - return - - log_tail = "" - with contextlib.suppress(OSError): - log_tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] - raise RuntimeError( - f"XTuner trace viewer failed to start: exited with code {exit_code}, url={self.url}, log_tail={log_tail!r}" - ) - - def restart_command(self) -> str: - return shlex.join(self.command) - - def close(self) -> None: - process = self._process - self._process = None - if process is None: - return - if process.poll() is None: - process.terminate() - with contextlib.suppress(subprocess.TimeoutExpired): - process.wait(timeout=5) - if process.poll() is None: - process.kill() - with contextlib.suppress(subprocess.TimeoutExpired): - process.wait(timeout=5) - - -def _ensure_trace_viewer_port_available(*, host: str, port: int, run_id: str) -> None: - existing_viewer = _find_existing_trace_viewer_process_on_port(port) - if existing_viewer is not None: - pid = existing_viewer.get("pid", "unknown") - cmdline = existing_viewer.get("cmdline", "") - raise RuntimeError( - f"XTuner trace viewer port {port} already has an existing XTuner trace viewer process: " - f"pid={pid}, cmdline={cmdline!r}. " - f"Stop the old viewer/training process or set XTUNER_TRACE_VIEWER_PORT to another port " - f"before starting run_id={run_id}." - ) - - bind_host = host or "0.0.0.0" - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind((bind_host, port)) - except OSError as exc: - if exc.errno == errno.EADDRINUSE: - raise RuntimeError( - f"XTuner trace viewer port {port} is already in use on host {host!r}. " - "Stop the process using that port or set XTUNER_TRACE_VIEWER_PORT to another port." - ) from exc - raise - - -def _find_existing_trace_viewer_process_on_port(port: int) -> Mapping[str, Any] | None: - listening_inodes = _listening_socket_inodes_on_port(port) - if not listening_inodes: - return None - - proc_root = Path("/proc") - with contextlib.suppress(OSError): - proc_dirs = list(proc_root.iterdir()) - for proc_dir in proc_dirs: - if not proc_dir.name.isdigit(): - continue - cmdline = _read_process_cmdline(proc_dir) - if _TRACE_VIEWER_SERVER_MODULE not in cmdline: - continue - if _process_has_socket_inode(proc_dir, listening_inodes): - return {"pid": int(proc_dir.name), "cmdline": cmdline} - return None - - -def _listening_socket_inodes_on_port(port: int) -> set[str]: - inodes: set[str] = set() - for proc_net_path in (Path("/proc/net/tcp"), Path("/proc/net/tcp6")): - with contextlib.suppress(OSError, ValueError): - for line in proc_net_path.read_text(encoding="utf-8", errors="replace").splitlines()[1:]: - parts = line.split() - if len(parts) < 10: - continue - local_address, state, inode = parts[1], parts[3], parts[9] - if state != "0A": - continue - _, port_hex = local_address.rsplit(":", 1) - if int(port_hex, 16) == port: - inodes.add(inode) - return inodes - - -def _read_process_cmdline(proc_dir: Path) -> str: - with contextlib.suppress(OSError): - raw_cmdline = (proc_dir / "cmdline").read_bytes() - return " ".join(part.decode("utf-8", errors="replace") for part in raw_cmdline.split(b"\0") if part) - return "" - - -def _process_has_socket_inode(proc_dir: Path, socket_inodes: set[str]) -> bool: - fd_dir = proc_dir / "fd" - with contextlib.suppress(OSError): - for fd_path in fd_dir.iterdir(): - with contextlib.suppress(OSError): - target = os.readlink(fd_path) - if target.startswith("socket:[") and target.endswith("]") and target[8:-1] in socket_inodes: - return True - return False - - -def _build_trace_viewer_command( +def _build_xtuner_viewer_command( *, trace_jsonl_path: Path, jaeger_query_url: str | None, - jaeger_link_url: str | None, service_name: str, run_id: str, host: str, @@ -465,7 +273,7 @@ def _build_trace_viewer_command( command = [ sys.executable, "-m", - "xtuner.tools.trace_viewer.server", + _TRACE_VIEWER_SERVER_MODULE, "--trace-jsonl", os.fspath(trace_jsonl_path), "--service", @@ -479,8 +287,6 @@ def _build_trace_viewer_command( ] if jaeger_query_url is not None: command.extend(["--jaeger-query-url", jaeger_query_url]) - if jaeger_link_url is not None: - command.extend(["--jaeger-link-url", jaeger_link_url]) return command @@ -492,11 +298,12 @@ class _TraceRuntimeHandle: collector_port: int | None = None collector: _OTelCollector | None = None provider: Any | None = None - viewer_host: str | None = None - viewer_port: int = 0 - viewer_jaeger_query_url: str | None = None - viewer_jaeger_link_url: str | None = None - viewer: _TraceViewerProcess | None = None + xtuner_viewer_host: str | None = None + xtuner_viewer_port: int = 0 + xtuner_viewer_jaeger_query_url: str | None = None + xtuner_viewer_process: subprocess.Popen | None = None + xtuner_viewer_command: list[str] | None = None + xtuner_viewer_url: str | None = None def start(self) -> None: apply_trace_env(self.env_vars) @@ -517,20 +324,32 @@ def start(self) -> None: endpoint=self.endpoint, protocol=self.env_vars["OTEL_EXPORTER_OTLP_PROTOCOL"], ) - if self.runtime.mode == "driver" and self.viewer_host is not None: - self.viewer = _TraceViewerProcess.start( + if self.runtime.mode == "driver" and self.xtuner_viewer_host is not None: + if self.xtuner_viewer_port == 0: + self.xtuner_viewer_port = find_free_ports(nums=1, host=self.xtuner_viewer_host)[0] + self.xtuner_viewer_command = _build_xtuner_viewer_command( trace_jsonl_path=self.runtime.trace_jsonl_path, - jaeger_query_url=self.viewer_jaeger_query_url, - jaeger_link_url=self.viewer_jaeger_link_url, + jaeger_query_url=self.xtuner_viewer_jaeger_query_url, service_name=self.runtime.service_name, run_id=self.runtime.run_id, - host=self.viewer_host, - port=self.viewer_port, + host=self.xtuner_viewer_host, + port=self.xtuner_viewer_port, ) + self.xtuner_viewer_process = subprocess.Popen(self.xtuner_viewer_command) + try: + exit_code = self.xtuner_viewer_process.wait(timeout=_TRACE_VIEWER_STARTUP_CHECK_TIMEOUT_S) + except subprocess.TimeoutExpired: + pass + else: + raise RuntimeError( + "XTuner trace viewer failed to start: " + f"exit_code={exit_code}. Manual command: {shlex.join(self.xtuner_viewer_command)}" + ) + self.xtuner_viewer_url = f"http://{self.xtuner_viewer_host}:{self.xtuner_viewer_port}" self.runtime = replace( self.runtime, - trace_viewer_url=self.viewer.url, - trace_viewer_port=self.viewer.port, + trace_viewer_url=self.xtuner_viewer_url, + trace_viewer_port=self.xtuner_viewer_port, ) except Exception: self.close(stop_viewer=True) @@ -540,19 +359,25 @@ def start(self) -> None: f"XTuner OTel tracing enabled: run_id={self.runtime.run_id}, endpoint={self.endpoint}, " f"traces={self.runtime.trace_jsonl_path}" ) - if self.viewer is not None: + if self.xtuner_viewer_process is not None: logger.info( - f"XTuner trace viewer enabled: url={self.viewer.url}, host={self.viewer.host}, " - f"port={self.viewer.port}. Restart with: {self.viewer.restart_command()}" + f"XTuner trace viewer enabled: url={self.xtuner_viewer_url}. " + f"Manual command: {shlex.join(self.xtuner_viewer_command or [])}" ) def close(self, *, stop_viewer: bool = True) -> None: - viewer = self.viewer - self.viewer = None - if viewer is not None and stop_viewer: + xtuner_viewer_process = self.xtuner_viewer_process + self.xtuner_viewer_process = None + if xtuner_viewer_process is not None and stop_viewer: logger.info("XTuner trace viewer stopped with training process.") - with contextlib.suppress(Exception): - viewer.close() + if xtuner_viewer_process.poll() is None: + xtuner_viewer_process.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + xtuner_viewer_process.wait(timeout=5) + if xtuner_viewer_process.poll() is None: + xtuner_viewer_process.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + xtuner_viewer_process.wait(timeout=5) provider = self.provider self.provider = None @@ -649,10 +474,9 @@ def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: endpoint=endpoint, env_vars=env_vars, collector_port=port, - viewer_host=config.viewer_host if config.viewer_enabled else None, - viewer_port=config.viewer_port, - viewer_jaeger_query_url=config.viewer_jaeger_query_url, - viewer_jaeger_link_url=config.viewer_jaeger_link_url, + xtuner_viewer_host=config.xtuner_viewer_host if config.xtuner_viewer_enabled else None, + xtuner_viewer_port=config.xtuner_viewer_port, + xtuner_viewer_jaeger_query_url=config.xtuner_viewer_jaeger_query_url, ) From b501bba62030f8fa33b74b6b6051bf8ccc4c0427 Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Wed, 15 Jul 2026 08:56:24 +0000 Subject: [PATCH 6/6] mv setup_trace.sh to recipe --- .claude/skills/xtuner-trace/SKILL.md | 9 ++++----- .../xtuner-trace/references/trace-patterns.md | 7 ++++--- recipe/{otle => trace}/README.md | 12 ++++++------ recipe/{otle => trace}/jaeger/jaeger-memory.yaml | 0 .../{otle => trace/scripts}/install_otel_tools.sh | 0 .../scripts}/restart_jaeger_memory.sh | 5 +++-- .../v1/scripts => recipe/trace}/setup_trace.sh | 15 +++++++-------- recipe/{trace_viewer => trace/viewer}/payload.py | 0 recipe/{trace_viewer => trace/viewer}/render.py | 0 recipe/{trace_viewer => trace/viewer}/server.py | 4 ++-- tests/rl/test_trace.py | 6 +++--- xtuner/v1/rl/trace/runtime.py | 2 +- 12 files changed, 30 insertions(+), 30 deletions(-) rename recipe/{otle => trace}/README.md (85%) rename recipe/{otle => trace}/jaeger/jaeger-memory.yaml (100%) rename recipe/{otle => trace/scripts}/install_otel_tools.sh (100%) rename recipe/{otle => trace/scripts}/restart_jaeger_memory.sh (92%) rename {examples/v1/scripts => recipe/trace}/setup_trace.sh (63%) rename recipe/{trace_viewer => trace/viewer}/payload.py (100%) rename recipe/{trace_viewer => trace/viewer}/render.py (100%) rename recipe/{trace_viewer => trace/viewer}/server.py (99%) diff --git a/.claude/skills/xtuner-trace/SKILL.md b/.claude/skills/xtuner-trace/SKILL.md index 1b56d7d2a9..15da6cb4a5 100644 --- a/.claude/skills/xtuner-trace/SKILL.md +++ b/.claude/skills/xtuner-trace/SKILL.md @@ -16,8 +16,7 @@ Keep this package infrastructure-only: - Public facade: `xtuner/v1/rl/trace/__init__.py` - Basic API: `xtuner/v1/rl/trace/api.py` - Rollout starter preset: `xtuner/v1/rl/trace/rollout_api.py` -- Viewer: `recipe/trace_viewer/` -- Local tooling: `recipe/otle/` and `examples/v1/scripts/setup_trace.sh` +- Local trace tooling and viewer: `recipe/trace/` Do not add rollout, agent, judger, Ray remote, HTTP proxy, reward, status, or session-server business semantics to `api.py`, `runtime.py`, or `otel_utils.py`. Keep rollout-specific starter behavior inside `rollout_api.py` and gated by `TraceConfig.enable_rollout_trace`. @@ -25,7 +24,7 @@ Do not add rollout, agent, judger, Ray remote, HTTP proxy, reward, status, or se Choose the viewer by the observation question: -- Use `recipe/trace_viewer` when the user wants rollout-oriented inspection, such as comparing samples within a training step by stage duration, status, reward/filter metadata, or recorded sample call chain. +- Use `recipe/trace/viewer` when the user wants rollout-oriented inspection, such as comparing samples within a training step by stage duration, status, reward/filter metadata, or recorded sample call chain. - Use Jaeger when the question is not sample/step-centric, such as inspecting raw spans, service boundaries, cross-process causality, backend/client calls, collector/export behavior, or custom instrumentation that does not emit the rollout/sample attributes expected by the recipe viewer. The recipe viewer reads `traces.jsonl` and adds XTuner rollout/sample aggregation. Jaeger remains the raw trace drill-down surface for the same traced run. @@ -78,7 +77,7 @@ the basic trace API. When adding trace instrumentation to an XTuner run: 1. Ask for or locate the launch script and training config before editing. -2. In the launch script, source `examples/v1/scripts/setup_trace.sh` when +2. In the launch script, source `recipe/trace/setup_trace.sh` when `XTUNER_TRACE_ENABLED=1`. 3. In the training config, add `TraceConfig` and set `enabled=True`; set `xtuner_viewer_enabled=True` when the user needs interactive sample/step inspection. Set @@ -112,6 +111,6 @@ For concrete patterns, read [references/trace-patterns.md](references/trace-patt Use focused checks: -- `PYTHONPATH=. python -m compileall -q xtuner/v1/rl/trace recipe/trace_viewer` +- `PYTHONPATH=. python -m compileall -q xtuner/v1/rl/trace recipe/trace` - `PYTHONPATH=. python -m unittest discover -s tests/rl -p 'test_trace*.py' -v` when the trace tests exist in the worktree. - `git diff --check` diff --git a/.claude/skills/xtuner-trace/references/trace-patterns.md b/.claude/skills/xtuner-trace/references/trace-patterns.md index ca62e8234b..6ecc5437fe 100644 --- a/.claude/skills/xtuner-trace/references/trace-patterns.md +++ b/.claude/skills/xtuner-trace/references/trace-patterns.md @@ -53,8 +53,9 @@ Enable trace bootstrap from the launcher, guarded by `XTUNER_TRACE_ENABLED`: ```bash SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then - source "${SCRIPT_DIR}/setup_trace.sh" + source "${REPO_ROOT}/recipe/trace/setup_trace.sh" fi ``` @@ -354,7 +355,7 @@ report them to the user after the run. Useful reference points from the full trace implementation: - `examples/v1/scripts/run_rl_run.sh` -- `examples/v1/scripts/setup_trace.sh` +- `recipe/trace/setup_trace.sh` - `examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py` ## Runtime @@ -404,7 +405,7 @@ Avoid recording prompts, responses, full configs, secrets, raw headers, or large ## Local Setup -`examples/v1/scripts/setup_trace.sh` and `recipe/otle/` are local helper tooling for installing OTel collector binaries, starting the local Jaeger dependency, and clearing stale `recipe.trace_viewer.server` processes on `XTUNER_TRACE_VIEWER_PORT` or the default viewer port `18080`. Keep these as setup assets; do not use them to add automatic trace behavior to training configs or launch scripts unless that integration is explicitly requested. +`recipe/trace/setup_trace.sh` and the rest of `recipe/trace/` are local helper tooling for installing OTel collector binaries, starting the local Jaeger dependency, and clearing stale `recipe.trace.viewer.server` processes on `XTUNER_TRACE_VIEWER_PORT` or the default viewer port `18080`. Keep these as setup assets; do not use them to add automatic trace behavior to training configs or launch scripts unless that integration is explicitly requested. ## Rollout Starter Preset diff --git a/recipe/otle/README.md b/recipe/trace/README.md similarity index 85% rename from recipe/otle/README.md rename to recipe/trace/README.md index 6dda038c98..4f5eef463a 100644 --- a/recipe/otle/README.md +++ b/recipe/trace/README.md @@ -1,4 +1,4 @@ -# XTuner OTel Trace +# XTuner Trace Tools XTuner exports rollout traces through OpenTelemetry. For local inspection, start Jaeger with `jaeger/jaeger-memory.yaml`, enable trace in the training config, and @@ -13,21 +13,21 @@ The reference Jaeger config exposes: Install local binaries: ```bash -bash recipe/otle/install_otel_tools.sh +bash recipe/trace/scripts/install_otel_tools.sh export PATH=/tmp/xtuner_otel/bin:$PATH ``` Start Jaeger: ```bash -jaeger --config recipe/otle/jaeger/jaeger-memory.yaml +jaeger --config recipe/trace/jaeger/jaeger-memory.yaml ``` For local smoke tests, restart the in-memory Jaeger before each experiment so old services, operations, and trace ids cannot be mixed with the new run: ```bash -bash recipe/otle/restart_jaeger_memory.sh +bash recipe/trace/scripts/restart_jaeger_memory.sh ``` Run XTuner with trace enabled: @@ -43,14 +43,14 @@ OTLP gRPC endpoint `http://127.0.0.1:14317`. Set `TraceConfig(xtuner_viewer_enabled=True, ...)` to start the XTuner rollout viewer with the trace runtime. The viewer output goes to the same terminal or training log as the training process. The default viewer port is `18080`. -`examples/v1/scripts/setup_trace.sh` clears stale `recipe.trace_viewer.server` +`recipe/trace/setup_trace.sh` clears stale `recipe.trace.viewer.server` processes on `XTUNER_TRACE_VIEWER_PORT` or `18080` before preparing the local trace dependencies. Open the rollout viewer: ```bash -python -m recipe.trace_viewer.server \ +python -m recipe.trace.viewer.server \ --trace-jsonl /traces/traces.jsonl \ --jaeger-query-url http://127.0.0.1:16686 \ --service xtuner-rollout diff --git a/recipe/otle/jaeger/jaeger-memory.yaml b/recipe/trace/jaeger/jaeger-memory.yaml similarity index 100% rename from recipe/otle/jaeger/jaeger-memory.yaml rename to recipe/trace/jaeger/jaeger-memory.yaml diff --git a/recipe/otle/install_otel_tools.sh b/recipe/trace/scripts/install_otel_tools.sh similarity index 100% rename from recipe/otle/install_otel_tools.sh rename to recipe/trace/scripts/install_otel_tools.sh diff --git a/recipe/otle/restart_jaeger_memory.sh b/recipe/trace/scripts/restart_jaeger_memory.sh similarity index 92% rename from recipe/otle/restart_jaeger_memory.sh rename to recipe/trace/scripts/restart_jaeger_memory.sh index b35eb8930f..03539ae567 100755 --- a/recipe/otle/restart_jaeger_memory.sh +++ b/recipe/trace/scripts/restart_jaeger_memory.sh @@ -6,7 +6,8 @@ set -euo pipefail ROOT="${XTUNER_OTEL_ROOT:-/tmp/xtuner_otel}" JAEGER_BIN="${JAEGER_BIN:-${ROOT}/bin/jaeger}" -CONFIG="${1:-recipe/otle/jaeger/jaeger-memory.yaml}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +CONFIG="${1:-${SCRIPT_DIR}/../jaeger/jaeger-memory.yaml}" PID_FILE="${XTUNER_JAEGER_PID_FILE:-/tmp/xtuner_jaeger_memory.pid}" LOG_FILE="${XTUNER_JAEGER_LOG_FILE:-/tmp/xtuner_jaeger_memory.log}" QUERY_URL="${XTUNER_JAEGER_QUERY_URL:-http://127.0.0.1:16686}" @@ -14,7 +15,7 @@ WAIT_TIMEOUT_S="${XTUNER_JAEGER_WAIT_TIMEOUT_S:-30}" if ! command -v "${JAEGER_BIN}" >/dev/null 2>&1; then echo "Jaeger binary not found: ${JAEGER_BIN}" >&2 - echo "Install it first: bash recipe/otle/install_otel_tools.sh" >&2 + echo "Install it first: bash recipe/trace/scripts/install_otel_tools.sh" >&2 exit 1 fi diff --git a/examples/v1/scripts/setup_trace.sh b/recipe/trace/setup_trace.sh similarity index 63% rename from examples/v1/scripts/setup_trace.sh rename to recipe/trace/setup_trace.sh index 3db641f384..08dae26a7c 100644 --- a/examples/v1/scripts/setup_trace.sh +++ b/recipe/trace/setup_trace.sh @@ -1,21 +1,20 @@ #!/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}" +OTEL_INSTALL_SCRIPT="${XTUNER_OTEL_INSTALL_SCRIPT:-${SETUP_TRACE_DIR}/scripts/install_otel_tools.sh}" +JAEGER_RESTART_SCRIPT="${XTUNER_JAEGER_RESTART_SCRIPT:-${SETUP_TRACE_DIR}/scripts/restart_jaeger_memory.sh}" +JAEGER_CONFIG="${XTUNER_JAEGER_CONFIG:-${SETUP_TRACE_DIR}/jaeger/jaeger-memory.yaml}" TRACE_VIEWER_PORT="${XTUNER_TRACE_VIEWER_PORT:-18080}" -if pgrep -f "recipe.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" >/dev/null 2>&1 || - pgrep -f "recipe.trace_viewer.server.*--port=${TRACE_VIEWER_PORT}" >/dev/null 2>&1; then +if pgrep -f "recipe.trace.viewer.server.*--port ${TRACE_VIEWER_PORT}" >/dev/null 2>&1 || + pgrep -f "recipe.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 "recipe.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" 2>/dev/null || true - pkill -f "recipe.trace_viewer.server.*--port=${TRACE_VIEWER_PORT}" 2>/dev/null || true + pkill -f "recipe.trace.viewer.server.*--port ${TRACE_VIEWER_PORT}" 2>/dev/null || true + pkill -f "recipe.trace.viewer.server.*--port=${TRACE_VIEWER_PORT}" 2>/dev/null || true sleep 1 fi diff --git a/recipe/trace_viewer/payload.py b/recipe/trace/viewer/payload.py similarity index 100% rename from recipe/trace_viewer/payload.py rename to recipe/trace/viewer/payload.py diff --git a/recipe/trace_viewer/render.py b/recipe/trace/viewer/render.py similarity index 100% rename from recipe/trace_viewer/render.py rename to recipe/trace/viewer/render.py diff --git a/recipe/trace_viewer/server.py b/recipe/trace/viewer/server.py similarity index 99% rename from recipe/trace_viewer/server.py rename to recipe/trace/viewer/server.py index 2ad5f07923..d2dbab2315 100644 --- a/recipe/trace_viewer/server.py +++ b/recipe/trace/viewer/server.py @@ -20,12 +20,12 @@ from urllib.parse import parse_qs, urlsplit from urllib.request import Request, urlopen -from recipe.trace_viewer.payload import ( +from recipe.trace.viewer.payload import ( build_rollout_view_payload_from_jaeger_traces, filter_rollout_view_payload_by_train_step, load_jaeger_traces_from_otel_jsonl, ) -from recipe.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html +from recipe.trace.viewer.render import render_rollout_trace_html, write_rollout_trace_html _JAEGER_PROXY_PREFIX = "/jaeger" diff --git a/tests/rl/test_trace.py b/tests/rl/test_trace.py index c2782fa546..1844daf307 100644 --- a/tests/rl/test_trace.py +++ b/tests/rl/test_trace.py @@ -50,7 +50,7 @@ def test_nested_trace_span_preserves_parent_to_child_order(self): self.assertEqual(output["span_name_paths"]["order.child"], ["order.parent", "order.child"]) def test_viewer_uses_span_name_path_for_display_chain(self): - from recipe.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces + from recipe.trace.viewer.payload import build_rollout_view_payload_from_jaeger_traces traces = [ { @@ -95,8 +95,8 @@ def test_viewer_uses_span_name_path_for_display_chain(self): self.assertEqual(payload["samples"][0]["chain"], "parent.phase -> child.phase") def test_viewer_filters_latest_train_step_and_renders_payload(self): - from recipe.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces - from recipe.trace_viewer.render import render_rollout_trace_html + from recipe.trace.viewer.payload import build_rollout_view_payload_from_jaeger_traces + from recipe.trace.viewer.render import render_rollout_trace_html traces = [ { diff --git a/xtuner/v1/rl/trace/runtime.py b/xtuner/v1/rl/trace/runtime.py index 893bbee61a..11ae9d4268 100644 --- a/xtuner/v1/rl/trace/runtime.py +++ b/xtuner/v1/rl/trace/runtime.py @@ -42,7 +42,7 @@ DEFAULT_JAEGER_OTLP_GRPC_ENDPOINT = "127.0.0.1:14317" -_TRACE_VIEWER_SERVER_MODULE = "recipe.trace_viewer.server" +_TRACE_VIEWER_SERVER_MODULE = "recipe.trace.viewer.server" _READY_TIMEOUT_S = 3.0 _READY_POLL_INTERVAL_S = 0.05 _LOG_TAIL_CHARS = 4000