diff --git a/.claude/skills/xtuner-trace/SKILL.md b/.claude/skills/xtuner-trace/SKILL.md new file mode 100644 index 0000000000..a1802685b4 --- /dev/null +++ b/.claude/skills/xtuner-trace/SKILL.md @@ -0,0 +1,62 @@ +--- +name: xtuner-trace +description: Use when adding, modifying, or reviewing XTuner RL trace instrumentation for rollout, worker, session server, localhost/sandbox agent loop, judger, infer engine, Ray remote calls, HTTP proxy calls, or trace viewer fields. Helps choose the correct XTuner trace API, preserve Ray/HTTP context propagation, and add focused verification for a requested observation phase. +--- + +# XTuner Trace + +Use this skill to instrument or review a specific XTuner RL phase with trace spans, events, attributes, and context propagation. + +## Core Workflow + +1. Identify the requested observation phase and its boundary: + - Local code block or function only. + - Stable XTuner endpoint around `RolloutState`, agent item, or judger. + - Ray `.remote(...)` boundary. + - HTTP request/proxy boundary. + - Viewer field or payload change only. + +2. Read the current code before editing: + - Public facade: `xtuner/v1/rl/trace/__init__.py` + - Basic and endpoint API: `xtuner/v1/rl/trace/api.py` + - Propagation helpers: `xtuner/v1/rl/trace/context_propagation.py` + - Span names and attribute builders: `xtuner/v1/rl/trace/trace_utils.py` + - Existing examples near the requested phase. + +3. Choose the narrowest API that matches the boundary: + - Use `trace_span(...)` or `trace_function(...)` for local work. + - Use `traced_rollout_endpoint(...)` for async endpoints whose lifecycle is a `RolloutState`. + - Use `traced_agent_item_endpoint(...)` for agent item runner entrypoints. + - Use `traced_judger_endpoint(...)` for agent-loop judger entrypoints with `owner.judger`. + - Use `trace_remote(...)` for Ray remote submission and keep the receiver as a traced rollout endpoint or a span with `parent_carrier`. + - Use `traced_aiohttp_request(...)` or `instrument_aiohttp_client()` for HTTP propagation. + +4. Keep semantics centralized: + - Add new stable span names as `TRACE_SPAN_*` in `trace_utils.py`. + - Add reusable field mapping to `trace_utils.py` instead of scattering key construction. + - Use `failure_attributes(...)`, `reward_trace_attributes(...)`, `rollout_state_*_attributes(...)`, `agent_item_*_attributes(...)`, and `agent_judger_initial_attributes(...)` when they fit. + +5. Verify propagation, not just local span creation: + - Ray: the outbound call must use `trace_remote(...)`; the inbound side must restore the carrier through `traced_rollout_endpoint(...)` or `trace_span(..., parent_carrier=...)`. + - HTTP: inbound request context must be extracted from headers or body carrier; outbound request context must be injected into headers or request body as appropriate. + - Do not call OpenTelemetry SDK directly from rollout, agent, judger, or session-server business code. + +For concrete patterns and snippets, read [references/trace-patterns.md](references/trace-patterns.md) before making code changes. + +## Guardrails + +- Do not invent public APIs that are not exported from `xtuner.v1.rl.trace`. +- Do not use `set_trace_attribute(...)`, `extract_trace_context(...)`, or `attach_trace_context(...)`; the current public facade does not expose them. +- Do not make a thin endpoint helper for a phase without a stable target and final-result lifecycle. Use explicit `trace_span(...)` instead. +- Do not manually keep `_xtuner_trace_carrier` on `RolloutState.extra_fields`; `trace_remote(...)` attaches it temporarily and restores local state. +- Do not pass multiple `RolloutState` objects, a collection of rollout states, or no rollout state to `trace_remote(...)`; pass exactly one or use `target=...`. +- Do not record prompts, responses, full configs, secrets, or large payloads as attributes. Prefer IDs, statuses, counts, timings, categories, and stable viewer fields. + +## Verification + +Prefer focused verification: + +- Run `git diff --check` after edits. +- Run `python -m compileall -q` on touched trace and caller files. +- Add or update behavior-oriented trace tests when changing propagation, endpoint decorators, viewer payload fields, or final-state recording. +- For viewer changes, verify both JSONL and Jaeger-source assumptions if the change touches source or payload selection. diff --git a/.claude/skills/xtuner-trace/references/trace-patterns.md b/.claude/skills/xtuner-trace/references/trace-patterns.md new file mode 100644 index 0000000000..371d5fc89f --- /dev/null +++ b/.claude/skills/xtuner-trace/references/trace-patterns.md @@ -0,0 +1,175 @@ +# XTuner Trace Patterns + +Use this reference after the `xtuner-trace` skill triggers and before editing trace instrumentation. + +## Public API Surface + +Import these from `xtuner.v1.rl.trace` unless noted: + +| API | Use | +| --- | --- | +| `trace_span(name, attributes=None, parent_carrier=None)` | Create a span around a local phase. Pass `parent_carrier` only when restoring an inbound context. | +| `trace_function(name=None, attributes=None)` | Wrap an entire sync or async function with one span. | +| `trace_event(name, attributes=None)` | Add an event to the current span. | +| `set_trace_attributes(attributes)` | Set one or more attributes on the current span. For one key, pass `{key: value}`. | +| `inject_trace_context(carrier=None)` | Inject current W3C context into a dict carrier. | +| `traced_rollout_endpoint(span_name, target_arg="rollout_state", initial_attributes=None)` | Wrap an async endpoint whose target and final result are a `RolloutState`. | +| `traced_agent_item_endpoint(span_name, item_arg="item", initial_attributes=None)` | Wrap an async agent item runner endpoint. | +| `traced_judger_endpoint(span_name="judger.run", target_arg="rollout_state")` | Wrap an async judger endpoint; owner must expose `self.judger`. | +| `trace_remote(remote_method, *args, target=None, **kwargs)` | Call Ray `.remote(...)` while temporarily attaching the current trace carrier to exactly one `RolloutState`. | + +Import these from `xtuner.v1.rl.trace.context_propagation` only when working on framework or transport code: + +| API | Use | +| --- | --- | +| `traced_aiohttp_request(...)` | Create a send span, inject headers, perform an aiohttp request, and record HTTP status/send latency. | +| `instrument_aiohttp_client()` | Patch generic aiohttp clients, such as lagent's OpenAI-compatible client, so they propagate context. | +| `extract_trace_carrier_from_mapping(...)` | Restore inbound body carrier at HTTP proxy boundaries. | +| `extract_trace_attributes_from_mapping(...)` | Recover upstream semantic attributes passed through request body. | +| `remove_trace_carrier_from_mapping(...)` | Remove internal trace carrier fields before forwarding user/backend payloads. | + +## Decision Matrix + +| Requested observation | Preferred pattern | +| --- | --- | +| Add timing around a small local block | `with trace_span(TRACE_SPAN_..., attributes=...)` | +| Trace a whole helper function with no custom final result | `@trace_function(...)` | +| Trace agent loop, rollout controller, or rollout worker entrypoint | `@traced_rollout_endpoint(TRACE_SPAN_...)` | +| Submit a traced Ray call | `await trace_remote(actor.method, rollout_state=state)` or store the returned object ref | +| Trace localhost/sandbox agent item runner entrypoint | `@traced_agent_item_endpoint(TRACE_SPAN_...)` | +| Trace agent infer/validate/invoke substage | Explicit `trace_span(...)` with `agent_item_initial_attributes(item)` | +| Trace judger information and final reward | `@traced_judger_endpoint(...)` for agent-loop judger entrypoints; explicit `trace_span(...)` inside agent runner validation stages | +| Trace HTTP proxy request phases | One parent request `trace_span(..., parent_carrier=...)`, then child spans for prepare/send/read/record | +| Propagate context through outbound HTTP | `traced_aiohttp_request(...)`; for third-party aiohttp clients call `instrument_aiohttp_client()` before the client issues requests | +| Add viewer-visible fields | Stable attribute builders or result recorders in `trace_utils.py`, then payload/render tests if viewer grouping changes | + +## Current Code Examples + +- Ray rollout chain: + - `xtuner/v1/rl/agent_loop/single_turn_agent_loop.py` + - `xtuner/v1/rl/rollout/controller.py` + - `xtuner/v1/rl/rollout/worker.py` + - Pattern: caller uses `trace_remote(...)`; receiver uses `@traced_rollout_endpoint(...)`. + +- Session server HTTP proxy: + - `xtuner/v1/rl/rollout/session_server.py` + - Pattern: `_handle_request()` restores parent context from headers/body, `_prepare_request()` records request metadata, `_forward_request()` uses `traced_aiohttp_request(...)`, `_read_upstream_response()` records stream/non-stream latency, `_record_response()` records tokenization/session trace result. + +- Localhost agent loop: + - `xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py` + - `xtuner/v1/rl/agent_loop/localhost_agent_loop/runner.py` + - `xtuner/v1/rl/agent_loop/localhost_agent_loop/stage.py` + - `xtuner/v1/rl/agent_loop/localhost_agent_loop/judger.py` + - Pattern: rollout entrypoint uses `@traced_rollout_endpoint(...)`; item runner uses `@traced_agent_item_endpoint(...)`; infer, agent invoke, validate, judger, and materialization use explicit `trace_span(...)`. + +## Templates + +### Local Phase + +```python +from xtuner.v1.rl.trace import set_trace_attributes, trace_span +from xtuner.v1.rl.trace.trace_utils import TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN, agent_item_initial_attributes + +attributes = agent_item_initial_attributes(item) +attributes["stage.name"] = "infer" +with trace_span(TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN, attributes=attributes): + result = await self.infer.run(item, item.infer) + set_trace_attributes({"agent.message_count": message_count}) +``` + +### Failure Attributes + +```python +from xtuner.v1.rl.trace import set_trace_attributes +from xtuner.v1.rl.trace.trace_utils import failure_attributes + +set_trace_attributes( + failure_attributes( + error.category, + message=error.message, + type=error.type, + ) +) +``` + +### Ray Boundary + +```python +from xtuner.v1.rl.trace import trace_remote, traced_rollout_endpoint +from xtuner.v1.rl.trace.trace_utils import TRACE_SPAN_ROLLOUT_CONTROLLER_GENERATE + +@traced_rollout_endpoint(TRACE_SPAN_ROLLOUT_CONTROLLER_GENERATE) +async def generate(self, rollout_state: RolloutState) -> RolloutState: + response_ref = trace_remote(worker.generate, rollout_state=rollout_state) + return await response_ref +``` + +Rules: + +- `trace_remote(...)` only propagates context; it does not create the receiver's business span. +- The receiver must create the span, usually through `@traced_rollout_endpoint(...)`. +- If the `RolloutState` is not directly in args/kwargs, pass `target=rollout_state`. + +### HTTP Proxy Boundary + +```python +parent_carrier, source = _request_trace_context(headers, request_data) +trace_attributes = extract_trace_attributes_from_mapping(request_data) + +with trace_span("session_server.request", attributes=trace_attributes, parent_carrier=parent_carrier): + prepared = await self._prepare_request(...) + forwarded = await self._forward_request(request, prepared) +``` + +Outbound aiohttp: + +```python +async with traced_aiohttp_request( + client, + span_name="session_server.send_request", + method=request.method, + url=target_url, + headers=forward_headers, + data=prepared.body, + attributes=trace_attributes, +) as resp: + return await self._read_upstream_response(request, resp, prepared) +``` + +Rules: + +- Prefer `traceparent` headers when present; otherwise use body carrier. +- Remove internal carrier fields from request JSON before forwarding to model backends unless the downstream service is expected to consume them. +- Keep request, prepare, send, read, and record as separate spans in proxy code; this makes latency attribution useful. + +## Attribute Guidance + +Prefer stable scalar attributes: + +- Identity: `xtuner.rollout_id`, `xtuner.group_id`, `xtuner.session_id`, `xtuner.task_name` +- Status: `xtuner.status`, `agent.status`, `session.stream` +- Timing/counts: `session.first_token_ms`, `session.response_total_ms`, `completion.tokens`, `agent.message_count` +- Reward: `reward.score`, `reward.pass` +- Error: `error.category`, `error.message`, `error.type` + +Avoid: + +- Full prompts, full responses, raw chat histories, stack traces, large tool payloads. +- Secrets, bearer tokens, API keys, raw HTTP headers. +- Per-call random field names that the viewer cannot aggregate. + +## Verification Targets + +Use the smallest meaningful checks: + +```bash +python -m compileall -q xtuner/v1/rl/trace xtuner/v1/rl/rollout/session_server.py +git diff --check +``` + +For propagation changes, add focused tests that assert: + +- Child spans share the parent `trace_id`. +- `trace_remote(...)` cleans the temporary carrier from local `RolloutState.extra_fields`. +- HTTP body carriers are removed before forwarding when required. +- Final status/reward/error attributes are written after the operation completes. diff --git a/examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py b/examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py index da270fd527..0608d51993 100644 --- a/examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py +++ b/examples/v1/config/agentic_rl_qwen3p5vl_mtp_ep_code.py @@ -15,6 +15,7 @@ from xtuner.v1.rl.replay_buffer import SyncReplayBufferConfig from xtuner.v1.rl.rollout.worker import RolloutConfig from xtuner.v1.rl.trainer import RolloutImportanceSampling, WorkerConfig +from xtuner.v1.rl.trace import TraceConfig from xtuner.v1.train.trainer import LoadCheckpointConfig from xtuner.v1.rl.utils import AcceleratorResourcesConfig from xtuner.v1.train.rl_trainer import RLColocateTrainerConfig @@ -298,6 +299,16 @@ def _build_dataset_cfg(specs, default_recipe): ), ) +trace_config = TraceConfig( + enabled=os.environ.get("XTUNER_TRACE_ENABLED") == "1", + output_dir=Path(work_dir) / "otel", + service_name="xtuner-agent-rollout", + viewer_enabled=True, + viewer_host="0.0.0.0", + viewer_port=18080, + viewer_jaeger_query_url="http://127.0.0.1:16686", +) + trainer = RLColocateTrainerConfig( resources=resources, train_worker_cfg=train_worker_cfg, @@ -326,4 +337,5 @@ def _build_dataset_cfg(specs, default_recipe): debug_rollout_dir=debug_rollout_dir, debug_train=debug_train, debug_rollout=debug_rollout, + trace_config=trace_config, ) diff --git a/examples/v1/config/rl_grpo_gsm8k_judge.py b/examples/v1/config/rl_grpo_gsm8k_judge.py index c7c3255f53..153744f99d 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 @@ -183,7 +184,18 @@ # 7. evaluator evaluator_config = EvaluatorConfig(compute_metric_func=None) -# 8. RL Colocate Trainer Config(CLI 通过 config["trainer"].build() 得到 Trainer) +# 8. Trace Config +trace_config = TraceConfig( + enabled=os.environ.get("XTUNER_TRACE_ENABLED") == "1", + output_dir=Path(work_dir) / "otel", + service_name="xtuner-rollout", + viewer_enabled=True, + viewer_host="0.0.0.0", + viewer_port=18080, + viewer_jaeger_query_url="http://127.0.0.1:16686", +) + +# 9. RL Colocate Trainer Config(CLI 通过 config["trainer"].build() 得到 Trainer) trainer = RLColocateTrainerConfig( resources=resources, train_worker_cfg=train_worker_cfg, # TODO: uniform naming of cfg and config @@ -203,4 +215,5 @@ work_dir=work_dir, seed=123, debug_rollout=False, + trace_config=trace_config, ) diff --git a/examples/v1/scripts/run_rl.sh b/examples/v1/scripts/run_rl.sh index a81de88e3d..66cca6d505 100644 --- a/examples/v1/scripts/run_rl.sh +++ b/examples/v1/scripts/run_rl.sh @@ -1,4 +1,8 @@ set -ex +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then + source "${SCRIPT_DIR}/setup_trace.sh" +fi ray stop --force # examples of usage: # qwen3_8B_grpo_gsm8k training: diff --git a/examples/v1/scripts/run_rl_submit.sh b/examples/v1/scripts/run_rl_submit.sh index 636bfa116f..0beea62129 100644 --- a/examples/v1/scripts/run_rl_submit.sh +++ b/examples/v1/scripts/run_rl_submit.sh @@ -1,4 +1,8 @@ set -ex +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +if [ "${XTUNER_TRACE_ENABLED:-0}" = "1" ]; then + source "${SCRIPT_DIR}/setup_trace.sh" +fi ray stop --force # examples of usage: # qwen3_8B_grpo_gsm8k training: bash examples/v1/scripts/run_rl.sh examples/v1/config/rl_qwen3_8B_grpo.py "sglang" $MODEL_PATH $DATA_PATH $EVAL_DATA_PATH diff --git a/examples/v1/scripts/setup_trace.sh b/examples/v1/scripts/setup_trace.sh new file mode 100644 index 0000000000..560f01eaa8 --- /dev/null +++ b/examples/v1/scripts/setup_trace.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash + +SETUP_TRACE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SETUP_TRACE_DIR}/../../.." && pwd)" + +export XTUNER_OTEL_ROOT="${XTUNER_OTEL_ROOT:-/tmp/xtuner_otel}" +export PATH="${XTUNER_OTEL_ROOT}/bin:${PATH}" + +OTEL_INSTALL_SCRIPT="${XTUNER_OTEL_INSTALL_SCRIPT:-${REPO_ROOT}/recipe/otle/install_otel_tools.sh}" +JAEGER_RESTART_SCRIPT="${XTUNER_JAEGER_RESTART_SCRIPT:-${REPO_ROOT}/recipe/otle/restart_jaeger_memory.sh}" +JAEGER_CONFIG="${XTUNER_JAEGER_CONFIG:-${REPO_ROOT}/recipe/otle/jaeger/jaeger-memory.yaml}" +TRACE_VIEWER_PORT="${XTUNER_TRACE_VIEWER_PORT:-18080}" + +if pgrep -f "xtuner.tools.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" >/dev/null 2>&1 || + pgrep -f "xtuner.tools.trace_viewer.server.*--port=${TRACE_VIEWER_PORT}" >/dev/null 2>&1; then + echo "Stopping previous XTuner trace viewer on port ${TRACE_VIEWER_PORT}" + pkill -f "xtuner.tools.trace_viewer.server.*--port ${TRACE_VIEWER_PORT}" 2>/dev/null || true + pkill -f "xtuner.tools.trace_viewer.server.*--port=${TRACE_VIEWER_PORT}" 2>/dev/null || true + sleep 1 +fi + +if [ ! -d "$XTUNER_OTEL_ROOT" ]; then + echo "Installing XTuner OTel tools to ${XTUNER_OTEL_ROOT}" + bash "$OTEL_INSTALL_SCRIPT" "$XTUNER_OTEL_ROOT" || return 1 2>/dev/null || exit 1 +fi + +otel_collector="$(command -v otelcol-contrib || command -v otelcol || true)" +if [ -z "$otel_collector" ]; then + echo "Error: XTuner trace collector not found after checking ${XTUNER_OTEL_ROOT}." >&2 + echo "Expected otelcol-contrib or otelcol under ${XTUNER_OTEL_ROOT}/bin." >&2 + return 1 2>/dev/null || exit 1 +fi + +echo "XTuner trace collector: ${otel_collector}" +bash "$JAEGER_RESTART_SCRIPT" "$JAEGER_CONFIG" || return 1 2>/dev/null || exit 1 diff --git a/recipe/otle/README.md b/recipe/otle/README.md new file mode 100644 index 0000000000..b3d432e87e --- /dev/null +++ b/recipe/otle/README.md @@ -0,0 +1,54 @@ +# XTuner OTel Trace + +XTuner exports rollout traces through OpenTelemetry. For local inspection, start +Jaeger with `jaeger/jaeger-memory.yaml`, enable trace in the training config, and +open the XTuner rollout viewer. + +The reference Jaeger config exposes: + +- Jaeger UI and Query API: `http://127.0.0.1:16686` +- OTLP gRPC receiver: `http://127.0.0.1:14317` +- OTLP HTTP receiver: `http://127.0.0.1:14318/v1/traces` + +Install local binaries: + +```bash +bash recipe/otle/install_otel_tools.sh +export PATH=/tmp/xtuner_otel/bin:$PATH +``` + +Start Jaeger: + +```bash +jaeger --config recipe/otle/jaeger/jaeger-memory.yaml +``` + +For local smoke tests, restart the in-memory Jaeger before each experiment so +old services, operations, and trace ids cannot be mixed with the new run: + +```bash +bash recipe/otle/restart_jaeger_memory.sh +``` + +Run XTuner with trace enabled: + +```bash +export XTUNER_TRACE_ENABLED=1 +export XTUNER_TRACE_SERVICE_NAME=xtuner-rollout +``` + +By default, XTuner starts a local collector that writes +`/traces/traces.jsonl` and forwards spans to the reference Jaeger +OTLP gRPC endpoint `http://127.0.0.1:14317`. + +Open the rollout viewer: + +```bash +python -m xtuner.tools.trace_viewer.server \ + --jaeger-query-url http://127.0.0.1:16686 \ + --service xtuner-rollout +``` + +The viewer is a thin Jaeger Query API adapter. Jaeger remains the trace backend; +XTuner only groups spans by rollout metadata such as `xtuner.rollout_id`, +`xtuner.group_id`, `xtuner.task_name`, and `xtuner.status`. diff --git a/recipe/otle/install_otel_tools.sh b/recipe/otle/install_otel_tools.sh new file mode 100755 index 0000000000..b8f3c0c86e --- /dev/null +++ b/recipe/otle/install_otel_tools.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="${1:-/tmp/xtuner_otel}" +BIN_DIR="${ROOT}/bin" +DOWNLOAD_DIR="${ROOT}/downloads" +OTEL_VERSION="${OTEL_VERSION:-0.128.0}" +JAEGER_VERSION="${JAEGER_VERSION:-2.19.0}" + +mkdir -p "${BIN_DIR}" "${DOWNLOAD_DIR}" + +download_and_extract() { + local url="$1" + local archive="$2" + + if [ ! -f "${DOWNLOAD_DIR}/${archive}" ]; then + curl -L --fail --retry 3 --output "${DOWNLOAD_DIR}/${archive}" "${url}" + fi + tar -xzf "${DOWNLOAD_DIR}/${archive}" -C "${BIN_DIR}" +} + +download_and_extract \ + "https://github.com/open-telemetry/opentelemetry-collector-releases/releases/download/v${OTEL_VERSION}/otelcol_${OTEL_VERSION}_linux_amd64.tar.gz" \ + "otelcol_${OTEL_VERSION}_linux_amd64.tar.gz" + +download_and_extract \ + "https://github.com/open-telemetry/opentelemetry-collector-releases/releases/download/v${OTEL_VERSION}/otelcol-contrib_${OTEL_VERSION}_linux_amd64.tar.gz" \ + "otelcol-contrib_${OTEL_VERSION}_linux_amd64.tar.gz" + +download_and_extract \ + "https://github.com/jaegertracing/jaeger/releases/download/v${JAEGER_VERSION}/jaeger-${JAEGER_VERSION}-linux-amd64.tar.gz" \ + "jaeger-${JAEGER_VERSION}-linux-amd64.tar.gz" + +if [ ! -f "${BIN_DIR}/jaeger" ] && [ -f "${BIN_DIR}/jaeger-${JAEGER_VERSION}-linux-amd64/jaeger" ]; then + ln -sfn "${BIN_DIR}/jaeger-${JAEGER_VERSION}-linux-amd64/jaeger" "${BIN_DIR}/jaeger" +fi + +chmod +x "${BIN_DIR}/otelcol" "${BIN_DIR}/otelcol-contrib" "${BIN_DIR}/jaeger" + +echo "Installed:" +"${BIN_DIR}/otelcol" --version +"${BIN_DIR}/otelcol-contrib" --version +"${BIN_DIR}/jaeger" version +echo "Add to PATH: export PATH=${BIN_DIR}:\$PATH" diff --git a/recipe/otle/jaeger/jaeger-memory.yaml b/recipe/otle/jaeger/jaeger-memory.yaml new file mode 100644 index 0000000000..533ad65925 --- /dev/null +++ b/recipe/otle/jaeger/jaeger-memory.yaml @@ -0,0 +1,40 @@ +receivers: + otlp: + protocols: + grpc: + endpoint: 0.0.0.0:14317 + http: + endpoint: 0.0.0.0:14318 + +processors: + batch: + +exporters: + jaeger_storage_exporter: + trace_storage: memstore + +extensions: + jaeger_storage: + backends: + memstore: + memory: + max_traces: 100000 + jaeger_query: + storage: + traces: memstore + base_path: / + http: + endpoint: 0.0.0.0:16686 + grpc: + endpoint: 0.0.0.0:16685 + +service: + telemetry: + metrics: + level: none + extensions: [jaeger_storage, jaeger_query] + pipelines: + traces: + receivers: [otlp] + processors: [batch] + exporters: [jaeger_storage_exporter] diff --git a/recipe/otle/restart_jaeger_memory.sh b/recipe/otle/restart_jaeger_memory.sh new file mode 100755 index 0000000000..b35eb8930f --- /dev/null +++ b/recipe/otle/restart_jaeger_memory.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Explicit local reset tool: restarting this process clears Jaeger in-memory traces. +# Do not use this against a shared Jaeger deployment. + +ROOT="${XTUNER_OTEL_ROOT:-/tmp/xtuner_otel}" +JAEGER_BIN="${JAEGER_BIN:-${ROOT}/bin/jaeger}" +CONFIG="${1:-recipe/otle/jaeger/jaeger-memory.yaml}" +PID_FILE="${XTUNER_JAEGER_PID_FILE:-/tmp/xtuner_jaeger_memory.pid}" +LOG_FILE="${XTUNER_JAEGER_LOG_FILE:-/tmp/xtuner_jaeger_memory.log}" +QUERY_URL="${XTUNER_JAEGER_QUERY_URL:-http://127.0.0.1:16686}" +WAIT_TIMEOUT_S="${XTUNER_JAEGER_WAIT_TIMEOUT_S:-30}" + +if ! command -v "${JAEGER_BIN}" >/dev/null 2>&1; then + echo "Jaeger binary not found: ${JAEGER_BIN}" >&2 + echo "Install it first: bash recipe/otle/install_otel_tools.sh" >&2 + exit 1 +fi + +if [ ! -f "${CONFIG}" ]; then + echo "Jaeger config not found: ${CONFIG}" >&2 + exit 1 +fi + +if ! command -v setsid >/dev/null 2>&1; then + echo "setsid is required to detach Jaeger from the launcher process group." >&2 + exit 1 +fi + +stop_pid() { + local pid="$1" + if ! kill -0 "${pid}" >/dev/null 2>&1; then + return + fi + kill "${pid}" >/dev/null 2>&1 || true + for _ in $(seq 1 50); do + if ! kill -0 "${pid}" >/dev/null 2>&1; then + return + fi + sleep 0.1 + done + kill -9 "${pid}" >/dev/null 2>&1 || true +} + +if [ -f "${PID_FILE}" ]; then + old_pid="$(cat "${PID_FILE}")" + if [ -n "${old_pid}" ]; then + stop_pid "${old_pid}" + fi + rm -f "${PID_FILE}" +fi + +if command -v pgrep >/dev/null 2>&1; then + while IFS= read -r pid; do + [ -n "${pid}" ] || continue + if [ "${pid}" = "$$" ]; then + continue + fi + stop_pid "${pid}" + done < <(pgrep -f "jaeger.*jaeger-memory.yaml" || true) +fi + +mkdir -p "$(dirname "${PID_FILE}")" "$(dirname "${LOG_FILE}")" +setsid "${JAEGER_BIN}" --config "${CONFIG}" >"${LOG_FILE}" 2>&1 "${PID_FILE}" + +deadline=$((SECONDS + WAIT_TIMEOUT_S)) +until curl -fsS "${QUERY_URL}/api/services" >/dev/null 2>&1; do + if ! kill -0 "${new_pid}" >/dev/null 2>&1; then + echo "Jaeger exited before becoming ready. Log: ${LOG_FILE}" >&2 + exit 1 + fi + if [ "${SECONDS}" -ge "${deadline}" ]; then + echo "Timed out waiting for Jaeger Query API at ${QUERY_URL}. Log: ${LOG_FILE}" >&2 + exit 1 + fi + sleep 0.5 +done + +echo "Jaeger in-memory storage restarted." +echo "PID file: ${PID_FILE}" +echo "Log file: ${LOG_FILE}" +echo "Jaeger UI: ${QUERY_URL}" +echo "OTLP gRPC: http://127.0.0.1:14317" +echo "OTLP HTTP: http://127.0.0.1:14318/v1/traces" diff --git a/xtuner/tools/trace_viewer/__init__.py b/xtuner/tools/trace_viewer/__init__.py new file mode 100644 index 0000000000..a2ef41bd45 --- /dev/null +++ b/xtuner/tools/trace_viewer/__init__.py @@ -0,0 +1,18 @@ +from xtuner.tools.trace_viewer.payload import build_rollout_view_payload_from_jaeger_traces +from xtuner.tools.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html +from xtuner.tools.trace_viewer.source import ( + JaegerQuerySource, + JsonlTraceSource, + fetch_jaeger_traces, + normalize_jaeger_query_url, +) + +__all__ = [ + "JaegerQuerySource", + "JsonlTraceSource", + "build_rollout_view_payload_from_jaeger_traces", + "fetch_jaeger_traces", + "normalize_jaeger_query_url", + "render_rollout_trace_html", + "write_rollout_trace_html", +] diff --git a/xtuner/tools/trace_viewer/payload.py b/xtuner/tools/trace_viewer/payload.py new file mode 100644 index 0000000000..22c83d5c13 --- /dev/null +++ b/xtuner/tools/trace_viewer/payload.py @@ -0,0 +1,1449 @@ +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 + + +_LOGICAL_PATH_ATTRIBUTE = "xtuner.logical_path" + +_INITIAL_SAMPLE_STATUSES = {"init"} +_NON_TERMINAL_SAMPLE_STATUSES = {"", "init", "pending", "queued", "running", "scheduled", "started", "unknown"} +_ERROR_SAMPLE_STATUSES = {"aborted", "error", "exception", "failed", "timeout", "timed_out"} +_TERMINAL_STAGE_STATUSES = { + "completed", + "failed", + "aborted", + "timeout", + "timed_out", + "expired", + "stale", + "filtered", +} + + +def build_rollout_view_payload_from_jaeger_traces( + traces: Iterable[dict[str, Any]], + *, + jaeger_query_url: str | None = None, + jaeger_link_url: str | None = None, + live_records: Iterable[dict[str, Any]] | None = None, + service_name: str | None = None, + run_id: str | None = None, + expected_global_batch_size: int | None = None, + expected_prompt_repeat_k: int | None = None, + stale_threshold: int | None = None, +) -> dict[str, Any]: + samples_by_key: dict[tuple[str, Any], dict[str, Any]] = {} + status_counts: Counter[str] = Counter() + stage_counts: Counter[str] = Counter() + group_ids: set[Any] = set() + missing_parents: list[dict[str, Any]] = [] + trace_context_sources: Counter[str] = Counter() + jaeger_trace_link_base_url = _jaeger_trace_link_base_url(jaeger_query_url, jaeger_link_url) + + for trace_data in traces: + trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") + if not trace_id: + continue + process_metadata = _process_metadata(trace_data) + span_entries = [] + entries_by_span_id: dict[str, dict[str, Any]] = {} + for span in trace_data.get("spans") or []: + process = process_metadata.get(str(span.get("processID") or ""), {}) + span_service_name = process.get("service_name") + if service_name is not None and span_service_name != service_name: + continue + tags = _tags_to_dict(span.get("tags") or []) + span_run_id = tags.get("run.id") or process.get("run_id") + if run_id is not None and span_run_id != run_id: + continue + context_source = tags.get("session.trace_context_source") + if context_source is not None: + trace_context_sources[str(context_source)] += 1 + entry = { + "span": span, + "tags": tags, + "service_name": span_service_name, + "run_id": span_run_id, + "span_id": _span_id(span), + "trace_id": trace_id, + } + span_entries.append(entry) + if entry["span_id"]: + entries_by_span_id[entry["span_id"]] = entry + + for entry in span_entries: + parent_span_id = _parent_span_id(entry["span"]) + if parent_span_id and parent_span_id not in entries_by_span_id: + missing_parents.append( + { + "trace_id": trace_id, + "span_id": entry["span_id"], + "parent_span_id": parent_span_id, + "span_name": str(entry["span"].get("operationName") or entry["span"].get("name") or ""), + } + ) + + for entry in span_entries: + span = entry["span"] + tags = entry["tags"] + rollout_id, sample_tags = _resolve_rollout_sample(entry, entries_by_span_id) + if rollout_id is None: + continue + span_service_name = entry["service_name"] + span_run_id = entry["run_id"] + sample_key = (trace_id, rollout_id) + sample = samples_by_key.setdefault( + sample_key, + { + "trace_id": trace_id, + "rollout_id": rollout_id, + "group_id": sample_tags.get("xtuner.group_id"), + "producer_future_step": _producer_future_step(sample_tags), + "task_name": sample_tags.get("xtuner.task_name"), + "status": sample_tags.get("xtuner.status"), + "service_name": span_service_name, + "run_id": span_run_id, + "jaeger_url": _jaeger_trace_url(jaeger_trace_link_base_url, trace_id), + "spans": [], + }, + ) + _merge_sample_fields(sample, sample_tags) + if span_run_id is not None: + sample["run_id"] = span_run_id + sample["spans"].append(_span_payload(span, tags, service_name=span_service_name, run_id=span_run_id)) + + _merge_live_records(samples_by_key, live_records or (), jaeger_trace_link_base_url) + + generated_at_s = time.time() + samples = [] + for sample in samples_by_key.values(): + sample["spans"].sort(key=lambda item: (item["start_time_us"], item["span_id"])) + sample["span_count"] = len(sample["spans"]) + _apply_live_state(sample, generated_at_s) + _apply_sample_display_status(sample) + _apply_sample_reward_filter(sample) + sample["stage"] = _sample_stage(sample) + status = str(sample.get("status") or "unknown") + status_counts[status] += 1 + stage_counts[str(sample["stage"])] += 1 + if sample.get("group_id") is not None: + group_ids.add(sample["group_id"]) + samples.append(sample) + + samples.sort(key=lambda item: (str(item.get("group_id")), str(item.get("rollout_id")), item["trace_id"])) + step_group_summaries = _build_step_group_summaries(samples) + debug = { + "expected_topology": _build_expected_topology( + expected_global_batch_size, + expected_prompt_repeat_k, + ), + "topology_summaries": _build_topology_summaries( + samples, + expected_global_batch_size, + expected_prompt_repeat_k, + stale_threshold, + ), + "trace_diagnostics": _build_trace_diagnostics(samples, missing_parents, trace_context_sources), + } + return { + "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, + "sample_count": len(samples), + "group_count": len(group_ids), + "step_count": len(step_group_summaries), + "step_group_summaries": step_group_summaries, + "stage_occupancy": _build_stage_occupancy(samples, generated_at_s), + "stage_duration_summaries": _build_stage_duration_summaries(samples), + "critical_paths": _build_critical_paths(samples), + "reward_filter_summaries": _build_reward_filter_summaries(samples), + "status_counts": dict(sorted(status_counts.items())), + "stage_counts": dict(sorted(stage_counts.items())), + "debug": debug, + "samples": samples, + } + + +def load_jaeger_traces_from_otel_jsonl(trace_jsonl_path: Path | str) -> list[dict[str, Any]]: + traces_by_id: dict[str, dict[str, Any]] = {} + process_ids: dict[tuple[str, str, tuple[tuple[str, str], ...]], str] = {} + path = Path(trace_jsonl_path).expanduser() + if not path.is_file(): + return [] + + for record in _iter_jsonl_records(path): + if not isinstance(record, dict): + continue + jaeger_traces = _jaeger_traces_from_json_record(record) + if jaeger_traces is not None: + for trace_data in jaeger_traces: + _merge_jaeger_trace(traces_by_id, trace_data) + continue + for resource_span in record.get("resourceSpans") or []: + if not isinstance(resource_span, dict): + continue + resource_attrs = _otel_attributes_to_dict( + (resource_span.get("resource") or {}).get("attributes") or [] + ) + service_name = str(resource_attrs.get("service.name") or "unknown") + process_tags = _dict_to_jaeger_tags(resource_attrs) + for scope_span in _otel_scope_spans(resource_span): + for otel_span in scope_span.get("spans") or []: + if not isinstance(otel_span, dict): + continue + trace_id = str( + otel_span.get("traceId") + or otel_span.get("traceID") + or otel_span.get("trace_id") + or "" + ) + span_id = str( + otel_span.get("spanId") + or otel_span.get("spanID") + or otel_span.get("span_id") + or "" + ) + if not trace_id or not span_id: + continue + trace_data = traces_by_id.setdefault( + trace_id, + { + "traceID": trace_id, + "processes": {}, + "spans": [], + }, + ) + process_key = ( + trace_id, + service_name, + tuple(sorted((str(key), str(value)) for key, value in resource_attrs.items())), + ) + process_id = process_ids.get(process_key) + if process_id is None: + process_id = f"p{len(trace_data['processes']) + 1}" + process_ids[process_key] = process_id + trace_data["processes"][process_id] = { + "serviceName": service_name, + "tags": process_tags, + } + trace_data["spans"].append( + _otel_span_to_jaeger_span(otel_span, trace_id, span_id, process_id) + ) + return list(traces_by_id.values()) + + +def load_live_trace_records(live_jsonl_path: Path | str | None) -> list[dict[str, Any]]: + if live_jsonl_path is None: + return [] + path = Path(live_jsonl_path).expanduser() + if not path.is_file(): + return [] + return list(_iter_jsonl_records(path)) + + +def _iter_jsonl_records(path: Path) -> Iterable[dict[str, Any]]: + with path.open("r", encoding="utf-8") as handle: + for line in handle: + stripped = line.strip() + if not stripped: + continue + try: + payload = json.loads(stripped) + except json.JSONDecodeError: + continue + if isinstance(payload, dict): + yield payload + + +def _jaeger_traces_from_json_record(record: dict[str, Any]) -> list[dict[str, Any]] | None: + data = record.get("data") + if isinstance(data, list): + return [trace_data for trace_data in data if isinstance(trace_data, dict)] + if record.get("traceID") is not None and isinstance(record.get("spans"), list): + return [record] + return None + + +def _merge_jaeger_trace(traces_by_id: dict[str, dict[str, Any]], trace_data: dict[str, Any]) -> None: + trace_id = str(trace_data.get("traceID") or trace_data.get("trace_id") or "") + if not trace_id: + return + target = traces_by_id.setdefault(trace_id, {"traceID": trace_id, "processes": {}, "spans": []}) + if isinstance(trace_data.get("processes"), dict): + target["processes"].update(trace_data["processes"]) + target["spans"].extend([span for span in trace_data.get("spans") or [] if isinstance(span, dict)]) + + +def _otel_scope_spans(resource_span: dict[str, Any]) -> list[dict[str, Any]]: + scope_spans = resource_span.get("scopeSpans") + if isinstance(scope_spans, list): + return [scope_span for scope_span in scope_spans if isinstance(scope_span, dict)] + legacy_scope_spans = resource_span.get("instrumentationLibrarySpans") + if isinstance(legacy_scope_spans, list): + return [scope_span for scope_span in legacy_scope_spans if isinstance(scope_span, dict)] + return [] + + +def _otel_span_to_jaeger_span( + otel_span: dict[str, Any], + trace_id: str, + span_id: str, + process_id: str, +) -> dict[str, Any]: + attributes = _otel_attributes_to_dict(otel_span.get("attributes") or []) + tags = _dict_to_jaeger_tags(attributes) + tags.extend(_otel_status_tags(otel_span.get("status") or {})) + start_ns = _int_from_otel_time( + otel_span.get("startTimeUnixNano") + or otel_span.get("start_time_unix_nano") + or otel_span.get("startTime") + or 0 + ) + end_ns = _int_from_otel_time( + otel_span.get("endTimeUnixNano") + or otel_span.get("end_time_unix_nano") + or otel_span.get("endTime") + or start_ns + ) + parent_span_id = otel_span.get("parentSpanId") or otel_span.get("parent_span_id") + references = [] + if parent_span_id is not None and str(parent_span_id): + references.append({"refType": "CHILD_OF", "traceID": trace_id, "spanID": str(parent_span_id)}) + return { + "traceID": trace_id, + "spanID": span_id, + "operationName": str(otel_span.get("name") or otel_span.get("operationName") or "unknown"), + "processID": process_id, + "startTime": start_ns // 1_000, + "duration": max(0, end_ns - start_ns) // 1_000, + "references": references, + "tags": tags, + } + + +def _otel_status_tags(status: dict[str, Any]) -> list[dict[str, Any]]: + code = str(status.get("code") or status.get("statusCode") or "STATUS_CODE_UNSET") + if code in {"STATUS_CODE_ERROR", "ERROR", "2"}: + normalized = "ERROR" + elif code in {"STATUS_CODE_OK", "OK", "1"}: + normalized = "OK" + else: + normalized = "UNSET" + tags = [{"key": "otel.status_code", "type": "string", "value": normalized}] + message = status.get("message") or status.get("description") + if message: + tags.append({"key": "otel.status_description", "type": "string", "value": str(message)}) + if normalized == "ERROR": + tags.append({"key": "error.message", "type": "string", "value": str(message)}) + return tags + + +def _otel_attributes_to_dict(attributes: Any) -> dict[str, Any]: + if isinstance(attributes, dict): + return dict(attributes) + result: dict[str, Any] = {} + for attribute in attributes or []: + if not isinstance(attribute, dict): + continue + key = attribute.get("key") + if key is None: + continue + result[str(key)] = _otel_any_value_to_python(attribute.get("value")) + return result + + +def _otel_any_value_to_python(value: Any) -> Any: + if not isinstance(value, dict): + return value + if "stringValue" in value: + return value["stringValue"] + if "boolValue" in value: + return bool(value["boolValue"]) + if "intValue" in value: + return _int_or_original(value["intValue"]) + if "doubleValue" in value: + return float(value["doubleValue"]) + if "bytesValue" in value: + return value["bytesValue"] + if "arrayValue" in value: + return [_otel_any_value_to_python(item) for item in (value["arrayValue"].get("values") or [])] + if "kvlistValue" in value: + return { + str(item.get("key")): _otel_any_value_to_python(item.get("value")) + for item in (value["kvlistValue"].get("values") or []) + if isinstance(item, dict) and item.get("key") is not None + } + return value + + +def _dict_to_jaeger_tags(attributes: dict[str, Any]) -> list[dict[str, Any]]: + return [ + {"key": str(key), "type": _jaeger_tag_type(value), "value": value} + for key, value in attributes.items() + if value is not None + ] + + +def _jaeger_tag_type(value: Any) -> str: + if isinstance(value, bool): + return "bool" + if isinstance(value, int): + return "int64" + if isinstance(value, float): + return "float64" + return "string" + + +def _int_from_otel_time(value: Any) -> int: + parsed = _int_or_original(value) + return parsed if isinstance(parsed, int) and not isinstance(parsed, bool) else 0 + + +def _int_or_original(value: Any) -> int | Any: + if isinstance(value, bool): + return value + if isinstance(value, int): + return value + try: + return int(str(value)) + except (TypeError, ValueError): + return value + + +def _merge_sample_fields(sample: dict[str, Any], tags: dict[str, Any]) -> None: + for sample_key, tag_key in ( + ("rollout_id", "xtuner.rollout_id"), + ("group_id", "xtuner.group_id"), + ("task_name", "xtuner.task_name"), + ): + if tags.get(tag_key) is not None: + sample[sample_key] = tags[tag_key] + _merge_sample_status(sample, tags.get("xtuner.status")) + producer_future_step = _producer_future_step(tags) + if producer_future_step is not None: + sample["producer_future_step"] = producer_future_step + + +def _merge_sample_status(sample: dict[str, Any], status: Any) -> None: + if status is None: + return + current_status = sample.get("status") + if current_status is None or _sample_status_priority(status) >= _sample_status_priority(current_status): + sample["status"] = status + + +def _apply_sample_display_status(sample: dict[str, Any]) -> None: + status = str(sample.get("status") or "").strip().lower() + if status not in _INITIAL_SAMPLE_STATUSES: + return + if _sample_has_observed_stage(sample): + sample["status"] = "running" + + +def _sample_has_observed_stage(sample: dict[str, Any]) -> bool: + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict) and current_stage.get("name"): + return True + return bool(sample.get("spans")) + + +def _sample_status_priority(status: Any) -> int: + normalized = str(status or "").strip().lower() + if normalized in _ERROR_SAMPLE_STATUSES: + return 3 + if normalized in _TERMINAL_STAGE_STATUSES: + return 2 + if normalized in _NON_TERMINAL_SAMPLE_STATUSES: + return 0 + return 1 + + +def _merge_live_records( + samples_by_key: dict[tuple[str, Any], dict[str, Any]], + live_records: Iterable[dict[str, Any]], + jaeger_trace_link_base_url: str | None, +) -> None: + for record in live_records: + if not isinstance(record, dict): + continue + trace_id = str(record.get("trace_id") or "") + attributes = record.get("attributes") + if not trace_id or not isinstance(attributes, dict): + continue + rollout_id = attributes.get("xtuner.rollout_id") + if rollout_id is None: + continue + sample_key = (trace_id, rollout_id) + sample = samples_by_key.setdefault( + sample_key, + { + "trace_id": trace_id, + "rollout_id": rollout_id, + "group_id": attributes.get("xtuner.group_id"), + "producer_future_step": _producer_future_step(attributes), + "task_name": attributes.get("xtuner.task_name"), + "status": attributes.get("xtuner.status"), + "service_name": None, + "run_id": None, + "jaeger_url": _jaeger_trace_url(jaeger_trace_link_base_url, trace_id), + "spans": [], + }, + ) + _merge_sample_fields(sample, attributes) + sample.setdefault("_live_records", []).append(record) + + +def _apply_live_state(sample: dict[str, Any], generated_at_s: float) -> None: + live_states = _live_span_states(sample.pop("_live_records", [])) + active_states = [state for state in live_states if state.get("status") == "running"] + active_states.sort(key=lambda state: (len(state.get("logical_path") or []), float(state.get("started_at_s") or 0.0))) + current_state = active_states[-1] if active_states else None + if current_state is not None: + started_at_s = float(current_state.get("started_at_s") or generated_at_s) + sample["current_stage"] = { + "name": current_state.get("span_name"), + "status": "running", + "elapsed_ms": round(max(0.0, generated_at_s - started_at_s) * 1000.0, 3), + "started_at_s": started_at_s, + } + else: + sample["current_stage"] = None + sample["live_spans"] = live_states + sample["display_path"] = _build_display_path(sample, live_states, current_state, generated_at_s) + sample["chain"] = " -> ".join(node["name"] for node in sample["display_path"]) + + +def _live_span_states(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + states: dict[str, dict[str, Any]] = {} + for index, record in enumerate(records): + span_name = str(record.get("span_name") or "") + span_id = str(record.get("span_id") or "") + if not span_id: + span_id = f"live:{span_name}:{index}" + state = states.setdefault( + span_id, + { + "span_id": span_id, + "span_name": span_name, + "logical_path": _logical_path_from_value(record.get("logical_path")), + "attributes": record.get("attributes") if isinstance(record.get("attributes"), dict) else {}, + "status": "running", + }, + ) + event = str(record.get("event") or "") + if event == "start": + state["started_at_s"] = _float_or_none(record.get("time_s")) + state["status"] = "running" + elif event == "end": + state["ended_at_s"] = _float_or_none(record.get("time_s")) + state["duration_ms"] = _float_or_none(record.get("duration_ms")) + state["status"] = str(record.get("status") or "completed") + if record.get("error_message"): + state["error_message"] = record["error_message"] + if record.get("trace_id") is not None: + state["trace_id"] = record["trace_id"] + return sorted(states.values(), key=lambda state: (float(state.get("started_at_s") or 0.0), str(state.get("span_id")))) + + +def _build_display_path( + sample: dict[str, Any], + live_states: list[dict[str, Any]], + current_state: dict[str, Any] | None, + generated_at_s: float, +) -> list[dict[str, Any]]: + spans = sample.get("spans") or [] + spans_by_name = {str(span.get("name") or ""): span for span in spans} + live_by_name = {str(state.get("span_name") or ""): state for state in live_states} + path = _display_path_names(spans, live_states, current_state) + nodes = [] + for name in path: + span = spans_by_name.get(name) + live_state = live_by_name.get(name) + if current_state is not None and name == current_state.get("span_name"): + started_at_s = float(current_state.get("started_at_s") or generated_at_s) + nodes.append( + { + "name": name, + "source": "live", + "status": "running", + "elapsed_ms": round(max(0.0, generated_at_s - started_at_s) * 1000.0, 3), + } + ) + elif span is not None: + nodes.append( + { + "name": name, + "source": "span", + "status": "done" if str(span.get("status") or "").upper() != "ERROR" else "error", + "duration_ms": span.get("duration_ms"), + } + ) + elif live_state is not None and live_state.get("status") == "running": + nodes.append({"name": name, "source": "live", "status": "active"}) + else: + nodes.append({"name": name, "source": "logical", "status": "inferred"}) + return nodes + + +def _display_path_names( + spans: list[dict[str, Any]], + live_states: list[dict[str, Any]], + current_state: dict[str, Any] | None, +) -> list[str]: + path = _logical_path_from_value(current_state.get("logical_path")) if current_state is not None else [] + if not path: + path = [str(span.get("name") or "") for span in spans if span.get("name")] + if not path: + path = [str(state.get("span_name") or "") for state in live_states if state.get("span_name")] + if not path: + span_paths = [ + _logical_path_from_value((span.get("attributes") or {}).get(_LOGICAL_PATH_ATTRIBUTE)) + for span in spans + ] + path = max(span_paths, key=len, default=[]) + if not path: + live_paths = [_logical_path_from_value(state.get("logical_path")) for state in live_states] + path = max(live_paths, key=len, default=[]) + return _unique_path(path) + + +def _logical_path_from_value(value: Any) -> list[str]: + if isinstance(value, str): + try: + decoded = json.loads(value) + except json.JSONDecodeError: + decoded = [part.strip() for part in value.split("->")] + value = decoded + if isinstance(value, (list, tuple)): + return [str(item).strip() for item in value if isinstance(item, str) and item.strip()] + return [] + + +def _unique_path(path: list[str]) -> list[str]: + result = [] + for name in path: + if not result or result[-1] != name: + result.append(name) + return result + + +def _producer_future_step(tags: dict[str, Any]) -> Any: + value = tags.get("xtuner.producer_future_step") + return value if value is not None else tags.get("xtuner.train_step") + + +def _build_step_group_summaries(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: + steps: dict[str, dict[str, Any]] = {} + for sample in samples: + producer_future_step = sample.get("producer_future_step") + group_id = sample.get("group_id") + step_key = _summary_key(producer_future_step) + group_key = _summary_key(group_id) + + step_summary = steps.setdefault( + step_key, + { + "producer_future_step": producer_future_step, + "sample_count": 0, + "groups": {}, + }, + ) + step_summary["sample_count"] += 1 + groups = step_summary["groups"] + group_summary = groups.setdefault( + group_key, + { + "group_id": group_id, + "sample_count": 0, + "statuses": Counter(), + "stages": Counter(), + "rollout_ids": [], + }, + ) + group_summary["sample_count"] += 1 + group_summary["statuses"][str(sample.get("status") or "unknown")] += 1 + if _should_show_group_stage(sample): + group_summary["stages"][str(sample.get("stage") or "unknown")] += 1 + group_summary["rollout_ids"].append(sample.get("rollout_id")) + + summaries = [] + for step_summary in steps.values(): + groups = [] + for group_summary in step_summary["groups"].values(): + group_summary["statuses"] = dict(sorted(group_summary["statuses"].items())) + group_summary["stages"] = dict(sorted(group_summary["stages"].items())) + group_summary["rollout_ids"].sort(key=lambda value: str(value)) + groups.append(group_summary) + groups.sort(key=lambda item: _sortable_summary_value(item["group_id"])) + summaries.append( + { + "producer_future_step": step_summary["producer_future_step"], + "group_count": sum(1 for group in groups if group["group_id"] is not None), + "sample_count": step_summary["sample_count"], + "groups": groups, + } + ) + summaries.sort(key=lambda item: _sortable_summary_value(item["producer_future_step"])) + return summaries + + +def _build_stage_duration_summaries(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: + durations_by_stage: dict[str, list[float]] = 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 + durations_by_stage[str(span.get("name") or "unknown")].append(duration_s) + error = _extract_sample_error(sample) + if error is not None: + errors_by_stage[str(error.get("span_name") 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)), + "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, + } + ) + 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, + "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, + } + ) + return summaries + + +def _summarize_stage_errors(errors: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[tuple[str, str, Any], dict[str, Any]] = {} + for error in errors: + key = ( + str(error.get("error_type") or "error"), + str(error.get("message") or ""), + error.get("http_status_code"), + ) + summary = grouped.setdefault( + key, + { + "error_type": key[0], + "message": key[1], + "http_status_code": key[2], + "sample_count": 0, + "rollout_ids": [], + "groups": [], + "steps": [], + "jaeger_urls": [], + }, + ) + summary["sample_count"] += 1 + _append_unique(summary["rollout_ids"], error.get("rollout_id")) + _append_unique(summary["groups"], error.get("group_id")) + _append_unique(summary["steps"], error.get("producer_future_step")) + _append_unique(summary["jaeger_urls"], error.get("jaeger_url")) + result = list(grouped.values()) + for summary in result: + summary["rollout_ids"].sort(key=lambda value: str(value)) + summary["groups"].sort(key=lambda value: str(value)) + summary["steps"].sort(key=lambda value: str(value)) + result.sort(key=lambda item: (-item["sample_count"], str(item["error_type"]), str(item["message"]))) + return result + + +def _build_stage_occupancy(samples: list[dict[str, Any]], generated_at_s: float) -> list[dict[str, Any]]: + buckets: dict[str, dict[str, Any]] = {} + for sample in samples: + stage = _stage_bucket_for_sample(sample) + latest_start_s = _latest_span_start_s(sample) + age_s = 0.0 if stage in _TERMINAL_STAGE_STATUSES else max(0.0, generated_at_s - latest_start_s) + bucket = buckets.setdefault( + stage, + { + "stage": stage, + "sample_count": 0, + "oldest_age_s": 0.0, + "oldest_rollout_id": None, + "oldest_group_id": None, + "oldest_producer_future_step": None, + }, + ) + bucket["sample_count"] += 1 + if age_s >= bucket["oldest_age_s"]: + bucket["oldest_age_s"] = _round_duration(age_s) + bucket["oldest_rollout_id"] = sample.get("rollout_id") + bucket["oldest_group_id"] = sample.get("group_id") + bucket["oldest_producer_future_step"] = sample.get("producer_future_step") + return list(buckets.values()) + + +def _stage_bucket_for_sample(sample: dict[str, Any]) -> str: + status = str(sample.get("status") or "").strip().lower() + if status in _TERMINAL_STAGE_STATUSES: + return status + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict) and current_stage.get("name"): + return str(current_stage["name"]) + spans = sample.get("spans") or [] + if not spans: + return status or "unknown" + name = str(spans[-1].get("name") or "").lower() + if name.startswith("session_server."): + return "session_server" + if "trajectory.materialize" in name: + return "materialize" + if name.startswith("judger."): + return "judger" + if spans[-1].get("rollout_backend") or any(token in name for token in ("backend", "lmdeploy", "sglang", "generate")): + return "backend" + if status in {"pending", "queued", "scheduled"}: + return "scheduled" + return "running" + + +def _latest_span_start_s(sample: dict[str, Any]) -> float: + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict) and current_stage.get("started_at_s") is not None: + try: + return float(current_stage["started_at_s"]) + except (TypeError, ValueError): + pass + spans = sample.get("spans") or [] + if not spans: + return time.time() + return max(float(span.get("start_time_us") or 0) / 1_000_000.0 for span in spans) + + +def _build_critical_paths(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: + groups: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list) + for sample in samples: + groups[(_summary_key(sample.get("producer_future_step")), _summary_key(sample.get("group_id")))].append(sample) + + summaries = [] + for group_samples in groups.values(): + group_samples.sort(key=lambda item: str(item.get("rollout_id"))) + sample_rows = [] + for sample in group_samples: + slowest_span = _slowest_span(sample) + sample_rows.append( + { + "rollout_id": sample.get("rollout_id"), + "duration_s": _sample_wall_time_s(sample), + "slowest_span_name": slowest_span.get("name") if slowest_span else None, + "slowest_span_duration_s": _span_duration_s(slowest_span) if slowest_span else 0.0, + "status": sample.get("status"), + "jaeger_url": sample.get("jaeger_url"), + } + ) + if not sample_rows: + continue + slowest = max(sample_rows, key=lambda item: (item["duration_s"], str(item["rollout_id"]))) + spans_start = [_span_start_us(span) for sample in group_samples for span in sample.get("spans") or []] + spans_end = [_span_end_time_us(span) for sample in group_samples for span in sample.get("spans") or []] + group_wall_time_s = _round_duration((max(spans_end) - min(spans_start)) / 1_000_000.0) if spans_start else 0.0 + first_sample = group_samples[0] + summaries.append( + { + "producer_future_step": first_sample.get("producer_future_step"), + "group_id": first_sample.get("group_id"), + "sample_count": len(group_samples), + "group_wall_time_s": group_wall_time_s, + "slowest_rollout_id": slowest["rollout_id"], + "slowest_sample_duration_s": slowest["duration_s"], + "slowest_span_name": slowest["slowest_span_name"], + "slowest_span_duration_s": slowest["slowest_span_duration_s"], + "jaeger_url": slowest["jaeger_url"], + "samples": sample_rows, + } + ) + summaries.sort(key=lambda item: (_sortable_summary_value(item["producer_future_step"]), _sortable_summary_value(item["group_id"]))) + return summaries + + +def _slowest_span(sample: dict[str, Any]) -> dict[str, Any] | None: + spans = sample.get("spans") or [] + if not spans: + return None + return max(spans, key=lambda span: (_span_duration_s(span), str(span.get("span_id")))) + + +def _span_duration_s(span: dict[str, Any] | None) -> float: + if span is None: + return 0.0 + return _round_duration(float(span.get("duration_ms") or 0.0) / 1000.0) + + +def _span_start_us(span: dict[str, Any]) -> int: + return int(span.get("start_time_us") or 0) + + +def _span_end_time_us(span: dict[str, Any]) -> int: + return _span_start_us(span) + int(float(span.get("duration_ms") or 0.0) * 1000) + + +def _sample_wall_time_s(sample: dict[str, Any]) -> float: + spans = sample.get("spans") or [] + if not spans: + return 0.0 + starts = [_span_start_us(span) for span in spans] + ends = [_span_end_time_us(span) for span in spans] + return _round_duration((max(ends) - min(starts)) / 1_000_000.0) + + +def _apply_sample_reward_filter(sample: dict[str, Any]) -> None: + values: dict[str, Any] = {} + for span in sample.get("spans") or []: + attrs = span.get("attributes") or {} + for target, keys in ( + ("reward_score", ("reward.score", "reward_score")), + ("reward_pass", ("reward.pass",)), + ("filter_decision", ("filter.decision",)), + ("filter_reason", ("filter.reason",)), + ("train_included", ("train.included",)), + ("oversample_source", ("oversample.source",)), + ("drop_reason", ("drop.reason",)), + ): + for key in keys: + if key in attrs: + values[target] = attrs[key] + if "reward_score" in values: + values["reward_score"] = _to_float(values["reward_score"]) + if "reward_pass" in values: + values["reward_pass"] = _to_bool(values["reward_pass"]) + if "train_included" in values: + values["train_included"] = _to_bool(values["train_included"]) + sample.update(values) + + +def _build_reward_filter_summaries(samples: list[dict[str, Any]]) -> list[dict[str, Any]]: + grouped: dict[tuple[str, str], dict[str, Any]] = {} + for sample in samples: + key = (_summary_key(sample.get("producer_future_step")), _summary_key(sample.get("group_id"))) + row = grouped.setdefault( + key, + { + "producer_future_step": sample.get("producer_future_step"), + "group_id": sample.get("group_id"), + "sample_count": 0, + "reward_count": 0, + "reward_avg": None, + "reward_min": None, + "reward_max": None, + "pass_count": 0, + "kept_count": 0, + "dropped_count": 0, + "filter_reasons": Counter(), + "drop_reasons": Counter(), + "_rewards": [], + }, + ) + row["sample_count"] += 1 + reward_score = sample.get("reward_score") + if isinstance(reward_score, (int, float)): + row["_rewards"].append(float(reward_score)) + row["reward_count"] += 1 + if sample.get("reward_pass") is True: + row["pass_count"] += 1 + decision = str(sample.get("filter_decision") or "").lower() + if decision == "kept" or sample.get("train_included") is True: + row["kept_count"] += 1 + if decision in {"filtered", "dropped"} or sample.get("train_included") is False: + row["dropped_count"] += 1 + if sample.get("filter_reason") is not None: + row["filter_reasons"][str(sample["filter_reason"])] += 1 + if sample.get("drop_reason") is not None: + row["drop_reasons"][str(sample["drop_reason"])] += 1 + + summaries = [] + for row in grouped.values(): + rewards = row.pop("_rewards") + if rewards: + row["reward_avg"] = _round_duration(sum(rewards) / len(rewards)) + row["reward_min"] = _round_duration(min(rewards)) + row["reward_max"] = _round_duration(max(rewards)) + row["filter_reasons"] = dict(sorted(row["filter_reasons"].items())) + row["drop_reasons"] = dict(sorted(row["drop_reasons"].items())) + summaries.append(row) + summaries.sort(key=lambda item: (_sortable_summary_value(item["producer_future_step"]), _sortable_summary_value(item["group_id"]))) + return summaries + + +def _build_expected_topology( + expected_global_batch_size: int | None, + expected_prompt_repeat_k: int | None, +) -> dict[str, Any]: + expected_samples = ( + expected_global_batch_size * expected_prompt_repeat_k + if expected_global_batch_size is not None and expected_prompt_repeat_k is not None + else None + ) + source = "cli" if expected_global_batch_size is not None or expected_prompt_repeat_k is not None else "unknown" + return { + "global_batch_size": expected_global_batch_size, + "prompt_repeat_k": expected_prompt_repeat_k, + "expected_groups_per_step": expected_global_batch_size, + "expected_samples_per_step": expected_samples, + "source": source, + } + + +def _build_topology_summaries( + samples: list[dict[str, Any]], + expected_global_batch_size: int | None, + expected_prompt_repeat_k: int | None, + stale_threshold: int | None, +) -> list[dict[str, Any]]: + steps: dict[str, dict[str, Any]] = {} + for sample in samples: + step_key = _summary_key(sample.get("producer_future_step")) + group_key = _summary_key(sample.get("group_id")) + step = steps.setdefault( + step_key, + {"producer_future_step": sample.get("producer_future_step"), "groups": defaultdict(list)}, + ) + step["groups"][group_key].append(sample) + + summaries = [] + for step in steps.values(): + group_rows = [] + actual_group_ids = set() + actual_group_indices = set() + missing_groups = [] + extra_groups = [] + missing_samples = 0 + extra_samples = 0 + stale_samples = 0 + expected_group_ids = set(range(expected_global_batch_size or 0)) + for group_key, group_samples in step["groups"].items(): + group_id = group_samples[0].get("group_id") + actual_group_ids.add(group_id) + group_index = _int_value(group_id) + if group_index is not None: + actual_group_indices.add(group_index) + actual_repeats = len(group_samples) + stale_count = sum(1 for sample in group_samples if _is_stale_sample(sample, stale_threshold)) + stale_samples += stale_count + expected_repeats = expected_prompt_repeat_k + group_missing = max(0, (expected_repeats or actual_repeats) - actual_repeats) if expected_repeats else 0 + group_extra = max(0, actual_repeats - expected_repeats) if expected_repeats else 0 + missing_samples += group_missing + extra_samples += group_extra + if ( + expected_global_batch_size is not None + and group_index is not None + and group_index >= expected_global_batch_size + ): + extra_groups.append(group_id) + group_status = _topology_status(group_missing, group_extra, stale_count) + group_rows.append( + { + "group_id": group_id, + "expected_repeats": expected_repeats, + "actual_repeats": actual_repeats, + "missing_samples": group_missing, + "extra_samples": group_extra, + "stale_samples": stale_count, + "status": group_status, + "rollout_ids": sorted([sample.get("rollout_id") for sample in group_samples], key=lambda value: str(value)), + } + ) + if expected_global_batch_size is not None: + for group_id in sorted(expected_group_ids - actual_group_indices): + missing_groups.append(group_id) + missing_count = expected_prompt_repeat_k or 0 + missing_samples += missing_count + group_rows.append( + { + "group_id": group_id, + "expected_repeats": expected_prompt_repeat_k, + "actual_repeats": 0, + "missing_samples": missing_count, + "extra_samples": 0, + "stale_samples": 0, + "status": "missing", + "rollout_ids": [], + } + ) + group_rows.sort(key=lambda item: _sortable_summary_value(item["group_id"])) + expected_samples = ( + expected_global_batch_size * expected_prompt_repeat_k + if expected_global_batch_size is not None and expected_prompt_repeat_k is not None + else None + ) + status = _topology_status(missing_samples, extra_samples + len(extra_groups), stale_samples) + summaries.append( + { + "producer_future_step": step["producer_future_step"], + "expected_groups": expected_global_batch_size, + "actual_groups": len([group for group in actual_group_ids if group is not None]), + "expected_samples": expected_samples, + "actual_samples": sum(len(group_samples) for group_samples in step["groups"].values()), + "missing_groups": missing_groups, + "extra_groups": sorted(extra_groups, key=lambda value: str(value)), + "missing_samples": missing_samples, + "extra_samples": extra_samples, + "stale_samples": stale_samples, + "status": status, + "groups": group_rows, + } + ) + summaries.sort(key=lambda item: _sortable_summary_value(item["producer_future_step"])) + return summaries + + +def _topology_status(missing_count: int, extra_count: int, stale_count: int) -> str: + states = [] + if missing_count: + states.append("missing") + if extra_count: + states.append("extra") + if stale_count: + states.append("stale") + if not states: + return "ok" + return states[0] if len(states) == 1 else "mixed" + + +def _is_stale_sample(sample: dict[str, Any], stale_threshold: int | None) -> bool: + status = str(sample.get("status") or "").lower() + if status in {"expired", "stale"}: + return True + for span in sample.get("spans") or []: + attrs = span.get("attributes") or {} + if _to_bool(attrs.get("xtuner.stale")) is True: + return True + if stale_threshold is None: + return False + try: + return float(sample.get("seq_staleness") or 0) >= stale_threshold + except (TypeError, ValueError): + return False + + +def _build_trace_diagnostics( + samples: list[dict[str, Any]], + missing_parents: list[dict[str, Any]], + trace_context_sources: Counter[str], +) -> dict[str, Any]: + samples_missing_session_server = [] + for sample in samples: + span_names = [str(span.get("name") or "") for span in sample.get("spans") or []] + has_agent_span = any(name.startswith("agent") or name == "agent_loop.run" for name in span_names) + has_session_span = any(name.startswith("session_server.") for name in span_names) + if has_agent_span and not has_session_span: + _append_unique(samples_missing_session_server, sample.get("rollout_id")) + samples_missing_session_server.sort(key=lambda value: str(value)) + return { + "missing_parent_count": len(missing_parents), + "missing_parents": missing_parents, + "samples_missing_session_server": samples_missing_session_server, + "trace_context_sources": dict(sorted(trace_context_sources.items())), + "orphan_span_count": 0, + } + + +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, + } + 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, + } + 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 _sample_stage(sample: dict[str, Any]) -> str: + status = str(sample.get("status") or "").strip().lower() + if status and status not in _NON_TERMINAL_SAMPLE_STATUSES: + return status + current_stage = sample.get("current_stage") + if isinstance(current_stage, dict) and current_stage.get("name"): + return str(current_stage["name"]) + + spans = sample.get("spans") or [] + for span in spans: + attributes = span.get("attributes") or {} + error_value = str(attributes.get("error") or "").strip().lower() + if str(span.get("status") or "").upper() == "ERROR" or error_value == "true": + return "error" + if spans: + return str(spans[-1].get("name") or "unknown") + return status or "unknown" + + +def _should_show_group_stage(sample: dict[str, Any]) -> bool: + status = str(sample.get("status") or "unknown").strip().lower() + stage = str(sample.get("stage") or "unknown").strip().lower() + return stage != status + + +def _summary_key(value: Any) -> str: + return "" if value is None else str(value) + + +def _sortable_summary_value(value: Any) -> tuple[int, float | str]: + if value is None: + return (1, "") + if isinstance(value, bool): + return (0, str(value)) + if isinstance(value, (int, float)): + return (0, float(value)) + text = str(value) + try: + return (0, float(text)) + except ValueError: + return (0, text) + + +def _span_payload( + span: dict[str, Any], + tags: dict[str, Any], + *, + service_name: str | None, + run_id: str | None, +) -> dict[str, Any]: + return { + "name": str(span.get("operationName") or span.get("name") or "unknown"), + "span_id": str(span.get("spanID") or span.get("span_id") or ""), + "parent_span_id": _parent_span_id(span), + "start_time_us": int(span.get("startTime") or 0), + "duration_ms": float(span.get("duration") or 0) / 1000.0, + "status": tags.get("otel.status_code") or tags.get("status.code") or "UNSET", + "service_name": service_name, + "run_id": run_id, + "rollout_backend": tags.get("rollout.backend"), + "attributes": { + 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"} + }, + } + + +def _span_id(span: dict[str, Any]) -> str: + return str(span.get("spanID") or span.get("span_id") or "") + + +def _resolve_rollout_sample( + entry: dict[str, Any], + entries_by_span_id: dict[str, dict[str, Any]], +) -> tuple[Any | None, dict[str, Any]]: + tags = entry["tags"] + + ancestor_sample: tuple[Any, dict[str, Any]] | None = None + visited: set[str] = set() + parent_span_id = _parent_span_id(entry["span"]) + while parent_span_id and parent_span_id not in visited: + visited.add(parent_span_id) + parent = entries_by_span_id.get(parent_span_id) + if parent is None: + break + parent_tags = parent["tags"] + parent_rollout_id = parent_tags.get("xtuner.rollout_id") + if parent_rollout_id is not None: + ancestor_sample = (parent_rollout_id, parent_tags) + parent_span_id = _parent_span_id(parent["span"]) + + if ancestor_sample is not None: + return ancestor_sample + + rollout_id = tags.get("xtuner.rollout_id") + if rollout_id is not None: + return rollout_id, tags + return None, {} + + +def _process_metadata(trace_data: dict[str, Any]) -> dict[str, dict[str, Any]]: + metadata: dict[str, dict[str, Any]] = {} + for process_id, process in (trace_data.get("processes") or {}).items(): + if not isinstance(process, dict): + continue + tags = _tags_to_dict(process.get("tags") or []) + metadata[str(process_id)] = { + "service_name": str(process["serviceName"]) if process.get("serviceName") is not None else None, + "run_id": tags.get("run.id"), + } + return metadata + + +def _tags_to_dict(tags: list[dict[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for tag in tags: + key = tag.get("key") + if key is None: + continue + result[str(key)] = tag.get("value") + return result + + +def _parent_span_id(span: dict[str, Any]) -> str | None: + for reference in span.get("references") or []: + if reference.get("refType") == "CHILD_OF" and reference.get("spanID") is not None: + return str(reference["spanID"]) + return None + + +def _normalize_jaeger_query_url(jaeger_query_url: str | None) -> str | None: + if jaeger_query_url is None: + return None + stripped = jaeger_query_url.strip() + return stripped.rstrip("/") if stripped else None + + +def _jaeger_trace_url(jaeger_query_url: str | None, trace_id: str) -> str | None: + base = _normalize_jaeger_query_url(jaeger_query_url) + if base is None: + return None + return f"{base}/trace/{trace_id}" + + +def _jaeger_trace_link_base_url(jaeger_query_url: str | None, jaeger_link_url: str | None) -> str | None: + return _normalize_jaeger_query_url(jaeger_link_url) or _normalize_jaeger_query_url(jaeger_query_url) + + +def _append_unique(values: list[Any], value: Any) -> None: + if value is not None and value not in values: + values.append(value) + + +def _to_float(value: Any) -> float | Any: + try: + return float(value) + except (TypeError, ValueError): + return value + + +def _float_or_none(value: Any) -> float | None: + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _to_bool(value: Any) -> bool | None: + if isinstance(value, bool): + return value + if value is None: + return None + text = str(value).strip().lower() + if text in {"1", "true", "yes", "y", "on"}: + return True + if text in {"0", "false", "no", "n", "off"}: + return False + return None + + +def _int_value(value: Any) -> int | None: + if isinstance(value, bool) or value is None: + return None + if isinstance(value, int): + return value + text = str(value).strip() + if not text: + return None + try: + parsed = float(text) + except ValueError: + return None + if parsed.is_integer(): + return int(parsed) + return None + + +__all__ = [ + "build_rollout_view_payload_from_jaeger_traces", + "load_jaeger_traces_from_otel_jsonl", + "load_live_trace_records", +] diff --git a/xtuner/tools/trace_viewer/render.py b/xtuner/tools/trace_viewer/render.py new file mode 100644 index 0000000000..f2ee6741cb --- /dev/null +++ b/xtuner/tools/trace_viewer/render.py @@ -0,0 +1,577 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + + +def render_rollout_trace_html( + payload: dict[str, Any], + *, + live: bool = False, + api_url: str = "/api/trace", + refresh_interval_s: float = 2.0, +) -> str: + data = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + return ( + _HTML_TEMPLATE.replace("__TRACE_DATA__", data) + .replace("__LIVE_MODE__", json.dumps(live)) + .replace("__TRACE_API_URL__", json.dumps(api_url)) + .replace("__REFRESH_INTERVAL_MS__", str(int(refresh_interval_s * 1000))) + ) + + +def write_rollout_trace_html(payload: dict[str, Any], output_path: Path) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(render_rollout_trace_html(payload), encoding="utf-8") + + +_HTML_TEMPLATE = r""" + + + + + XTuner Rollout Trace Viewer + + + +
+
+

