diff --git a/.claude/skills/xtuner-trace/SKILL.md b/.claude/skills/xtuner-trace/SKILL.md new file mode 100644 index 000000000..15da6cb4a --- /dev/null +++ b/.claude/skills/xtuner-trace/SKILL.md @@ -0,0 +1,116 @@ +--- +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` +- Rollout starter preset: `xtuner/v1/rl/trace/rollout_api.py` +- 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`. + +## 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 +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()` + +## 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. + +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: + +1. Ask for or locate the launch script and training config before editing. +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 + `enable_rollout_trace=True` only when the user wants the built-in rollout + starter trace. +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 + 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, XTuner viewer URL, + and equivalent manual viewer command when the XTuner viewer is enabled. + +## Guardrails + +- 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 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 new file mode 100644 index 000000000..6ecc5437f --- /dev/null +++ b/.claude/skills/xtuner-trace/references/trace-patterns.md @@ -0,0 +1,441 @@ +# 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)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then + source "${REPO_ROOT}/recipe/trace/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 xtuner.v1.rl.trace import TraceConfig + + +trace_config = TraceConfig( + enabled=os.environ.get("XTUNER_TRACE_ENABLED") == "1", + service_name="xtuner-agent-rollout", + 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, +) + +trainer = RLColocateTrainerConfig( + ..., + trace_config=trace_config, +) +``` + +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 +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 `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: + +- `examples/v1/scripts/run_rl_run.sh` +- `recipe/trace/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", + xtuner_viewer_enabled=True, + ) +) + +try: + ... +finally: + close_trace() +``` + +The runtime owns OTel collector setup, trace JSONL output, and optional lightweight XTuner viewer 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 + +`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 + +The built-in rollout starter trace is opt-in: + +```python +trace_config = TraceConfig( + enabled=True, + xtuner_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: + +- `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/config/rl_grpo_gsm8k_judge.py b/examples/v1/config/rl_grpo_gsm8k_judge.py index c7c3255f5..1f68946a8 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", + 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, +) + # 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/trace/README.md b/recipe/trace/README.md new file mode 100644 index 000000000..4f5eef463 --- /dev/null +++ b/recipe/trace/README.md @@ -0,0 +1,62 @@ +# 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 +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/trace/scripts/install_otel_tools.sh +export PATH=/tmp/xtuner_otel/bin:$PATH +``` + +Start Jaeger: + +```bash +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/trace/scripts/restart_jaeger_memory.sh +``` + +Run XTuner with trace enabled: + +```bash +export XTUNER_TRACE_ENABLED=1 +``` + +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`. +`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 \ + --trace-jsonl /traces/traces.jsonl \ + --jaeger-query-url http://127.0.0.1:16686 \ + --service xtuner-rollout +``` + +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/recipe/trace/jaeger/jaeger-memory.yaml b/recipe/trace/jaeger/jaeger-memory.yaml new file mode 100644 index 000000000..533ad6592 --- /dev/null +++ b/recipe/trace/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/trace/scripts/install_otel_tools.sh b/recipe/trace/scripts/install_otel_tools.sh new file mode 100755 index 000000000..b8f3c0c86 --- /dev/null +++ b/recipe/trace/scripts/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/trace/scripts/restart_jaeger_memory.sh b/recipe/trace/scripts/restart_jaeger_memory.sh new file mode 100755 index 000000000..03539ae56 --- /dev/null +++ b/recipe/trace/scripts/restart_jaeger_memory.sh @@ -0,0 +1,88 @@ +#!/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}" +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}" +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/trace/scripts/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/recipe/trace/setup_trace.sh b/recipe/trace/setup_trace.sh new file mode 100644 index 000000000..08dae26a7 --- /dev/null +++ b/recipe/trace/setup_trace.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash + +SETUP_TRACE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && 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:-${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 + 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 + 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/trace/viewer/payload.py b/recipe/trace/viewer/payload.py new file mode 100644 index 000000000..b80a42170 --- /dev/null +++ b/recipe/trace/viewer/payload.py @@ -0,0 +1,987 @@ +"""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 +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, + 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 = _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: + raise ValueError("Jaeger trace is missing traceID") + process_metadata = _process_metadata(trace_data) + span_entries = [] + entries_by_span_id: dict[str, dict[str, Any]] = {} + 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 + 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, + "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_status_from_span(span, tags), + "service_name": span_service_name, + "run_id": span_run_id, + "jaeger_url": f"{jaeger_trace_link_base_url}/trace/{trace_id}" + if jaeger_trace_link_base_url is not None + else None, + "spans": [], + }, + ) + 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)) + + 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"]) + 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) + + 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(): + raise FileNotFoundError(f"trace_jsonl path does not exist: {path}") + + 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, context=context) + continue + 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) + 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") + 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: + 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, + { + "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 _iter_jsonl_records(path: Path) -> Iterable[tuple[int, dict[str, Any]]]: + with path.open("r", encoding="utf-8") as 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 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], *, context: str) -> list[dict[str, Any]] | None: + data = record.get("data") + 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], *, context: str) -> None: + trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") + if not trace_id: + raise ValueError(f"{context}: Jaeger trace is missing traceID") + target = traces_by_id.setdefault(trace_id, {"traceID": trace_id, "processes": {}, "spans": []}) + 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( + 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_time = ( + otel_span.get("startTimeUnixNano") + or otel_span.get("start_time_unix_nano") + or otel_span.get("startTime") + or 0 + ) + 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): + 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: 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 []: + key = attribute.get("key") + 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: + return value["stringValue"] + if "boolValue" in value: + return bool(value["boolValue"]) + if "intValue" in value: + return int(str(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]]: + 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: + 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 _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} + 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) + if 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"), + } + ) + else: + nodes.append({"name": name, "source": "logical", "status": "inferred"}) + return nodes + + +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")] + 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(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: + value = json.loads(value) + except json.JSONDecodeError: + 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 _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 _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_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) + 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() + 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 + 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"]) + + 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, + "sample_count": len(visible_samples), + "group_count": len(group_ids), + "step_count": len(visible_steps), + "stage_duration_summaries": _build_stage_duration_summaries(visible_samples), + "status_counts": dict(sorted(status_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)) + 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 _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: + 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: + 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 in _ERROR_SAMPLE_STATUSES: + return status + + 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 _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]: + 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() + 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": span_id, + "parent_span_id": _parent_span_id(span), + "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, + "rollout_backend": tags.get("rollout.backend"), + "attributes": attributes, + } + + +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]] = {} + 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): + 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"]), + "run_id": tags.get("run.id"), + } + return metadata + + +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 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: + 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: + 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 + + +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 _append_unique(values: list[Any], value: Any) -> None: + if value is not None and value not in values: + values.append(value) + + +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 + + +__all__ = [ + "build_rollout_view_payload_from_jaeger_traces", + "filter_rollout_view_payload_by_train_step", + "load_jaeger_traces_from_otel_jsonl", +] diff --git a/recipe/trace/viewer/render.py b/recipe/trace/viewer/render.py new file mode 100644 index 000000000..acfa73078 --- /dev/null +++ b/recipe/trace/viewer/render.py @@ -0,0 +1,532 @@ +"""Render the XTuner trace viewer HTML from an already-built payload.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + + +def render_rollout_trace_html( + payload: dict[str, Any], + *, + 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("__AUTO_REFRESH__", json.dumps(auto_refresh)) + .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 Durations

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

Samples

+
+ + +
+
+ + + + + + + + + + + + + +
SampleStatusGroupStepRewardPathJaeger
+
+
+ +
+ + + +""" + + +__all__ = ["render_rollout_trace_html", "write_rollout_trace_html"] diff --git a/recipe/trace/viewer/server.py b/recipe/trace/viewer/server.py new file mode 100644 index 000000000..d2dbab231 --- /dev/null +++ b/recipe/trace/viewer/server.py @@ -0,0 +1,364 @@ +"""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 +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 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 + + +_JAEGER_PROXY_PREFIX = "/jaeger" +JAEGER_DEFAULT_QUERY_URL = "http://127.0.0.1:16686" +_PROXY_TIMEOUT_S = 10.0 +_HOP_BY_HOP_HEADERS = { + "connection", + "content-length", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +} + + +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, + 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_rollout_view_payload_from_trace_jsonl( + trace_jsonl_path: Path | str, + *, + jaeger_query_url: str | None = None, + jaeger_link_url: str | None = None, + service_name: str | None = None, + run_id: str | None = None, + train_step: str | int | None = "latest", +) -> dict[str, Any]: + 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, + 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, + host: str = "127.0.0.1", + port: int = 0, + refresh_interval_s: float = 2.0, + 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) + + def load_base_payload() -> dict[str, Any]: + 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: + return _source_signature(trace_jsonl_path) + + payload_cache = _TraceViewerPayloadCache( + load_base_payload, + source_signature=current_source_signature, + ) + payload_cache.get(train_step) + + 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)), + auto_refresh=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]: + 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") + + 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 _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: + 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(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + 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, 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("--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(argv) + + +def main() -> None: + args = _parse_args() + + if args.output 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, + service_name=args.service, + run_id=args.run_id, + train_step=args.train_step, + ) + write_rollout_trace_html(payload, args.output) + print(args.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, + host=args.host, + port=args.port, + train_step=args.train_step, + ) + print(f"XTuner Rollout Trace Viewer: {handle.url}", 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) + 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) + 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_rollout_view_payload_from_trace_jsonl", + "render_rollout_trace_html", + "start_rollout_trace_viewer", + "write_rollout_trace_html", +] diff --git a/tests/rl/test_trace.py b/tests/rl/test_trace.py new file mode 100644 index 000000000..1844daf30 --- /dev/null +++ b/tests/rl/test_trace.py @@ -0,0 +1,155 @@ +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"]) + + def test_viewer_uses_span_name_path_for_display_chain(self): + from recipe.trace.viewer.payload import build_rollout_view_payload_from_jaeger_traces + + 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"]}, + ], + }, + { + "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"]], + ["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 recipe.trace.viewer.payload import build_rollout_view_payload_from_jaeger_traces + from recipe.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 000000000..b92fc20a0 --- /dev/null +++ b/tests/rl/trace_utils.py @@ -0,0 +1,169 @@ +import json +import os +import subprocess +import sys +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 ( + 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"] + 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() + }, + } + ) + + +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/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index fff3658bb..6a69d0f80 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 86e1539da..2368d1c3b 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/agent_loop_manager/produce_utils.py b/xtuner/v1/rl/agent_loop_manager/produce_utils.py index 33f1e115b..b80a2f9e9 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -144,6 +144,11 @@ 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 + 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/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index ae4572334..028abb1ac 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 c50ffa6bf..3a79c9309 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/__init__.py b/xtuner/v1/rl/trace/__init__.py new file mode 100644 index 000000000..41b356263 --- /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 000000000..c0d17c809 --- /dev/null +++ b/xtuner/v1/rl/trace/api.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +import contextvars +import inspect +import json +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from functools import wraps +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" +# 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, 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) + try: + with otel_utils.start_span(span_name, attributes=normalized_attributes): + try: + yield + except Exception as exc: + otel_utils.record_failure(exc) + raise + 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__}>" + + +__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 000000000..2677cc909 --- /dev/null +++ b/xtuner/v1/rl/trace/otel_utils.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Mapping, MutableMapping + + +if TYPE_CHECKING: + from opentelemetry.sdk.trace import TracerProvider + + +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 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. " + "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.""" + + from opentelemetry import propagate + + 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.""" + + 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, + record_exception=False, + set_status_on_exception=False, + ) + + +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: + 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: + from opentelemetry import trace + + 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: + from opentelemetry import trace + + 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: + from opentelemetry import trace + from opentelemetry.trace import Status, StatusCode + + 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: + from opentelemetry import trace + from opentelemetry.trace import Status, StatusCode + + 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/rollout_api.py b/xtuner/v1/rl/trace/rollout_api.py new file mode 100644 index 000000000..e20a4cccd --- /dev/null +++ b/xtuner/v1/rl/trace/rollout_api.py @@ -0,0 +1,235 @@ +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 . import otel_utils +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_CALL_CHAIN_EXTRA_FIELD = "_xtuner_trace_call_chain" + + +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 _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] + + 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]: + 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 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) + 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} + + +@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( + 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_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 new file mode 100644 index 000000000..11ae9d426 --- /dev/null +++ b/xtuner/v1/rl/trace/runtime.py @@ -0,0 +1,583 @@ +from __future__ import annotations + +import atexit +import contextlib +import os +import shlex +import shutil +import socket +import subprocess +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 + +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_TRACE_ENABLE_ROLLOUT", + "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 = "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: + 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" + 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") + @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 + 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) + + +def _build_xtuner_viewer_command( + *, + trace_jsonl_path: Path, + jaeger_query_url: str | None, + service_name: str, + run_id: str, + host: str, + port: int, +) -> list[str]: + command = [ + sys.executable, + "-m", + _TRACE_VIEWER_SERVER_MODULE, + "--trace-jsonl", + os.fspath(trace_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]) + 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 + 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) + 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.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.xtuner_viewer_jaeger_query_url, + service_name=self.runtime.service_name, + run_id=self.runtime.run_id, + 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.xtuner_viewer_url, + trace_viewer_port=self.xtuner_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.xtuner_viewer_process is not None: + logger.info( + 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: + 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.") + 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 + 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(), + 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" + 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] + 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_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, + "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, + service_name=config.service_name, + trace_viewer_url=None, + trace_viewer_port=None, + ), + endpoint=endpoint, + env_vars=env_vars, + collector_port=port, + 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, + ) + + +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() + 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, + 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.""" + + if _RUNTIME is None: + ensure_trace_runtime_from_env() + 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 diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index eecb796dd..d3a989a98 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,