XTuner Rollout Trace Viewer

+
+
+
+
+
+
+ +
+

Stage Occupancy

+
+ + + +
StageSamples
+
+
+ +
+

Stage Durations

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

Samples

+
+ + + + +
+
+ + + + + + + + + + + + + +
SampleStatusGroupStepRewardPath / CurrentJaeger
+
+
+ + + +
+ + + +""" + + +__all__ = ["render_rollout_trace_html", "write_rollout_trace_html"] diff --git a/xtuner/tools/trace_viewer/server.py b/xtuner/tools/trace_viewer/server.py new file mode 100644 index 0000000000..00459e95b1 --- /dev/null +++ b/xtuner/tools/trace_viewer/server.py @@ -0,0 +1,430 @@ +from __future__ import annotations + +import argparse +import http.server +import json +import os +import threading +from pathlib import Path +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.parse import urlsplit +from urllib.request import Request, urlopen + +from xtuner.tools.trace_viewer.payload import ( + build_rollout_view_payload_from_jaeger_traces, + load_live_trace_records, +) +from xtuner.tools.trace_viewer.render import render_rollout_trace_html, write_rollout_trace_html +from xtuner.tools.trace_viewer.source import ( + JAEGER_DEFAULT_LIMIT, + JAEGER_DEFAULT_LOOKBACK_S, + JAEGER_DEFAULT_QUERY_URL, + JaegerQuerySource, + JsonlTraceSource, + normalize_jaeger_query_url, + require_jaeger_query_url, +) + + +_JAEGER_PROXY_PREFIX = "/jaeger" +_PROXY_TIMEOUT_S = 10.0 +_HOP_BY_HOP_HEADERS = { + "connection", + "content-length", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +} + + +class TraceViewerHandle: + def __init__( + self, + server: http.server.ThreadingHTTPServer, + thread: threading.Thread, + *, + host: str, + port: int, + url: str, + ) -> None: + self.server = server + self.thread = thread + self.host = host + self.port = port + self.url = url + + def close(self) -> None: + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + + +def fetch_jaeger_traces( + jaeger_query_url: str, + *, + service_name: str, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + timeout_s: float = 5.0, +) -> list[dict[str, Any]]: + return JaegerQuerySource( + query_url=jaeger_query_url, + service_name=service_name, + lookback_s=lookback_s, + limit=limit, + timeout_s=timeout_s, + ).load() + + +def fetch_rollout_view_payload( + jaeger_query_url: str, + *, + jaeger_link_url: str | None = None, + live_jsonl_path: Path | str | None = None, + service_name: str, + run_id: str | None = None, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + expected_global_batch_size: int | None = None, + expected_prompt_repeat_k: int | None = None, + stale_threshold: int | None = None, +) -> dict[str, Any]: + traces = JaegerQuerySource( + query_url=jaeger_query_url, + service_name=service_name, + lookback_s=lookback_s, + limit=limit, + ).load() + payload = build_rollout_view_payload_from_jaeger_traces( + traces, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=jaeger_link_url, + live_records=load_live_trace_records(live_jsonl_path), + service_name=service_name, + run_id=run_id, + expected_global_batch_size=expected_global_batch_size, + expected_prompt_repeat_k=expected_prompt_repeat_k, + stale_threshold=stale_threshold, + ) + payload["service_name"] = service_name + payload["run_id"] = run_id + payload["lookback_s"] = lookback_s + payload["limit"] = limit + return payload + + +def fetch_rollout_view_payload_from_trace_jsonl( + trace_jsonl_path: Path | str, + *, + jaeger_query_url: str | None = None, + jaeger_link_url: str | None = None, + live_jsonl_path: Path | str | None = None, + service_name: str | None = None, + run_id: str | None = None, + expected_global_batch_size: int | None = None, + expected_prompt_repeat_k: int | None = None, + stale_threshold: int | None = None, +) -> dict[str, Any]: + traces = JsonlTraceSource(trace_jsonl_path).load() + payload = build_rollout_view_payload_from_jaeger_traces( + traces, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=jaeger_link_url, + live_records=load_live_trace_records(live_jsonl_path), + service_name=service_name, + run_id=run_id, + expected_global_batch_size=expected_global_batch_size, + expected_prompt_repeat_k=expected_prompt_repeat_k, + stale_threshold=stale_threshold, + ) + payload["source"] = "trace_jsonl" + payload["trace_jsonl_path"] = os.fspath(Path(trace_jsonl_path).expanduser()) + payload["service_name"] = service_name + payload["run_id"] = run_id + return payload + + +def start_rollout_trace_viewer( + jaeger_query_url: str | None = JAEGER_DEFAULT_QUERY_URL, + *, + jaeger_link_url: str | None = None, + service_name: str, + run_id: str | None = None, + trace_jsonl_path: Path | str | None = None, + live_jsonl_path: Path | str | None = None, + payload_output_path: Path | str | None = None, + host: str = "127.0.0.1", + port: int = 0, + refresh_interval_s: float = 2.0, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + expected_global_batch_size: int | None = None, + expected_prompt_repeat_k: int | None = None, + stale_threshold: int | None = None, +) -> TraceViewerHandle: + jaeger_query_url = normalize_jaeger_query_url(jaeger_query_url) + jaeger_link_url = normalize_jaeger_query_url(jaeger_link_url) + viewer_jaeger_link_url = jaeger_link_url or (_JAEGER_PROXY_PREFIX if jaeger_query_url is not None else None) + if trace_jsonl_path is None: + jaeger_query_url = require_jaeger_query_url(jaeger_query_url) + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self) -> None: + path = self.path.split("?", 1)[0] + 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(), + live=True, + api_url="/api/trace", + refresh_interval_s=refresh_interval_s, + ) + self._send_bytes(html_body.encode("utf-8"), "text/html; charset=utf-8") + return + if path == "/api/trace": + self._send_json(self._payload()) + return + self.send_error(404) + + def _payload(self) -> dict[str, Any]: + if trace_jsonl_path is not None: + payload = fetch_rollout_view_payload_from_trace_jsonl( + trace_jsonl_path, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=viewer_jaeger_link_url, + live_jsonl_path=live_jsonl_path, + service_name=service_name, + run_id=run_id, + expected_global_batch_size=expected_global_batch_size, + expected_prompt_repeat_k=expected_prompt_repeat_k, + stale_threshold=stale_threshold, + ) + else: + payload = fetch_rollout_view_payload( + jaeger_query_url, + jaeger_link_url=viewer_jaeger_link_url, + live_jsonl_path=live_jsonl_path, + service_name=service_name, + run_id=run_id, + lookback_s=lookback_s, + limit=limit, + expected_global_batch_size=expected_global_batch_size, + expected_prompt_repeat_k=expected_prompt_repeat_k, + stale_threshold=stale_threshold, + ) + if payload_output_path is not None: + _write_payload_json(payload, payload_output_path) + return payload + + def _send_json(self, payload: dict[str, Any]) -> None: + self._send_bytes(json.dumps(payload, ensure_ascii=False).encode("utf-8"), "application/json") + + def _proxy_jaeger(self) -> None: + if jaeger_query_url is None: + self.send_error(502, "Jaeger query URL is not configured") + return + try: + target_url = _jaeger_proxy_target_url(jaeger_query_url, self.path) + request = Request( + target_url, + headers={ + "Accept": self.headers.get("Accept", "*/*"), + "User-Agent": self.headers.get("User-Agent", "XTunerTraceViewer"), + }, + ) + with urlopen(request, timeout=_PROXY_TIMEOUT_S) as response: + self._send_proxy_response(response.status, response.headers.items(), response.read()) + except HTTPError as exc: + self._send_proxy_response(exc.code, exc.headers.items(), exc.read()) + except (OSError, URLError) as exc: + self.send_error(502, f"Failed to proxy Jaeger request: {exc}") + + def _send_proxy_response(self, status: int, headers: Any, body: bytes) -> None: + self.send_response(status) + for key, value in headers: + if key.lower() in _HOP_BY_HOP_HEADERS: + continue + self.send_header(key, value) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _send_bytes(self, body: bytes, content_type: str) -> None: + self.send_response(200) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: Any) -> None: + return + + server = http.server.ThreadingHTTPServer((host, port), Handler) + server_host, server_port = server.server_address + display_host = server_host or host + thread = threading.Thread(target=server.serve_forever, name="XTunerRolloutTraceViewer", daemon=True) + thread.start() + return TraceViewerHandle( + server=server, + thread=thread, + host=display_host, + port=server_port, + url=f"http://{display_host}:{server_port}", + ) + + +def _write_payload_json(payload: dict[str, Any], output_path: Path | str) -> None: + path = Path(output_path).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + +def _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 _env_int(name: str) -> int | None: + value = os.environ.get(name) + if value is None or value == "": + return None + return int(value) + + +def _first_int(*values: int | None) -> int | None: + for value in values: + if value is not None: + return value + return None + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Serve or render an XTuner rollout trace viewer backed by Jaeger or JSONL." + ) + parser.add_argument("--jaeger-query-url", default=JAEGER_DEFAULT_QUERY_URL) + parser.add_argument("--jaeger-link-url", default=None) + parser.add_argument("--trace-jsonl", type=Path, default=None) + parser.add_argument("--live-jsonl", type=Path, default=None) + parser.add_argument("--service", "--service-name", dest="service", default="xtuner-rollout") + parser.add_argument("--run-id", default=None) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=0) + parser.add_argument("--lookback", type=int, default=JAEGER_DEFAULT_LOOKBACK_S) + parser.add_argument("--limit", type=int, default=JAEGER_DEFAULT_LIMIT) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument("--payload-output", type=Path, default=None) + parser.add_argument("--global-batch-size", type=int, default=None) + parser.add_argument("--prompt-repeat-k", type=int, default=None) + parser.add_argument("--stale-threshold", type=int, default=None) + return parser.parse_args() + + +def main() -> None: + args = _parse_args() + expected_global_batch_size = _first_int(args.global_batch_size, _env_int("XTUNER_TRACE_GLOBAL_BATCH_SIZE")) + expected_prompt_repeat_k = _first_int(args.prompt_repeat_k, _env_int("XTUNER_TRACE_PROMPT_REPEAT_K")) + stale_threshold = _first_int(args.stale_threshold, _env_int("XTUNER_TRACE_STALE_THRESHOLD")) + + if args.output is not None: + if args.trace_jsonl is not None: + payload = fetch_rollout_view_payload_from_trace_jsonl( + args.trace_jsonl, + jaeger_query_url=args.jaeger_query_url, + jaeger_link_url=args.jaeger_link_url, + live_jsonl_path=args.live_jsonl, + service_name=args.service, + run_id=args.run_id, + expected_global_batch_size=expected_global_batch_size, + expected_prompt_repeat_k=expected_prompt_repeat_k, + stale_threshold=stale_threshold, + ) + else: + payload = fetch_rollout_view_payload( + args.jaeger_query_url, + jaeger_link_url=args.jaeger_link_url, + live_jsonl_path=args.live_jsonl, + service_name=args.service, + run_id=args.run_id, + lookback_s=args.lookback, + limit=args.limit, + expected_global_batch_size=expected_global_batch_size, + expected_prompt_repeat_k=expected_prompt_repeat_k, + stale_threshold=stale_threshold, + ) + if args.payload_output is not None: + _write_payload_json(payload, args.payload_output) + write_rollout_trace_html(payload, args.output) + print(args.output) + if args.payload_output is not None: + print(args.payload_output) + return + + handle = start_rollout_trace_viewer( + args.jaeger_query_url, + jaeger_link_url=args.jaeger_link_url, + service_name=args.service, + run_id=args.run_id, + trace_jsonl_path=args.trace_jsonl, + live_jsonl_path=args.live_jsonl, + payload_output_path=args.payload_output, + host=args.host, + port=args.port, + lookback_s=args.lookback, + limit=args.limit, + expected_global_batch_size=expected_global_batch_size, + expected_prompt_repeat_k=expected_prompt_repeat_k, + stale_threshold=stale_threshold, + ) + print(f"XTuner Rollout Trace Viewer: {handle.url}", flush=True) + if args.trace_jsonl is not None: + print(f"Trace JSONL: {args.trace_jsonl}", flush=True) + if args.live_jsonl is not None: + print(f"Live JSONL: {args.live_jsonl}", flush=True) + jaeger_query_url = normalize_jaeger_query_url(args.jaeger_query_url) + if jaeger_query_url is not None: + print(f"Jaeger Trace Viewer: {jaeger_query_url}", flush=True) + print(f"Jaeger Same-Origin Proxy: {handle.url}{_JAEGER_PROXY_PREFIX}/", flush=True) + jaeger_link_url = normalize_jaeger_query_url(args.jaeger_link_url) + if jaeger_link_url is not None: + print(f"Jaeger Open Links: {jaeger_link_url}", flush=True) + if args.payload_output is not None: + print(f"Viewer Payload JSON: {args.payload_output}", flush=True) + try: + handle.thread.join() + except KeyboardInterrupt: + pass + finally: + handle.close() + + +if __name__ == "__main__": + main() + + +__all__ = [ + "TraceViewerHandle", + "build_rollout_view_payload_from_jaeger_traces", + "fetch_jaeger_traces", + "fetch_rollout_view_payload", + "fetch_rollout_view_payload_from_trace_jsonl", + "render_rollout_trace_html", + "start_rollout_trace_viewer", + "write_rollout_trace_html", +] diff --git a/xtuner/tools/trace_viewer/source.py b/xtuner/tools/trace_viewer/source.py new file mode 100644 index 0000000000..3c82ef641c --- /dev/null +++ b/xtuner/tools/trace_viewer/source.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.parse import urlencode +from urllib.request import Request, urlopen + +from xtuner.tools.trace_viewer.payload import load_jaeger_traces_from_otel_jsonl + + +JAEGER_DEFAULT_QUERY_URL = "http://127.0.0.1:16686" +JAEGER_DEFAULT_LOOKBACK_S = 60 * 60 +JAEGER_DEFAULT_LIMIT = 500 + + +@dataclass(frozen=True) +class JaegerQuerySource: + query_url: str + service_name: str + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S + limit: int = JAEGER_DEFAULT_LIMIT + timeout_s: float = 5.0 + + def load(self) -> list[dict[str, Any]]: + return fetch_jaeger_traces( + self.query_url, + service_name=self.service_name, + lookback_s=self.lookback_s, + limit=self.limit, + timeout_s=self.timeout_s, + ) + + +@dataclass(frozen=True) +class JsonlTraceSource: + trace_jsonl_path: Path | str + + def load(self) -> list[dict[str, Any]]: + return load_jaeger_traces_from_otel_jsonl(self.trace_jsonl_path) + + +def fetch_jaeger_traces( + jaeger_query_url: str, + *, + service_name: str, + lookback_s: int = JAEGER_DEFAULT_LOOKBACK_S, + limit: int = JAEGER_DEFAULT_LIMIT, + timeout_s: float = 5.0, +) -> list[dict[str, Any]]: + base_url = require_jaeger_query_url(jaeger_query_url) + query = urlencode( + { + "service": service_name, + "lookback": f"{max(1, lookback_s)}s", + "limit": max(1, limit), + } + ) + request = Request(f"{base_url}/api/traces?{query}", headers={"Accept": "application/json"}) + with urlopen(request, timeout=timeout_s) as response: + payload = json.loads(response.read().decode("utf-8")) + traces = payload.get("data") if isinstance(payload, dict) else None + return traces if isinstance(traces, list) else [] + + +def normalize_jaeger_query_url(jaeger_query_url: str | None) -> str | None: + if jaeger_query_url is None: + return None + stripped = jaeger_query_url.strip() + return stripped.rstrip("/") if stripped else None + + +def require_jaeger_query_url(jaeger_query_url: str | None) -> str: + normalized = normalize_jaeger_query_url(jaeger_query_url) + if normalized is None: + raise ValueError("jaeger_query_url is required") + return normalized + + +__all__ = [ + "JAEGER_DEFAULT_LIMIT", + "JAEGER_DEFAULT_LOOKBACK_S", + "JAEGER_DEFAULT_QUERY_URL", + "JaegerQuerySource", + "JsonlTraceSource", + "fetch_jaeger_traces", + "normalize_jaeger_query_url", + "require_jaeger_query_url", +] diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index fff3658bb0..bf8d37fbc8 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -13,6 +13,7 @@ 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 import traced_judger_endpoint from xtuner.v1.rl.utils import ( JUDGER_PAUSE_JUDGE_TASK_TIMEOUT_S, CPUActorLauncher, @@ -183,7 +184,11 @@ def __init__( self._judger_pause_event = asyncio.Event() @abstractmethod - async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: ... + async def generate_sample( + self, + rollout_state: RolloutState, + **kwargs, + ) -> RolloutState: ... async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: pending_tasks = [] @@ -204,6 +209,7 @@ async def run_judger(self, rollout_state: RolloutState) -> RolloutState: ... @overload async def run_judger(self, rollout_state: list[RolloutState]) -> list[RolloutState]: ... + @traced_judger_endpoint() async def run_judger(self, rollout_state: RolloutState | list[RolloutState]) -> RolloutState | list[RolloutState]: assert self.judger is not None if isinstance(rollout_state, list): @@ -211,32 +217,36 @@ async def run_judger(self, rollout_state: RolloutState | list[RolloutState]) -> else: judge_task = create_task(self.judger.judge(rollout_state)) pause_task = create_task(self._judger_pause_event.wait()) + result: RolloutState | list[RolloutState] = rollout_state try: done, _ = await asyncio.wait({judge_task, pause_task}, return_when=asyncio.FIRST_COMPLETED) if judge_task in done: - return await judge_task - try: - return await asyncio.wait_for( - asyncio.shield(judge_task), - timeout=JUDGER_PAUSE_JUDGE_TASK_TIMEOUT_S, - ) - except asyncio.TimeoutError: - await cancel_and_drain([judge_task]) - for sample in rollout_state if isinstance(rollout_state, list) else [rollout_state]: - sample.status = Status.ABORTED - sample.finish_reason = "abort" - sample.reward = None - return rollout_state + result = await judge_task + else: + try: + result = await asyncio.wait_for( + asyncio.shield(judge_task), + timeout=JUDGER_PAUSE_JUDGE_TASK_TIMEOUT_S, + ) + except asyncio.TimeoutError: + await cancel_and_drain([judge_task]) + for sample in rollout_state if isinstance(rollout_state, list) else [rollout_state]: + sample.status = Status.ABORTED + sample.finish_reason = "abort" + sample.reward = None + result = rollout_state except asyncio.CancelledError: await cancel_and_drain([judge_task]) for sample in rollout_state if isinstance(rollout_state, list) else [rollout_state]: sample.status = Status.ABORTED sample.finish_reason = "abort" sample.reward = None - return rollout_state + result = rollout_state finally: await cancel_and_drain([pause_task]) + return result + async def pause(self) -> None: self._judger_pause_event.set() try: diff --git a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py index 13044d4f35..2b1c0d48d4 100644 --- a/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py +++ b/xtuner/v1/rl/agent_loop/gsm8k_with_tool.py @@ -9,6 +9,9 @@ from xtuner.v1.rl.agent_loop import AgentLoop, AgentLoopConfig from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController +from xtuner.v1.rl.trace import traced_rollout_endpoint +from xtuner.v1.rl.trace.context_propagation import trace_remote +from xtuner.v1.rl.trace.trace_utils import TRACE_SPAN_AGENT_LOOP_RUN from xtuner.v1.utils import get_logger @@ -53,6 +56,7 @@ def __init__( judger=judger, enable_batch_judge=enable_batch_judge, ) + self.rollout_ctl: RolloutController = rollout_ctl self.max_turns = max_turns self.tool_call_pattern = re.compile(r"\n*(.*?)", re.DOTALL) self.tool_call_start_token: str = "" @@ -86,7 +90,12 @@ def extract_tool_calls(self, rollout_state: RolloutState) -> tuple[str, list[Fun content = self.tool_call_pattern.sub("", text) return content, function_calls - async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: + @traced_rollout_endpoint(TRACE_SPAN_AGENT_LOOP_RUN) + async def generate_sample( + self, + rollout_state: RolloutState, + **kwargs, + ) -> RolloutState: # Respect state passed from preprocess for partial rollout continuation. base_sample_params = copy.deepcopy(rollout_state.sample_params or self.sample_params) final_response_mask: list[int] = [] @@ -105,8 +114,10 @@ async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> Rollou rollout_state.sample_params = copy.deepcopy(base_sample_params) rollout_state.sample_params.max_tokens = remaining_max_tokens - 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_remote( + self.rollout_ctl.generate, # type: ignore[attr-defined] + rollout_state=rollout_state, + ) cur_turn += 1 response_ids = cast(list[int], rollout_state.response_ids) cur_turn_tokens.extend(response_ids) diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py index b70edb888a..90879e5d69 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/agent_in_localhost_loop.py @@ -18,6 +18,14 @@ from xtuner.v1.rl.rollout import RolloutController from xtuner.v1.rl.rollout.chat_template import canonicalize_messages_for_chat_template from xtuner.v1.rl.rollout.trace_store import get_store +from xtuner.v1.rl.trace import set_trace_attributes, trace_span, traced_rollout_endpoint +from xtuner.v1.rl.trace.context_propagation import instrument_aiohttp_client +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_AGENT_LOCALHOST_TRAJECTORY_MATERIALIZE, + TRACE_SPAN_AGENT_LOOP_RUN, + agent_item_initial_attributes, + rollout_state_initial_attributes, +) from xtuner.v1.rl.utils import create_task from ..agent_loop import AgentLoop, AgentLoopConfig @@ -141,6 +149,7 @@ async def generate_one(state: RolloutState) -> RolloutState: return await asyncio.gather(*tasks) + @traced_rollout_endpoint(TRACE_SPAN_AGENT_LOOP_RUN) async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: try: if self.sample_timeout_s is not None and self.sample_timeout_s > 0: @@ -202,57 +211,81 @@ async def _run_item(self, item: AgentRolloutItem) -> AgentRolloutItem: runner = _resolve_runner(item.pipeline) if runner is None: raise ValueError("AgentRolloutItem.pipeline is required.") + instrument_aiohttp_client() with ctx_session_id.set(str(item.uid)): return await runner.run(item) async def _fill_rollout_state(self, rollout_state: RolloutState, item: AgentRolloutItem) -> None: - if self.mode == "eval": - self._fill_eval_rollout_state(rollout_state, item) - return - - response_message = item.artifacts.get("response_message") or {} - rollout_state.status = Status.COMPLETED if item.status == RolloutStatus.COMPLETED else Status.FAILED - rollout_state.finish_reason = str( - response_message.get("finish_reason") or ("stop" if item.status == RolloutStatus.COMPLETED else "error") - ) - rollout_state.reward = _extract_reward_payload(item) - rollout_state.extra_fields["agent_status"] = item.status.value - rollout_state.extra_fields["agent_artifacts"] = item.artifacts - rollout_state.extra_fields["agent_judgers"] = { - name: record.model_dump(mode="json") for name, record in item.judgers.items() - } - messages, tools = _load_eval_trace_segment(item.artifacts) - if messages: - rollout_state.extra_fields["agent_messages"] = messages - rollout_state.extra_fields["agent_tools"] = tools - rollout_state.extra_fields["agent_tool_turns"] = _count_tool_turns(messages) - finish_info = response_message.get("finish_info") - if isinstance(finish_info, dict) and finish_info: - rollout_state.extra_fields["agent_finish_info"] = finish_info - if item.error is not None: - rollout_state.error_msg = f"{item.error.stage}/{item.error.category}: {item.error.message}" - if item.status != RolloutStatus.COMPLETED: - return - - segment = item.artifacts["messages"][-1] - text = self.tokenizer.apply_chat_template( - canonicalize_messages_for_chat_template(segment["messages"]), - tools=segment["tools"], - tokenize=False, - add_generation_prompt=False, - ) - prompt_text = text[:-1] if text.endswith("\n") else text - data = await get_store().export_training_trace.remote(str(rollout_state.session_id), prompt_text) - - rollout_state.input_ids = data["input_ids"] - rollout_state.labels = data["labels"] - rollout_state.response_ids = [ - token_id for token_id, label in zip(data["input_ids"][1:], data["labels"][1:]) if label != -100 - ] - rollout_state.logprobs = data["logprobs"] - rollout_state.routed_experts = data["routed_experts"] - content = response_message.get("content") - rollout_state.response = content if isinstance(content, str) else (str(content) if content is not None else "") + trace_attributes = agent_item_initial_attributes(item) + trace_attributes.update(rollout_state_initial_attributes(rollout_state)) + with trace_span( + TRACE_SPAN_AGENT_LOCALHOST_TRAJECTORY_MATERIALIZE, + attributes=trace_attributes, + ): + if self.mode == "eval": + self._fill_eval_rollout_state(rollout_state, item) + set_trace_attributes({"rollout.mode": self.mode}) + return + + response_message = item.artifacts.get("response_message") or {} + rollout_state.status = Status.COMPLETED if item.status == RolloutStatus.COMPLETED else Status.FAILED + rollout_state.finish_reason = str( + response_message.get("finish_reason") + or ("stop" if item.status == RolloutStatus.COMPLETED else "error") + ) + rollout_state.reward = _extract_reward_payload(item) + rollout_state.extra_fields["agent_status"] = item.status.value + rollout_state.extra_fields["agent_artifacts"] = item.artifacts + rollout_state.extra_fields["agent_judgers"] = { + name: record.model_dump(mode="json") for name, record in item.judgers.items() + } + messages, tools = _load_eval_trace_segment(item.artifacts) + if messages: + rollout_state.extra_fields["agent_messages"] = messages + rollout_state.extra_fields["agent_tools"] = tools + rollout_state.extra_fields["agent_tool_turns"] = _count_tool_turns(messages) + set_trace_attributes( + { + "agent.message_count": len(messages), + "agent.tool_turns": rollout_state.extra_fields["agent_tool_turns"], + } + ) + finish_info = response_message.get("finish_info") + if isinstance(finish_info, dict) and finish_info: + rollout_state.extra_fields["agent_finish_info"] = finish_info + if item.error is not None: + rollout_state.error_msg = f"{item.error.stage}/{item.error.category}: {item.error.message}" + if item.status != RolloutStatus.COMPLETED: + return + + segment = item.artifacts["messages"][-1] + text = self.tokenizer.apply_chat_template( + canonicalize_messages_for_chat_template(segment["messages"]), + tools=segment["tools"], + tokenize=False, + add_generation_prompt=False, + ) + prompt_text = text[:-1] if text.endswith("\n") else text + data = await get_store().export_training_trace.remote(str(rollout_state.session_id), prompt_text) + + rollout_state.input_ids = data["input_ids"] + rollout_state.labels = data["labels"] + rollout_state.response_ids = [ + token_id for token_id, label in zip(data["input_ids"][1:], data["labels"][1:]) if label != -100 + ] + rollout_state.logprobs = data["logprobs"] + rollout_state.routed_experts = data["routed_experts"] + content = response_message.get("content") + rollout_state.response = ( + content if isinstance(content, str) else (str(content) if content is not None else "") + ) + set_trace_attributes( + { + "training.input_tokens": len(rollout_state.input_ids or []), + "training.label_tokens": len([label for label in rollout_state.labels or [] if label != -100]), + "training.response_tokens": len(rollout_state.response_ids or []), + } + ) def _fill_eval_rollout_state(self, rollout_state: RolloutState, item: AgentRolloutItem) -> None: is_success = item.status == RolloutStatus.COMPLETED diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/judger.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/judger.py index 1979139ab1..0a96d61d8a 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/judger.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/judger.py @@ -16,6 +16,13 @@ StageStatus, ) from xtuner.v1.rl.judger.native import Judger +from xtuner.v1.rl.trace import set_trace_attributes, trace_span +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_AGENT_LOCALHOST_JUDGER_RUN, + agent_judger_initial_attributes, + failure_attributes, + reward_trace_attributes, +) class LocalhostJudgerStage: @@ -40,46 +47,54 @@ def __init__( self.weight = weight async def run(self, item: AgentRolloutItem, record: StageRecord) -> float: - record.status = StageStatus.RUNNING - record.started_at = record.started_at or time.monotonic() - try: - # reward_model stays as-is (dataset-provided ground_truth/style etc.). - # Per-rollout artifacts (response message, agent trace) flow through extra_fields. - reward_model = dict(item.reward_model or {}) - segment = item.artifacts["messages"][-1] - response_message = item.artifacts.get("response_message") or {} - content = response_message.get("content") - response = content if isinstance(content, str) else (str(content) if content is not None else "") + with trace_span( + TRACE_SPAN_AGENT_LOCALHOST_JUDGER_RUN, + attributes=agent_judger_initial_attributes(self, item, reward_key=self.reward_key), + ): + record.status = StageStatus.RUNNING + record.started_at = record.started_at or time.monotonic() + try: + # reward_model stays as-is (dataset-provided ground_truth/style etc.). + # Per-rollout artifacts (response message, agent trace) flow through extra_fields. + reward_model = dict(item.reward_model or {}) + segment = item.artifacts["messages"][-1] + response_message = item.artifacts.get("response_message") or {} + content = response_message.get("content") + response = content if isinstance(content, str) else (str(content) if content is not None else "") - rollout_state = RolloutState( - message=[{"role": "user", "content": item.instruction}], - response=response, - reward_model=reward_model, - extra_fields={ - "agent_messages": segment["messages"], - "response_message": response_message, - }, - status=Status.COMPLETED, - ) - judged = await self.judger.judge(rollout_state) - reward_payload = judged.reward or {} - if self.reward_key not in reward_payload: - raise KeyError(f"judger reward payload has no {self.reward_key!r}: {reward_payload!r}") - record.metadata["reward"] = reward_payload - record.score = float(reward_payload[self.reward_key]) - record.status = StageStatus.COMPLETED - return record.score - except Exception as exc: - record.status = StageStatus.FAILED - record.error = record.error or RolloutError( - stage=self.name, - category="judger", - type=type(exc).__name__, - message=str(exc), - ) - raise - finally: - record.finished_at = time.monotonic() + rollout_state = RolloutState( + message=[{"role": "user", "content": item.instruction}], + response=response, + reward_model=reward_model, + extra_fields={ + "agent_messages": segment["messages"], + "response_message": response_message, + }, + status=Status.COMPLETED, + ) + judged = await self.judger.judge(rollout_state) + reward_payload = judged.reward or {} + if self.reward_key not in reward_payload: + raise KeyError(f"judger reward payload has no {self.reward_key!r}: {reward_payload!r}") + record.metadata["reward"] = reward_payload + record.score = float(reward_payload[self.reward_key]) + set_trace_attributes(reward_trace_attributes(record.score)) + record.status = StageStatus.COMPLETED + return record.score + except Exception as exc: + record.status = StageStatus.FAILED + record.error = record.error or RolloutError( + stage=self.name, + category="judger", + type=type(exc).__name__, + message=str(exc), + ) + set_trace_attributes( + failure_attributes(record.error.category, message=record.error.message, type=record.error.type) + ) + raise + finally: + record.finished_at = time.monotonic() __all__ = ["LocalhostJudgerStage"] diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/runner.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/runner.py index 79a17e2266..26d7ed5d03 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/runner.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/runner.py @@ -16,6 +16,15 @@ StageRecord, StageStatus, ) +from xtuner.v1.rl.trace import set_trace_attributes, trace_span, traced_agent_item_endpoint +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN, + TRACE_SPAN_AGENT_LOCALHOST_RUNNER_RUN, + TRACE_SPAN_AGENT_LOCALHOST_VALIDATE_RUN, + agent_item_initial_attributes, + failure_attributes, + reward_trace_attributes, +) from xtuner.v1.utils import get_logger @@ -30,6 +39,7 @@ def __init__( self.infer: LocalhostStage = create_object(infer) if isinstance(infer, dict) else infer self.validate = create_object(validate) if isinstance(validate, dict) else validate + @traced_agent_item_endpoint(TRACE_SPAN_AGENT_LOCALHOST_RUNNER_RUN) async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: if not item.instruction: raise ValueError("AgentRolloutItem.instruction is required by LocalhostRunner.run") @@ -38,9 +48,22 @@ async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: t_validate: float | None = None try: - await self.infer.run(item, item.infer) - if item.infer.status != StageStatus.COMPLETED: - return self._fail(item, item.infer.error) + infer_attributes = agent_item_initial_attributes(item) + infer_attributes["stage.name"] = getattr(self.infer, "name", "infer") + with trace_span( + TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN, + attributes=infer_attributes, + ): + await self.infer.run(item, item.infer) + if item.infer.status != StageStatus.COMPLETED: + error = item.infer.error + set_trace_attributes( + failure_attributes( + error.category if error is not None else "agent_exception", + message=error.message if error is not None else None, + ) + ) + return self._fail(item, error) if self.validate is not None: t0 = time.monotonic() @@ -50,10 +73,16 @@ async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: StageRecord(), ) try: - score = float(await self.validate.run(item, validate_record)) + validate_attributes = agent_item_initial_attributes(item) + validate_attributes["stage.name"] = validate_name + with trace_span( + TRACE_SPAN_AGENT_LOCALHOST_VALIDATE_RUN, + attributes=validate_attributes, + ): + score = float(await self.validate.run(item, validate_record)) + set_trace_attributes(reward_trace_attributes(score)) except Exception: - return self._fail( - item, + error = ( validate_record.error or _first_judger_error(item) or RolloutError( @@ -61,8 +90,10 @@ async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: category="validate_failed", type=type(self.validate).__name__, message="validate failed", - ), + ) ) + set_trace_attributes(failure_attributes(error.category, message=error.message)) + return self._fail(item, error) t_validate = time.monotonic() - t0 item.reward = score @@ -80,6 +111,7 @@ async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: ) ) get_logger().error(f"[{tid}] traceback:\n{traceback.format_exc()}") + set_trace_attributes(failure_attributes(promoted.category, message=promoted.message)) return self._fail(item, promoted) finally: self._log_final(tid, item, t_validate) diff --git a/xtuner/v1/rl/agent_loop/localhost_agent_loop/stage.py b/xtuner/v1/rl/agent_loop/localhost_agent_loop/stage.py index 42c621a7db..4267860f7e 100644 --- a/xtuner/v1/rl/agent_loop/localhost_agent_loop/stage.py +++ b/xtuner/v1/rl/agent_loop/localhost_agent_loop/stage.py @@ -19,6 +19,12 @@ StageResult, StageStatus, ) +from xtuner.v1.rl.trace import set_trace_attributes, trace_span +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_AGENT_LOCALHOST_AGENT_INVOKE, + agent_item_initial_attributes, + failure_attributes, +) from xtuner.v1.utils import get_logger @@ -41,47 +47,63 @@ def __init__( self.name = name async def run(self, item: AgentRolloutItem, record: StageRecord) -> StageResult: - record.status = StageStatus.RUNNING - record.started_at = record.started_at or time.monotonic() - agent = None - try: - spec = self._pick_agent(item, record) - agent = create_object(deepcopy(_resolve_agent_config(spec.config))) - output = await agent(item.instruction) - response_message = output.model_dump(mode="json") if hasattr(output, "model_dump") else None - if response_message is None: - raise TypeError("Agent forward must return an AgentMessage-like object.") - item.artifacts["response_message"] = response_message - messages = agent.get_messages() - if not isinstance(messages, list) or not messages: - raise ValueError("Agent messages artifact must be a non-empty list.") - segment = messages[-1] - if not isinstance(segment, dict) or "messages" not in segment or "tools" not in segment: - raise ValueError("Agent messages trace segment must contain messages and tools.") - if not isinstance(segment["messages"], list): - raise TypeError("Agent messages trace segment.messages must be a list.") - item.artifacts["messages"] = messages - content = response_message.get("content") - stdout = content if isinstance(content, str) else (str(content) if content is not None else "") - result = StageResult(stdout=stdout, return_code=0) - record.entry_result = result - record.status = StageStatus.COMPLETED - return result - except Exception as exc: - record.status = StageStatus.FAILED - record.error = record.error or RolloutError( - stage=self.name, - category="agent_exception", - type=type(exc).__name__, - message=str(exc), - ) - result = StageResult(return_code=None, error=str(exc), stderr=str(exc)) - record.entry_result = result - return result - finally: - record.finished_at = time.monotonic() - if agent is not None: - await _close_agent(agent) + attributes = agent_item_initial_attributes(item) + attributes["stage.name"] = self.name + with trace_span( + TRACE_SPAN_AGENT_LOCALHOST_AGENT_INVOKE, + attributes=attributes, + ): + record.status = StageStatus.RUNNING + record.started_at = record.started_at or time.monotonic() + agent = None + try: + spec = self._pick_agent(item, record) + set_trace_attributes({"agent.name": spec.name, "agent.config": spec.config}) + agent = create_object(deepcopy(_resolve_agent_config(spec.config))) + output = await agent(item.instruction) + response_message = output.model_dump(mode="json") if hasattr(output, "model_dump") else None + if response_message is None: + raise TypeError("Agent forward must return an AgentMessage-like object.") + item.artifacts["response_message"] = response_message + messages = agent.get_messages() + if not isinstance(messages, list) or not messages: + raise ValueError("Agent messages artifact must be a non-empty list.") + segment = messages[-1] + if not isinstance(segment, dict) or "messages" not in segment or "tools" not in segment: + raise ValueError("Agent messages trace segment must contain messages and tools.") + if not isinstance(segment["messages"], list): + raise TypeError("Agent messages trace segment.messages must be a list.") + item.artifacts["messages"] = messages + set_trace_attributes( + { + "agent.message_count": len(segment["messages"]), + "agent.has_tools": bool(segment["tools"]), + } + ) + content = response_message.get("content") + stdout = content if isinstance(content, str) else (str(content) if content is not None else "") + result = StageResult(stdout=stdout, return_code=0) + record.entry_result = result + record.status = StageStatus.COMPLETED + return result + except Exception as exc: + record.status = StageStatus.FAILED + record.error = record.error or RolloutError( + stage=self.name, + category="agent_exception", + type=type(exc).__name__, + message=str(exc), + ) + set_trace_attributes( + failure_attributes(record.error.category, message=record.error.message, type=record.error.type) + ) + result = StageResult(return_code=None, error=str(exc), stderr=str(exc)) + record.entry_result = result + return result + finally: + record.finished_at = time.monotonic() + if agent is not None: + await _close_agent(agent) def _pick_agent(self, item: AgentRolloutItem, record: StageRecord) -> LocalhostAgentSpec: group_id = item.group_id or 0 diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/judger.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/judger.py index ed1d1af8a7..f90caba7b4 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/judger.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/judger.py @@ -13,6 +13,13 @@ StageRecord, StageStatus, ) +from xtuner.v1.rl.trace import set_trace_attributes, trace_span +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_AGENT_SANDBOX_JUDGER_RUN, + agent_judger_initial_attributes, + failure_attributes, + reward_trace_attributes, +) class Judger: @@ -28,29 +35,38 @@ async def run( pool: SandboxPool, record: StageRecord, ) -> float: - try: - sandbox_name = _judger_sandbox_name(self.stage, pool) - client = await pool.get(sandbox_name, record=record) - spec = pool.spec(sandbox_name) - record.sandbox_name = sandbox_name - record.sandbox_image = spec.image - record.workspace = spec.workspace_path - record.judger_name = self.name - await self.stage.run(client, item, record) - if record.status == StageStatus.FAILED: - raise RuntimeError(record.error.message if record.error is not None else "judger stage failed") - if record.score is None: - raise RuntimeError("judger stage did not produce record.score") - return float(record.score) - except Exception as exc: - record.status = StageStatus.FAILED - record.error = record.error or RolloutError( - stage=self.name, - category="judger", - type=type(exc).__name__, - message=str(exc), - ) - raise + with trace_span( + TRACE_SPAN_AGENT_SANDBOX_JUDGER_RUN, + attributes=agent_judger_initial_attributes(self, item), + ): + try: + sandbox_name = _judger_sandbox_name(self.stage, pool) + client = await pool.get(sandbox_name, record=record) + spec = pool.spec(sandbox_name) + record.sandbox_name = sandbox_name + record.sandbox_image = spec.image + record.workspace = spec.workspace_path + record.judger_name = self.name + set_trace_attributes({"sandbox.name": sandbox_name, "sandbox.image": spec.image}) + await self.stage.run(client, item, record) + if record.status == StageStatus.FAILED: + raise RuntimeError(record.error.message if record.error is not None else "judger stage failed") + if record.score is None: + raise RuntimeError("judger stage did not produce record.score") + set_trace_attributes(reward_trace_attributes(record.score)) + return float(record.score) + except Exception as exc: + record.status = StageStatus.FAILED + record.error = record.error or RolloutError( + stage=self.name, + category="judger", + type=type(exc).__name__, + message=str(exc), + ) + set_trace_attributes( + failure_attributes(record.error.category, message=record.error.message, type=record.error.type) + ) + raise def _judger_sandbox_name(stage: SandboxStage, pool: SandboxPool) -> str: diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/runner.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/runner.py index b442586cf2..6911afc676 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/runner.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/runner.py @@ -27,7 +27,16 @@ RolloutStatus, StageRecord, ) -from xtuner.v1.rl.agent_loop.sandbox_agent_loop.trace import span +from xtuner.v1.rl.trace import set_trace_attributes, trace_span, traced_agent_item_endpoint +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_AGENT_SANDBOX_ACQUIRE, + TRACE_SPAN_AGENT_SANDBOX_INFER_RUN, + TRACE_SPAN_AGENT_SANDBOX_RUNNER_RUN, + TRACE_SPAN_AGENT_SANDBOX_VALIDATE_RUN, + agent_item_initial_attributes, + failure_attributes, + reward_trace_attributes, +) from xtuner.v1.utils import get_logger @@ -49,6 +58,7 @@ def __init__( self.infer = create_object(infer) self.validate = create_object(validate) + @traced_agent_item_endpoint(TRACE_SPAN_AGENT_SANDBOX_RUNNER_RUN) async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: """Run one rollout sample and return the same item with result fields filled. @@ -78,53 +88,81 @@ async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: item.infer.workspace = infer_spec.workspace_path tid = item.id - uid_obs = str(item.uid) if item.uid is not None else "" t_acquire: float | None = None t_infer: float | None = None t_validate: float | None = None try: - with span(uid_obs, "run_total", task_id=tid) as total_span: - # ─── acquire infer sandbox ─────────────────────────────── - t0 = time.monotonic() - with span(uid_obs, "acquire", task_id=tid) as acquire_span: - infer_client = await pool.get(infer_sandbox, record=item.infer) - item.infer.sandbox_env_id = pool.env_id(infer_sandbox) - sandbox_url = pool.url(infer_sandbox) - if sandbox_url is not None: - item.infer.sandbox_url = sandbox_url - acquire_span.annotate( - sandbox_name=infer_sandbox, - sandbox_env_id=item.infer.sandbox_env_id, - sandbox_url=sandbox_url, - sandbox_image=infer_spec.image, - ) - t_acquire = time.monotonic() - t0 - - # ─── infer ────────────────────────────────────────────── - t1 = time.monotonic() - with span(uid_obs, "infer", task_id=tid) as infer_span: - infer_result = await self.infer.run(infer_client, item, item.infer) - if not infer_result.ok: - infer_span.mark_error(_format_error(item.infer.error)) - t_infer = time.monotonic() - t1 + # ─── acquire infer sandbox ─────────────────────────────── + t0 = time.monotonic() + acquire_attributes = agent_item_initial_attributes(item) + acquire_attributes["sandbox.name"] = infer_sandbox + with trace_span( + TRACE_SPAN_AGENT_SANDBOX_ACQUIRE, + attributes=acquire_attributes, + ): + infer_client = await pool.get(infer_sandbox, record=item.infer) + item.infer.sandbox_env_id = pool.env_id(infer_sandbox) + sandbox_url = pool.url(infer_sandbox) + if sandbox_url is not None: + item.infer.sandbox_url = sandbox_url + set_trace_attributes( + { + "sandbox.name": infer_sandbox, + "sandbox.env_id": item.infer.sandbox_env_id, + "sandbox.url": sandbox_url, + "sandbox.image": infer_spec.image, + } + ) + t_acquire = time.monotonic() - t0 + + # ─── infer ────────────────────────────────────────────── + t1 = time.monotonic() + infer_attributes = agent_item_initial_attributes(item) + infer_attributes.update({"stage.name": "infer", "sandbox.name": infer_sandbox}) + with trace_span( + TRACE_SPAN_AGENT_SANDBOX_INFER_RUN, + attributes=infer_attributes, + ): + infer_result = await self.infer.run(infer_client, item, item.infer) if not infer_result.ok: - total_span.mark_error(_format_error(item.infer.error)) - return self._fail(item, item.infer.error) - - # ─── validate ─────────────────────────────────────────── - t2 = time.monotonic() - with span(uid_obs, "validate", task_id=tid): - validate_name = self.validate.name - validate_record = item.judgers.setdefault( - validate_name, - StageRecord(judger_name=validate_name), + error = item.infer.error + set_trace_attributes( + failure_attributes( + error.category if error is not None else "runner_exception", + message=_format_error(error), + ) ) - score = float(await self.validate.run(item, pool, validate_record)) - t_validate = time.monotonic() - t2 - item.reward = score + t_infer = time.monotonic() - t1 + if not infer_result.ok: + error = item.infer.error + set_trace_attributes( + failure_attributes( + error.category if error is not None else "runner_exception", + message=_format_error(error), + ) + ) + return self._fail(item, error) + + # ─── validate ─────────────────────────────────────────── + t2 = time.monotonic() + validate_name = self.validate.name + validate_record = item.judgers.setdefault( + validate_name, + StageRecord(judger_name=validate_name), + ) + validate_attributes = agent_item_initial_attributes(item) + validate_attributes["stage.name"] = validate_name + with trace_span( + TRACE_SPAN_AGENT_SANDBOX_VALIDATE_RUN, + attributes=validate_attributes, + ): + score = float(await self.validate.run(item, pool, validate_record)) + set_trace_attributes(reward_trace_attributes(score)) + t_validate = time.monotonic() - t2 + item.reward = score - item.status = RolloutStatus.COMPLETED - return item + item.status = RolloutStatus.COMPLETED + return item except Exception as exc: promoted = ( item.infer.error @@ -137,6 +175,7 @@ async def run(self, item: AgentRolloutItem) -> AgentRolloutItem: ) ) get_logger().error(f"[{tid}] traceback:\n{traceback.format_exc()}") + set_trace_attributes(failure_attributes(promoted.category, message=_format_error(promoted))) return self._fail(item, promoted) finally: self._log_final(tid, item, t_acquire, t_infer, t_validate) diff --git a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py index a3fd579c13..cfc4bb1072 100644 --- a/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py +++ b/xtuner/v1/rl/agent_loop/sandbox_agent_loop/sandbox.py @@ -47,7 +47,15 @@ async def hook(client, item, record) -> None StageResult, StageStatus, ) -from xtuner.v1.rl.agent_loop.sandbox_agent_loop.trace import span +from xtuner.v1.rl.trace import set_trace_attributes, trace_span +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_AGENT_SANDBOX_ENTRY_RUN, + TRACE_SPAN_AGENT_SANDBOX_HOOK_RUN, + agent_item_initial_attributes, + sandbox_entry_attributes, + sandbox_entry_error_attributes, + sandbox_entry_result_attributes, +) from xtuner.v1.utils import get_logger @@ -237,25 +245,35 @@ async def run( record.entries.append(entry) entry.status = StageStatus.RUNNING entry.started_at = time.monotonic() - uid_obs = str(item.uid) if item.uid is not None else "" - with span(uid_obs, f"entry:{self.name}", entry_kind="ShellEntry"): - try: - outcome = await self._execute(client, self.env) - if not outcome.ok and self.failure is not None: - outcome = await self.failure.handle(client, item, record, entry, outcome) - self._finish_record(entry, outcome) - return outcome - except Exception as exc: - error = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip() - outcome = EntryOutcome( - source=type(self).__name__, - reason="exception", - result=StageResult(return_code=None, stderr=error, error=error), - ) - if self.failure is not None: - outcome = await self.failure.handle(client, item, record, entry, outcome) - self._finish_record(entry, outcome, exc=exc) - raise + trace_attributes = agent_item_initial_attributes(item) + trace_attributes.update(sandbox_entry_attributes(self)) + with trace_span(TRACE_SPAN_AGENT_SANDBOX_ENTRY_RUN, attributes=trace_attributes): + + async def run_entry() -> EntryOutcome: + try: + outcome = await self._execute(client, self.env) + if not outcome.ok and self.failure is not None: + outcome = await self.failure.handle(client, item, record, entry, outcome) + self._finish_record(entry, outcome) + return outcome + except Exception as exc: + error = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip() + outcome = EntryOutcome( + source=type(self).__name__, + reason="exception", + result=StageResult(return_code=None, stderr=error, error=error), + ) + if self.failure is not None: + outcome = await self.failure.handle(client, item, record, entry, outcome) + self._finish_record(entry, outcome, exc=exc) + raise + + outcome = await run_entry() + set_trace_attributes(sandbox_entry_result_attributes(outcome)) + error_attributes = sandbox_entry_error_attributes(outcome) + if error_attributes: + set_trace_attributes(error_attributes) + return outcome async def _execute(self, client: Any, env: dict[str, str]) -> EntryOutcome: exec_res = await exec_in( @@ -363,26 +381,36 @@ async def run( record.entries.append(entry) entry.status = StageStatus.RUNNING entry.started_at = time.monotonic() - uid_obs = str(item.uid) if item.uid is not None else "" - with span(uid_obs, f"entry:{self.name}", entry_kind="DetachedShellEntry"): - try: - outcome = await self._run_detached(client, item, entry, self.env) - await self._fill_output_files(client, entry, outcome.result) - if not outcome.ok and self.failure is not None: - outcome = await self.failure.handle(client, item, record, entry, outcome) - self._finish_record(entry, outcome) - return outcome - except Exception as exc: - error = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip() - outcome = EntryOutcome( - source=type(self).__name__, - reason="exception", - result=StageResult(return_code=None, stderr=error, error=error), - ) - if self.failure is not None: - outcome = await self.failure.handle(client, item, record, entry, outcome) - self._finish_record(entry, outcome, exc=exc) - raise + trace_attributes = agent_item_initial_attributes(item) + trace_attributes.update(sandbox_entry_attributes(self)) + with trace_span(TRACE_SPAN_AGENT_SANDBOX_ENTRY_RUN, attributes=trace_attributes): + + async def run_entry() -> EntryOutcome: + try: + outcome = await self._run_detached(client, item, entry, self.env) + await self._fill_output_files(client, entry, outcome.result) + if not outcome.ok and self.failure is not None: + outcome = await self.failure.handle(client, item, record, entry, outcome) + self._finish_record(entry, outcome) + return outcome + except Exception as exc: + error = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)).rstrip() + outcome = EntryOutcome( + source=type(self).__name__, + reason="exception", + result=StageResult(return_code=None, stderr=error, error=error), + ) + if self.failure is not None: + outcome = await self.failure.handle(client, item, record, entry, outcome) + self._finish_record(entry, outcome, exc=exc) + raise + + outcome = await run_entry() + set_trace_attributes(sandbox_entry_result_attributes(outcome)) + error_attributes = sandbox_entry_error_attributes(outcome) + if error_attributes: + set_trace_attributes(error_attributes) + return outcome def _new_record(self) -> EntryRecord: suffix = uuid.uuid4().hex[:12] @@ -798,7 +826,15 @@ async def warn_if_stuck() -> None: warn_task = asyncio.create_task(warn_if_stuck()) if stuck_warn_sec > 0 else None try: - await hook(client, item, record) + trace_attributes = agent_item_initial_attributes(item) + trace_attributes.update( + {"stage.name": record.judger_name or "infer", "hook.phase": phase, "hook.name": hook_name} + ) + with trace_span( + TRACE_SPAN_AGENT_SANDBOX_HOOK_RUN, + attributes=trace_attributes, + ): + await hook(client, item, record) finally: done = True if warn_task is not None: diff --git a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py index 86e1539da8..e848936203 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,9 @@ 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 import traced_rollout_endpoint +from xtuner.v1.rl.trace.context_propagation import trace_remote +from xtuner.v1.rl.trace.trace_utils import TRACE_SPAN_AGENT_LOOP_RUN from .agent_loop import AgentLoop, AgentLoopConfig @@ -63,7 +66,9 @@ def __init__( logger=logger, enable_batch_judge=enable_batch_judge, ) + self.rollout_ctl: RolloutController = rollout_ctl + @traced_rollout_endpoint(TRACE_SPAN_AGENT_LOOP_RUN) async def generate_sample( self, rollout_state: RolloutState, @@ -73,8 +78,10 @@ async def generate_sample( rollout_state.tokens = rollout_state.prompt_ids # 推理引擎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_remote( + self.rollout_ctl.generate, # type: ignore[attr-defined] + rollout_state=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 33f1e115b2..2df92e521e 100644 --- a/xtuner/v1/rl/agent_loop_manager/produce_utils.py +++ b/xtuner/v1/rl/agent_loop_manager/produce_utils.py @@ -135,7 +135,14 @@ async def expired_count(self) -> int: async def sample_group(self, *, from_expired_pool: bool) -> list[RolloutState]: group_status = [Status.EXPIRED, Status.ABORTED] if from_expired_pool else [Status.ABORTED] - return await self.sampler.sample(task_name=self.task_name, group_status=group_status) + group = await self.sampler.sample(task_name=self.task_name, group_status=group_status) + for item in group: + extra_fields = getattr(item, "extra_fields", None) + if extra_fields is None: + extra_fields = {} + setattr(item, "extra_fields", extra_fields) + extra_fields["producer_future_step"] = self.train_step + return group async def generate_group( self, diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index 5620be7772..73ede95a45 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -268,6 +268,12 @@ async def produce_batch(self, ctx: ProduceContext) -> None: for _ in range(ctx.task_batch_size): rollout_state = await ctx.sampler.sample(task_name=ctx.task_name) + for item in rollout_state: + extra_fields = getattr(item, "extra_fields", None) + if extra_fields is None: + extra_fields = {} + setattr(item, "extra_fields", extra_fields) + extra_fields["producer_future_step"] = ctx.train_step task = create_task(ctx.generate_group(rollout_state)) pending_tasks.add(task) @@ -301,6 +307,12 @@ async def produce_batch(self, ctx: ProduceContext) -> None: completed_sample_count, ctx.task_batch_size ): rollout_state = await ctx.sampler.sample(task_name=ctx.task_name) + for item in rollout_state: + extra_fields = getattr(item, "extra_fields", None) + if extra_fields is None: + extra_fields = {} + setattr(item, "extra_fields", extra_fields) + extra_fields["producer_future_step"] = ctx.train_step task = create_task(ctx.generate_group(rollout_state)) pending_tasks.add(task) progress_displayer.close() diff --git a/xtuner/v1/rl/replay_buffer.py b/xtuner/v1/rl/replay_buffer.py index d5f9f1336a..80034fa0c0 100644 --- a/xtuner/v1/rl/replay_buffer.py +++ b/xtuner/v1/rl/replay_buffer.py @@ -464,6 +464,8 @@ async def put( if model_step is not None: for item in items: update_sample_version(item, model_step) + for item in items: + item.task_name = task_name if current_train_step is not None: refresh_seq_staleness(items, current_train_step) if stale_threshold is not None: diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index ae4572334c..03aef39a71 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -7,6 +7,9 @@ from ray.util.placement_group import PlacementGroup from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.trace import traced_rollout_endpoint +from xtuner.v1.rl.trace.context_propagation import trace_remote +from xtuner.v1.rl.trace.trace_utils import TRACE_SPAN_ROLLOUT_CONTROLLER_GENERATE from xtuner.v1.rl.utils import AutoAcceleratorWorkers from xtuner.v1.utils import XTUNER_DETERMINISTIC, get_logger @@ -89,6 +92,7 @@ def validate_registered_workers_to_proxy(self) -> None: self.proxy_manager.validate_registered_session_urls() @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) + @traced_rollout_endpoint(TRACE_SPAN_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 +108,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_remote( + worker.generate, # type: ignore[attr-defined] + rollout_state=rollout_state, + ) try: response_rollout_state = await asyncio.wait_for( response_ref, @@ -197,11 +204,12 @@ 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." ) - return ray.remote( - concurrency_groups={ + ray_kwargs = { + "concurrency_groups": { ROLLOUT_CONCURRENCY_GROUP_GENERATE: ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY, - }, - )(worker_base_cls) + } + } + return ray.remote(**ray_kwargs)(worker_base_cls) def _init_workers(self, placement_group: PlacementGroup) -> RolloutWorkerRegistry: """Initializes and configures the pool of RolloutWorker actors. diff --git a/xtuner/v1/rl/rollout/session_server.py b/xtuner/v1/rl/rollout/session_server.py index 666cfc740b..5f90282772 100644 --- a/xtuner/v1/rl/rollout/session_server.py +++ b/xtuner/v1/rl/rollout/session_server.py @@ -1,4 +1,6 @@ import json +import time +from dataclasses import dataclass from functools import reduce from operator import add from typing import Any, Optional @@ -8,6 +10,23 @@ from aiohttp import ClientConnectionResetError, ClientSession, ClientTimeout, web from transformers import AutoTokenizer +from xtuner.v1.rl.trace import ( + set_trace_attributes, + trace_span, +) +from xtuner.v1.rl.trace.context_propagation import ( + extract_trace_attributes_from_mapping, + extract_trace_carrier_from_mapping, + remove_trace_carrier_from_mapping, + traced_aiohttp_request, +) +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_SESSION_SERVER_PREPARE_REQUEST, + TRACE_SPAN_SESSION_SERVER_READ_RESPONSE, + TRACE_SPAN_SESSION_SERVER_RECORD_RESPONSE, + TRACE_SPAN_SESSION_SERVER_REQUEST, + TRACE_SPAN_SESSION_SERVER_SEND_REQUEST, +) from xtuner.v1.utils import get_logger from .chat_template import canonicalize_messages_for_chat_template @@ -73,7 +92,33 @@ def _extract_output_logprobs(choice: dict, output_token_ids: list[int]) -> list[ return [item[0] for item in output_token_logprobs] -_SESSION_SERVER_ONLY_KEYS = {"session_id"} +def _worker_payload_from_session_request(req_body: dict) -> dict: + worker_req = dict(req_body) + worker_req.pop("session_id", None) + remove_trace_carrier_from_mapping(worker_req) + return worker_req + + +def _request_body_json(request_body: bytes) -> dict[str, Any] | None: + if not request_body: + return None + try: + request_data = json.loads(request_body) + except json.JSONDecodeError: + return None + return request_data if isinstance(request_data, dict) else None + + +def _request_trace_context( + headers: dict[str, str], + request_data: dict[str, Any] | None, +) -> tuple[dict[str, str], str]: + if any(key.lower() == "traceparent" for key in headers): + return headers, "headers" + body_carrier = extract_trace_carrier_from_mapping(request_data) + if body_carrier: + return body_carrier, "body" + return headers, "none" def _bool_request_value(value: Any, default: bool = False) -> bool: @@ -88,6 +133,38 @@ def _request_uses_trace_store(req_body: dict) -> bool: return _bool_request_value(req_body.get("return_token_ids"), True) +@dataclass(slots=True) +class _PreparedSessionRequest: + body: bytes + data: dict | None + session_id: Any | None + messages: Any | None + tools: Any | None + trace_attributes: dict[str, Any] + trace_enabled: bool + orig_return_logprob: bool + orig_return_token_ids: bool + orig_return_routed_experts: bool + + @property + def is_stream(self) -> bool: + return self.data.get("stream", False) if self.data else False + + +@dataclass(slots=True) +class _ForwardedSessionResponse: + response: web.StreamResponse + raw_response: bytes + is_stream: bool + + +class _PrepareRequestError(RuntimeError): + def __init__(self, message: str, *, status: int = 500): + super().__init__(message) + self.message = message + self.status = status + + class SessionServer: """SessionServer intercepts and records requests sent to a remote LLM API worker. @@ -136,7 +213,7 @@ async def on_request(self, req_body: dict, *, trace_enabled: bool = True) -> dic """Hook for processing/modifying the request before forwarding.""" if not trace_enabled: - worker_req = {k: v for k, v in req_body.items() if k not in _SESSION_SERVER_ONLY_KEYS} + worker_req = _worker_payload_from_session_request(req_body) if "logprobs" in worker_req: worker_req.setdefault("return_logprob", worker_req.pop("logprobs")) if not _bool_request_value(worker_req.get("return_logprob"), False): @@ -166,19 +243,19 @@ async def on_request(self, req_body: dict, *, trace_enabled: bool = True) -> dic input_ids = reduce(add, [node.value.token_ids for node in nodes] + [delta_ids]) # 3. 组装 OpenAI chat completions 请求。 - worker_req = { - **{ - k: v - for k, v in req_body.items() - if k not in _SESSION_SERVER_ONLY_KEYS | {"messages", "logprobs", "top_logprobs"} - }, - "messages": [], - "input_ids": input_ids, - "return_token_ids": True, - "return_routed_experts": True, - "return_logprob": True, - "include_stop_str_in_output": True, - } + worker_req = _worker_payload_from_session_request(req_body) + for key in ("messages", "logprobs", "top_logprobs"): + worker_req.pop(key, None) + worker_req.update( + { + "messages": [], + "input_ids": input_ids, + "return_token_ids": True, + "return_routed_experts": True, + "return_logprob": True, + "include_stop_str_in_output": True, + } + ) return worker_req async def on_response(self, worker_resp: dict, *, trace_enabled: bool = True) -> dict: @@ -292,214 +369,367 @@ async def stop(self): self._app = None get_logger().info("SessionServer stopped.") - async def _handle_request(self, request: web.Request) -> web.Response: + async def _handle_request(self, request: web.Request) -> web.StreamResponse: """Proxy handler for the worker API.""" - # Read the request body + headers = dict(request.headers) request_body = await request.read() - request_data = session_id = messages = None + request_data = _request_body_json(request_body) + parent_carrier, trace_context_source = _request_trace_context(headers, request_data) + trace_attributes = extract_trace_attributes_from_mapping(request_data) + with trace_span( + TRACE_SPAN_SESSION_SERVER_REQUEST, + attributes=trace_attributes, + parent_carrier=parent_carrier, + ): + try: + prepared = await self._prepare_request( + request, + request_body=request_body, + request_data=request_data, + trace_context_source=trace_context_source, + trace_attributes=trace_attributes, + ) + except _PrepareRequestError as exc: + return web.json_response( + _lmdeploy_error_payload(exc.message, status=exc.status), + status=exc.status, + ) + forwarded = await self._forward_request(request, prepared) + skip_done, error_msg = await self._record_response(prepared, forwarded) + return await self._finalize_response(forwarded, skip_done=skip_done, error_msg=error_msg) + + async def _prepare_request( + self, + request: web.Request, + *, + request_body: bytes | None = None, + request_data: dict[str, Any] | None = None, + trace_context_source: str | None = None, + trace_attributes: dict[str, Any] | None = None, + ) -> _PreparedSessionRequest: + if request_body is None: + request_body = await request.read() + if request_data is None: + request_data = _request_body_json(request_body) + if trace_context_source is None: + _, trace_context_source = _request_trace_context(dict(request.headers), request_data) + if trace_attributes is None: + trace_attributes = extract_trace_attributes_from_mapping(request_data) + session_id = messages = tools = None trace_enabled = False orig_return_logprob = orig_return_token_ids = orig_return_routed_experts = False - if request_body: - try: - request_data = json.loads(request_body) + prepare_trace_attributes = dict(trace_attributes) + prepare_trace_attributes.update( + { + "session.request_bytes": len(request_body), + "session.trace_context_source": trace_context_source, + } + ) - trace_enabled = _request_uses_trace_store(request_data) - orig_return_logprob = _bool_request_value( - request_data.get("return_logprob", request_data.get("logprobs")), False - ) - orig_return_token_ids = _bool_request_value(request_data.get("return_token_ids"), False) - orig_return_routed_experts = _bool_request_value(request_data.get("return_routed_experts"), True) + with trace_span( + TRACE_SPAN_SESSION_SERVER_PREPARE_REQUEST, + attributes=prepare_trace_attributes, + ): + if request_body: + try: + if request_data is None: + request_data = json.loads(request_body) - session_id = request_data.get("session_id") - messages = request_data.get("messages") - tools = request_data.get("tools", None) + trace_enabled = _request_uses_trace_store(request_data) + orig_return_logprob = _bool_request_value( + request_data.get("return_logprob", request_data.get("logprobs")), False + ) + orig_return_token_ids = _bool_request_value(request_data.get("return_token_ids"), False) + orig_return_routed_experts = _bool_request_value(request_data.get("return_routed_experts"), True) - # Apply purely abstract on_request processing - request_data = await self.on_request(request_data, trace_enabled=trace_enabled) - # Re-serialize the modified payload back to bytes - request_body = json.dumps(request_data).encode("utf-8") - except json.JSONDecodeError: - pass - except Exception as exc: - message = f"SessionServer request hook failed: {type(exc).__name__}: {exc}" - get_logger().error(message) - return web.json_response(_lmdeploy_error_payload(message), status=500) + session_id = request_data["session_id"] + set_trace_attributes({"xtuner.session_id": str(session_id)}) + messages = request_data.get("messages") + tools = request_data.get("tools", None) + + request_data = await self.on_request( + request_data, + trace_enabled=trace_enabled, + ) + remove_trace_carrier_from_mapping(request_data) + set_trace_attributes({"session.stream": bool(request_data.get("stream", False))}) + request_body = json.dumps(request_data).encode("utf-8") + except json.JSONDecodeError: + pass + except Exception as exc: + message = f"SessionServer request hook failed: {type(exc).__name__}: {exc}" + get_logger().error(message) + raise _PrepareRequestError(message) from exc + + return _PreparedSessionRequest( + body=request_body, + data=request_data, + session_id=session_id, + messages=messages, + tools=tools, + trace_attributes=trace_attributes, + trace_enabled=trace_enabled, + orig_return_logprob=orig_return_logprob, + orig_return_token_ids=orig_return_token_ids, + orig_return_routed_experts=orig_return_routed_experts, + ) - # Build forwarding headers, dropping original Host + async def _forward_request( + self, + request: web.Request, + prepared: _PreparedSessionRequest, + ) -> _ForwardedSessionResponse: forward_headers = dict(request.headers) forward_headers.pop("Host", None) forward_headers.pop("host", None) forward_headers.pop("Content-Length", None) forward_headers.pop("content-length", None) - # Re-build Path req_path = request.match_info["path"] target_url = f"{self.worker_base_url}/{req_path.lstrip('/')}" if request.query_string: target_url += f"?{request.query_string}" - is_stream = request_data.get("stream", False) if request_data else False - - def _clean_data(data: dict) -> bool: - modified = False - for key, drop in [ - ("output_token_logprobs", not orig_return_logprob), - ("output_ids", not orig_return_token_ids), - ("routed_experts", not orig_return_routed_experts), - ]: - if drop and key in data: - data.pop(key) - modified = True - if drop: - for c in data.get("choices", []): - if key in c: - c.pop(key) - modified = True - - for c in data.get("choices", []): - if "logprobs" in c: - c.pop("logprobs") - modified = True - - for c in data.get("choices", []): - if c.get("message") and isinstance(c["message"].get("content"), str): - if self.stop_word in c["message"]["content"]: - c["message"]["content"] = c["message"]["content"].replace(self.stop_word, "") - modified = True - if c.get("delta") and isinstance(c["delta"].get("content"), str): - if self.stop_word in c["delta"]["content"]: - c["delta"]["content"] = c["delta"]["content"].replace(self.stop_word, "") - modified = True - - return modified + trace_attributes: dict[str, Any] = dict(prepared.trace_attributes) + trace_attributes.update( + { + "target_url": target_url, + "session.stream": bool(prepared.is_stream), + } + ) + if prepared.session_id is not None: + trace_attributes.setdefault("xtuner.session_id", str(prepared.session_id)) - # Forward the request to the upstream worker - # read_bufsize controls StreamReader's line buffer limit; SSE events with large - # tool_calls/reasoning_content payloads can exceed the 64KB default and trigger - # "Chunk too big" from readuntil(b"\n"). timeout = ClientTimeout(total=self.request_timeout, sock_connect=30) async with ClientSession(read_bufsize=self.read_bufsize, timeout=timeout) as client: - async with client.request( - method=request.method, url=target_url, headers=forward_headers, data=request_body + async with traced_aiohttp_request( + client, + span_name=TRACE_SPAN_SESSION_SERVER_SEND_REQUEST, + method=request.method, + url=target_url, + headers=forward_headers, + data=prepared.body, + attributes=trace_attributes, ) as resp: - # Setup proper stream vs sync response objects - if is_stream: - response_chunks = [] - response = web.StreamResponse( - status=resp.status, - headers={ - k: v - for k, v in resp.headers.items() - if k.lower() not in ("transfer-encoding", "content-length", "content-encoding") - }, - ) - await response.prepare(request) - # If the downstream client closes the socket mid-stream - # (e.g. AsyncAPIClient bails out on a finish_reason=='error' - # chunk after the prompt overflowed the session window), - # keep draining the upstream so the trace is still recorded - # in full but stop attempting to write to the closed socket. - client_alive = True - async for line in resp.content: - # Keep unmodified line for trace store parsing - if trace_enabled: - response_chunks.append(line) - - # Dynamically prune added fields before writing to client - if request_data is not None and line.startswith(b"data: ") and line.strip() != b"data: [DONE]": - try: - text = line.decode("utf-8") - data = json.loads(text[6:]) - if _clean_data(data): - line = ("data: " + json.dumps(data) + "\n").encode("utf-8") - except Exception: - pass - - # Delay [DONE] only while a training trace still needs to be exported. - if client_alive and (not trace_enabled or line.strip() != b"data: [DONE]"): - try: - await response.write(line) - except (ConnectionError, ClientConnectionResetError): - client_alive = False - - raw_response = b"".join(response_chunks) if trace_enabled else b"" - else: - raw_response = await resp.read() - final_raw_response = raw_response + return await self._read_upstream_response(request, resp, prepared) - if request_data is not None: + async def _read_upstream_response( + self, + request: web.Request, + resp: Any, + prepared: _PreparedSessionRequest, + ) -> _ForwardedSessionResponse: + is_stream = prepared.is_stream + if is_stream: + response_chunks = [] + response = web.StreamResponse( + status=resp.status, + headers={ + k: v + for k, v in resp.headers.items() + if k.lower() not in ("transfer-encoding", "content-length", "content-encoding") + }, + ) + await response.prepare(request) + client_alive = True + response_started_at = time.monotonic() + chunk_count = 0 + stream_trace_attributes: dict[str, Any] = dict(prepared.trace_attributes) + stream_trace_attributes["session.stream"] = True + if prepared.session_id is not None: + stream_trace_attributes.setdefault("xtuner.session_id", str(prepared.session_id)) + with trace_span( + TRACE_SPAN_SESSION_SERVER_READ_RESPONSE, + attributes=stream_trace_attributes, + ): + async for line in resp.content: + chunk_count += 1 + if chunk_count == 1: + set_trace_attributes( + {"session.first_token_ms": int((time.monotonic() - response_started_at) * 1000)} + ) + + if prepared.trace_enabled: + response_chunks.append(line) + + if prepared.data is not None and line.startswith(b"data: ") and line.strip() != b"data: [DONE]": try: - clean_data = json.loads(raw_response) - if _clean_data(clean_data): - final_raw_response = json.dumps(clean_data).encode("utf-8") + text = line.decode("utf-8") + data = json.loads(text[6:]) + if self._clean_response_data(data, prepared): + line = ("data: " + json.dumps(data) + "\n").encode("utf-8") except Exception: pass - response = web.Response( - status=resp.status, - headers={ - k: v - for k, v in resp.headers.items() - if k.lower() not in ("transfer-encoding", "content-length", "content-encoding") - }, - body=final_raw_response, # Modified raw response without our injected trace params - ) + if client_alive and (not prepared.trace_enabled or line.strip() != b"data: [DONE]"): + try: + await response.write(line) + except (ConnectionError, ClientConnectionResetError): + client_alive = False + set_trace_attributes( + { + "session.chunks": chunk_count, + "session.response_total_ms": int((time.monotonic() - response_started_at) * 1000), + } + ) + + raw_response = b"".join(response_chunks) if prepared.trace_enabled else b"" + return _ForwardedSessionResponse(response=response, raw_response=raw_response, is_stream=True) + + response_started_at = time.monotonic() + response_trace_attributes: dict[str, Any] = dict(prepared.trace_attributes) + response_trace_attributes["session.stream"] = False + if prepared.session_id is not None: + response_trace_attributes.setdefault("xtuner.session_id", str(prepared.session_id)) + with trace_span( + TRACE_SPAN_SESSION_SERVER_READ_RESPONSE, + attributes=response_trace_attributes, + ): + raw_response = await resp.read() + set_trace_attributes( + { + "session.response_bytes": len(raw_response), + "session.response_total_ms": int((time.monotonic() - response_started_at) * 1000), + } + ) + final_raw_response = raw_response - # Apply abstract on_response processing + if prepared.data is not None: + try: + clean_data = json.loads(raw_response) + if self._clean_response_data(clean_data, prepared): + final_raw_response = json.dumps(clean_data).encode("utf-8") + except Exception: + pass + + response = web.Response( + status=resp.status, + headers={ + k: v + for k, v in resp.headers.items() + if k.lower() not in ("transfer-encoding", "content-length", "content-encoding") + }, + body=final_raw_response, + ) + return _ForwardedSessionResponse(response=response, raw_response=raw_response, is_stream=False) + + async def _record_response( + self, + prepared: _PreparedSessionRequest, + forwarded: _ForwardedSessionResponse, + ) -> tuple[bool, str | None]: response_data = None - skip_done = bool(is_stream and not trace_enabled) + skip_done = bool(forwarded.is_stream and not prepared.trace_enabled) session_error_msg = None - if request_data and trace_enabled: - if is_stream: - skip_done = not _stream_has_traceable_choices(raw_response) - if not skip_done: + + trace_attributes: dict[str, Any] = dict(prepared.trace_attributes) + trace_attributes.update( + { + "session.stream": bool(forwarded.is_stream), + "session.response_bytes": len(forwarded.raw_response) if forwarded.raw_response else 0, + } + ) + if prepared.session_id is not None: + trace_attributes.setdefault("xtuner.session_id", str(prepared.session_id)) + with trace_span( + TRACE_SPAN_SESSION_SERVER_RECORD_RESPONSE, + attributes=trace_attributes, + ): + if prepared.data and prepared.trace_enabled: + if forwarded.is_stream: + skip_done = not _stream_has_traceable_choices(forwarded.raw_response) + if not skip_done: + try: + response_data = self._parse_stream_response(forwarded.raw_response) + except Exception as exc: + session_error_msg = f"SessionServer stream trace failed: {type(exc).__name__}: {exc}" + else: try: - response_data = self._parse_stream_response(raw_response) - except Exception as exc: - session_error_msg = f"SessionServer stream trace failed: {type(exc).__name__}: {exc}" - else: - try: - response_data = json.loads(raw_response) - except json.JSONDecodeError: - pass - if isinstance(response_data, dict) and _is_error_payload(response_data): - response_data = None + response_data = json.loads(forwarded.raw_response) + except json.JSONDecodeError: + pass + if isinstance(response_data, dict) and _is_error_payload(response_data): + response_data = None - if response_data is not None: - try: - for c in response_data.get("choices", []): - if c.get("message") and isinstance(c["message"].get("content"), str): - c["message"]["content"] = c["message"]["content"].replace(self.stop_word, "") - - response_data["session_id"] = session_id - response_data["messages"] = messages - response_data["tools"] = tools - await self.on_response(response_data, trace_enabled=trace_enabled) - except Exception as exc: - session_error_msg = f"SessionServer response hook failed: {type(exc).__name__}: {exc}" + if response_data is not None: + try: + for c in response_data.get("choices", []): + if c.get("message") and isinstance(c["message"].get("content"), str): + c["message"]["content"] = c["message"]["content"].replace(self.stop_word, "") + + response_data["session_id"] = prepared.session_id + response_data["messages"] = prepared.messages + response_data["tools"] = prepared.tools + await self.on_response( + response_data, + trace_enabled=prepared.trace_enabled, + ) + except Exception as exc: + session_error_msg = f"SessionServer response hook failed: {type(exc).__name__}: {exc}" if session_error_msg: get_logger().error(session_error_msg) - if is_stream: + return skip_done, session_error_msg + + async def _finalize_response( + self, + forwarded: _ForwardedSessionResponse, + *, + skip_done: bool, + error_msg: str | None, + ) -> web.StreamResponse: + if forwarded.is_stream: try: - if session_error_msg: - error_payload = _lmdeploy_error_payload(session_error_msg) - await response.write( + if error_msg: + error_payload = _lmdeploy_error_payload(error_msg) + await forwarded.response.write( ("data: " + json.dumps(error_payload, ensure_ascii=False) + "\n\n").encode("utf-8") ) skip_done = True if not skip_done: - await response.write(b"data: [DONE]\n\n") - await response.write_eof() + await forwarded.response.write(b"data: [DONE]\n\n") + await forwarded.response.write_eof() except (ConnectionError, ClientConnectionResetError): - # Client already gone; trace was still recorded above. pass - elif session_error_msg: - return web.json_response(_lmdeploy_error_payload(session_error_msg), status=500) + elif error_msg: + return web.json_response(_lmdeploy_error_payload(error_msg), status=500) + + return forwarded.response + + def _clean_response_data(self, data: dict, prepared: _PreparedSessionRequest) -> bool: + modified = False + for key, drop in [ + ("output_token_logprobs", not prepared.orig_return_logprob), + ("output_ids", not prepared.orig_return_token_ids), + ("routed_experts", not prepared.orig_return_routed_experts), + ]: + if drop and key in data: + data.pop(key) + modified = True + if drop: + for c in data.get("choices", []): + if key in c: + c.pop(key) + modified = True + + for c in data.get("choices", []): + if "logprobs" in c: + c.pop("logprobs") + modified = True + + for c in data.get("choices", []): + if c.get("message") and isinstance(c["message"].get("content"), str): + if self.stop_word in c["message"]["content"]: + c["message"]["content"] = c["message"]["content"].replace(self.stop_word, "") + modified = True + if c.get("delta") and isinstance(c["delta"].get("content"), str): + if self.stop_word in c["delta"]["content"]: + c["delta"]["content"] = c["delta"]["content"].replace(self.stop_word, "") + modified = True - return response + return modified async def _decode_routed_experts(self, routed_experts: Any) -> np.ndarray: if isinstance(routed_experts, str): diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index c50ffa6bfb..d97b1e8afd 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -28,12 +28,18 @@ reset_rollout_response, update_status_from_finish_reason, ) +from xtuner.v1.rl.trace import inject_trace_context, trace_span, traced_rollout_endpoint +from xtuner.v1.rl.trace.trace_utils import ( + TRACE_SPAN_ROLLOUT_WORKER_GENERATE, + rollout_state_final_attributes, +) from xtuner.v1.rl.utils import ( AutoAcceleratorWorkers, CPUResourcesConfig, SingleAcceleratorWorker, get_eos_token, register_cpu_resources, + with_trace_runtime_env, ) from xtuner.v1.utils import get_logger from xtuner.v1.utils.httpx_utils import HttpRequestErrorType, HttpRequestResult @@ -531,13 +537,14 @@ def build(self, placement_group: "PlacementGroup"): name="rollout_controller", cpu_resources=CPUResourcesConfig(num_workers=num_workers), ) + ray_kwargs = { + "concurrency_groups": { + ROLLOUT_CONCURRENCY_GROUP_GENERATE: ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY, + }, + } return ( - ray.remote( - concurrency_groups={ - ROLLOUT_CONCURRENCY_GROUP_GENERATE: ROLLOUT_RAY_GENERATE_MAX_CONCURRENCY, - }, - )(RolloutController) - .options(num_cpus=num_workers) + ray.remote(**ray_kwargs)(RolloutController) + .options(**with_trace_runtime_env({"num_cpus": num_workers})) .remote(self, placement_group) ) @@ -882,8 +889,12 @@ def _start_session_server(self) -> None: self.session_server_actor = ( ray.remote(SessionServerActor) .options( - scheduling_strategy=scheduling_strategy, - num_cpus=0, + **with_trace_runtime_env( + { + "scheduling_strategy": scheduling_strategy, + "num_cpus": 0, + } + ) ) .remote( worker_base_url=self.server_url, @@ -968,6 +979,7 @@ async def _decode_routed_experts(self, routed_experts: Any) -> Any: return routed_experts @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) + @traced_rollout_endpoint(TRACE_SPAN_ROLLOUT_WORKER_GENERATE) async def generate(self, rollout_state: RolloutState) -> RolloutState: request_max_tokens = rollout_state.sample_params.max_tokens try: @@ -1021,7 +1033,12 @@ async def generate(self, rollout_state: RolloutState) -> RolloutState: for attempt in range(max_retries + 1): is_last_attempt = attempt == max_retries - http_result = await self._safe_post_request(endpoint_url, headers=headers, payload=payload) + http_result = await self._safe_post_request( + endpoint_url, + headers=headers, + payload=payload, + rollout_state=rollout_state, + ) # Case 1: HTTP Request is Successful if http_result.response: @@ -1225,15 +1242,18 @@ def _launch_server(self): placement_group_bundle_index=self.engine_bundle_idxs[0], ) assert ray.is_initialized() - ray_kwargs = ( - {"runtime_env": server_configs.ray_runtime_env} if hasattr(server_configs, "ray_runtime_env") else {} + runtime_env = copy.deepcopy(getattr(server_configs, "ray_runtime_env", None)) or {} + ray_options = with_trace_runtime_env( + { + "scheduling_strategy": scheduling_strategy, + **AutoAcceleratorWorkers.get_pg_options(current_pg), + "runtime_env": runtime_env, + } ) self.server_task = ( ray.remote(self.server_func) .options( - scheduling_strategy=scheduling_strategy, - **AutoAcceleratorWorkers.get_pg_options(current_pg), - **ray_kwargs, + **ray_options, ) .remote(server_configs) ) @@ -1267,18 +1287,35 @@ def _launch_server(self): ray.cancel(self.server_task) raise TimeoutError("Server failed to start within the timeout period.") - async def _safe_post_request(self, url, headers, payload) -> HttpRequestResult: + async def _safe_post_request( + self, + url, + headers, + payload, + *, + rollout_state: RolloutState | None = None, + ) -> HttpRequestResult: try: if self.receive_abort_request.is_set(): return HttpRequestResult(error_type=HttpRequestErrorType.REQUEST_ABORTED, url=url, payload=payload) - req = self.client.build_request( - "POST", - url, - headers=headers, - json=payload, - ) - r = await self.client.send(req) - r.raise_for_status() + trace_attributes = { + "rollout.backend": self.config.rollout_backend, + "http.method": "POST", + "http.url": url, + } + if rollout_state is not None: + trace_attributes.update(rollout_state_final_attributes(rollout_state)) + request_headers = dict(headers) + with trace_span("infer_engine.generate", attributes=trace_attributes): + inject_trace_context(request_headers) + req = self.client.build_request( + "POST", + url, + headers=request_headers, + json=payload, + ) + r = await self.client.send(req) + r.raise_for_status() return HttpRequestResult(response=r) except Exception as e: error_type = HttpRequestErrorType.from_exception(e) diff --git a/xtuner/v1/rl/trace/__init__.py b/xtuner/v1/rl/trace/__init__.py new file mode 100644 index 0000000000..d18728b1cf --- /dev/null +++ b/xtuner/v1/rl/trace/__init__.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from .api import ( + inject_trace_context, + set_trace_attributes, + trace_event, + trace_function, + trace_span, + traced_agent_item_endpoint, + traced_judger_endpoint, + traced_rollout_endpoint, +) +from .context_propagation import trace_remote +from .runtime import ( + TraceConfig, + TraceRuntime, + close_trace, + configure_trace, + configure_trace_runtime, + current_trace_runtime, + get_trace_env_vars, +) + + +__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_remote", + "trace_span", + "traced_agent_item_endpoint", + "traced_judger_endpoint", + "traced_rollout_endpoint", +] diff --git a/xtuner/v1/rl/trace/api.py b/xtuner/v1/rl/trace/api.py new file mode 100644 index 0000000000..640fbfd1d9 --- /dev/null +++ b/xtuner/v1/rl/trace/api.py @@ -0,0 +1,417 @@ +from __future__ import annotations + +import contextlib +import contextvars +import inspect +import json +import os +import time +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from functools import wraps +from pathlib import Path +from typing import Any, TypeVar + +from xtuner.v1.data_proto.rl_data import RolloutState + +from .runtime import ( + ensure_trace_runtime_from_env as _ensure_trace_runtime_from_env, +) +from .runtime import is_trace_enabled +from .trace_utils import _normalize_trace_attributes, _validate_trace_name + + +F = TypeVar("F", bound=Callable[..., Any]) + + +_LOGICAL_PATH_CARRIER_KEY = "xtuner-trace-logical-path" +_LOGICAL_PATH_ATTRIBUTE = "xtuner.logical_path" +_LIVE_TRACE_PATH_ENV = "XTUNER_OTEL_LIVE_JSONL_PATH" +_CURRENT_LOGICAL_PATH: contextvars.ContextVar[tuple[str, ...]] = contextvars.ContextVar( + "xtuner_trace_logical_path", + default=(), +) +_LIVE_ATTRIBUTE_PREFIXES = ( + "xtuner.", + "agent.", + "session.", + "judger.", + "http.", + "error.", + "exception.", + "reward.", + "filter.", + "drop.", + "train.", + "rollout.", +) + + +def _otel_utils(): + try: + from . import otel_utils + 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 otel_utils + + +@contextmanager +def trace_span( + name: str, + attributes: Mapping[str, Any] | None = None, + *, + parent_carrier: Mapping[str, str] | None = None, +): + """Create a current trace span using XTuner trace semantics.""" + + span_name = _validate_trace_name(name, kind="span name") + normalized_attributes = _normalize_trace_attributes(attributes) + _ensure_trace_runtime_from_env() + if not is_trace_enabled(): + yield + return + + with _attach_parent_carrier(parent_carrier): + logical_path = (*_CURRENT_LOGICAL_PATH.get(), span_name) + normalized_attributes = dict(normalized_attributes) + normalized_attributes.setdefault(_LOGICAL_PATH_ATTRIBUTE, logical_path) + path_token = _CURRENT_LOGICAL_PATH.set(logical_path) + started_at_s = time.time() + status = "completed" + error_message: str | None = None + otel_utils = _otel_utils() + try: + with otel_utils.start_span(span_name, attributes=normalized_attributes): + span_ids = otel_utils.current_span_ids() + _record_live_span( + "start", + span_name=span_name, + logical_path=logical_path, + attributes=normalized_attributes, + span_ids=span_ids, + started_at_s=started_at_s, + ) + try: + yield + except Exception as exc: + status = "error" + error_message = str(exc) + otel_utils.record_failure(exc) + raise + finally: + _record_live_span( + "end", + span_name=span_name, + logical_path=logical_path, + attributes=normalized_attributes, + span_ids=span_ids, + started_at_s=started_at_s, + ended_at_s=time.time(), + status=status, + error_message=error_message, + ) + finally: + _CURRENT_LOGICAL_PATH.reset(path_token) + + +def trace_function( + name: str | None = None, + attributes: Mapping[str, Any] | None = None, +) -> Callable[[F], F]: + """Decorate a function so each call is wrapped in a trace span.""" + + def decorator(func: F) -> F: + span_name = _validate_trace_name( + name or f"{func.__module__}.{func.__qualname__}", + kind="span name", + ) + + 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: + event_name = _validate_trace_name(name, kind="event name") + normalized_attributes = _normalize_trace_attributes(attributes) + if not is_trace_enabled(): + return + otel_utils = _otel_utils() + otel_utils.add_event(event_name, attributes=normalized_attributes) + + +def set_trace_attributes(attributes: Mapping[str, Any] | None) -> None: + normalized_attributes = _normalize_trace_attributes(attributes) + if not is_trace_enabled(): + return + otel_utils = _otel_utils() + otel_utils.set_attributes(normalized_attributes) + + +def inject_trace_context(carrier: dict[str, str] | None = None) -> dict[str, str]: + target = carrier if carrier is not None else {} + _ensure_trace_runtime_from_env() + if not is_trace_enabled(): + return target + otel_utils = _otel_utils() + otel_utils.inject_otel_context(carrier=target) + logical_path = _CURRENT_LOGICAL_PATH.get() + if logical_path: + target[_LOGICAL_PATH_CARRIER_KEY] = json.dumps(logical_path, separators=(",", ":")) + return target + + +def traced_rollout_endpoint( + span_name: str, + *, + target_arg: str = "rollout_state", + initial_attributes: Callable[[Any, RolloutState], Mapping[str, Any]] | None = None, +) -> Callable[[F], F]: + """Decorate a framework-owned rollout endpoint with trace lifecycle + logic.""" + + def decorator(func: F) -> F: + if not inspect.iscoroutinefunction(func): + raise TypeError("traced_rollout_endpoint() only supports async functions") + + signature = inspect.signature(func) + + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + if target_arg not in bound.arguments: + raise TypeError(f"traced rollout endpoint target argument {target_arg!r} was not bound") + rollout_state = bound.arguments[target_arg] + if not isinstance(rollout_state, RolloutState): + raise TypeError("traced_rollout_endpoint target must be a RolloutState") + owner = bound.arguments.get("self", args[0] if args else None) + if initial_attributes is not None: + attributes = dict(initial_attributes(owner, rollout_state)) + else: + from xtuner.v1.rl.trace.trace_utils import rollout_state_initial_attributes + + attributes = rollout_state_initial_attributes(rollout_state) + + from xtuner.v1.rl.trace.context_propagation import ( + attach_trace_attributes, + extract_rollout_trace_parent_carrier, + ) + + parent_carrier = extract_rollout_trace_parent_carrier(rollout_state) + with attach_trace_attributes(attributes): + with trace_span(span_name, attributes=attributes, parent_carrier=parent_carrier): + result = await func(*args, **kwargs) + + from xtuner.v1.rl.trace.trace_utils import record_rollout_state_result + + record_rollout_state_result(result if isinstance(result, RolloutState) else rollout_state) + return result + + return wrapper # type: ignore[return-value] + + return decorator + + +def traced_agent_item_endpoint( + span_name: str, + *, + item_arg: str = "item", + initial_attributes: Callable[[Any, Any], Mapping[str, Any]] | None = None, +) -> Callable[[F], F]: + """Decorate an agent item endpoint with trace lifecycle logic.""" + + def decorator(func: F) -> F: + if not inspect.iscoroutinefunction(func): + raise TypeError("traced_agent_item_endpoint() only supports async functions") + + signature = inspect.signature(func) + + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + if item_arg not in bound.arguments: + raise TypeError(f"traced agent item endpoint target argument {item_arg!r} was not bound") + item = bound.arguments[item_arg] + owner = bound.arguments.get("self", args[0] if args else None) + if initial_attributes is not None: + attributes = dict(initial_attributes(owner, item)) + else: + from xtuner.v1.rl.trace.trace_utils import agent_item_initial_attributes + + attributes = agent_item_initial_attributes(item) + + with trace_span(span_name, attributes=attributes): + result = await func(*args, **kwargs) + + from xtuner.v1.rl.trace.trace_utils import record_agent_item_result + + record_agent_item_result(item, result) + return result + + return wrapper # type: ignore[return-value] + + return decorator + + +def traced_judger_endpoint( + span_name: str = "judger.run", + *, + target_arg: str = "rollout_state", +) -> Callable[[F], F]: + """Decorate an agent-loop judger endpoint with trace lifecycle logic.""" + + def decorator(func: F) -> F: + if not inspect.iscoroutinefunction(func): + raise TypeError("traced_judger_endpoint() only supports async functions") + + signature = inspect.signature(func) + + @wraps(func) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + if target_arg not in bound.arguments: + raise TypeError(f"traced judger endpoint target argument {target_arg!r} was not bound") + owner = bound.arguments.get("self", args[0] if args else None) + judger = getattr(owner, "judger", None) + if judger is None: + raise RuntimeError("traced_judger_endpoint() requires owner.judger") + target_obj = bound.arguments[target_arg] + + from xtuner.v1.rl.trace.trace_utils import judger_trace_attributes + + with trace_span(span_name, attributes=judger_trace_attributes(judger, target_obj)): + result = await func(*args, **kwargs) + set_trace_attributes(judger_trace_attributes(judger, result)) + return result + + return wrapper # type: ignore[return-value] + + return decorator + + +@contextmanager +def _attach_parent_carrier(parent_carrier: Mapping[str, str] | None): + if parent_carrier is None: + yield + return + + otel_utils = _otel_utils() + parent_context = otel_utils.extract_otel_context(parent_carrier) + token = otel_utils.attach_otel_context(parent_context) + logical_path = _extract_logical_path(parent_carrier) + path_token = _CURRENT_LOGICAL_PATH.set(logical_path) if logical_path else None + try: + yield + finally: + if path_token is not None: + _CURRENT_LOGICAL_PATH.reset(path_token) + otel_utils.detach_otel_context(token) + + +def _extract_logical_path(carrier: Mapping[str, str]) -> tuple[str, ...]: + payload: str | None = None + for key, value in carrier.items(): + if str(key).lower() == _LOGICAL_PATH_CARRIER_KEY: + payload = str(value) + break + if not payload: + return () + try: + value = json.loads(payload) + except json.JSONDecodeError: + return () + if not isinstance(value, list): + return () + path = tuple(str(item).strip() for item in value if isinstance(item, str) and item.strip()) + return path + + +def _record_live_span( + event: str, + *, + span_name: str, + logical_path: tuple[str, ...], + attributes: Mapping[str, Any], + span_ids: Mapping[str, str] | None, + started_at_s: float, + ended_at_s: float | None = None, + status: str | None = None, + error_message: str | None = None, +) -> None: + live_path = os.environ.get(_LIVE_TRACE_PATH_ENV) + if not live_path: + return + record: dict[str, Any] = { + "event": event, + "time_s": ended_at_s if ended_at_s is not None else started_at_s, + "span_name": span_name, + "logical_path": list(logical_path), + "attributes": _live_trace_attributes(attributes), + } + if span_ids is not None: + record.update(span_ids) + if status is not None: + record["status"] = status + if ended_at_s is not None: + record["duration_ms"] = round(max(0.0, ended_at_s - started_at_s) * 1000.0, 3) + if error_message: + record["error_message"] = error_message + + with contextlib.suppress(Exception): + path = Path(live_path).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def _live_trace_attributes(attributes: Mapping[str, Any]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in attributes.items(): + if key == "agent.config": + continue + if key == _LOGICAL_PATH_ATTRIBUTE or any(key.startswith(prefix) for prefix in _LIVE_ATTRIBUTE_PREFIXES): + result[key] = _json_safe_trace_value(value) + return result + + +def _json_safe_trace_value(value: Any) -> Any: + if isinstance(value, (str, bool, int, float)) or value is None: + return value + if isinstance(value, (list, tuple)): + return [_json_safe_trace_value(item) for item in value] + return str(value) + + +__all__ = [ + "inject_trace_context", + "set_trace_attributes", + "trace_event", + "trace_function", + "trace_span", + "traced_agent_item_endpoint", + "traced_judger_endpoint", + "traced_rollout_endpoint", +] diff --git a/xtuner/v1/rl/trace/context_propagation.py b/xtuner/v1/rl/trace/context_propagation.py new file mode 100644 index 0000000000..a2b74aff2a --- /dev/null +++ b/xtuner/v1/rl/trace/context_propagation.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +import contextvars +import time +from collections.abc import Mapping +from contextlib import asynccontextmanager, contextmanager +from typing import Any, Protocol + +from xtuner.v1.data_proto.rl_data import RolloutState +from xtuner.v1.rl.trace import inject_trace_context, set_trace_attributes, trace_span + + +_TRACE_CARRIER_EXTRA_FIELD = "_xtuner_trace_carrier" +_TRACE_ATTRIBUTES_EXTRA_FIELD = "_xtuner_trace_attributes" +_AIOHTTP_INSTRUMENTED_ATTR = "_xtuner_trace_instrumented" +_AIOHTTP_ORIGINAL_REQUEST_ATTR = "_xtuner_trace_original_request" +_CURRENT_TRACE_ATTRIBUTES: contextvars.ContextVar[dict[str, Any] | None] = contextvars.ContextVar( + "xtuner_trace_attributes", + default=None, +) + + +class _RayRemoteMethod(Protocol): + def remote(self, *args: Any, **kwargs: Any) -> Any: ... + + +class _RolloutTraceCarrier: + field_name = _TRACE_CARRIER_EXTRA_FIELD + + @classmethod + @contextmanager + def attach_temporarily(cls, rollout_state: RolloutState, carrier: Mapping[str, str]): + if not carrier: + yield + return + + extra_fields = rollout_state.extra_fields + if extra_fields is None: + extra_fields = {} + rollout_state.extra_fields = extra_fields + had_previous_carrier = cls.field_name in extra_fields + previous_carrier = extra_fields.get(cls.field_name) + extra_fields[cls.field_name] = dict(carrier) + try: + yield + finally: + if had_previous_carrier: + extra_fields[cls.field_name] = previous_carrier + else: + extra_fields.pop(cls.field_name, None) + + @classmethod + def pop_from_rollout_state(cls, rollout_state: RolloutState) -> dict[str, str] | None: + if not isinstance(rollout_state, RolloutState): + raise TypeError("extract_rollout_trace_parent_carrier() requires a RolloutState") + + extra_fields = rollout_state.extra_fields + if not extra_fields: + return None + return cls.extract_from_mapping(extra_fields, remove=True) + + @classmethod + def extract_from_mapping(cls, payload: Mapping[str, Any] | None, *, remove: bool = False) -> dict[str, str] | None: + if not isinstance(payload, Mapping): + return None + carrier = payload.get(cls.field_name) + if remove and isinstance(payload, dict): + payload.pop(cls.field_name, None) + if not isinstance(carrier, Mapping): + return None + return {str(key): str(value) for key, value in carrier.items()} + + @classmethod + def extract_attributes_from_mapping(cls, payload: Mapping[str, Any] | None) -> dict[str, Any]: + if not isinstance(payload, Mapping): + return {} + attributes = payload.get(_TRACE_ATTRIBUTES_EXTRA_FIELD) + if not isinstance(attributes, Mapping): + return {} + return {str(key): _trace_payload_value(value) for key, value in attributes.items() if value is not None} + + @classmethod + def inject_into_json_body(cls, kwargs: dict[str, Any], carrier: Mapping[str, str]) -> None: + if not carrier: + return + payload = kwargs.get("json") + if not isinstance(payload, dict): + return + if cls.field_name in payload: + return + payload = dict(payload) + payload[cls.field_name] = dict(carrier) + attributes = _CURRENT_TRACE_ATTRIBUTES.get() + if attributes: + payload.setdefault(_TRACE_ATTRIBUTES_EXTRA_FIELD, dict(attributes)) + kwargs["json"] = payload + + +@contextmanager +def attach_trace_attributes(attributes: Mapping[str, Any] | None): + normalized = _normalize_trace_payload_attributes(attributes) + token = _CURRENT_TRACE_ATTRIBUTES.set(normalized or None) + try: + yield + finally: + _CURRENT_TRACE_ATTRIBUTES.reset(token) + + +@asynccontextmanager +async def traced_aiohttp_request( + client: Any, + *, + span_name: str, + method: str, + url: str, + headers: dict[str, str], + attributes: Mapping[str, Any] | None = None, + inject_context: bool = True, + **request_kwargs: Any, +): + started_at = time.monotonic() + span_attributes: dict[str, Any] = {"http.method": method, "http.url": url} + span_attributes.update(dict(attributes or {})) + with trace_span(span_name, attributes=span_attributes): + if inject_context: + inject_trace_context(headers) + request_context = client.request(method=method, url=url, headers=headers, **request_kwargs) + resp = await request_context.__aenter__() + set_trace_attributes( + { + "http.status_code": getattr(resp, "status", None), + "send_request_ms": int((time.monotonic() - started_at) * 1000), + } + ) + try: + yield resp + except BaseException as exc: + await request_context.__aexit__(type(exc), exc, exc.__traceback__) + raise + else: + await request_context.__aexit__(None, None, None) + + +def trace_remote( + remote_method: _RayRemoteMethod, + *args: Any, + target: Any | None = None, + **kwargs: Any, +) -> Any: + """Call a Ray remote method and propagate the current trace context.""" + + rollout_state = _resolve_rollout_state_target(args, kwargs, target=target, owner="trace_remote") + carrier: dict[str, str] = {} + inject_trace_context(carrier) + with _RolloutTraceCarrier.attach_temporarily(rollout_state, carrier): + return remote_method.remote(*args, **kwargs) + + +def instrument_aiohttp_client() -> bool: + """Inject the current trace context into generic aiohttp client requests. + + This covers third-party clients, such as lagent's OpenAI-compatible client, + that build their own aiohttp requests instead of going through + ``traced_aiohttp_request``. + """ + + try: + import aiohttp + except ImportError: + return False + + client_session_cls = aiohttp.ClientSession + if getattr(client_session_cls, _AIOHTTP_INSTRUMENTED_ATTR, False): + return False + + original_request = client_session_cls._request + + async def _xtuner_trace_request(self, method: str, str_or_url: Any, **kwargs: Any): + headers = _copy_headers(kwargs.get("headers")) + carrier: dict[str, str] = {} + inject_trace_context(carrier) + headers.update(carrier) + kwargs["headers"] = headers + _RolloutTraceCarrier.inject_into_json_body(kwargs, carrier) + return await original_request(self, method, str_or_url, **kwargs) + + setattr(client_session_cls, _AIOHTTP_ORIGINAL_REQUEST_ATTR, original_request) + setattr(client_session_cls, "_request", _xtuner_trace_request) + setattr(client_session_cls, _AIOHTTP_INSTRUMENTED_ATTR, True) + return True + + +def reset_aiohttp_client_instrumentation() -> None: + try: + import aiohttp + except ImportError: + return + + client_session_cls = aiohttp.ClientSession + original_request = getattr(client_session_cls, _AIOHTTP_ORIGINAL_REQUEST_ATTR, None) + if original_request is not None: + setattr(client_session_cls, "_request", original_request) + for attr in (_AIOHTTP_INSTRUMENTED_ATTR, _AIOHTTP_ORIGINAL_REQUEST_ATTR): + if hasattr(client_session_cls, attr): + delattr(client_session_cls, attr) + + +def extract_rollout_trace_parent_carrier(rollout_state: RolloutState) -> dict[str, str] | None: + """Remove and return the parent trace carrier attached to a rollout + call.""" + + return _RolloutTraceCarrier.pop_from_rollout_state(rollout_state) + + +def extract_trace_carrier_from_mapping(payload: Mapping[str, Any] | None) -> dict[str, str] | None: + return _RolloutTraceCarrier.extract_from_mapping(payload) + + +def extract_trace_attributes_from_mapping(payload: Mapping[str, Any] | None) -> dict[str, Any]: + return _RolloutTraceCarrier.extract_attributes_from_mapping(payload) + + +def remove_trace_carrier_from_mapping(payload: dict[str, Any]) -> None: + payload.pop(_RolloutTraceCarrier.field_name, None) + payload.pop(_TRACE_ATTRIBUTES_EXTRA_FIELD, None) + + +def _resolve_rollout_state_target( + args: tuple[Any, ...], + kwargs: Mapping[str, Any], + *, + target: Any | 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) + continue + if _contains_rollout_state_collection(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] + + +def _contains_rollout_state_collection(value: Any) -> bool: + if not isinstance(value, (list, tuple, set, frozenset)): + return False + return any(isinstance(item, RolloutState) for item in value) + + +__all__ = [ + "trace_remote", +] + + +def _normalize_trace_payload_attributes(attributes: Mapping[str, Any] | None) -> dict[str, Any]: + if not isinstance(attributes, Mapping): + return {} + return {str(key): _trace_payload_value(value) for key, value in attributes.items() if value is not None} + + +def _trace_payload_value(value: Any) -> str | bool | int | float: + if isinstance(value, bool): + return value + if isinstance(value, (str, int, float)): + return value + return str(value) + + +def _copy_headers(headers: Any) -> dict[str, str]: + if headers is None: + return {} + if isinstance(headers, Mapping): + return {str(key): str(value) for key, value in headers.items()} + return {str(key): str(value) for key, value in dict(headers).items()} diff --git a/xtuner/v1/rl/trace/otel_utils.py b/xtuner/v1/rl/trace/otel_utils.py new file mode 100644 index 0000000000..896b808ea2 --- /dev/null +++ b/xtuner/v1/rl/trace/otel_utils.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from typing import Any, Mapping, MutableMapping + +from opentelemetry import context as otel_context +from opentelemetry import propagate, trace +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.trace import Status, StatusCode + + +def configure_tracer_provider( + *, + service_name: str, + run_id: str, + endpoint: str, + protocol: str = "grpc", +) -> TracerProvider: + if protocol != "grpc": + raise ValueError(f"Unsupported OTel trace export protocol: {protocol!r}") + try: + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter + except ImportError as exc: + raise RuntimeError( + "XTuner OTel tracing requires the official OpenTelemetry OTLP gRPC trace exporter. " + "Install `opentelemetry-exporter-otlp-proto-grpc` before enabling trace." + ) from exc + resource = Resource.create({"service.name": service_name, "run.id": run_id}) + provider = TracerProvider(resource=resource) + provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter(endpoint=endpoint, insecure=True))) + trace.set_tracer_provider(provider) + return provider + + +def inject_otel_context( + context: Any | None = None, + carrier: MutableMapping[str, str] | None = None, +) -> MutableMapping[str, str]: + """Inject the current or provided OTel context into a W3C carrier.""" + + carrier = carrier if carrier is not None else {} + propagate.inject(carrier, context=context) + return carrier + + +def extract_otel_context( + carrier: Mapping[str, str], + context: Any | None = None, +) -> Any: + """Extract W3C TraceContext from a carrier into an OTel context.""" + + return propagate.extract(carrier, context=context) + + +def attach_otel_context(context: Any) -> object: + """Attach extracted context to the current execution scope.""" + + return otel_context.attach(context) + + +def detach_otel_context(token: object) -> None: + """Detach a previously attached OTel context token.""" + + otel_context.detach(token) + + +def start_span(name: str, *, attributes: Mapping[str, Any] | None = None): + """Start a current OTel span with XTuner-managed exception handling.""" + + return trace.get_tracer("xtuner").start_as_current_span( + name, + attributes=attributes, + record_exception=False, + set_status_on_exception=False, + ) + + +def current_span_ids() -> dict[str, str] | None: + span = trace.get_current_span() + span_context = span.get_span_context() + if not span_context.is_valid: + return None + return { + "trace_id": f"{span_context.trace_id:032x}", + "span_id": f"{span_context.span_id:016x}", + } + + +def add_event(name: str, *, attributes: Mapping[str, Any] | None = None) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + span.add_event(name, attributes=attributes) + + +def set_attributes(attributes: Mapping[str, Any]) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + for key, value in attributes.items(): + span.set_attribute(key, value) + + +def set_error_status(message: str | None = None) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + span.set_attribute("error", True) + description = message or "error" + span.set_attribute("error.message", description) + span.set_status(Status(StatusCode.ERROR, description)) + + +def record_failure(exc: BaseException) -> None: + span = trace.get_current_span() + if not span.is_recording(): + return + error_attributes = { + "error.type": type(exc).__name__, + "error.message": str(exc), + "error": True, + } + for key, value in error_attributes.items(): + span.set_attribute(key, value) + span.record_exception(exc, attributes=error_attributes) + span.set_status(Status(StatusCode.ERROR, str(exc))) diff --git a/xtuner/v1/rl/trace/runtime.py b/xtuner/v1/rl/trace/runtime.py new file mode 100644 index 0000000000..77e48f9285 --- /dev/null +++ b/xtuner/v1/rl/trace/runtime.py @@ -0,0 +1,769 @@ +from __future__ import annotations + +import atexit +import contextlib +import errno +import os +import shlex +import shutil +import socket +import subprocess +import sys +import time +import uuid +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Literal, Mapping + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from xtuner.v1.rl.utils.misc import find_free_ports +from xtuner.v1.utils import get_logger + + +logger = get_logger() + +TraceRuntimeMode = Literal["disabled", "driver", "inherited"] + + +TRACE_ENV_KEYS = ( + "XTUNER_OTEL_ENABLED", + "XTUNER_OTEL_OUTPUT_DIR", + "XTUNER_OTEL_RUN_ID", + "XTUNER_OTEL_RUN_DIR", + "XTUNER_OTEL_JSONL_PATH", + "XTUNER_OTEL_LIVE_JSONL_PATH", + "OTEL_TRACES_EXPORTER", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_EXPORTER_OTLP_PROTOCOL", + "OTEL_SERVICE_NAME", +) + + +DEFAULT_JAEGER_OTLP_GRPC_ENDPOINT = "127.0.0.1:14317" +_TRACE_VIEWER_SERVER_MODULE = "xtuner.tools.trace_viewer.server" +_READY_TIMEOUT_S = 3.0 +_READY_POLL_INTERVAL_S = 0.05 +_LOG_TAIL_CHARS = 4000 + +_OTELCOL_CONFIG_YAML_TEMPLATE = """ +receivers: + otlp: + protocols: + grpc: + endpoint: 0.0.0.0:{port} + +exporters: + file: + path: {output_path} + rotation: + max_megabytes: 64 + max_days: 0 + max_backups: 0 + +service: + telemetry: + metrics: + level: none + pipelines: + traces: + receivers: [otlp] + exporters: [{exporters}] +""".lstrip() + +_OTELCOL_OTLP_GRPC_EXPORTER_YAML_TEMPLATE = """ + otlp/jaeger: + endpoint: {endpoint} + tls: + insecure: true +""".rstrip() + + +def _configure_tracer_provider( + *, + service_name: str, + run_id: str, + endpoint: str, + protocol: str, +) -> Any: + try: + from xtuner.v1.rl.trace.otel_utils import configure_tracer_provider + except ModuleNotFoundError as exc: + if exc.name == "opentelemetry" or (exc.name or "").startswith("opentelemetry."): + raise RuntimeError( + "XTuner OTel tracing requires OpenTelemetry packages. " + "Install `opentelemetry-sdk` and `opentelemetry-exporter-otlp-proto-grpc` " + "before enabling trace." + ) from exc + raise + return configure_tracer_provider( + service_name=service_name, + run_id=run_id, + endpoint=endpoint, + protocol=protocol, + ) + + +class TraceConfig(BaseModel): + """Public rollout tracing configuration. + + The interface is intentionally XTuner-level. OTel endpoint, exporter, and collector choices are runtime + implementation details. + """ + + model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") + + enabled: bool = False + output_dir: Path | str | None = Field(default=None) + service_name: str = "xtuner-rollout" + viewer_enabled: bool = False + viewer_host: str = "127.0.0.1" + viewer_port: int = Field(default=0, ge=0, le=65535) + viewer_jaeger_query_url: str | None = None + viewer_jaeger_link_url: str | None = None + + @field_validator("output_dir") + @classmethod + def _expand_output_dir(cls, value: Path | str | None) -> Path | None: + if value is None: + return None + return Path(value).expanduser() + + +@dataclass(frozen=True) +class TraceRuntime: + enabled: bool + mode: TraceRuntimeMode + run_id: str + run_dir: Path + trace_jsonl_path: Path + live_jsonl_path: Path + service_name: str + trace_viewer_url: str | None = None + trace_viewer_port: int | None = None + + +@dataclass +class _OTelCollector: + port: int + output_path: Path + _binary_path: str = field(repr=False) + _config_path: Path = field(repr=False) + _stdout_path: Path = field(repr=False) + _stderr_path: Path = field(repr=False) + _process: subprocess.Popen | None = field(repr=False) + + @classmethod + def start( + cls, + *, + port: int, + output_path: Path, + ) -> _OTelCollector: + otelcol = shutil.which("otelcol-contrib") or shutil.which("otelcol") + if otelcol is None: + raise RuntimeError( + "XTuner OTel tracing requires `otelcol-contrib` or `otelcol` on PATH. " + "Install an official OpenTelemetry Collector binary before enabling trace." + ) + + jaeger_exporter = "\n" + _OTELCOL_OTLP_GRPC_EXPORTER_YAML_TEMPLATE.format( + endpoint=DEFAULT_JAEGER_OTLP_GRPC_ENDPOINT + ) + + config_path = output_path.parent / "otelcol.yaml" + config_yaml = _OTELCOL_CONFIG_YAML_TEMPLATE.format( + port=port, + output_path=output_path, + exporters="file, otlp/jaeger", + ).replace("\nservice:", f"{jaeger_exporter}\n\nservice:") + config_path.write_text(config_yaml, encoding="utf-8") + stdout_path = output_path.parent / "otelcol.stdout.log" + stderr_path = output_path.parent / "otelcol.stderr.log" + with stdout_path.open("wb") as stdout_file, stderr_path.open("wb") as stderr_file: + process = subprocess.Popen( + [otelcol, "--config", os.fspath(config_path)], + stdout=stdout_file, + stderr=stderr_file, + ) + collector = cls( + port=port, + output_path=output_path, + _binary_path=otelcol, + _config_path=config_path, + _stdout_path=stdout_path, + _stderr_path=stderr_path, + _process=process, + ) + try: + collector._wait_until_ready() + except Exception: + collector.close() + raise + return collector + + def _wait_until_ready(self) -> None: + deadline = time.monotonic() + _READY_TIMEOUT_S + last_error: OSError | None = None + while time.monotonic() < deadline: + process = self._process + if process is None: + stderr_tail = "" + with contextlib.suppress(OSError): + stderr_tail = self._stderr_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + "OpenTelemetry collector failed to start: collector process is not available. " + f"binary={self._binary_path}, config={self._config_path}, port={self.port}, " + f"stderr_tail={stderr_tail!r}" + ) + exit_code = process.poll() + if exit_code is not None: + stderr_tail = "" + with contextlib.suppress(OSError): + stderr_tail = self._stderr_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + f"OpenTelemetry collector failed to start: collector exited with code {exit_code}. " + f"binary={self._binary_path}, config={self._config_path}, port={self.port}, " + f"stderr_tail={stderr_tail!r}" + ) + try: + with socket.create_connection(("127.0.0.1", self.port), timeout=0.1): + return + except OSError as exc: + last_error = exc + time.sleep(_READY_POLL_INTERVAL_S) + + detail = f"collector did not become ready within {_READY_TIMEOUT_S:.1f}s" + if last_error is not None: + detail += f"; last connection error: {last_error}" + stderr_tail = "" + with contextlib.suppress(OSError): + stderr_tail = self._stderr_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + f"OpenTelemetry collector failed to start: {detail}. " + f"binary={self._binary_path}, config={self._config_path}, port={self.port}, " + f"stderr_tail={stderr_tail!r}" + ) + + def close(self) -> None: + process = self._process + self._process = None + if process is None: + return + if process.poll() is None: + process.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + if process.poll() is None: + process.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + + +@dataclass +class _TraceViewerProcess: + host: str + port: int + url: str + log_path: Path + _process: subprocess.Popen | None = field(repr=False) + command: list[str] = field(default_factory=list, repr=False) + + @classmethod + def start( + cls, + *, + trace_jsonl_path: Path, + live_jsonl_path: Path, + jaeger_query_url: str | None, + jaeger_link_url: str | None, + service_name: str, + run_id: str, + host: str, + port: int, + ) -> _TraceViewerProcess: + if port == 0: + port = find_free_ports(nums=1, host=host)[0] + _ensure_trace_viewer_port_available(host=host, port=port, run_id=run_id) + command = _build_trace_viewer_command( + trace_jsonl_path=trace_jsonl_path, + live_jsonl_path=live_jsonl_path, + jaeger_query_url=jaeger_query_url, + jaeger_link_url=jaeger_link_url, + service_name=service_name, + run_id=run_id, + host=host, + port=port, + ) + log_path = trace_jsonl_path.parent / "viewer.log" + with log_path.open("ab") as log_file: + process = subprocess.Popen( + command, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + viewer = cls( + host=host, + port=port, + url=f"http://{host}:{port}", + log_path=log_path, + command=command, + _process=process, + ) + try: + viewer._wait_until_ready() + except Exception: + viewer.close() + raise + return viewer + + def _wait_until_ready(self) -> None: + deadline = time.monotonic() + _READY_TIMEOUT_S + last_error: OSError | None = None + while time.monotonic() < deadline: + self._raise_if_process_exited() + connect_host = "127.0.0.1" if self.host in {"", "0.0.0.0"} else self.host + try: + with socket.create_connection((connect_host, self.port), timeout=0.1): + time.sleep(_READY_POLL_INTERVAL_S) + self._raise_if_process_exited() + return + except OSError as exc: + last_error = exc + time.sleep(_READY_POLL_INTERVAL_S) + + detail = f"viewer did not become ready within {_READY_TIMEOUT_S:.1f}s" + if last_error is not None: + detail += f"; last connection error: {last_error}" + log_tail = "" + with contextlib.suppress(OSError): + log_tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError(f"XTuner trace viewer failed to start: {detail}, url={self.url}, log_tail={log_tail!r}") + + def _raise_if_process_exited(self) -> None: + process = self._process + if process is None: + raise RuntimeError(f"XTuner trace viewer failed to start: process is not available, log={self.log_path}") + exit_code = process.poll() + if exit_code is None: + return + + log_tail = "" + with contextlib.suppress(OSError): + log_tail = self.log_path.read_text(encoding="utf-8", errors="replace")[-_LOG_TAIL_CHARS:] + raise RuntimeError( + f"XTuner trace viewer failed to start: exited with code {exit_code}, url={self.url}, log_tail={log_tail!r}" + ) + + def restart_command(self) -> str: + return shlex.join(self.command) + + def close(self) -> None: + process = self._process + self._process = None + if process is None: + return + if process.poll() is None: + process.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + if process.poll() is None: + process.kill() + with contextlib.suppress(subprocess.TimeoutExpired): + process.wait(timeout=5) + + +def _ensure_trace_viewer_port_available(*, host: str, port: int, run_id: str) -> None: + existing_viewer = _find_existing_trace_viewer_process_on_port(port) + if existing_viewer is not None: + pid = existing_viewer.get("pid", "unknown") + cmdline = existing_viewer.get("cmdline", "") + raise RuntimeError( + f"XTuner trace viewer port {port} already has an existing XTuner trace viewer process: " + f"pid={pid}, cmdline={cmdline!r}. " + f"Stop the old viewer/training process or set XTUNER_TRACE_VIEWER_PORT to another port " + f"before starting run_id={run_id}." + ) + + bind_host = host or "0.0.0.0" + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((bind_host, port)) + except OSError as exc: + if exc.errno == errno.EADDRINUSE: + raise RuntimeError( + f"XTuner trace viewer port {port} is already in use on host {host!r}. " + "Stop the process using that port or set XTUNER_TRACE_VIEWER_PORT to another port." + ) from exc + raise + + +def _find_existing_trace_viewer_process_on_port(port: int) -> Mapping[str, Any] | None: + listening_inodes = _listening_socket_inodes_on_port(port) + if not listening_inodes: + return None + + proc_root = Path("/proc") + with contextlib.suppress(OSError): + proc_dirs = list(proc_root.iterdir()) + for proc_dir in proc_dirs: + if not proc_dir.name.isdigit(): + continue + cmdline = _read_process_cmdline(proc_dir) + if _TRACE_VIEWER_SERVER_MODULE not in cmdline: + continue + if _process_has_socket_inode(proc_dir, listening_inodes): + return {"pid": int(proc_dir.name), "cmdline": cmdline} + return None + + +def _listening_socket_inodes_on_port(port: int) -> set[str]: + inodes: set[str] = set() + for proc_net_path in (Path("/proc/net/tcp"), Path("/proc/net/tcp6")): + with contextlib.suppress(OSError, ValueError): + for line in proc_net_path.read_text(encoding="utf-8", errors="replace").splitlines()[1:]: + parts = line.split() + if len(parts) < 10: + continue + local_address, state, inode = parts[1], parts[3], parts[9] + if state != "0A": + continue + _, port_hex = local_address.rsplit(":", 1) + if int(port_hex, 16) == port: + inodes.add(inode) + return inodes + + +def _read_process_cmdline(proc_dir: Path) -> str: + with contextlib.suppress(OSError): + raw_cmdline = (proc_dir / "cmdline").read_bytes() + return " ".join(part.decode("utf-8", errors="replace") for part in raw_cmdline.split(b"\0") if part) + return "" + + +def _process_has_socket_inode(proc_dir: Path, socket_inodes: set[str]) -> bool: + fd_dir = proc_dir / "fd" + with contextlib.suppress(OSError): + for fd_path in fd_dir.iterdir(): + with contextlib.suppress(OSError): + target = os.readlink(fd_path) + if target.startswith("socket:[") and target.endswith("]") and target[8:-1] in socket_inodes: + return True + return False + + +def _build_trace_viewer_command( + *, + trace_jsonl_path: Path, + live_jsonl_path: Path, + jaeger_query_url: str | None, + jaeger_link_url: str | None, + service_name: str, + run_id: str, + host: str, + port: int, +) -> list[str]: + command = [ + sys.executable, + "-m", + "xtuner.tools.trace_viewer.server", + "--trace-jsonl", + os.fspath(trace_jsonl_path), + "--live-jsonl", + os.fspath(live_jsonl_path), + "--service", + service_name, + "--run-id", + run_id, + "--host", + host, + "--port", + str(port), + ] + if jaeger_query_url is not None: + command.extend(["--jaeger-query-url", jaeger_query_url]) + if jaeger_link_url is not None: + command.extend(["--jaeger-link-url", jaeger_link_url]) + return command + + +@dataclass +class _TraceRuntimeHandle: + runtime: TraceRuntime + endpoint: str + env_vars: dict[str, str] + collector_port: int | None = None + collector: _OTelCollector | None = None + provider: Any | None = None + viewer_host: str | None = None + viewer_port: int = 0 + viewer_jaeger_query_url: str | None = None + viewer_jaeger_link_url: str | None = None + viewer: _TraceViewerProcess | None = None + + def start(self) -> None: + apply_trace_env(self.env_vars) + if not self.runtime.enabled: + logger.info("XTuner OTel tracing disabled.") + return + try: + if self.runtime.mode == "driver": + if self.collector_port is None: + raise RuntimeError("driver trace runtime requires a collector port") + self.collector = _OTelCollector.start( + port=self.collector_port, + output_path=self.runtime.trace_jsonl_path, + ) + self.provider = _configure_tracer_provider( + service_name=self.runtime.service_name, + run_id=self.runtime.run_id, + endpoint=self.endpoint, + protocol=self.env_vars["OTEL_EXPORTER_OTLP_PROTOCOL"], + ) + if self.runtime.mode == "driver" and self.viewer_host is not None: + self.viewer = _TraceViewerProcess.start( + trace_jsonl_path=self.runtime.trace_jsonl_path, + live_jsonl_path=self.runtime.live_jsonl_path, + jaeger_query_url=self.viewer_jaeger_query_url, + jaeger_link_url=self.viewer_jaeger_link_url, + service_name=self.runtime.service_name, + run_id=self.runtime.run_id, + host=self.viewer_host, + port=self.viewer_port, + ) + self.runtime = replace( + self.runtime, + trace_viewer_url=self.viewer.url, + trace_viewer_port=self.viewer.port, + ) + except Exception: + self.close(stop_viewer=True) + clear_trace_env() + raise + logger.info( + f"XTuner OTel tracing enabled: run_id={self.runtime.run_id}, endpoint={self.endpoint}, " + f"traces={self.runtime.trace_jsonl_path}" + ) + if self.viewer is not None: + logger.info( + f"XTuner trace viewer enabled: url={self.viewer.url}, host={self.viewer.host}, " + f"port={self.viewer.port}. Restart with: {self.viewer.restart_command()}" + ) + + def close(self, *, stop_viewer: bool = True) -> None: + viewer = self.viewer + self.viewer = None + if viewer is not None and stop_viewer: + logger.info("XTuner trace viewer stopped with training process.") + with contextlib.suppress(Exception): + viewer.close() + + provider = self.provider + self.provider = None + if provider is not None: + with contextlib.suppress(Exception): + provider.shutdown() + + collector = self.collector + self.collector = None + if collector is not None: + with contextlib.suppress(Exception): + collector.close() + + +_RUNTIME: _TraceRuntimeHandle | None = None +_ATEXIT_REGISTERED = False + + +def configure_trace(config: TraceConfig | None = None) -> TraceRuntime: + return configure_trace_runtime(config or TraceConfig()) + + +def configure_trace_runtime(config: TraceConfig) -> TraceRuntime: + """Configure Layer1 OTel runtime for the current process.""" + + global _RUNTIME + + close_trace() + runtime_handle = _build_trace_runtime_handle(config) + runtime_handle.start() + _RUNTIME = runtime_handle + if runtime_handle.runtime.enabled: + register_atexit_once(close_trace) + return runtime_handle.runtime + + +def _build_trace_runtime_handle(config: TraceConfig) -> _TraceRuntimeHandle: + if not config.enabled: + return _TraceRuntimeHandle( + runtime=TraceRuntime( + enabled=False, + mode="disabled", + run_id="", + run_dir=Path(), + trace_jsonl_path=Path(), + live_jsonl_path=Path(), + service_name=config.service_name, + trace_viewer_url=None, + trace_viewer_port=None, + ), + endpoint="", + env_vars={}, + ) + + output_dir = Path(config.output_dir or Path.cwd() / "otel_traces").expanduser() + timestamp = time.strftime("%Y%m%d-%H%M%S", time.localtime()) + run_id = f"{timestamp}-{os.getpid()}-{uuid.uuid4().hex[:8]}" + run_dir = output_dir / run_id + traces_dir = run_dir / "traces" + traces_dir.mkdir(parents=True, exist_ok=True) + trace_jsonl_path = traces_dir / "traces.jsonl" + live_jsonl_path = traces_dir / "live.jsonl" + + try: + port = find_free_ports(nums=1, host="127.0.0.1", start_port=4317, end_port=4318)[0] + except RuntimeError: + port = find_free_ports(nums=1, host="127.0.0.1")[0] + endpoint = f"http://127.0.0.1:{port}" + protocol = "grpc" + + env_vars = { + "XTUNER_OTEL_ENABLED": "1", + "XTUNER_OTEL_OUTPUT_DIR": os.fspath(output_dir), + "XTUNER_OTEL_RUN_ID": run_id, + "XTUNER_OTEL_RUN_DIR": os.fspath(run_dir), + "XTUNER_OTEL_JSONL_PATH": os.fspath(trace_jsonl_path), + "XTUNER_OTEL_LIVE_JSONL_PATH": os.fspath(live_jsonl_path), + "OTEL_TRACES_EXPORTER": "otlp", + "OTEL_EXPORTER_OTLP_ENDPOINT": endpoint, + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": endpoint, + "OTEL_EXPORTER_OTLP_PROTOCOL": protocol, + "OTEL_SERVICE_NAME": config.service_name, + } + return _TraceRuntimeHandle( + runtime=TraceRuntime( + enabled=True, + mode="driver", + run_id=run_id, + run_dir=run_dir, + trace_jsonl_path=trace_jsonl_path, + live_jsonl_path=live_jsonl_path, + service_name=config.service_name, + trace_viewer_url=None, + trace_viewer_port=None, + ), + endpoint=endpoint, + env_vars=env_vars, + collector_port=port, + viewer_host=config.viewer_host if config.viewer_enabled else None, + viewer_port=config.viewer_port, + viewer_jaeger_query_url=config.viewer_jaeger_query_url, + viewer_jaeger_link_url=config.viewer_jaeger_link_url, + ) + + +def get_trace_env_vars() -> dict[str, str]: + """Return the env vars that should be injected into child Ray processes.""" + + if _RUNTIME is None or not _RUNTIME.runtime.enabled: + return get_trace_env_vars_from_env() + return dict(_RUNTIME.env_vars) + + +def current_trace_runtime() -> TraceRuntime | None: + """Return the active trace runtime owned by this process, if any.""" + + if _RUNTIME is None: + return None + return _RUNTIME.runtime + + +def get_trace_env_vars_from_env() -> dict[str, str]: + """Return inherited trace env vars before process-local runtime exists.""" + + if os.environ.get("XTUNER_OTEL_ENABLED") != "1": + return {} + return {key: os.environ[key] for key in TRACE_ENV_KEYS if key in os.environ} + + +def ensure_trace_runtime_from_env() -> bool: + """Lazily configure trace runtime in Ray child processes from inherited + env.""" + + global _RUNTIME + + env_vars = get_trace_env_vars_from_env() + if _RUNTIME is not None and _RUNTIME.runtime.enabled: + return True + if not env_vars: + return False + + env_vars = {key: str(env_vars[key]) for key in TRACE_ENV_KEYS if key in env_vars} + if "OTEL_EXPORTER_OTLP_ENDPOINT" not in env_vars: + return False + env_vars.setdefault("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", env_vars["OTEL_EXPORTER_OTLP_ENDPOINT"]) + env_vars.setdefault("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc") + env_vars.setdefault("OTEL_TRACES_EXPORTER", "otlp") + + run_dir = Path(env_vars.get("XTUNER_OTEL_RUN_DIR") or Path.cwd()).expanduser() + trace_jsonl_path = Path(env_vars.get("XTUNER_OTEL_JSONL_PATH") or run_dir / "traces" / "traces.jsonl").expanduser() + live_jsonl_path = Path( + env_vars.get("XTUNER_OTEL_LIVE_JSONL_PATH") or run_dir / "traces" / "live.jsonl" + ).expanduser() + runtime_handle = _TraceRuntimeHandle( + runtime=TraceRuntime( + enabled=True, + mode="inherited", + run_id=env_vars.get("XTUNER_OTEL_RUN_ID", ""), + run_dir=run_dir, + trace_jsonl_path=trace_jsonl_path, + live_jsonl_path=live_jsonl_path, + service_name=env_vars.get("OTEL_SERVICE_NAME", "xtuner-rollout"), + trace_viewer_url=None, + trace_viewer_port=None, + ), + endpoint=env_vars["OTEL_EXPORTER_OTLP_ENDPOINT"], + env_vars=env_vars, + ) + runtime_handle.start() + _RUNTIME = runtime_handle + register_atexit_once(close_trace) + return True + + +def is_trace_enabled() -> bool: + """Return whether XTuner trace runtime is enabled in this process.""" + + return _RUNTIME is not None and _RUNTIME.runtime.enabled + + +def close_trace() -> None: + """Flush provider state and stop local trace processes owned by this + process.""" + + global _RUNTIME + + runtime = _RUNTIME + _RUNTIME = None + if runtime is not None: + runtime.close() + clear_trace_env() + + +def apply_trace_env(env_vars: Mapping[str, str]) -> None: + clear_trace_env() + os.environ.update(env_vars) + + +def clear_trace_env() -> None: + for key in TRACE_ENV_KEYS: + os.environ.pop(key, None) + + +def register_atexit_once(close_fn) -> None: + global _ATEXIT_REGISTERED + if not _ATEXIT_REGISTERED: + atexit.register(close_fn) + _ATEXIT_REGISTERED = True diff --git a/xtuner/v1/rl/trace/trace_utils.py b/xtuner/v1/rl/trace/trace_utils.py new file mode 100644 index 0000000000..d482615333 --- /dev/null +++ b/xtuner/v1/rl/trace/trace_utils.py @@ -0,0 +1,451 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from xtuner.v1.data_proto.rl_data import RolloutState +from xtuner.v1.rl.trace.runtime import is_trace_enabled + + +# #################################################################### +# ***** Trace Span Names ***** +# #################################################################### +TRACE_SPAN_ROLLOUT_CONTROLLER_GENERATE = "rollout_controller.generate" +TRACE_SPAN_ROLLOUT_WORKER_GENERATE = "rollout_worker.generate" +TRACE_SPAN_AGENT_LOOP_RUN = "agent_loop.run" + +TRACE_SPAN_AGENT_LOCALHOST_RUNNER_RUN = "agent.localhost.runner.run" +TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN = "agent.localhost.infer.run" +TRACE_SPAN_AGENT_LOCALHOST_AGENT_INVOKE = "agent.localhost.agent.invoke" +TRACE_SPAN_AGENT_LOCALHOST_VALIDATE_RUN = "agent.localhost.validate.run" +TRACE_SPAN_AGENT_LOCALHOST_JUDGER_RUN = "agent.localhost.judger.run" +TRACE_SPAN_AGENT_LOCALHOST_TRAJECTORY_MATERIALIZE = "agent.localhost.trajectory.materialize" + +TRACE_SPAN_AGENT_SANDBOX_RUNNER_RUN = "agent.sandbox.runner.run" +TRACE_SPAN_AGENT_SANDBOX_ACQUIRE = "agent.sandbox.acquire" +TRACE_SPAN_AGENT_SANDBOX_INFER_RUN = "agent.sandbox.infer.run" +TRACE_SPAN_AGENT_SANDBOX_HOOK_RUN = "agent.sandbox.hook.run" +TRACE_SPAN_AGENT_SANDBOX_ENTRY_RUN = "agent.sandbox.entry.run" +TRACE_SPAN_AGENT_SANDBOX_VALIDATE_RUN = "agent.sandbox.validate.run" +TRACE_SPAN_AGENT_SANDBOX_JUDGER_RUN = "agent.sandbox.judger.run" + +TRACE_SPAN_SESSION_SERVER_REQUEST = "session_server.request" +TRACE_SPAN_SESSION_SERVER_PREPARE_REQUEST = "session_server.prepare_request" +TRACE_SPAN_SESSION_SERVER_SEND_REQUEST = "session_server.send_request" +TRACE_SPAN_SESSION_SERVER_READ_RESPONSE = "session_server.read_response" +TRACE_SPAN_SESSION_SERVER_RECORD_RESPONSE = "session_server.record_response" + +TraceAttributes = Mapping[str, Any] + + +_ERROR_STATUS_VALUES = {"error", "exception", "failed", "timeout", "timed_out"} + + +# #################################################################### +# ***** Attribute Builders ***** +# #################################################################### +def rollout_state_initial_attributes( + rollout_state: RolloutState, + *, + task_name: str | None = None, +) -> dict[str, Any]: + attributes: dict[str, Any] = { + "xtuner.status": rollout_state.status.value, + "xtuner.seq_staleness": rollout_state.seq_staleness, + } + trace_task_name = task_name or rollout_state.task_name + producer_future_step = rollout_state.extra_fields.get("producer_future_step") + _set_if_not_none(attributes, "xtuner.rollout_id", rollout_state.rollout_id) + _set_if_not_none(attributes, "xtuner.group_id", rollout_state.group_id) + _set_if_not_none(attributes, "xtuner.session_id", rollout_state.session_id) + _set_if_not_none(attributes, "xtuner.task_name", trace_task_name) + _set_if_not_none(attributes, "xtuner.producer_future_step", producer_future_step) + if rollout_state.prompt_ids is not None: + attributes["prompt.tokens"] = len(rollout_state.prompt_ids) + return attributes + + +def rollout_state_final_attributes( + rollout_state: RolloutState, + *, + task_name: str | None = None, +) -> dict[str, Any]: + attributes = rollout_state_initial_attributes(rollout_state, task_name=task_name) + _set_if_not_none(attributes, "xtuner.finish_reason", rollout_state.finish_reason) + if rollout_state.response_ids is not None: + attributes["completion.tokens"] = len(rollout_state.response_ids) + _set_if_not_none(attributes, "xtuner.error_msg", rollout_state.error_msg) + if rollout_state.reward is not None: + assert isinstance(rollout_state.reward, Mapping) + attributes.update( + reward_trace_attributes( + rollout_state.reward["score"], + passed=rollout_state.reward.get("pass"), + ) + ) + return attributes + + +def judger_trace_attributes( + judger: Any, + rollout_state: RolloutState | list[RolloutState], +) -> dict[str, Any]: + judger_name = judger.get_judger_name() + assert judger_name is not None + attributes: dict[str, Any] = {"judger.name": judger_name} + if isinstance(rollout_state, list): + attributes["judger.batch_size"] = len(rollout_state) + return attributes + attributes.update(rollout_state_final_attributes(rollout_state)) + return attributes + + +def agent_item_initial_attributes(item: Any) -> dict[str, Any]: + if isinstance(item, RolloutState): + return rollout_state_initial_attributes(item) + + attributes: dict[str, Any] = {} + session_id = getattr(item, "uid", None) + group_id = getattr(item, "group_id", None) + task_name = getattr(item, "data_source", None) + item_id = getattr(item, "id", None) + ability = getattr(item, "ability", None) + status = getattr(item, "status", None) + status_value = getattr(status, "value", status) + _set_if_not_none(attributes, "xtuner.session_id", session_id) + _set_if_not_none(attributes, "xtuner.group_id", group_id) + _set_if_not_none(attributes, "xtuner.task_name", task_name) + _set_if_not_none(attributes, "agent.item_id", item_id) + _set_if_not_none(attributes, "agent.data_source", task_name) + _set_if_not_none(attributes, "agent.ability", ability) + _set_if_not_none(attributes, "agent.status", status_value) + _set_if_not_none(attributes, "xtuner.status", status_value) + return attributes + + +def agent_judger_initial_attributes( + judger: Any, + item: Any, + *, + reward_key: str | None = None, +) -> dict[str, Any]: + attributes = agent_item_initial_attributes(item) + judger_name = _judger_name(judger) + assert judger_name is not None + attributes["judger.name"] = judger_name + if reward_key is not None: + attributes["reward.key"] = reward_key + return attributes + + +def agent_item_final_attributes(item: Any) -> dict[str, Any]: + if isinstance(item, RolloutState): + return rollout_state_final_attributes(item) + + status = getattr(item, "status", None) + assert status is not None + status_value = status.value + attributes: dict[str, Any] = { + "agent.status": status_value, + "xtuner.status": status_value, + } + reward = getattr(item, "reward", None) + if reward is not None: + attributes.update(reward_trace_attributes(reward)) + + artifacts = getattr(item, "artifacts", None) + assert isinstance(artifacts, Mapping) + messages = artifacts.get("messages") + if messages is not None: + assert isinstance(messages, list) + if messages: + last_segment = messages[-1] + assert isinstance(last_segment, Mapping) + segment_messages = last_segment.get("messages") + assert isinstance(segment_messages, list) + attributes["agent.message_count"] = len(segment_messages) + attributes["agent.tool_turns"] = _count_tool_turns(segment_messages) + segment_tools = last_segment.get("tools") + if segment_tools is not None: + attributes["agent.has_tools"] = bool(segment_tools) + response_message = artifacts.get("response_message") + if response_message is not None: + assert isinstance(response_message, Mapping) + finish_info = response_message.get("finish_info") + if finish_info is not None: + attributes["agent.finish_info"] = str(finish_info) + return attributes + + +def sandbox_entry_attributes(entry: Any) -> dict[str, Any]: + entry_name = getattr(entry, "name", None) + assert entry_name is not None + entry_kind = getattr(entry, "kind", type(entry).__name__) + assert entry_kind is not None + return { + "entry.name": str(entry_name), + "entry.kind": str(entry_kind), + } + + +def sandbox_entry_result_attributes(result: Any) -> dict[str, Any]: + payload = getattr(result, "result", result) + return_code = getattr(payload, "return_code", None) + assert return_code is not None + attributes: dict[str, Any] = {"entry.return_code": return_code} + retryable = getattr(result, "retryable", None) + if retryable is not None: + attributes["entry.retryable"] = retryable + return attributes + + +def sandbox_entry_error_attributes(result: Any) -> dict[str, Any]: + payload = getattr(result, "result", result) + return_code = getattr(payload, "return_code", None) + error = getattr(payload, "error", None) + stderr = getattr(payload, "stderr", None) + + if return_code not in (None, 0): + return failure_attributes("entry_return_code", message=f"return_code={return_code}", return_code=return_code) + if error: + return failure_attributes("entry_exception", message=str(error)) + if return_code is None and stderr: + return failure_attributes("entry_exception", message=str(stderr)) + return {} + + +def failure_attributes( + category: str, + *, + message: str | None = None, + attributes: Mapping[str, Any] | None = None, + **kwargs: Any, +) -> dict[str, Any]: + if not isinstance(category, str) or not category.strip(): + raise ValueError("failure category cannot be empty") + failure: dict[str, Any] = {"error.category": category.strip()} + if message is not None: + failure["error.message"] = message + merged = dict(attributes or {}) + merged.update(kwargs) + for key, value in merged.items(): + failure[f"error.{key}"] = value + return failure + + +def reward_trace_attributes(score: Any, *, passed: Any | None = None) -> dict[str, Any]: + assert score is not None + assert not isinstance(score, bool) + normalized_score = float(score) + if passed is None: + passed = normalized_score > 0 + assert isinstance(passed, bool) + return { + "reward.score": normalized_score, + "reward_score": normalized_score, + "reward.pass": passed, + } + + +# #################################################################### +# ***** Span And Event Recorders ***** +# #################################################################### +def add_rollout_status_event( + rollout_state: RolloutState, + *, + task_name: str | None = None, +) -> None: + from xtuner.v1.rl.trace import trace_event + + trace_event("rollout.status", rollout_state_final_attributes(rollout_state, task_name=task_name)) + + +def record_rollout_state_result( + rollout_state: Any, + *, + task_name: str | None = None, +) -> None: + if not isinstance(rollout_state, RolloutState): + return + from xtuner.v1.rl.trace import set_trace_attributes + + set_trace_attributes(rollout_state_final_attributes(rollout_state, task_name=task_name)) + add_rollout_status_event(rollout_state, task_name=task_name) + if _is_error_status(rollout_state.status): + _set_trace_error(_rollout_error_message(rollout_state)) + + +def add_agent_status_event(item: Any) -> None: + from xtuner.v1.rl.trace import trace_event + + trace_event("agent.status", agent_item_final_attributes(item)) + + +def record_agent_item_result(item: Any, result: Any | None = None) -> None: + target = result if result is not None and result is item else item + from xtuner.v1.rl.trace import set_trace_attributes + + set_trace_attributes(agent_item_final_attributes(target)) + add_agent_status_event(target) + if _is_error_status(getattr(target, "status", None)): + _set_trace_error(_agent_item_error_message(target)) + + +# #################################################################### +# ***** Internal Helpers ***** +# #################################################################### +def _validate_trace_name(name: str, *, kind: str) -> str: + if not isinstance(name, str): + raise TypeError(f"{kind} must be a string") + normalized_name = name.strip() + if not normalized_name: + raise ValueError(f"{kind} cannot be empty") + return normalized_name + + +def _normalize_trace_attributes(attributes: Mapping[str, Any] | None) -> dict[str, Any]: + if attributes is None: + return {} + 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 + + +def _normalize_trace_attribute_value(value: Any) -> Any: + def normalize_scalar(item: Any) -> str | bool | int | float: + if isinstance(item, bool): + return item + if isinstance(item, int): + return item if -(2**63) <= item <= 2**63 - 1 else str(item) + if isinstance(item, (str, float)): + return item + return f"<{type(item).__name__}>" + + if isinstance(value, bool): + return value + if isinstance(value, int): + return value if -(2**63) <= value <= 2**63 - 1 else str(value) + if isinstance(value, (str, float)): + return value + if isinstance(value, (list, tuple)): + items = tuple(normalize_scalar(item) for item in value) + if not items: + return items + if all(isinstance(item, str) for item in items): + return items + if all(isinstance(item, bool) for item in items): + return items + if all(isinstance(item, int) and not isinstance(item, bool) for item in items): + return items + if all(isinstance(item, float) for item in items): + return items + return tuple(str(item) for item in items) + return f"<{type(value).__name__}>" + + +def _set_trace_error(message: str | None = None) -> None: + if not is_trace_enabled(): + return + from xtuner.v1.rl.trace import otel_utils + + otel_utils.set_error_status(message) + + +def _set_if_not_none(attributes: dict[str, Any], key: str, value: Any) -> None: + if value is not None: + attributes[key] = value + + +def _is_error_status(status: Any) -> bool: + value = getattr(status, "value", status) + return isinstance(value, str) and value.strip().lower() in _ERROR_STATUS_VALUES + + +def _rollout_error_message(rollout_state: RolloutState) -> str: + if rollout_state.error_msg: + return str(rollout_state.error_msg) + status = getattr(rollout_state.status, "value", rollout_state.status) + return f"rollout status={status}" + + +def _agent_item_error_message(item: Any) -> str: + error = getattr(item, "error", None) + if error is not None: + message = getattr(error, "message", None) + stage = getattr(error, "stage", None) + category = getattr(error, "category", None) + parts = [str(value) for value in (stage, category, message) if value] + if parts: + return ": ".join(parts) + return str(error) + status = getattr(item, "status", None) + return f"agent item status={status.value if status is not None else None}" + + +def _judger_name(judger: Any) -> str | None: + value = getattr(judger, "name", None) + if value is not None: + return str(value) + get_judger_name = getattr(judger, "get_judger_name", None) + if callable(get_judger_name): + return str(get_judger_name()) + return None + + +def _count_tool_turns(messages: list[Any]) -> int: + count = 0 + for message in messages: + if not isinstance(message, Mapping): + continue + if message.get("tool_calls") or message.get("role") == "tool": + count += 1 + return count + + +__all__ = [ + "TRACE_SPAN_AGENT_LOCALHOST_AGENT_INVOKE", + "TRACE_SPAN_AGENT_LOCALHOST_INFER_RUN", + "TRACE_SPAN_AGENT_LOCALHOST_JUDGER_RUN", + "TRACE_SPAN_AGENT_LOCALHOST_RUNNER_RUN", + "TRACE_SPAN_AGENT_LOCALHOST_TRAJECTORY_MATERIALIZE", + "TRACE_SPAN_AGENT_LOCALHOST_VALIDATE_RUN", + "TRACE_SPAN_AGENT_LOOP_RUN", + "TRACE_SPAN_AGENT_SANDBOX_ACQUIRE", + "TRACE_SPAN_AGENT_SANDBOX_ENTRY_RUN", + "TRACE_SPAN_AGENT_SANDBOX_HOOK_RUN", + "TRACE_SPAN_AGENT_SANDBOX_INFER_RUN", + "TRACE_SPAN_AGENT_SANDBOX_JUDGER_RUN", + "TRACE_SPAN_AGENT_SANDBOX_RUNNER_RUN", + "TRACE_SPAN_AGENT_SANDBOX_VALIDATE_RUN", + "TRACE_SPAN_ROLLOUT_CONTROLLER_GENERATE", + "TRACE_SPAN_ROLLOUT_WORKER_GENERATE", + "TRACE_SPAN_SESSION_SERVER_PREPARE_REQUEST", + "TRACE_SPAN_SESSION_SERVER_READ_RESPONSE", + "TRACE_SPAN_SESSION_SERVER_RECORD_RESPONSE", + "TRACE_SPAN_SESSION_SERVER_REQUEST", + "TRACE_SPAN_SESSION_SERVER_SEND_REQUEST", + "TraceAttributes", + # Attribute builders. + "agent_judger_initial_attributes", + "agent_item_final_attributes", + "agent_item_initial_attributes", + "failure_attributes", + "judger_trace_attributes", + "reward_trace_attributes", + "rollout_state_final_attributes", + "rollout_state_initial_attributes", + "sandbox_entry_attributes", + "sandbox_entry_error_attributes", + "sandbox_entry_result_attributes", + "record_agent_item_result", + "record_rollout_state_result", +] diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index c86545bbcd..7cdb8fe4e0 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -166,16 +166,20 @@ def build(self, placement_group: "PlacementGroup"): from xtuner.v1.rl.trainer.controller import TrainingController from xtuner.v1.rl.utils import AutoAcceleratorWorkers - TrainingWorkerCls = ray.remote( - runtime_env={ - "env_vars": { - "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", - "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", - "HCCL_NPU_SOCKET_PORT_RANGE": "auto", - } + training_runtime_env = { + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + "HCCL_NPU_SOCKET_PORT_RANGE": "auto", } - )(TrainingWorker) - train_workers, _ = AutoAcceleratorWorkers.from_placement_group(TrainingWorkerCls, self, placement_group) + } + TrainingWorkerCls = ray.remote(TrainingWorker) + train_workers, _ = AutoAcceleratorWorkers.from_placement_group( + TrainingWorkerCls, + self, + placement_group, + runtime_env=training_runtime_env, + ) ray.wait([w.ready.remote() for w in train_workers]) return TrainingController(workers=train_workers) diff --git a/xtuner/v1/rl/utils/__init__.py b/xtuner/v1/rl/utils/__init__.py index d2410ba795..2d2282774e 100644 --- a/xtuner/v1/rl/utils/__init__.py +++ b/xtuner/v1/rl/utils/__init__.py @@ -47,7 +47,9 @@ free_object_refs, get_accelerator_ids, get_ray_accelerator, + merge_trace_runtime_env, register_cleanup, + with_trace_runtime_env, ) @@ -83,6 +85,8 @@ "get_accelerator_ids", "free_object_refs", "bind_train_rollout", + "merge_trace_runtime_env", + "with_trace_runtime_env", "handle_task_exception", "create_task", "cancel_and_drain", diff --git a/xtuner/v1/rl/utils/ray_accelerator_worker.py b/xtuner/v1/rl/utils/ray_accelerator_worker.py index a7d52d981a..71fd51b300 100644 --- a/xtuner/v1/rl/utils/ray_accelerator_worker.py +++ b/xtuner/v1/rl/utils/ray_accelerator_worker.py @@ -15,7 +15,7 @@ ) from typing_extensions import Annotated -from .ray_utils import find_master_addr_and_port, get_accelerator_ids +from .ray_utils import find_master_addr_and_port, get_accelerator_ids, with_trace_runtime_env PG_READY_TIMEOUT = os.getenv("XTUNER_PG_READY_TIMEOUT", 30) # default 30 seconds @@ -443,7 +443,12 @@ def from_config(cls, worker_cls, worker_config, accelerator_config: AcceleratorR @classmethod def from_placement_group( - cls, worker_cls: ActorClass[T], worker_config, pg: PlacementGroup + cls, + worker_cls: ActorClass[T], + worker_config, + pg: PlacementGroup, + *, + runtime_env: dict[str, Any] | None = None, ) -> tuple[list[ActorProxy[T]], list[tuple[int, int]]]: """Create workers from an existing placement group. @@ -464,11 +469,22 @@ def from_placement_group( workers_list: list[ActorProxy[T]] = [] rank_bundle_idx_list: list[tuple[int, int]] = [] for rank, bundle_idx in enumerate(sorted_bundle_idxs): - worker = worker_cls.options( - placement_group=pg, - placement_group_bundle_index=bundle_idx, - **pg_options, - ).remote(worker_config, rank, master_addr, master_port, world_size, device_type) + ray_options = with_trace_runtime_env( + { + "placement_group": pg, + "placement_group_bundle_index": bundle_idx, + "runtime_env": runtime_env, + **pg_options, + } + ) + worker = worker_cls.options(**ray_options).remote( + worker_config, + rank, + master_addr, + master_port, + world_size, + device_type, + ) workers_list.append(worker) rank_bundle_idx_list.append((rank, bundle_idx)) diff --git a/xtuner/v1/rl/utils/ray_cpu_worker.py b/xtuner/v1/rl/utils/ray_cpu_worker.py index ed9ebc3783..535db06c81 100644 --- a/xtuner/v1/rl/utils/ray_cpu_worker.py +++ b/xtuner/v1/rl/utils/ray_cpu_worker.py @@ -19,6 +19,8 @@ from xtuner.v1.utils.logger import get_logger +from .ray_utils import with_trace_runtime_env + PG_READY_TIMEOUT = os.getenv("XTUNER_PG_READY_TIMEOUT", 30) # default 30 seconds PlacementGroups: TypeAlias = PlacementGroup | list[PlacementGroup] | tuple[PlacementGroup, ...] | None @@ -187,7 +189,7 @@ def build_actor( actor_options["memory"] = resolved_memory if pg is None: - return actor_cls.options(**actor_options).remote(*init_args, **init_kwargs) + return actor_cls.options(**with_trace_runtime_env(actor_options)).remote(*init_args, **init_kwargs) resolved_num_cpus, resolved_memory = cls._resolve_actor_resources( pg=pg, @@ -203,7 +205,7 @@ def build_actor( placement_group_bundle_index=bundle_idx, placement_group_capture_child_tasks=capture_child_tasks, ) - return actor_cls.options(**actor_options).remote(*init_args, **init_kwargs) + return actor_cls.options(**with_trace_runtime_env(actor_options)).remote(*init_args, **init_kwargs) @classmethod def build_actors( diff --git a/xtuner/v1/rl/utils/ray_utils.py b/xtuner/v1/rl/utils/ray_utils.py index ad6b1dcc9c..55f6f2728a 100644 --- a/xtuner/v1/rl/utils/ray_utils.py +++ b/xtuner/v1/rl/utils/ray_utils.py @@ -1,8 +1,8 @@ import atexit import signal import subprocess -from collections.abc import Iterable -from typing import TYPE_CHECKING, Optional, cast +from collections.abc import Iterable, Mapping +from typing import TYPE_CHECKING, Any, Optional, cast import ray from ray import ObjectRef @@ -18,6 +18,51 @@ logger = get_logger() +_PROTECTED_TRACE_ENV_KEYS = frozenset( + { + "XTUNER_OTEL_ENABLED", + "XTUNER_OTEL_OUTPUT_DIR", + "XTUNER_OTEL_RUN_ID", + "XTUNER_OTEL_RUN_DIR", + "XTUNER_OTEL_JSONL_PATH", + "OTEL_TRACES_EXPORTER", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_EXPORTER_OTLP_PROTOCOL", + } +) + + +def merge_trace_runtime_env(runtime_env: Mapping[str, Any] | None = None) -> dict[str, Any]: + """Merge active trace runtime env into a Ray runtime_env dict.""" + + merged_runtime_env = dict(runtime_env or {}) + + from xtuner.v1.rl.trace.runtime import get_trace_env_vars + + trace_env = get_trace_env_vars() + if not trace_env: + return merged_runtime_env + + env_vars = dict(merged_runtime_env.get("env_vars") or {}) + for key, value in trace_env.items(): + if key in _PROTECTED_TRACE_ENV_KEYS or key not in env_vars: + env_vars[key] = value + merged_runtime_env["env_vars"] = env_vars + return merged_runtime_env + + +def with_trace_runtime_env(ray_options: Mapping[str, Any] | None = None) -> dict[str, Any]: + """Return Ray options with active trace runtime env merged in.""" + + options = dict(ray_options or {}) + runtime_env = merge_trace_runtime_env(options.get("runtime_env")) + if runtime_env: + options["runtime_env"] = runtime_env + else: + options.pop("runtime_env", None) + return options + @ray.remote def find_master_addr_and_port( diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index eecb796ddd..d3a989a988 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -41,6 +41,7 @@ ) from xtuner.v1.rl.rollout.controller import RolloutControllerProxy from xtuner.v1.rl.rollout.worker import RolloutConfig +from xtuner.v1.rl.trace import TraceConfig, close_trace, configure_trace from xtuner.v1.rl.trainer.controller import TrainingController from xtuner.v1.rl.trainer.worker import WorkerConfig, WorkerLogItem from xtuner.v1.rl.utils import ( @@ -360,6 +361,7 @@ class BaseRLTrainerConfig(BaseModel): debug_train: bool = False skip_checkpoint_validation: bool = False exp_tracker: Literal["tensorboard", "jsonl"] = "tensorboard" + trace_config: TraceConfig = Field(default_factory=TraceConfig) @model_validator(mode="after") def _validate_sync_intervals(self): @@ -583,6 +585,7 @@ def _init_common(self, cfg: BaseRLTrainerConfig, *, meta_path: str, logger_tag: self._init_load_source(cfg) self._init_save_config(cfg) log_dir = self._init_logger(cfg, logger_tag) + self._init_trace(cfg) self._save_runtime_environment(log_dir) self._init_train_state(cfg) self._init_train_worker_config(cfg, log_dir) @@ -633,6 +636,12 @@ def _init_logger(self, cfg: BaseRLTrainerConfig, logger_tag: str) -> Path: patch_default_save_plan() return log_dir + def _init_trace(self, cfg: BaseRLTrainerConfig) -> None: + trace_config = cfg.trace_config + if trace_config.output_dir is None: + trace_config = trace_config.model_copy(update={"output_dir": self.exp_dir / "otel"}) + self._trace_runtime = configure_trace(trace_config) + def _save_runtime_environment(self, log_dir: Path) -> None: if get_rank() != 0: return @@ -1627,6 +1636,7 @@ def fit(self): self._fit() finally: self._exp_tracker.close() + close_trace() def _fit(self): self.logger.info("Start RL training") @@ -1855,6 +1865,7 @@ def fit(self): return asyncio_run(self._fit()) finally: self._exp_tracker.close() + close_trace() async def _get_batch_or_raise_producer_failure( self,