From d9221f56bbf57ae27854249631f45bc14fadefbb Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Mon, 8 Jun 2026 11:43:00 +0000 Subject: [PATCH 1/3] [Feat] Add producer task trace and viewer --- xtuner/tools/producer_trace_analysis.py | 252 ++++ xtuner/tools/producer_trace_hotspots.py | 656 ++++++++++ xtuner/tools/producer_trace_viewer.py | 813 ++++++++++++ xtuner/v1/rl/agent_loop/agent_loop.py | 11 + .../rl/agent_loop/single_turn_agent_loop.py | 6 + .../agent_loop_manager/agent_loop_manager.py | 7 +- xtuner/v1/rl/agent_loop_manager/producer.py | 46 + xtuner/v1/rl/agent_loop_manager/sampler.py | 11 +- xtuner/v1/rl/rollout/controller.py | 6 + xtuner/v1/rl/rollout/worker.py | 18 +- xtuner/v1/rl/trace.py | 1104 +++++++++++++++++ xtuner/v1/rl/utils/ray_accelerator_worker.py | 3 + xtuner/v1/rl/utils/ray_cpu_worker.py | 2 + xtuner/v1/train/rl_trainer.py | 48 +- 14 files changed, 2977 insertions(+), 6 deletions(-) create mode 100644 xtuner/tools/producer_trace_analysis.py create mode 100644 xtuner/tools/producer_trace_hotspots.py create mode 100644 xtuner/tools/producer_trace_viewer.py create mode 100644 xtuner/v1/rl/trace.py diff --git a/xtuner/tools/producer_trace_analysis.py b/xtuner/tools/producer_trace_analysis.py new file mode 100644 index 0000000000..a85bd59040 --- /dev/null +++ b/xtuner/tools/producer_trace_analysis.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import dataclasses +import json +import time +from collections import defaultdict +from pathlib import Path +from typing import Iterable + +from xtuner.v1.rl.trace import ( + TRACE_JSONL_BASENAME, + TRACE_VIEWER_SCOPE_ALL, + TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, + TraceEvent, + TraceViewerScope, +) + + +TRACE_STAGE_LABELS = { + "xtuner.producer.sample_group": "sampler", + "xtuner.producer.generate_group": "producer.generate", + "xtuner.producer.put_generated_group": "producer.put", + "xtuner.agent_loop.generate_group": "agent_loop.generate_group", + "xtuner.agent_loop.generate_sample": "agent_loop.generate_sample", + "xtuner.rollout_controller.generate": "rollout.generate", + "xtuner.rollout_worker.generate": "rollout_worker.generate", + "xtuner.rollout_engine.generate": "engine.generate", + "xtuner.judger.judge": "judger", +} + + +@dataclasses.dataclass(frozen=True) +class TraceViewerRow: + trace_id: str + task_name: str | None + uid: int | str | None + status: str | None + latest_stage: str + latest_timestamp_s: float + event_count: int + open_span: str | None = None + open_age_s: float | None = None + + +@dataclasses.dataclass(frozen=True) +class OpenSpanSummary: + span: str + open_count: int + oldest_age_s: float + p50_age_s: float + p95_age_s: float + oldest_trace_id: str + oldest_task_name: str | None + oldest_uid: int | str | None + + +def display_trace_stage(span: str | None) -> str: + if not span: + return "unknown" + if span in TRACE_STAGE_LABELS: + return TRACE_STAGE_LABELS[span] + if span.startswith("xtuner.") and span.endswith(".request"): + return span.removeprefix("xtuner.") + return span.removeprefix("xtuner.") + + +def load_trace_jsonl(path: str | Path) -> list[TraceEvent]: + path = Path(path) + if path.is_file(): + files = [path] + else: + files = sorted(path.glob(f"{TRACE_JSONL_BASENAME}_*.jsonl")) + events: list[TraceEvent] = [] + for file in files: + with file.open("r", encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + events.append(TraceEvent.from_dict(json.loads(line))) + return events + + +def events_to_timelines(events: Iterable[TraceEvent]) -> dict[str, list[TraceEvent]]: + timelines: dict[str, list[TraceEvent]] = defaultdict(list) + for event in events: + timelines[event.trace_id].append(event) + for trace_id in timelines: + timelines[trace_id].sort(key=lambda event: event.timestamp_s) + return dict(timelines) + + +def filter_trace_events_by_scope( + events: Iterable[TraceEvent], + scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, +) -> list[TraceEvent]: + event_list = list(events) + if scope == TRACE_VIEWER_SCOPE_ALL: + return event_list + if scope != TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH: + raise ValueError(f"Unsupported trace viewer scope: {scope!r}") + + latest_batch_id = _latest_produce_batch_id(event_list) + if latest_batch_id is not None: + return [event for event in event_list if event.produce_batch_id == latest_batch_id] + + latest_key = _latest_produce_batch_key(event_list) + if latest_key is None: + return event_list + return [event for event in event_list if _produce_batch_key(event) == latest_key] + + +def build_viewer_rows( + timelines: dict[str, list[TraceEvent]] | Iterable[TraceEvent], + *, + now_s: float | None = None, +) -> list[TraceViewerRow]: + if not isinstance(timelines, dict): + timelines = events_to_timelines(timelines) + now_s = time.time() if now_s is None else now_s + rows: list[TraceViewerRow] = [] + for trace_id, events in timelines.items(): + if not events: + continue + sorted_events = sorted(events, key=lambda event: event.timestamp_s) + latest = sorted_events[-1] + open_spans = _get_open_spans(sorted_events) + newest_open = max(open_spans, key=lambda span: span[1].timestamp_s) if open_spans else None + open_span = newest_open[0] if newest_open is not None else None + open_age_s = now_s - newest_open[1].timestamp_s if newest_open is not None else None + rows.append( + TraceViewerRow( + trace_id=trace_id, + task_name=latest.task_name, + uid=latest.uid, + status=latest.status, + latest_stage=latest.stage, + latest_timestamp_s=latest.timestamp_s, + event_count=len(sorted_events), + open_span=open_span, + open_age_s=open_age_s, + ) + ) + return sorted(rows, key=lambda row: ((row.open_age_s or 0.0), row.latest_timestamp_s), reverse=True) + + +def build_open_span_summaries( + timelines: dict[str, list[TraceEvent]] | Iterable[TraceEvent], + *, + now_s: float | None = None, +) -> list[OpenSpanSummary]: + if not isinstance(timelines, dict): + timelines = events_to_timelines(timelines) + now_s = time.time() if now_s is None else now_s + grouped: dict[str, list[tuple[float, str, TraceEvent]]] = defaultdict(list) + for trace_id, events in timelines.items(): + for span, start_event in _get_open_spans(sorted(events, key=lambda event: event.timestamp_s)): + grouped[span].append((now_s - start_event.timestamp_s, trace_id, start_event)) + + summaries: list[OpenSpanSummary] = [] + for span, entries in grouped.items(): + entries.sort(key=lambda item: item[0], reverse=True) + ages = sorted(age for age, _, _ in entries) + oldest_age, oldest_trace_id, oldest_event = entries[0] + summaries.append( + OpenSpanSummary( + span=span, + open_count=len(entries), + oldest_age_s=oldest_age, + p50_age_s=_percentile(ages, 0.50), + p95_age_s=_percentile(ages, 0.95), + oldest_trace_id=oldest_trace_id, + oldest_task_name=oldest_event.task_name, + oldest_uid=oldest_event.uid, + ) + ) + return sorted(summaries, key=lambda summary: (summary.oldest_age_s, summary.open_count), reverse=True) + + +def _latest_produce_batch_key(events: list[TraceEvent]) -> tuple[int, int, int] | None: + keys: list[tuple[int, int, int]] = [] + for event in events: + key = _produce_batch_key(event) + if key is not None: + keys.append(key) + return max(keys) if keys else None + + +def _latest_produce_batch_id(events: list[TraceEvent]) -> str | None: + latest_event: TraceEvent | None = None + latest_sort_key: tuple[int, int, int, float] | None = None + for event in events: + if event.produce_batch_id is None: + continue + batch_key = _produce_batch_key(event) + if batch_key is None: + sort_key = (-1, -1, -1, event.timestamp_s) + else: + sort_key = (*batch_key, event.timestamp_s) + if latest_sort_key is None or sort_key > latest_sort_key: + latest_event = event + latest_sort_key = sort_key + if latest_event is None: + return None + return latest_event.produce_batch_id + + +def _produce_batch_key(event: TraceEvent) -> tuple[int, int, int] | None: + if event.train_step is None: + return None + return ( + event.train_step, + -1 if event.model_step is None else event.model_step, + -1 if event.producer_future_step is None else event.producer_future_step, + ) + + +def _get_open_spans(events: list[TraceEvent]) -> list[tuple[str, TraceEvent]]: + stacks: dict[str, list[TraceEvent]] = defaultdict(list) + for event in events: + span, suffix = _split_span_stage(event.stage) + if span is None or suffix is None: + continue + if suffix == "start": + stacks[span].append(event) + elif suffix in {"end", "error"} and stacks.get(span): + stacks[span].pop() + open_spans: list[tuple[str, TraceEvent]] = [] + for span, stack in stacks.items(): + open_spans.extend((span, event) for event in stack) + return open_spans + + +def _split_span_stage(stage: str) -> tuple[str | None, str | None]: + for suffix in (".start", ".end", ".error"): + if stage.endswith(suffix): + return stage[: -len(suffix)], suffix[1:] + return None, None + + +def _percentile(sorted_values: list[float], percentile: float) -> float: + if not sorted_values: + return 0.0 + if len(sorted_values) == 1: + return sorted_values[0] + position = (len(sorted_values) - 1) * percentile + lower = int(position) + upper = min(lower + 1, len(sorted_values) - 1) + if lower == upper: + return sorted_values[lower] + fraction = position - lower + return sorted_values[lower] * (1 - fraction) + sorted_values[upper] * fraction diff --git a/xtuner/tools/producer_trace_hotspots.py b/xtuner/tools/producer_trace_hotspots.py new file mode 100644 index 0000000000..b8418d2d0c --- /dev/null +++ b/xtuner/tools/producer_trace_hotspots.py @@ -0,0 +1,656 @@ +from __future__ import annotations + +import argparse +import dataclasses +import json +from collections import defaultdict +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable + +from xtuner.v1.rl.trace import ( + TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, + TraceEvent, + TraceViewerScope, +) +from xtuner.tools.producer_trace_analysis import ( + display_trace_stage, + filter_trace_events_by_scope, + load_trace_jsonl, +) + + +_PALETTE = [ + "#2563eb", + "#059669", + "#d97706", + "#7c3aed", + "#dc2626", + "#0891b2", + "#9333ea", + "#4d7c0f", + "#be123c", + "#0f766e", + "#b45309", + "#475569", +] + + +@dataclass +class TraceSpanRecord: + trace_id: str + span: str + display_stage: str + start_s: float + end_s: float + duration_s: float + outcome: str + depth: int = 0 + task_name: str | None = None + uid: int | str | None = None + status: str | None = None + train_step: int | None = None + worker_rank: int | None = None + error_msg: str | None = None + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Build an offline Producer Trace hotspot timeline HTML.") + parser.add_argument("trace_dir", type=Path, help="Directory containing producer_trace_*.jsonl files.") + parser.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="Output HTML file. Defaults to /producer_trace_hotspots.html.", + ) + parser.add_argument( + "--scope", + choices=("latest-produce-batch", "all"), + default=TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, + help="Trace scope shown by the hotspot viewer. Defaults to the latest produce batch.", + ) + return parser.parse_args() + + +def build_trace_span_records( + events: Iterable[TraceEvent], + *, + include_open: bool = True, + now_s: float | None = None, +) -> list[TraceSpanRecord]: + events_by_trace: dict[str, list[TraceEvent]] = defaultdict(list) + latest_event_s = 0.0 + for event in events: + events_by_trace[event.trace_id].append(event) + latest_event_s = max(latest_event_s, event.timestamp_s) + now_s = latest_event_s if now_s is None else now_s + + records: list[TraceSpanRecord] = [] + for trace_id, trace_events in events_by_trace.items(): + stacks: dict[str, list[TraceEvent]] = defaultdict(list) + for event in sorted(trace_events, key=lambda item: item.timestamp_s): + span, suffix = _split_span_stage(event.stage) + if span is None or suffix is None: + continue + if suffix == "start": + stacks[span].append(event) + continue + if suffix in {"end", "error"} and stacks.get(span): + start_event = stacks[span].pop() + records.append(_build_span_record(start_event, event, outcome=suffix)) + continue + if suffix in {"end", "error"} and event.elapsed_s is not None: + records.append(_build_elapsed_only_span_record(event, span=span, outcome=suffix)) + + if include_open: + for span_starts in stacks.values(): + for start_event in span_starts: + span, _ = _split_span_stage(start_event.stage) + if span is None: + continue + records.append(_build_open_span_record(start_event, span=span, now_s=now_s)) + + _assign_depths(records) + return sorted(records, key=lambda record: (record.trace_id, record.start_s, record.depth, record.end_s)) + + +def build_timeline_stage_records(records: Iterable[TraceSpanRecord]) -> list[TraceSpanRecord]: + return sorted(records, key=lambda record: (record.trace_id, record.start_s, record.depth, record.end_s)) + + +def build_hotspot_payload( + trace_dir: Path, + *, + scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, +) -> dict[str, Any]: + return build_hotspot_payload_from_events(load_trace_jsonl(trace_dir), trace_source=str(trace_dir), scope=scope) + + +def build_hotspot_payload_from_events( + events: Iterable[TraceEvent], + *, + trace_source: str, + scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, +) -> dict[str, Any]: + events = filter_trace_events_by_scope(events, scope) + raw_records = build_trace_span_records(events) + records = build_timeline_stage_records(raw_records) + if records: + global_start_s = min(record.start_s for record in records) + global_end_s = max(record.end_s for record in records) + else: + global_start_s = 0.0 + global_end_s = 1.0 + global_window_s = max(0.001, global_end_s - global_start_s) + stage_colors = _build_stage_colors(records) + + rows = [] + for trace_id, trace_records in _group_records(records).items(): + trace_start = min(record.start_s for record in trace_records) + trace_end = max(record.end_s for record in trace_records) + trace_window_s = max(0.001, trace_end - trace_start) + max_depth = max(record.depth for record in trace_records) if trace_records else 0 + rows.append( + { + "trace_id": trace_id, + "task_name": trace_records[-1].task_name, + "uid": trace_records[-1].uid, + "duration_s": trace_end - trace_start, + "start_s": trace_start, + "end_s": trace_end, + "row_height_px": 56 + (max_depth + 1) * 22, + "spans": [ + _span_record_to_payload(record, trace_start, trace_window_s, stage_colors) + for record in trace_records + ], + } + ) + rows.sort(key=lambda row: row["duration_s"], reverse=True) + + return { + "title": "Producer Trace Hotspots", + "trace_source": trace_source, + "scope": scope, + "task_count": len(rows), + "span_count": len(records), + "raw_span_count": len(raw_records), + "timeline_start_s": global_start_s, + "timeline_end_s": global_end_s, + "timeline_duration_s": global_window_s, + "max_task_duration_s": max((row["duration_s"] for row in rows), default=0.0), + "scale_mode": "task_relative", + "stage_colors": stage_colors, + "stage_stats": _build_stage_stats(records), + "rows": rows, + } + + +def write_hotspot_html(payload: dict[str, Any], output_path: Path) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(render_hotspot_html(payload), encoding="utf-8") + + +def render_hotspot_html(payload: dict[str, Any]) -> str: + data_json = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).replace(" None: + args = parse_args() + trace_dir = args.trace_dir + output_path = args.output or trace_dir / "producer_trace_hotspots.html" + payload = build_hotspot_payload(trace_dir, scope=args.scope) + write_hotspot_html(payload, output_path) + print(output_path) + + +def _build_span_record(start_event: TraceEvent, end_event: TraceEvent, *, outcome: str) -> TraceSpanRecord: + span, _ = _split_span_stage(start_event.stage) + assert span is not None + duration_s = max(0.0, end_event.timestamp_s - start_event.timestamp_s) + return TraceSpanRecord( + trace_id=start_event.trace_id, + span=span, + display_stage=display_trace_stage(span), + start_s=start_event.timestamp_s, + end_s=end_event.timestamp_s, + duration_s=duration_s, + outcome=outcome, + task_name=end_event.task_name or start_event.task_name, + uid=end_event.uid if end_event.uid is not None else start_event.uid, + status=end_event.status or start_event.status, + train_step=end_event.train_step if end_event.train_step is not None else start_event.train_step, + worker_rank=end_event.worker_rank if end_event.worker_rank is not None else start_event.worker_rank, + error_msg=end_event.error_msg, + ) + + +def _build_elapsed_only_span_record(end_event: TraceEvent, *, span: str, outcome: str) -> TraceSpanRecord: + elapsed_s = max(0.0, end_event.elapsed_s or 0.0) + return TraceSpanRecord( + trace_id=end_event.trace_id, + span=span, + display_stage=display_trace_stage(span), + start_s=end_event.timestamp_s - elapsed_s, + end_s=end_event.timestamp_s, + duration_s=elapsed_s, + outcome=outcome, + task_name=end_event.task_name, + uid=end_event.uid, + status=end_event.status, + train_step=end_event.train_step, + worker_rank=end_event.worker_rank, + error_msg=end_event.error_msg, + ) + + +def _build_open_span_record(start_event: TraceEvent, *, span: str, now_s: float) -> TraceSpanRecord: + duration_s = max(0.0, now_s - start_event.timestamp_s) + return TraceSpanRecord( + trace_id=start_event.trace_id, + span=span, + display_stage=display_trace_stage(span), + start_s=start_event.timestamp_s, + end_s=now_s, + duration_s=duration_s, + outcome="open", + task_name=start_event.task_name, + uid=start_event.uid, + status=start_event.status, + train_step=start_event.train_step, + worker_rank=start_event.worker_rank, + error_msg=start_event.error_msg, + ) + + +def _assign_depths(records: list[TraceSpanRecord]) -> None: + records_by_trace = _group_records(records) + for trace_records in records_by_trace.values(): + active_end_times: list[float] = [] + for record in sorted(trace_records, key=lambda item: (item.start_s, -item.end_s)): + active_end_times = [end_s for end_s in active_end_times if end_s > record.start_s] + record.depth = len(active_end_times) + active_end_times.append(record.end_s) + + +def _group_records(records: Iterable[TraceSpanRecord]) -> dict[str, list[TraceSpanRecord]]: + grouped: dict[str, list[TraceSpanRecord]] = defaultdict(list) + for record in records: + grouped[record.trace_id].append(record) + for trace_id in grouped: + grouped[trace_id].sort(key=lambda record: (record.start_s, record.depth, record.end_s)) + return dict(grouped) + + +def _span_record_to_payload( + record: TraceSpanRecord, + timeline_start_s: float, + timeline_duration_s: float, + stage_colors: dict[str, str], +) -> dict[str, Any]: + left_pct = (record.start_s - timeline_start_s) / timeline_duration_s * 100.0 + width_pct = max(0.2, record.duration_s / timeline_duration_s * 100.0) + return { + **dataclasses.asdict(record), + "left_pct": left_pct, + "width_pct": min(width_pct, max(0.2, 100.0 - left_pct)), + "top_px": record.depth * 22, + "color": stage_colors[record.display_stage], + } + + +def _build_stage_colors(records: Iterable[TraceSpanRecord]) -> dict[str, str]: + stages = sorted({record.display_stage for record in records}) + return {stage: _PALETTE[index % len(_PALETTE)] for index, stage in enumerate(stages)} + + +def _build_stage_stats(records: Iterable[TraceSpanRecord]) -> list[dict[str, Any]]: + grouped: dict[str, list[TraceSpanRecord]] = defaultdict(list) + for record in records: + grouped[record.display_stage].append(record) + + stats = [] + for stage, stage_records in grouped.items(): + durations = sorted(record.duration_s for record in stage_records) + stats.append( + { + "stage": stage, + "count": len(stage_records), + "open_count": sum(1 for record in stage_records if record.outcome == "open"), + "error_count": sum(1 for record in stage_records if record.outcome == "error"), + "total_s": sum(durations), + "avg_s": sum(durations) / len(durations), + "p50_s": _percentile(durations, 0.50), + "p95_s": _percentile(durations, 0.95), + "max_s": durations[-1], + } + ) + return sorted(stats, key=lambda item: (item["p95_s"], item["total_s"]), reverse=True) + + +def _split_span_stage(stage: str) -> tuple[str | None, str | None]: + for suffix in (".start", ".end", ".error"): + if stage.endswith(suffix): + return stage[: -len(suffix)], suffix[1:] + return None, None + + +def _percentile(sorted_values: list[float], percentile: float) -> float: + if not sorted_values: + return 0.0 + if len(sorted_values) == 1: + return sorted_values[0] + position = (len(sorted_values) - 1) * percentile + lower = int(position) + upper = min(lower + 1, len(sorted_values) - 1) + if lower == upper: + return sorted_values[lower] + fraction = position - lower + return sorted_values[lower] * (1 - fraction) + sorted_values[upper] * fraction + + +_HTML_TEMPLATE = """ + + + + + Producer Trace Hotspots + + + +
+
+

Producer Trace Hotspots

+
+
+
+
+ +
+
Tasks0
+
Spans0
+
Max Task Duration0s
+
+ +
+
+
+

Stage Hotspots

+
+ + + +
StageCountAvgP95Max
+
+
+
+
+ +
+
+

Task Timeline

+
Each row is one task. The x-axis is normalized to that task's own duration.
+
+
+
+
+ + + + + +""" + + +if __name__ == "__main__": + main() diff --git a/xtuner/tools/producer_trace_viewer.py b/xtuner/tools/producer_trace_viewer.py new file mode 100644 index 0000000000..ac89ea5c77 --- /dev/null +++ b/xtuner/tools/producer_trace_viewer.py @@ -0,0 +1,813 @@ +from __future__ import annotations + +import argparse +import dataclasses +import http.server +import json +import threading +import time +from collections import Counter, defaultdict +from pathlib import Path +from typing import Any + +from xtuner.v1.rl.trace import ( + TRACE_JSONL_BASENAME, + TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, + TraceEvent, + TraceViewerScope, +) +from xtuner.tools.producer_trace_analysis import ( + build_open_span_summaries, + build_viewer_rows, + display_trace_stage, + events_to_timelines, + filter_trace_events_by_scope, + load_trace_jsonl, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Inspect producer trace JSONL shards.") + parser.add_argument("trace_dir", type=Path, help="Directory containing producer_trace_*.jsonl files.") + parser.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="Output HTML file. Defaults to /producer_trace_viewer.html.", + ) + parser.add_argument( + "--serve", + action="store_true", + help="Serve a live viewer that polls the trace directory instead of writing a static HTML snapshot.", + ) + parser.add_argument("--host", default="127.0.0.1", help="Host used by --serve.") + parser.add_argument("--port", type=int, default=0, help="Port used by --serve. Defaults to an available port.") + parser.add_argument( + "--refresh-interval", + type=float, + default=1.0, + help="Live viewer refresh interval in seconds.", + ) + parser.add_argument( + "--scope", + choices=("latest-produce-batch", "all"), + default=TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, + help="Trace scope shown by the viewer. Defaults to the latest produce batch.", + ) + return parser.parse_args() + + +@dataclasses.dataclass +class ProducerTraceViewerHandle: + server: http.server.ThreadingHTTPServer + thread: threading.Thread + url: str + closed: bool = False + + def close(self) -> None: + if self.closed: + return + self.closed = True + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=5) + + +class TraceJsonlIndex: + def __init__(self, trace_dir: Path) -> None: + self.trace_dir = trace_dir.absolute() + self._offsets: dict[Path, int] = {} + self._events: list[TraceEvent] = [] + self._events_by_batch_id: defaultdict[str, list[TraceEvent]] = defaultdict(list) + self._latest_batch_id: str | None = None + self._latest_batch_sort_key: tuple[int, int, int, float] | None = None + self._lock = threading.RLock() + + def build_payload( + self, + *, + scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, + ) -> dict[str, Any]: + with self._lock: + self.refresh() + events = self._events_for_scope(scope) + return build_viewer_payload_from_events(events, trace_source=str(self.trace_dir), scope=scope) + + def refresh(self) -> None: + with self._lock: + self.trace_dir.mkdir(parents=True, exist_ok=True) + for file in sorted(self.trace_dir.glob(f"{TRACE_JSONL_BASENAME}_*.jsonl")): + self._tail_file(file) + + def _tail_file(self, file: Path) -> None: + offset = self._offsets.get(file, 0) + try: + size = file.stat().st_size + except OSError: + return + if size < offset: + offset = 0 + try: + with file.open("r", encoding="utf-8") as f: + f.seek(offset) + while True: + line_start = f.tell() + line = f.readline() + if not line: + break + if not line.endswith("\n"): + offset = line_start + break + offset = f.tell() + line = line.strip() + if not line: + continue + try: + self._add_event(TraceEvent.from_dict(json.loads(line))) + except (json.JSONDecodeError, KeyError, TypeError, ValueError): + continue + except OSError: + return + self._offsets[file] = offset + + def _add_event(self, event: TraceEvent) -> None: + self._events.append(event) + if event.produce_batch_id is None: + return + self._events_by_batch_id[event.produce_batch_id].append(event) + sort_key = self._batch_sort_key(event) + if self._latest_batch_sort_key is None or sort_key > self._latest_batch_sort_key: + self._latest_batch_id = event.produce_batch_id + self._latest_batch_sort_key = sort_key + + def _events_for_scope(self, scope: TraceViewerScope) -> list[TraceEvent]: + if scope == "all": + return list(self._events) + if self._latest_batch_id is not None: + return list(self._events_by_batch_id[self._latest_batch_id]) + return filter_trace_events_by_scope(self._events, scope) + + @staticmethod + def _batch_sort_key(event: TraceEvent) -> tuple[int, int, int, float]: + if event.train_step is None: + return (-1, -1, -1, event.timestamp_s) + model_step = -1 if event.model_step is None else event.model_step + producer_future_step = -1 if event.producer_future_step is None else event.producer_future_step + return (event.train_step, model_step, producer_future_step, event.timestamp_s) + + +def build_viewer_payload( + trace_dir: Path, + *, + scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, +) -> dict[str, Any]: + return build_viewer_payload_from_events(load_trace_jsonl(trace_dir), trace_source=str(trace_dir), scope=scope) + + +def build_viewer_payload_from_events( + events: list[TraceEvent], + *, + trace_source: str, + scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, +) -> dict[str, Any]: + raw_event_count = len(events) + events = filter_trace_events_by_scope(events, scope) + timelines = events_to_timelines(events) + rows = build_viewer_rows(timelines) + summaries = build_open_span_summaries(timelines) + latest_stage_counts = Counter(row.latest_stage for row in rows) + + return { + "generated_at_s": time.time(), + "trace_dir": trace_source, + "scope": scope, + "event_count": len(events), + "raw_event_count": raw_event_count, + "trace_count": len(timelines), + "open_trace_count": sum(1 for row in rows if row.open_span is not None), + "task_summary": build_task_summary(rows), + "latest_stage_counts": dict(latest_stage_counts.most_common()), + "open_span_summaries": [dataclasses.asdict(summary) for summary in summaries], + "rows": [dataclasses.asdict(row) for row in rows], + "timelines": { + trace_id: [event.to_dict() for event in trace_events] for trace_id, trace_events in timelines.items() + }, + } + + +def build_task_summary(rows: list[Any]) -> dict[str, Any]: + running_tasks = sum(1 for row in rows if row.open_span is not None) + current_stage_counts = Counter(display_trace_stage(row.open_span) for row in rows if row.open_span is not None) + return { + "total_tasks": len(rows), + "running_tasks": running_tasks, + "completed_tasks": max(0, len(rows) - running_tasks), + "current_stage_counts": dict(current_stage_counts.most_common()), + } + + +def write_viewer_html(payload: dict[str, Any], output_path: Path) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(render_viewer_html(payload), encoding="utf-8") + + +def render_viewer_html( + payload: dict[str, Any], + *, + live: bool = False, + api_url: str = "/api/trace", + refresh_interval_s: float = 1.0, +) -> str: + data_json = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).replace(" None: + handle = start_trace_viewer( + trace_dir, + host=host, + port=port, + refresh_interval_s=refresh_interval_s, + scope=scope, + ) + print(f"Serving Producer Trace Viewer on {handle.url}", flush=True) + print(f"Trace dir: {trace_dir.absolute()}", flush=True) + try: + handle.thread.join() + except KeyboardInterrupt: + pass + finally: + handle.close() + + +def start_trace_viewer( + trace_dir: Path, + *, + host: str = "127.0.0.1", + port: int = 0, + refresh_interval_s: float = 1.0, + scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH, +) -> ProducerTraceViewerHandle: + trace_dir = trace_dir.absolute() + trace_index = TraceJsonlIndex(trace_dir) + + class TraceViewerHandler(http.server.BaseHTTPRequestHandler): + def do_GET(self) -> None: + path = self.path.split("?", 1)[0] + if path in {"/", "/index.html"}: + payload = trace_index.build_payload(scope=scope) + html = render_viewer_html( + payload, + live=True, + api_url="/api/trace", + refresh_interval_s=refresh_interval_s, + ) + self._send_bytes(html.encode("utf-8"), "text/html; charset=utf-8") + return + + if path == "/api/trace": + payload = trace_index.build_payload(scope=scope) + body = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + self._send_bytes(body, "application/json; charset=utf-8") + return + + self.send_error(404) + + 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), TraceViewerHandler) + server_host, server_port = server.server_address + display_host = "127.0.0.1" if server_host in {"", "0.0.0.0"} else server_host + thread = threading.Thread(target=server.serve_forever, name="ProducerTraceViewer", daemon=True) + thread.start() + return ProducerTraceViewerHandle( + server=server, + thread=thread, + url=f"http://{display_host}:{server_port}", + ) + + +def main() -> None: + args = parse_args() + trace_dir = args.trace_dir + if args.serve: + serve_trace_viewer( + trace_dir, + host=args.host, + port=args.port, + refresh_interval_s=args.refresh_interval, + scope=args.scope, + ) + return + + output_path = args.output or trace_dir / "producer_trace_viewer.html" + payload = build_viewer_payload(trace_dir, scope=args.scope) + write_viewer_html(payload, output_path) + print(output_path) + + +_HTML_TEMPLATE = """ + + + + + Producer Trace Viewer + + + +
+
+

Producer Trace Viewer

+
+
+
+
+
+
+
+ +
+
Total tasks0
+
Running tasks0
+
Completed tasks0
+
+ Current Trace Function Stages +
+
+
+ +
+
+
+
+
+

Suspect Open Spans

+
+ + + +
SpanOpenOldestOldest Trace
+
+
+
+
+
+

Latest Stage Distribution

+
+ + + +
StageTasks
+
+
+
+
+ +
+
+ + +
+
+ + + + + + + + + + + +
TraceTaskLatest StageOpen SpanOpen Age
+
+
+
+ +
+
+
+

Task Timeline

+
+
+
+
+
+
+ + + + + +""" + + +if __name__ == "__main__": + main() diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 835683149a..55feaaeb2e 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.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 trace_function from xtuner.v1.rl.utils import ( JUDGER_PAUSE_JUDGE_TASK_TIMEOUT_S, CPUActorLauncher, @@ -189,6 +190,11 @@ def __init__( @abstractmethod async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: ... + @trace_function( + "xtuner.agent_loop.generate_group", + target="rollout_state", + result="return", + ) async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: pending_tasks = [] for state in rollout_state: @@ -208,6 +214,11 @@ async def run_judger(self, rollout_state: RolloutState) -> RolloutState: ... @overload async def run_judger(self, rollout_state: list[RolloutState]) -> list[RolloutState]: ... + @trace_function( + "xtuner.judger.judge", + target="rollout_state", + result="return", + ) async def run_judger(self, rollout_state: RolloutState | list[RolloutState]) -> RolloutState | list[RolloutState]: assert self.judger is not None if isinstance(rollout_state, list): diff --git a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py index 86e1539da8..203a065758 100644 --- a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py +++ b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py @@ -1,6 +1,7 @@ from xtuner.v1.data_proto.rl_data import RolloutState, SampleParams, Status from xtuner.v1.rl.judger import Judger from xtuner.v1.rl.rollout import RolloutController +from xtuner.v1.rl.trace import trace_function from .agent_loop import AgentLoop, AgentLoopConfig @@ -64,6 +65,11 @@ def __init__( enable_batch_judge=enable_batch_judge, ) + @trace_function( + "xtuner.agent_loop.generate_sample", + target="rollout_state", + result="return", + ) async def generate_sample( self, rollout_state: RolloutState, diff --git a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py index 9da3b3a3a8..4c9830dc45 100644 --- a/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py +++ b/xtuner/v1/rl/agent_loop_manager/agent_loop_manager.py @@ -14,6 +14,7 @@ from xtuner.v1.rl.judger import ComposedJudgerConfig, JudgerConfig, build_judger from xtuner.v1.rl.replay_buffer import ReplayBuffer from xtuner.v1.rl.rollout import RolloutController +from xtuner.v1.rl.trace import flush_trace from xtuner.v1.utils import get_logger from .producer import ( @@ -589,6 +590,7 @@ async def _produce_batch_to_buffer( ) finally: progress.add_produce_time(time.perf_counter() - produce_start) + flush_trace() return _aggregate_status(statuses) async def pause_produce( @@ -727,13 +729,15 @@ async def _get_batch_from_buffer( consume_progress.mark_consumed(consumed_counts) leftover_counts = await self.replay_buffer.count_statuses(self.task_names, _LEFTOVER_STATUSES) self._log_buffer_counts(task_batch_sizes, batch_by_task, leftover_counts) - return self._build_result_from_batch( + result = self._build_result_from_batch( task_batch_sizes, batch_by_task, leftover_counts, progress=consume_progress, pause_time_s=pause_time_s, ) + flush_trace() + return result async def continue_produce(self, model_step: int) -> None: # @@ -755,6 +759,7 @@ def shutdown(self) -> None: self._status = AgentLoopManagerStatus.FINISH self._update_event.set() self._finish_event.set() + flush_trace() async def _wait_for_status_exit(self, blocked_status: AgentLoopManagerStatus) -> None: while not self._finish_event.is_set() and self._status == blocked_status: diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index bf6dafe54d..bdad231c72 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -23,6 +23,14 @@ ) from xtuner.v1.rl.agent_loop import AgentLoopSpec from xtuner.v1.rl.replay_buffer import ReplayBuffer +from xtuner.v1.rl.trace import ( + TRACE_EXTRA_MODEL_STEP, + TRACE_EXTRA_PRODUCE_BATCH_ID, + TRACE_EXTRA_PRODUCER_FUTURE_STEP, + TRACE_EXTRA_TRAIN_STEP, + build_produce_batch_id, + trace_function, +) from xtuner.v1.rl.utils import ( AGENT_LOOP_PAUSE_REQUEST_TIMEOUT_S, PRODUCER_PAUSE_PENDING_TASK_TIMEOUT_S, @@ -326,6 +334,20 @@ def target_abs(self) -> int: def should_abort(self) -> bool: return self.update_event.is_set() + def trace_kwargs(self) -> dict[str, Any]: + produce_batch_id = build_produce_batch_id( + self.train_step, + self.model_step, + self.progress.producer_future_step, + ) + return { + "task_name": self.task_name, + "train_step": self.train_step, + "model_step": self.model_step, + "producer_future_step": self.progress.producer_future_step, + "produce_batch_id": produce_batch_id, + } + async def expired_count(self) -> int: return await self.replay_buffer.count(task_name=self.task_name, group_status=Status.EXPIRED) @@ -333,10 +355,21 @@ async def available_count(self) -> int: completed_count = await self.replay_buffer.count(task_name=self.task_name, group_status=Status.COMPLETED) return self.progress.consumed_samples[self.task_name] + completed_count + @trace_function( + "xtuner.producer.sample_group", + result="return", + trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs(), + ) 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) + @trace_function( + "xtuner.producer.generate_group", + target="rollout_state", + result="return", + trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs(), + ) async def generate_group( self, rollout_state: list[RolloutState], @@ -344,6 +377,14 @@ async def generate_group( enable_partial_rollout: bool = False, ) -> list[RolloutState]: # strategy 只表达“要生成”，不关心 agent_loop 是 ray actor 还是本地对象。 + trace_kwargs = self.trace_kwargs() + for state in rollout_state: + extra_fields = dict(state.extra_fields or {}) + extra_fields[TRACE_EXTRA_TRAIN_STEP] = trace_kwargs["train_step"] + extra_fields[TRACE_EXTRA_MODEL_STEP] = trace_kwargs["model_step"] + extra_fields[TRACE_EXTRA_PRODUCER_FUTURE_STEP] = trace_kwargs["producer_future_step"] + extra_fields[TRACE_EXTRA_PRODUCE_BATCH_ID] = trace_kwargs["produce_batch_id"] + state.extra_fields = extra_fields start = time.perf_counter() if isinstance(self.agent_loop, ray.actor.ActorHandle): result = await self.agent_loop.generate_group.remote( @@ -364,6 +405,11 @@ async def generate_group( extra_fields[GROUP_GENERATE_TIME_KEY] = elapsed return result + @trace_function( + "xtuner.producer.put_generated_group", + target="group", + trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs(), + ) async def put_generated_group(self, group: list[RolloutState]) -> bool: # 只有完整生成的 group 才需要业务有效性过滤；ABORTED / EXPIRED 保留原状态供重试或统计。 is_completed = get_group_status(group) == Status.COMPLETED diff --git a/xtuner/v1/rl/agent_loop_manager/sampler.py b/xtuner/v1/rl/agent_loop_manager/sampler.py index e5b0170bbe..e05bdf283d 100644 --- a/xtuner/v1/rl/agent_loop_manager/sampler.py +++ b/xtuner/v1/rl/agent_loop_manager/sampler.py @@ -128,11 +128,18 @@ def __init__( self.replay_buffer = replay_buffer async def sample(self, task_name: str, group_status: list[Status] | None = None) -> list[RolloutState]: + group = None for status in group_status or []: buffer_data = await self.replay_buffer.get(1, task_name=task_name, group_status=status) if buffer_data: - return buffer_data[0] - return self.sample_from_dataloader() + group = buffer_data[0] + break + if group is None: + group = self.sample_from_dataloader() + for state in group: + if state.task_name is None: + state.task_name = task_name + return group def save(self, checkpoint_path: Path | str) -> None: """Save the sampler's dataloader state to checkpoint.""" diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index efd49d62d6..de2f3b8746 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -12,6 +12,7 @@ from transformers import AutoTokenizer from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.trace import trace_function from xtuner.v1.rl.utils import AutoAcceleratorWorkers from xtuner.v1.utils import XTUNER_DETERMINISTIC, get_logger @@ -173,6 +174,11 @@ def get_generate_concurrency(self) -> int: return active_workers * concurrency_per_worker @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) + @trace_function( + "xtuner.rollout_controller.generate", + target="rollout_state", + result="return", + ) async def generate(self, rollout_state: RolloutState) -> RolloutState: if XTUNER_DETERMINISTIC: sample_params = rollout_state.sample_params.model_copy(deep=True) diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index d5f2c7391b..8899d15f56 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -28,6 +28,7 @@ reset_rollout_response, update_status_from_finish_reason, ) +from xtuner.v1.rl.trace import merge_trace_runtime_env, trace_function, trace_span from xtuner.v1.rl.utils import ( AutoAcceleratorWorkers, CPUResourcesConfig, @@ -445,13 +446,15 @@ def build(self, placement_group: "PlacementGroup"): ) generate_max_concurrency = self.get_controller_generate_concurrency(placement_group) get_logger().info(f"Calculated RolloutController generate concurrency: {generate_max_concurrency}") + actor_options = {"num_cpus": num_workers} + merge_trace_runtime_env(actor_options) return ( ray.remote( concurrency_groups={ ROLLOUT_CONCURRENCY_GROUP_GENERATE: generate_max_concurrency, }, )(RolloutController) - .options(num_cpus=num_workers) + .options(**actor_options) .remote(self, placement_group) ) @@ -678,6 +681,12 @@ async def _decode_routed_experts(self, routed_experts: Any) -> Any: return routed_experts @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) + @trace_function( + "xtuner.rollout_worker.generate", + target="rollout_state", + result="return", + trace_kwargs_getter=lambda self, *args, **kwargs: {"worker_rank": self.rank}, + ) async def generate(self, rollout_state: RolloutState) -> RolloutState: try: # TODO(@duanyanhui): @@ -739,7 +748,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) + async with trace_span( + rollout_state, + "xtuner.rollout_engine.generate", + worker_rank=self.rank, + ): + http_result = await self._safe_post_request(endpoint_url, headers=headers, payload=payload) # Case 1: HTTP Request is Successful if http_result.response: diff --git a/xtuner/v1/rl/trace.py b/xtuner/v1/rl/trace.py new file mode 100644 index 0000000000..f8d57af194 --- /dev/null +++ b/xtuner/v1/rl/trace.py @@ -0,0 +1,1104 @@ +from __future__ import annotations + +import atexit +import contextlib +import contextvars +import dataclasses +import functools +import inspect +import json +import os +import queue +import threading +import time +from collections import OrderedDict, deque +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, ClassVar, Iterable, Literal, Sequence, TextIO, cast + +from pydantic import BaseModel, ConfigDict + +from xtuner.v1.data_proto.rl_data import RolloutState +from xtuner.v1.utils import get_logger + + +logger = get_logger() + +TRACE_JSONL_FLUSH_INTERVAL_S = 1.0 +TRACE_JSONL_FLUSH_EVENTS = 1024 +TRACE_JSONL_FLUSH_BYTES = 1 * 1024 * 1024 +TRACE_JSONL_SHARD_BYTES = 256 * 1024 * 1024 +TRACE_JSONL_QUEUE_MAX_EVENTS = 100_000 +TRACE_JSONL_FLUSH_TIMEOUT_S = 0.2 +TRACE_JSONL_CLOSE_TIMEOUT_S = 2.0 +TRACE_JSONL_BASENAME = "producer_trace" +TRACE_ENV_ENABLED = "XTUNER_TRACE_ENABLED" +TRACE_ENV_OUTPUT_DIR = "XTUNER_TRACE_OUTPUT_DIR" +TRACE_ENV_MAX_EVENTS = "XTUNER_TRACE_MAX_EVENTS" +TRACE_ENV_MAX_EVENTS_PER_TRACE = "XTUNER_TRACE_MAX_EVENTS_PER_TRACE" +TraceViewerScope = Literal["latest-produce-batch", "all"] +TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH: TraceViewerScope = "latest-produce-batch" +TRACE_VIEWER_SCOPE_ALL: TraceViewerScope = "all" +TRACE_EXTRA_TRAIN_STEP = "_trace_train_step" +TRACE_EXTRA_MODEL_STEP = "_trace_model_step" +TRACE_EXTRA_PRODUCER_FUTURE_STEP = "_trace_producer_future_step" +TRACE_EXTRA_PRODUCE_BATCH_ID = "_trace_produce_batch_id" + + +# Config and event schema. +class TraceConfig(BaseModel): + model_config = ConfigDict(extra="forbid") + + # Whether producer task tracing is enabled. + enabled: bool = False + # Directory that receives producer_trace_*.jsonl shards. If None, only memory tracing is used. + output_dir: str | Path | None = None + # Max events retained in memory across all traces. + max_events: int = 100_000 + # Max events retained in memory for one trace_id. + max_events_per_trace: int = 256 + # Whether rank0 should start the live producer trace viewer when tracing is enabled. + viewer_enabled: bool = True + # Host for the live producer trace viewer. Use 0.0.0.0 explicitly when remote access is needed. + viewer_host: str = "127.0.0.1" + # Port for the live producer trace viewer. 0 asks the OS to pick an available port. + viewer_port: int = 0 + # Browser polling interval for the live producer trace viewer. + viewer_refresh_interval_s: float = 1.0 + # Default viewer data scope. The latest produce batch keeps the UI focused during training. + viewer_scope: TraceViewerScope = TRACE_VIEWER_SCOPE_LATEST_PRODUCE_BATCH + + +@dataclass(frozen=True) +class TraceEvent: + trace_id: str + stage: str + timestamp_s: float + status: str | None = None + task_name: str | None = None + uid: int | str | None = None + session_uid: int | str | None = None + train_step: int | None = None + model_step: int | None = None + producer_future_step: int | None = None + produce_batch_id: str | None = None + worker_rank: int | None = None + elapsed_s: float | None = None + error_msg: str | None = None + + def to_dict(self) -> dict[str, Any]: + return dataclasses.asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> TraceEvent: + return cls( + trace_id=str(data["trace_id"]), + stage=str(data["stage"]), + timestamp_s=float(data["timestamp_s"]), + status=data.get("status"), + task_name=data.get("task_name"), + uid=data.get("uid"), + session_uid=data.get("session_uid"), + train_step=data.get("train_step"), + model_step=data.get("model_step"), + producer_future_step=data.get("producer_future_step"), + produce_batch_id=data.get("produce_batch_id"), + worker_rank=data.get("worker_rank"), + elapsed_s=data.get("elapsed_s"), + error_msg=data.get("error_msg"), + ) + + +class TraceEventBuilder: + @classmethod + def trace_id(cls, task_name: str | None, uid: int | str | None) -> str | None: + if uid is None: + return None + return f"{task_name or 'unknown'}:{uid}" + + @classmethod + def produce_batch_id( + cls, + train_step: int | None, + model_step: int | None, + producer_future_step: int | None, + ) -> str | None: + if train_step is None and model_step is None and producer_future_step is None: + return None + return ( + f"train_step={cls._format_batch_value(train_step)}/" + f"model_step={cls._format_batch_value(model_step)}/" + f"producer_future_step={cls._format_batch_value(producer_future_step)}" + ) + + @classmethod + def build( + cls, + target: RolloutState | None, + stage: str, + *, + task_name: str | None = None, + uid: int | str | None = None, + session_uid: int | str | None = None, + status: Any | None = None, + train_step: int | None = None, + model_step: int | None = None, + producer_future_step: int | None = None, + produce_batch_id: str | None = None, + worker_rank: int | None = None, + elapsed_s: float | None = None, + error_msg: str | None = None, + timestamp_s: float | None = None, + ) -> TraceEvent | None: + resolved_task_name = task_name if task_name is not None else getattr(target, "task_name", None) + resolved_uid = uid if uid is not None else getattr(target, "uid", None) + trace_id = cls.trace_id(resolved_task_name, resolved_uid) + if trace_id is None: + return None + + resolved_status = status if status is not None else getattr(target, "status", None) + resolved_session_uid = session_uid if session_uid is not None else getattr(target, "session_uid", None) + trace_extra_fields = cls._extra_fields(target) + if train_step is None: + train_step = trace_extra_fields.get(TRACE_EXTRA_TRAIN_STEP) + if model_step is None: + model_step = trace_extra_fields.get(TRACE_EXTRA_MODEL_STEP, getattr(target, "model_step", None)) + if producer_future_step is None: + producer_future_step = trace_extra_fields.get(TRACE_EXTRA_PRODUCER_FUTURE_STEP) + if produce_batch_id is None: + produce_batch_id = trace_extra_fields.get(TRACE_EXTRA_PRODUCE_BATCH_ID) + if produce_batch_id is None: + produce_batch_id = cls.produce_batch_id(train_step, model_step, producer_future_step) + + return TraceEvent( + trace_id=trace_id, + stage=stage, + timestamp_s=time.time() if timestamp_s is None else timestamp_s, + status=cls._stringify_status(resolved_status), + task_name=resolved_task_name, + uid=resolved_uid, + session_uid=resolved_session_uid, + train_step=train_step, + model_step=model_step, + producer_future_step=producer_future_step, + produce_batch_id=produce_batch_id, + worker_rank=worker_rank, + elapsed_s=elapsed_s, + error_msg=error_msg, + ) + + @staticmethod + def short_error(exc: BaseException, max_len: int = 500) -> str: + message = f"{type(exc).__name__}: {exc}" + if len(message) <= max_len: + return message + return message[: max_len - 3] + "..." + + @staticmethod + def _extra_fields(target: RolloutState | None) -> dict[str, Any]: + extra_fields = getattr(target, "extra_fields", None) + if isinstance(extra_fields, dict): + return cast(dict[str, Any], extra_fields) + return {} + + @staticmethod + def _format_batch_value(value: int | None) -> str: + return "none" if value is None else str(value) + + @staticmethod + def _stringify_status(status: Any) -> str | None: + if status is None: + return None + value = getattr(status, "value", None) + if value is not None: + return str(value) + return str(status) + + +@dataclass(frozen=True) +class _StoredEvent: + seq: int + event: TraceEvent + + +@dataclass +class _FlushRequest: + done: threading.Event + + +@dataclass +class _CloseRequest: + done: threading.Event + + +# Durable JSONL writer. All writes are best-effort and must not block training. +class BufferedTraceJsonlWriter: + _writer_seq: ClassVar[int] = 0 + _writer_seq_lock: ClassVar[threading.Lock] = threading.Lock() + + def __init__(self, output_dir: str | Path) -> None: + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + self._writer_id = self._next_writer_id() + self._queue: queue.Queue[str | _FlushRequest | _CloseRequest] = queue.Queue( + maxsize=TRACE_JSONL_QUEUE_MAX_EVENTS + ) + self._thread = threading.Thread(target=self._run, name="TraceJsonlWriter", daemon=True) + self._closed = False + self._failed = False + self._dropped_events = 0 + self._closed_lock = threading.Lock() + self._file: TextIO | None = None + self._shard_idx = 0 + self._bytes_written = 0 + self._thread.start() + + def append(self, event: TraceEvent) -> bool: + with self._closed_lock: + if self._closed or self._failed: + return False + line = json.dumps(event.to_dict(), ensure_ascii=False, separators=(",", ":")) + "\n" + try: + self._queue.put_nowait(line) + except queue.Full: + self._record_drop("queue full") + return False + except Exception: + self._record_drop("queue append failed") + return False + return True + + def flush(self, timeout_s: float = TRACE_JSONL_FLUSH_TIMEOUT_S) -> bool: + with self._closed_lock: + if self._closed or self._failed: + return False + request = _FlushRequest(threading.Event()) + try: + self._queue.put(request, timeout=max(0.0, min(timeout_s, TRACE_JSONL_FLUSH_TIMEOUT_S))) + except Exception: + return False + return request.done.wait(timeout=max(0.0, timeout_s)) + + def close(self, timeout_s: float = TRACE_JSONL_CLOSE_TIMEOUT_S) -> bool: + with self._closed_lock: + if self._closed: + return True + self._closed = True + request = _CloseRequest(threading.Event()) + try: + self._queue.put(request, timeout=max(0.0, min(timeout_s, TRACE_JSONL_FLUSH_TIMEOUT_S))) + except Exception: + return False + closed = request.done.wait(timeout=max(0.0, timeout_s)) + self._thread.join(timeout=max(0.0, timeout_s)) + return closed + + def _run(self) -> None: + batch: list[str] = [] + batch_bytes = 0 + last_flush = time.monotonic() + while True: + timeout = TRACE_JSONL_FLUSH_INTERVAL_S if batch else None + try: + item = self._queue.get(timeout=timeout) + except queue.Empty: + self._safe_write_batch(batch) + batch = [] + batch_bytes = 0 + last_flush = time.monotonic() + continue + + if isinstance(item, str): + batch.append(item) + batch_bytes += len(item.encode("utf-8")) + self._queue.task_done() + should_flush = ( + len(batch) >= TRACE_JSONL_FLUSH_EVENTS + or batch_bytes >= TRACE_JSONL_FLUSH_BYTES + or time.monotonic() - last_flush >= TRACE_JSONL_FLUSH_INTERVAL_S + ) + if should_flush: + self._safe_write_batch(batch) + batch = [] + batch_bytes = 0 + last_flush = time.monotonic() + continue + + if isinstance(item, _FlushRequest): + self._safe_write_batch(batch) + batch = [] + batch_bytes = 0 + self._safe_flush_file() + last_flush = time.monotonic() + item.done.set() + self._queue.task_done() + continue + + if isinstance(item, _CloseRequest): + self._safe_write_batch(batch) + self._safe_flush_file() + self._safe_close_file() + item.done.set() + self._queue.task_done() + return + + def _safe_write_batch(self, batch: list[str]) -> bool: + try: + self._write_batch(batch) + except Exception: + self._mark_failed("Failed to write producer trace JSONL batch") + return False + return True + + def _write_batch(self, batch: list[str]) -> None: + if not batch: + return + for line in batch: + line_bytes = len(line.encode("utf-8")) + self._ensure_file(line_bytes) + assert self._file is not None + self._file.write(line) + self._bytes_written += line_bytes + self._flush_file() + + def _ensure_file(self, next_bytes: int) -> None: + if self._file is None: + self._open_file() + return + if self._bytes_written > 0 and self._bytes_written + next_bytes > TRACE_JSONL_SHARD_BYTES: + self._close_file() + self._shard_idx += 1 + self._open_file() + + def _open_file(self) -> None: + path = self.output_dir / f"{TRACE_JSONL_BASENAME}_{self._writer_id}_{self._shard_idx:06d}.jsonl" + self._file = path.open("a", encoding="utf-8") + self._bytes_written = path.stat().st_size if path.exists() else 0 + + def _flush_file(self) -> None: + if self._file is not None: + self._file.flush() + + def _close_file(self) -> None: + if self._file is not None: + self._file.close() + self._file = None + + def _safe_flush_file(self) -> bool: + try: + self._flush_file() + except Exception: + self._mark_failed("Failed to flush producer trace JSONL file") + return False + return True + + def _safe_close_file(self) -> bool: + try: + self._close_file() + except Exception: + self._mark_failed("Failed to close producer trace JSONL file") + return False + return True + + def _mark_failed(self, message: str) -> None: + with self._closed_lock: + if self._failed: + return + self._failed = True + logger.exception(message) + + def _record_drop(self, reason: str) -> None: + self._dropped_events += 1 + if self._dropped_events == 1 or self._dropped_events % 1000 == 0: + logger.warning( + "Dropped producer trace events because %s. dropped_events=%s", + reason, + self._dropped_events, + ) + + @classmethod + def _next_writer_id(cls) -> str: + with cls._writer_seq_lock: + cls._writer_seq += 1 + seq = cls._writer_seq + return f"{os.getpid()}_{seq:04d}" + + +# In-process state store used by the local tracer and tests. +class InMemoryTraceStore: + def __init__(self, config: TraceConfig | None = None) -> None: + self.config = config or TraceConfig() + self._timelines: dict[str, deque[_StoredEvent]] = {} + self._global_events: OrderedDict[int, _StoredEvent] = OrderedDict() + self._latest: dict[str, TraceEvent] = {} + self._seq = 0 + self._lock = threading.RLock() + self._jsonl_writer = None + if self.config.enabled and self.config.output_dir is not None: + self._jsonl_writer = BufferedTraceJsonlWriter(self.config.output_dir) + + def append(self, event: TraceEvent) -> None: + if not self.config.enabled: + return + with self._lock: + self._seq += 1 + stored = _StoredEvent(seq=self._seq, event=event) + timeline = self._timelines.setdefault(event.trace_id, deque()) + timeline.append(stored) + self._global_events[stored.seq] = stored + self._latest[event.trace_id] = event + self._enforce_per_trace_limit(event.trace_id) + self._enforce_global_limit() + if self._jsonl_writer is not None: + try: + self._jsonl_writer.append(event) + except Exception: + logger.exception("Failed to enqueue producer trace JSONL event") + + def get_timeline(self, trace_id: str) -> list[TraceEvent]: + with self._lock: + return [stored.event for stored in self._timelines.get(trace_id, ())] + + def get_all_timelines(self) -> dict[str, list[TraceEvent]]: + with self._lock: + return {trace_id: [stored.event for stored in timeline] for trace_id, timeline in self._timelines.items()} + + def get_latest(self, trace_id: str) -> TraceEvent | None: + with self._lock: + return self._latest.get(trace_id) + + def query_latest(self) -> dict[str, TraceEvent]: + with self._lock: + return dict(self._latest) + + def has_stage(self, trace_id: str, stage: str) -> bool: + with self._lock: + return any(stored.event.stage == stage for stored in self._timelines.get(trace_id, ())) + + def flush_jsonl(self) -> bool: + if self._jsonl_writer is not None: + try: + return self._jsonl_writer.flush() + except Exception: + logger.exception("Failed to flush producer trace JSONL writer") + return False + return True + + def close(self) -> bool: + if self._jsonl_writer is not None: + try: + closed = self._jsonl_writer.close() + except Exception: + logger.exception("Failed to close producer trace JSONL writer") + closed = False + self._jsonl_writer = None + return closed + return True + + def _enforce_per_trace_limit(self, trace_id: str) -> None: + limit = max(1, self.config.max_events_per_trace) + timeline = self._timelines.get(trace_id) + if timeline is None: + return + while len(timeline) > limit: + removed = timeline.popleft() + self._global_events.pop(removed.seq, None) + if not timeline: + self._timelines.pop(trace_id, None) + self._latest.pop(trace_id, None) + else: + self._latest[trace_id] = timeline[-1].event + + def _enforce_global_limit(self) -> None: + limit = max(1, self.config.max_events) + while len(self._global_events) > limit: + _, removed = self._global_events.popitem(last=False) + timeline = self._timelines.get(removed.event.trace_id) + if timeline is not None: + try: + timeline.remove(removed) + except ValueError: + pass + if not timeline: + self._timelines.pop(removed.event.trace_id, None) + self._latest.pop(removed.event.trace_id, None) + else: + self._latest[removed.event.trace_id] = timeline[-1].event + + +# Recorder API used by trace_event, trace_span, and trace_function. +class TraceRecorder: + def __init__(self, store: InMemoryTraceStore) -> None: + self.store = store + + def record( + self, + target: RolloutState | None, + stage: str, + *, + task_name: str | None = None, + uid: int | str | None = None, + session_uid: int | str | None = None, + status: Any | None = None, + train_step: int | None = None, + model_step: int | None = None, + producer_future_step: int | None = None, + produce_batch_id: str | None = None, + worker_rank: int | None = None, + elapsed_s: float | None = None, + error_msg: str | None = None, + timestamp_s: float | None = None, + ) -> TraceEvent | None: + event = build_trace_event( + target, + stage, + task_name=task_name, + uid=uid, + session_uid=session_uid, + status=status, + train_step=train_step, + model_step=model_step, + producer_future_step=producer_future_step, + produce_batch_id=produce_batch_id, + worker_rank=worker_rank, + elapsed_s=elapsed_s, + error_msg=error_msg, + timestamp_s=timestamp_s, + ) + if event is None: + return None + try: + self.store.append(event) + except Exception: + logger.exception("Failed to append trace event stage=%s trace_id=%s", stage, event.trace_id) + return None + return event + + def record_many(self, targets: Iterable[RolloutState], stage: str, **kwargs: Any) -> list[TraceEvent]: + events: list[TraceEvent] = [] + for target in targets: + event = self.record(target, stage, **kwargs) + if event is not None: + events.append(event) + return events + + async def mark(self, target: RolloutState | None, stage: str, **kwargs: Any) -> TraceEvent | None: + return self.record(target, stage, **kwargs) + + async def mark_many(self, targets: Iterable[RolloutState], stage: str, **kwargs: Any) -> list[TraceEvent]: + return self.record_many(targets, stage, **kwargs) + + +class NoopTraceRecorder: + def record(self, *args: Any, **kwargs: Any) -> None: + return None + + def record_many(self, *args: Any, **kwargs: Any) -> list[TraceEvent]: + return [] + + async def mark(self, *args: Any, **kwargs: Any) -> None: + return None + + async def mark_many(self, *args: Any, **kwargs: Any) -> list[TraceEvent]: + return [] + + +class TraceTargetResolver: + @classmethod + def as_rollout_state_list(cls, value: Any) -> list[RolloutState]: + if value is None: + return [] + if isinstance(value, RolloutState): + return [value] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + states: list[RolloutState] = [] + for item in value: + if isinstance(item, RolloutState): + states.append(item) + return states + return [] + + @classmethod + def resolve( + cls, + bound_arguments: dict[str, Any], + *, + target: str | RolloutState | Sequence[RolloutState] | Callable[..., Any] | None, + target_getter: Callable[..., Any] | None, + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> Any: + if target_getter is not None: + return target_getter(*args, **kwargs) + if isinstance(target, str): + return bound_arguments.get(target) + if callable(target): + return target(*args, **kwargs) + if target is not None: + return target + for value in bound_arguments.values(): + states = cls.as_rollout_state_list(value) + if states: + return states if len(states) > 1 else states[0] + return None + + @classmethod + def record_event(cls, target: Any, name: str, **kwargs: Any) -> TraceEvent | list[TraceEvent] | None: + targets = cls.as_rollout_state_list(target) + recorder = current_trace_recorder() + if targets: + if len(targets) == 1: + return recorder.record(targets[0], name, **kwargs) + return recorder.record_many(targets, name, **kwargs) + return recorder.record(None, name, **kwargs) + + @classmethod + async def mark_event(cls, target: Any, name: str, **kwargs: Any) -> TraceEvent | list[TraceEvent] | None: + targets = cls.as_rollout_state_list(target) + recorder = current_trace_recorder() + if targets: + if len(targets) == 1: + return await recorder.mark(targets[0], name, **kwargs) + return await recorder.mark_many(targets, name, **kwargs) + return await recorder.mark(None, name, **kwargs) + + +_NOOP_TRACE_RECORDER = NoopTraceRecorder() +_CURRENT_TRACE_RECORDER: contextvars.ContextVar[TraceRecorder | NoopTraceRecorder | None] = contextvars.ContextVar( + "xtuner_current_trace_recorder", + default=None, +) + + +# Global trace runtime. Each process owns one local runtime, propagated through Ray env vars. +def current_trace_recorder() -> TraceRecorder | NoopTraceRecorder: + recorder = _CURRENT_TRACE_RECORDER.get() + if recorder is not None: + return recorder + return get_tracer() + + +@contextlib.contextmanager +def use_trace_recorder(recorder: TraceRecorder | NoopTraceRecorder): + token = _CURRENT_TRACE_RECORDER.set(recorder) + try: + yield + finally: + _CURRENT_TRACE_RECORDER.reset(token) + + +def configure_trace( + config: TraceConfig | None, *, output_dir: str | Path | None = None +) -> TraceRecorder | NoopTraceRecorder: + return _TRACE_RUNTIME_MANAGER.configure(config, output_dir=output_dir) + + +def get_tracer() -> TraceRecorder | NoopTraceRecorder: + return _TRACE_RUNTIME_MANAGER.get_tracer() + + +def flush_trace() -> bool: + return _TRACE_RUNTIME_MANAGER.flush() + + +def close_trace() -> None: + _TRACE_RUNTIME_MANAGER.close() + + +def reset_trace_for_test() -> None: + close_trace() + _CURRENT_TRACE_RECORDER.set(None) + + +def get_trace_env_vars() -> dict[str, str]: + return _TRACE_RUNTIME_MANAGER.env_vars() + + +def merge_trace_runtime_env(actor_options: dict[str, Any]) -> dict[str, Any]: + return _TRACE_RUNTIME_MANAGER.merge_runtime_env(actor_options) + + +def build_trace_id(task_name: str | None, uid: int | str | None) -> str | None: + return TraceEventBuilder.trace_id(task_name, uid) + + +def build_produce_batch_id( + train_step: int | None, + model_step: int | None, + producer_future_step: int | None, +) -> str | None: + return TraceEventBuilder.produce_batch_id(train_step, model_step, producer_future_step) + + +def build_trace_event( + target: RolloutState | None, + stage: str, + *, + task_name: str | None = None, + uid: int | str | None = None, + session_uid: int | str | None = None, + status: Any | None = None, + train_step: int | None = None, + model_step: int | None = None, + producer_future_step: int | None = None, + produce_batch_id: str | None = None, + worker_rank: int | None = None, + elapsed_s: float | None = None, + error_msg: str | None = None, + timestamp_s: float | None = None, +) -> TraceEvent | None: + return TraceEventBuilder.build( + target, + stage, + task_name=task_name, + uid=uid, + session_uid=session_uid, + status=status, + train_step=train_step, + model_step=model_step, + producer_future_step=producer_future_step, + produce_batch_id=produce_batch_id, + worker_rank=worker_rank, + elapsed_s=elapsed_s, + error_msg=error_msg, + timestamp_s=timestamp_s, + ) + + +async def trace_event(target: Any, name: str, **kwargs: Any) -> TraceEvent | list[TraceEvent] | None: + return await TraceTargetResolver.mark_event(target, name, **kwargs) + + +@contextlib.asynccontextmanager +async def trace_span(target: Any, name: str, **kwargs: Any): + start_time = time.monotonic() + await trace_event(target, f"{name}.start", **kwargs) + try: + yield + except Exception as exc: + await trace_event( + target, + f"{name}.error", + elapsed_s=time.monotonic() - start_time, + error_msg=TraceEventBuilder.short_error(exc), + **kwargs, + ) + raise + else: + await trace_event(target, f"{name}.end", elapsed_s=time.monotonic() - start_time, **kwargs) + + +class TraceFunctionDecorator: + def __init__( + self, + name: str | Callable[..., str], + *, + target: str | RolloutState | Sequence[RolloutState] | Callable[..., Any] | None = None, + target_getter: Callable[..., Any] | None = None, + trace_kwargs_getter: Callable[..., dict[str, Any] | None] | None = None, + result: str | Callable[[Any], Any] | None = None, + trace_kwargs: dict[str, Any] | None = None, + ) -> None: + self.name = name + self.target = target + self.target_getter = target_getter + self.trace_kwargs_getter = trace_kwargs_getter + self.result = result + self.trace_kwargs = trace_kwargs or {} + + def decorate(self, func: Callable[..., Any]): + signature = inspect.signature(func) + if inspect.iscoroutinefunction(func): + return self._decorate_async(func, signature) + return self._decorate_sync(func, signature) + + def _start_target(self, signature: inspect.Signature, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any: + bound = signature.bind_partial(*args, **kwargs) + bound.apply_defaults() + return TraceTargetResolver.resolve( + bound.arguments, + target=self.target, + target_getter=self.target_getter, + args=args, + kwargs=kwargs, + ) + + def _end_target(self, start_target: Any, return_value: Any) -> Any: + if callable(self.result): + return self.result(return_value) + if self.result == "return": + return return_value + return start_target + + def _name(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> str: + if callable(self.name): + return self.name(*args, **kwargs) + return self.name + + def _kwargs(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: + resolved = dict(self.trace_kwargs) + if self.trace_kwargs_getter is None: + return resolved + dynamic_kwargs = self.trace_kwargs_getter(*args, **kwargs) + if dynamic_kwargs: + resolved.update(dynamic_kwargs) + return resolved + + def _decorate_async(self, func: Callable[..., Any], signature: inspect.Signature): + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any): + start_target = self._start_target(signature, args, kwargs) + trace_name = self._name(args, kwargs) + trace_kwargs = self._kwargs(args, kwargs) + start_time = time.monotonic() + await trace_event(start_target, f"{trace_name}.start", **trace_kwargs) + try: + return_value = await func(*args, **kwargs) + except Exception as exc: + await trace_event( + start_target, + f"{trace_name}.error", + elapsed_s=time.monotonic() - start_time, + error_msg=TraceEventBuilder.short_error(exc), + **trace_kwargs, + ) + raise + end_target = self._end_target(start_target, return_value) + await trace_event( + end_target, + f"{trace_name}.end", + elapsed_s=time.monotonic() - start_time, + **trace_kwargs, + ) + return return_value + + return async_wrapper + + def _decorate_sync(self, func: Callable[..., Any], signature: inspect.Signature): + @functools.wraps(func) + def sync_wrapper(*args: Any, **kwargs: Any): + start_target = self._start_target(signature, args, kwargs) + trace_name = self._name(args, kwargs) + trace_kwargs = self._kwargs(args, kwargs) + start_time = time.monotonic() + TraceTargetResolver.record_event(start_target, f"{trace_name}.start", **trace_kwargs) + try: + return_value = func(*args, **kwargs) + except Exception as exc: + TraceTargetResolver.record_event( + start_target, + f"{trace_name}.error", + elapsed_s=time.monotonic() - start_time, + error_msg=TraceEventBuilder.short_error(exc), + **trace_kwargs, + ) + raise + end_target = self._end_target(start_target, return_value) + TraceTargetResolver.record_event( + end_target, + f"{trace_name}.end", + elapsed_s=time.monotonic() - start_time, + **trace_kwargs, + ) + return return_value + + return sync_wrapper + + +def trace_function( + name: str | Callable[..., str], + *, + target: str | RolloutState | Sequence[RolloutState] | Callable[..., Any] | None = None, + target_getter: Callable[..., Any] | None = None, + trace_kwargs_getter: Callable[..., dict[str, Any] | None] | None = None, + result: str | Callable[[Any], Any] | None = None, + **trace_kwargs: Any, +): + return TraceFunctionDecorator( + name, + target=target, + target_getter=target_getter, + trace_kwargs_getter=trace_kwargs_getter, + result=result, + trace_kwargs=trace_kwargs, + ).decorate + + +# Runtime wrapper that owns store lifecycle. +@dataclass +class TraceRuntime: + config: TraceConfig + recorder: TraceRecorder | NoopTraceRecorder + store: InMemoryTraceStore | None = None + + def flush(self) -> bool: + if self.store is not None: + return self.store.flush_jsonl() + return True + + def close(self) -> bool: + if self.store is not None: + return self.store.close() + return True + + +def build_trace_runtime(config: TraceConfig | None, *, output_dir: str | Path | None = None) -> TraceRuntime: + if config is None: + config = TraceConfig() + if output_dir is not None and config.output_dir is None: + config = config.model_copy(update={"output_dir": output_dir}) + if not config.enabled: + return TraceRuntime(config=config, recorder=_NOOP_TRACE_RECORDER, store=None) + store = InMemoryTraceStore(config) + return TraceRuntime(config=config, recorder=TraceRecorder(store), store=store) + + +class TraceRuntimeManager: + def __init__(self) -> None: + self._runtime: TraceRuntime | None = None + self._identity: tuple[Any, ...] | None = None + self._lock = threading.RLock() + self._atexit_registered = False + + def configure( + self, config: TraceConfig | None, *, output_dir: str | Path | None = None + ) -> TraceRecorder | NoopTraceRecorder: + config = self._normalize_config(config or TraceConfig(), output_dir=output_dir) + self._export_env(config) + return self._replace_runtime(config).recorder + + def get_tracer(self) -> TraceRecorder | NoopTraceRecorder: + config = self._load_config_from_env() + identity = self._config_identity(config) + with self._lock: + if self._runtime is not None and self._identity == identity: + return self._runtime.recorder + return self._replace_runtime(config).recorder + + def flush(self) -> bool: + with self._lock: + runtime = self._runtime + if runtime is not None: + return runtime.flush() + return True + + def close(self) -> None: + with self._lock: + runtime = self._runtime + self._runtime = None + self._identity = None + self._clear_env() + if runtime is not None: + runtime.close() + + def env_vars(self) -> dict[str, str]: + if os.environ.get(TRACE_ENV_ENABLED) != "1": + return {} + env_vars: dict[str, str] = {} + for name in ( + TRACE_ENV_ENABLED, + TRACE_ENV_OUTPUT_DIR, + TRACE_ENV_MAX_EVENTS, + TRACE_ENV_MAX_EVENTS_PER_TRACE, + ): + value = os.environ.get(name) + if value is not None: + env_vars[name] = value + return env_vars + + def merge_runtime_env(self, actor_options: dict[str, Any]) -> dict[str, Any]: + trace_env_vars = self.env_vars() + if not trace_env_vars: + return actor_options + runtime_env = dict(actor_options.get("runtime_env") or {}) + env_vars = dict(runtime_env.get("env_vars") or {}) + env_vars.update(trace_env_vars) + runtime_env["env_vars"] = env_vars + actor_options["runtime_env"] = runtime_env + return actor_options + + def _replace_runtime(self, config: TraceConfig) -> TraceRuntime: + config = self._normalize_config(config) + identity = self._config_identity(config) + with self._lock: + if self._runtime is not None and self._identity == identity: + return self._runtime + old_runtime = self._runtime + self._runtime = build_trace_runtime(config) + self._identity = identity + self._register_atexit() + runtime = self._runtime + if old_runtime is not None: + old_runtime.close() + return runtime + + def _register_atexit(self) -> None: + if self._atexit_registered: + return + atexit.register(close_trace) + self._atexit_registered = True + + @staticmethod + def _normalize_config(config: TraceConfig, *, output_dir: str | Path | None = None) -> TraceConfig: + updates: dict[str, Any] = {} + resolved_output_dir = output_dir if output_dir is not None else config.output_dir + if resolved_output_dir is not None: + updates["output_dir"] = str(Path(resolved_output_dir).absolute()) + if updates: + config = config.model_copy(update=updates) + return config + + @staticmethod + def _export_env(config: TraceConfig) -> None: + if not config.enabled: + TraceRuntimeManager._clear_env() + os.environ[TRACE_ENV_ENABLED] = "0" + return + + os.environ[TRACE_ENV_ENABLED] = "1" + if config.output_dir is not None: + os.environ[TRACE_ENV_OUTPUT_DIR] = str(Path(config.output_dir).absolute()) + else: + os.environ.pop(TRACE_ENV_OUTPUT_DIR, None) + os.environ[TRACE_ENV_MAX_EVENTS] = str(config.max_events) + os.environ[TRACE_ENV_MAX_EVENTS_PER_TRACE] = str(config.max_events_per_trace) + + @staticmethod + def _clear_env() -> None: + for name in ( + TRACE_ENV_ENABLED, + TRACE_ENV_OUTPUT_DIR, + TRACE_ENV_MAX_EVENTS, + TRACE_ENV_MAX_EVENTS_PER_TRACE, + ): + os.environ.pop(name, None) + + @staticmethod + def _load_config_from_env() -> TraceConfig: + if os.environ.get(TRACE_ENV_ENABLED) != "1": + return TraceConfig() + + output_dir = os.environ.get(TRACE_ENV_OUTPUT_DIR) + max_events = int(os.environ.get(TRACE_ENV_MAX_EVENTS, TraceConfig.model_fields["max_events"].default)) + max_events_per_trace = int( + os.environ.get(TRACE_ENV_MAX_EVENTS_PER_TRACE, TraceConfig.model_fields["max_events_per_trace"].default) + ) + return TraceConfig( + enabled=True, + output_dir=output_dir, + max_events=max_events, + max_events_per_trace=max_events_per_trace, + ) + + @staticmethod + def _config_identity(config: TraceConfig) -> tuple[Any, ...]: + output_dir = str(Path(config.output_dir).absolute()) if config.output_dir is not None else None + return ( + config.enabled, + output_dir, + config.max_events, + config.max_events_per_trace, + ) + + +_TRACE_RUNTIME_MANAGER = TraceRuntimeManager() diff --git a/xtuner/v1/rl/utils/ray_accelerator_worker.py b/xtuner/v1/rl/utils/ray_accelerator_worker.py index a7d52d981a..8b9d40a1b2 100644 --- a/xtuner/v1/rl/utils/ray_accelerator_worker.py +++ b/xtuner/v1/rl/utils/ray_accelerator_worker.py @@ -15,6 +15,8 @@ ) from typing_extensions import Annotated +from xtuner.v1.rl.trace import merge_trace_runtime_env + from .ray_utils import find_master_addr_and_port, get_accelerator_ids @@ -458,6 +460,7 @@ def from_placement_group( (rank, bundle_index). """ pg_options = cls.get_pg_options(pg) + merge_trace_runtime_env(pg_options) device_type = cls.get_device_type(pg) sorted_bundle_idxs, master_addr, master_port, world_size = cls.get_spmd_info(pg) diff --git a/xtuner/v1/rl/utils/ray_cpu_worker.py b/xtuner/v1/rl/utils/ray_cpu_worker.py index ed9ebc3783..e69349f69b 100644 --- a/xtuner/v1/rl/utils/ray_cpu_worker.py +++ b/xtuner/v1/rl/utils/ray_cpu_worker.py @@ -17,6 +17,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from typing_extensions import Annotated +from xtuner.v1.rl.trace import merge_trace_runtime_env from xtuner.v1.utils.logger import get_logger @@ -185,6 +186,7 @@ def build_actor( } if resolved_memory is not None and resolved_memory > 0: actor_options["memory"] = resolved_memory + merge_trace_runtime_env(actor_options) if pg is None: return actor_cls.options(**actor_options).remote(*init_args, **init_kwargs) diff --git a/xtuner/v1/train/rl_trainer.py b/xtuner/v1/train/rl_trainer.py index 73c1abec19..3df1a6519c 100644 --- a/xtuner/v1/train/rl_trainer.py +++ b/xtuner/v1/train/rl_trainer.py @@ -39,6 +39,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 ( @@ -321,6 +322,7 @@ class BaseRLTrainerConfig(BaseModel): train_batch_size: int advantage_estimator_config: BaseAdvantageConfig = Field(default_factory=GRPOAdvantageConfig) sync_weights_interval: int = 1 + trace_config: TraceConfig = Field(default_factory=TraceConfig) enable_evaluate: bool = True enable_initial_evaluate: bool = False @@ -555,6 +557,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) @@ -595,6 +598,40 @@ def _init_save_config(self, cfg: BaseRLTrainerConfig) -> None: self._checkpoint_no_save_optimizer = cfg.checkpoint_no_save_optimizer self._load_checkpoint_cfg = self._resolve_load_checkpoint_cfg(cfg.auto_resume, cfg.load_checkpoint_cfg) + def _init_trace(self, cfg: BaseRLTrainerConfig) -> None: + trace_config = cfg.trace_config + if trace_config.enabled and trace_config.output_dir is None: + trace_config = trace_config.model_copy(update={"output_dir": self.exp_dir / "producer_trace"}) + self._trace_viewer_handle: Any | None = None + self._trace_config = trace_config + configure_trace(trace_config) + self._maybe_start_trace_viewer(trace_config) + + def _maybe_start_trace_viewer(self, trace_config: TraceConfig) -> None: + if not trace_config.enabled or not trace_config.viewer_enabled or trace_config.output_dir is None: + return + if get_rank() != 0: + return + + from xtuner.tools.producer_trace_viewer import start_trace_viewer + + handle = start_trace_viewer( + Path(trace_config.output_dir), + host=trace_config.viewer_host, + port=trace_config.viewer_port, + refresh_interval_s=trace_config.viewer_refresh_interval_s, + scope=trace_config.viewer_scope, + ) + self._trace_viewer_handle = handle + self.logger.info(f"Producer Trace Viewer: {handle.url}") + + def _close_trace(self) -> None: + handle = getattr(self, "_trace_viewer_handle", None) + if handle is not None: + handle.close() + self._trace_viewer_handle = None + close_trace() + def _init_logger(self, cfg: BaseRLTrainerConfig, logger_tag: str) -> Path: log_dir = self.exp_dir / "logs" log_dir.mkdir(parents=True, exist_ok=True) @@ -1599,6 +1636,12 @@ def _sync_weights_from_train_workers(self) -> None: self.logger.info("Rollout workers updated weights from train workers.") def fit(self): + try: + return self._fit() + finally: + self._close_trace() + + def _fit(self): self.logger.info("Start RL training") if self._cur_step >= self._total_train_steps: self.logger.info(f"Train steps {self._total_train_steps} reached, stop training") @@ -1796,7 +1839,10 @@ def _resume_from_checkpoint(self, checkpoint_path: Path | str) -> None: def fit(self): # 对外保留同步 fit 接口，内部用 async loop 组织 producer/consumer。 - return asyncio_run(self._fit()) + try: + return asyncio_run(self._fit()) + finally: + self._close_trace() async def _fit(self): self.logger.info("Start RL disaggregated training") From 62335a6207424c459c7b4270843bdc41f8586b4f Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Tue, 9 Jun 2026 08:17:40 +0000 Subject: [PATCH 2/3] simplify param in trace_function --- tests/rl/test_trace.py | 508 ++++++++++++++++++ xtuner/v1/rl/agent_loop/agent_loop.py | 12 +- .../rl/agent_loop/single_turn_agent_loop.py | 6 +- xtuner/v1/rl/agent_loop_manager/producer.py | 8 +- xtuner/v1/rl/rollout/controller.py | 6 +- xtuner/v1/rl/rollout/utils.py | 4 +- xtuner/v1/rl/rollout/worker.py | 7 +- xtuner/v1/rl/trace.py | 32 +- 8 files changed, 542 insertions(+), 41 deletions(-) create mode 100644 tests/rl/test_trace.py diff --git a/tests/rl/test_trace.py b/tests/rl/test_trace.py new file mode 100644 index 0000000000..da55943501 --- /dev/null +++ b/tests/rl/test_trace.py @@ -0,0 +1,508 @@ +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from xtuner.tools.producer_trace_analysis import load_trace_jsonl +from xtuner.tools.producer_trace_hotspots import build_hotspot_payload_from_events +from xtuner.tools.producer_trace_viewer import build_viewer_payload_from_events +from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.trace import ( + InMemoryTraceStore, + TraceConfig, + TraceEvent, + TraceRecorder, + close_trace, + configure_trace, + get_trace_env_vars, + merge_trace_runtime_env, + reset_trace_for_test, + trace_event, + trace_function, + trace_span, + use_trace_recorder, +) + + +def make_state(uid: int = 1, task_name: str = "gsm8k", status: Status = Status.INIT) -> RolloutState: + return RolloutState( + message=[{"role": "user", "content": "What is 1 + 1?"}], + uid=uid, + task_name=task_name, + session_uid=uid + 1000, + status=status, + ) + + +def make_event( + trace_id: str, + stage: str, + timestamp_s: float, + *, + task_name: str = "gsm8k", + uid: int = 1, + status: str = "init", + elapsed_s: float | None = None, + train_step: int | None = None, + model_step: int | None = None, + producer_future_step: int | None = None, + produce_batch_id: str | None = None, +) -> TraceEvent: + return TraceEvent( + trace_id=trace_id, + stage=stage, + timestamp_s=timestamp_s, + status=status, + task_name=task_name, + uid=uid, + session_uid=uid + 1000, + train_step=train_step, + model_step=model_step, + producer_future_step=producer_future_step, + produce_batch_id=produce_batch_id, + worker_rank=None, + elapsed_s=elapsed_s, + error_msg=None, + ) + + +def make_batch_event( + trace_id: str, + stage: str, + timestamp_s: float, + *, + uid: int, + batch: str, + train_step: int, + model_step: int = 0, + producer_future_step: int = 0, + elapsed_s: float | None = None, + status: str = "init", +) -> TraceEvent: + return make_event( + trace_id, + stage, + timestamp_s, + uid=uid, + status=status, + elapsed_s=elapsed_s, + train_step=train_step, + model_step=model_step, + producer_future_step=producer_future_step, + produce_batch_id=batch, + ) + + +class TraceCoreBehaviorTest(unittest.IsolatedAsyncioTestCase): + def tearDown(self): + reset_trace_for_test() + + async def test_trace_api_records_event_span_function_to_jsonl(self): + with tempfile.TemporaryDirectory() as tmp_dir: + store = InMemoryTraceStore( + TraceConfig(enabled=True, output_dir=tmp_dir, max_events=20, max_events_per_trace=20) + ) + recorder = TraceRecorder(store) + state = make_state(uid=123) + + with use_trace_recorder(recorder): + await trace_event(state, "custom.prepare") + + async with trace_span(state, "custom.work"): + state.status = Status.COMPLETED + + @trace_function("custom.fn") + async def traced_fn(rollout_state: RolloutState) -> RolloutState: + await trace_event(rollout_state, "custom.fn.inner") + return rollout_state + + await traced_fn(state) + + store.flush_jsonl() + store.close() + timeline = store.get_timeline("gsm8k:123") + disk_events = load_trace_jsonl(tmp_dir) + + expected_stages = [ + "custom.prepare", + "custom.work.start", + "custom.work.end", + "custom.fn.start", + "custom.fn.inner", + "custom.fn.end", + ] + self.assertEqual([event.stage for event in timeline], expected_stages) + self.assertEqual([event.stage for event in disk_events], expected_stages) + self.assertEqual(timeline[-1].status, "completed") + + async def test_trace_function_resolves_target_and_dynamic_kwargs(self): + store = InMemoryTraceStore(TraceConfig(enabled=True, max_events=10, max_events_per_trace=10)) + batch_id = "train_step=1/model_step=2/producer_future_step=3" + + class Worker: + def __init__(self): + self.worker_rank = 7 + + @trace_function( + "custom.worker", + target="state", + trace_kwargs_getter=lambda self, *args, **kwargs: { + "produce_batch_id": batch_id, + "worker_rank": self.worker_rank, + }, + ) + async def run(self, state: RolloutState) -> RolloutState: + return state.model_copy(update={"status": Status.COMPLETED}, deep=True) + + state = make_state(uid=456) + with use_trace_recorder(TraceRecorder(store)): + await Worker().run(state) + timeline = store.get_timeline("gsm8k:456") + + self.assertEqual([event.stage for event in timeline], ["custom.worker.start", "custom.worker.end"]) + self.assertEqual([event.produce_batch_id for event in timeline], [batch_id, batch_id]) + self.assertEqual([event.worker_rank for event in timeline], [7, 7]) + self.assertEqual(timeline[0].status, "init") + self.assertEqual(timeline[-1].status, "completed") + + async def test_trace_function_records_error_event_and_reraises(self): + store = InMemoryTraceStore(TraceConfig(enabled=True, max_events=10, max_events_per_trace=10)) + state = make_state(uid=789) + + @trace_function("custom.failure", target="state") + async def failing_fn(state: RolloutState) -> None: + raise ValueError("rollout failed") + + with use_trace_recorder(TraceRecorder(store)): + with self.assertRaisesRegex(ValueError, "rollout failed"): + await failing_fn(state) + + timeline = store.get_timeline("gsm8k:789") + self.assertEqual([event.stage for event in timeline], ["custom.failure.start", "custom.failure.error"]) + self.assertIsNotNone(timeline[-1].elapsed_s) + self.assertTrue(timeline[-1].error_msg.startswith("ValueError: rollout failed")) + + def test_trace_runtime_env_is_propagated_to_ray_actor_options(self): + with tempfile.TemporaryDirectory() as tmp_dir: + configure_trace(TraceConfig(enabled=True, output_dir=tmp_dir, max_events=20, max_events_per_trace=30)) + actor_options = {"num_cpus": 1, "runtime_env": {"env_vars": {"EXISTING": "1"}}} + + env_vars = get_trace_env_vars() + merged = merge_trace_runtime_env(actor_options) + close_trace() + + self.assertIs(merged, actor_options) + self.assertEqual(env_vars["XTUNER_TRACE_ENABLED"], "1") + self.assertEqual(env_vars["XTUNER_TRACE_OUTPUT_DIR"], str(Path(tmp_dir).absolute())) + self.assertEqual(env_vars["XTUNER_TRACE_MAX_EVENTS"], "20") + self.assertEqual(env_vars["XTUNER_TRACE_MAX_EVENTS_PER_TRACE"], "30") + self.assertEqual(actor_options["runtime_env"]["env_vars"]["EXISTING"], "1") + self.assertEqual(actor_options["runtime_env"]["env_vars"]["XTUNER_TRACE_ENABLED"], "1") + self.assertEqual(actor_options["runtime_env"]["env_vars"]["XTUNER_TRACE_OUTPUT_DIR"], str(Path(tmp_dir).absolute())) + self.assertEqual(get_trace_env_vars(), {}) + + def test_online_viewer_payload_reports_latest_batch_and_current_stalls(self): + old_batch = "train_step=1/model_step=0/producer_future_step=0" + latest_batch = "train_step=2/model_step=0/producer_future_step=0" + events = [ + make_batch_event( + "gsm8k:1", + "xtuner.rollout_controller.generate.start", + 5.0, + uid=1, + batch=old_batch, + train_step=1, + ), + make_batch_event( + "gsm8k:2", + "xtuner.rollout_controller.generate.start", + 10.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:3", + "xtuner.judger.judge.start", + 12.0, + uid=3, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:4", + "xtuner.producer.put_generated_group.start", + 13.0, + uid=4, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:4", + "xtuner.producer.put_generated_group.end", + 15.0, + uid=4, + batch=latest_batch, + train_step=2, + elapsed_s=2.0, + status="completed", + ), + ] + + with patch("xtuner.tools.producer_trace_analysis.time.time", return_value=30.0): + payload = build_viewer_payload_from_events(events, trace_source="/tmp/trace") + + summary = payload["task_summary"] + self.assertEqual(payload["raw_event_count"], 5) + self.assertEqual(payload["event_count"], 4) + self.assertEqual(summary["total_tasks"], 3) + self.assertEqual(summary["running_tasks"], 2) + self.assertEqual(summary["completed_tasks"], 1) + self.assertEqual(summary["current_stage_counts"]["rollout.generate"], 1) + self.assertEqual(summary["current_stage_counts"]["judger"], 1) + self.assertEqual({row["trace_id"] for row in payload["rows"]}, {"gsm8k:2", "gsm8k:3", "gsm8k:4"}) + + open_summary_by_span = {item["span"]: item for item in payload["open_span_summaries"]} + self.assertEqual(open_summary_by_span["xtuner.rollout_controller.generate"]["oldest_trace_id"], "gsm8k:2") + self.assertEqual(open_summary_by_span["xtuner.rollout_controller.generate"]["oldest_age_s"], 20.0) + self.assertEqual(open_summary_by_span["xtuner.judger.judge"]["oldest_age_s"], 18.0) + + def test_hotspot_payload_builds_nested_spans_and_stage_stats(self): + old_batch = "train_step=1/model_step=0/producer_future_step=0" + latest_batch = "train_step=2/model_step=0/producer_future_step=0" + events = [ + make_batch_event( + "gsm8k:1", + "xtuner.producer.generate_group.start", + 0.0, + uid=1, + batch=old_batch, + train_step=1, + ), + make_batch_event( + "gsm8k:1", + "xtuner.producer.generate_group.end", + 1.0, + uid=1, + batch=old_batch, + train_step=1, + elapsed_s=1.0, + ), + make_batch_event( + "gsm8k:2", + "xtuner.producer.generate_group.start", + 100.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:2", + "xtuner.agent_loop.generate_group.start", + 101.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:2", + "xtuner.agent_loop.generate_sample.start", + 102.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:2", + "xtuner.rollout_controller.generate.start", + 103.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:2", + "xtuner.rollout_worker.generate.start", + 104.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:2", + "xtuner.rollout_engine.generate.start", + 105.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:2", + "xtuner.rollout_engine.generate.end", + 108.0, + uid=2, + batch=latest_batch, + train_step=2, + elapsed_s=3.0, + ), + make_batch_event( + "gsm8k:2", + "xtuner.rollout_worker.generate.end", + 109.0, + uid=2, + batch=latest_batch, + train_step=2, + elapsed_s=5.0, + ), + make_batch_event( + "gsm8k:2", + "xtuner.rollout_controller.generate.end", + 110.0, + uid=2, + batch=latest_batch, + train_step=2, + elapsed_s=7.0, + ), + make_batch_event( + "gsm8k:2", + "xtuner.judger.judge.start", + 111.0, + uid=2, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:2", + "xtuner.judger.judge.end", + 115.0, + uid=2, + batch=latest_batch, + train_step=2, + elapsed_s=4.0, + ), + make_batch_event( + "gsm8k:2", + "xtuner.agent_loop.generate_sample.end", + 116.0, + uid=2, + batch=latest_batch, + train_step=2, + elapsed_s=14.0, + ), + make_batch_event( + "gsm8k:2", + "xtuner.agent_loop.generate_group.end", + 117.0, + uid=2, + batch=latest_batch, + train_step=2, + elapsed_s=16.0, + ), + make_batch_event( + "gsm8k:2", + "xtuner.producer.generate_group.end", + 118.0, + uid=2, + batch=latest_batch, + train_step=2, + elapsed_s=18.0, + ), + make_batch_event( + "gsm8k:3", + "xtuner.rollout_controller.generate.start", + 1000.0, + uid=3, + batch=latest_batch, + train_step=2, + ), + make_batch_event( + "gsm8k:3", + "xtuner.rollout_controller.generate.end", + 1010.0, + uid=3, + batch=latest_batch, + train_step=2, + elapsed_s=10.0, + ), + ] + + payload = build_hotspot_payload_from_events(events, trace_source="/tmp/trace") + all_payload = build_hotspot_payload_from_events(events, trace_source="/tmp/trace", scope="all") + + self.assertEqual(payload["task_count"], 2) + self.assertEqual(all_payload["task_count"], 3) + self.assertEqual(payload["scale_mode"], "task_relative") + + row_by_trace_id = {row["trace_id"]: row for row in payload["rows"]} + nested_spans = row_by_trace_id["gsm8k:2"]["spans"] + self.assertEqual( + [span["display_stage"] for span in nested_spans], + [ + "producer.generate", + "agent_loop.generate_group", + "agent_loop.generate_sample", + "rollout.generate", + "rollout_worker.generate", + "engine.generate", + "judger", + ], + ) + self.assertEqual([span["depth"] for span in nested_spans], [0, 1, 2, 3, 4, 5, 3]) + self.assertEqual(nested_spans[0]["left_pct"], 0.0) + self.assertEqual(row_by_trace_id["gsm8k:3"]["spans"][0]["left_pct"], 0.0) + + stats_by_stage = {stat["stage"]: stat for stat in payload["stage_stats"]} + self.assertEqual(stats_by_stage["engine.generate"]["avg_s"], 3.0) + self.assertEqual(stats_by_stage["engine.generate"]["p95_s"], 3.0) + self.assertEqual(stats_by_stage["engine.generate"]["max_s"], 3.0) + self.assertEqual(stats_by_stage["rollout.generate"]["count"], 2) + self.assertEqual(stats_by_stage["rollout.generate"]["avg_s"], 8.5) + self.assertEqual(stats_by_stage["rollout.generate"]["max_s"], 10.0) + + +class TraceTrainerIntegrationTest(unittest.TestCase): + def tearDown(self): + reset_trace_for_test() + + def test_trainer_starts_live_viewer_on_rank0(self): + with patch.dict("sys.modules", {"causal_conv1d_cuda": MagicMock()}): + from xtuner.v1.train import rl_trainer + + BaseRLTrainer = rl_trainer.BaseRLTrainer + + with tempfile.TemporaryDirectory() as tmp_dir: + trainer = object.__new__(BaseRLTrainer) + trainer._meta = SimpleNamespace(latest_exp=SimpleNamespace(exp_dir=tmp_dir)) + trainer.logger = MagicMock() + handle = MagicMock() + handle.url = "http://127.0.0.1:39563" + cfg = SimpleNamespace( + trace_config=TraceConfig( + enabled=True, + output_dir=None, + viewer_host="127.0.0.1", + viewer_port=39563, + viewer_refresh_interval_s=2.5, + viewer_scope="all", + ) + ) + + with ( + patch.object(rl_trainer, "get_rank", return_value=0), + patch("xtuner.tools.producer_trace_viewer.start_trace_viewer", return_value=handle) as start_viewer, + ): + trainer._init_trace(cfg) + trace_dir = Path(tmp_dir) / "producer_trace" + self.assertTrue(trace_dir.exists()) + trainer._close_trace() + + start_viewer.assert_called_once_with( + trace_dir, + host="127.0.0.1", + port=39563, + refresh_interval_s=2.5, + scope="all", + ) + trainer.logger.info.assert_any_call("Producer Trace Viewer: http://127.0.0.1:39563") + handle.close.assert_called_once() diff --git a/xtuner/v1/rl/agent_loop/agent_loop.py b/xtuner/v1/rl/agent_loop/agent_loop.py index 55feaaeb2e..c260b92bb3 100644 --- a/xtuner/v1/rl/agent_loop/agent_loop.py +++ b/xtuner/v1/rl/agent_loop/agent_loop.py @@ -190,11 +190,7 @@ def __init__( @abstractmethod async def generate_sample(self, rollout_state: RolloutState, **kwargs) -> RolloutState: ... - @trace_function( - "xtuner.agent_loop.generate_group", - target="rollout_state", - result="return", - ) + @trace_function("xtuner.agent_loop.generate_group") async def generate_group(self, rollout_state: list[RolloutState], **kwargs) -> list[RolloutState]: pending_tasks = [] for state in rollout_state: @@ -214,11 +210,7 @@ async def run_judger(self, rollout_state: RolloutState) -> RolloutState: ... @overload async def run_judger(self, rollout_state: list[RolloutState]) -> list[RolloutState]: ... - @trace_function( - "xtuner.judger.judge", - target="rollout_state", - result="return", - ) + @trace_function("xtuner.judger.judge") async def run_judger(self, rollout_state: RolloutState | list[RolloutState]) -> RolloutState | list[RolloutState]: assert self.judger is not None if isinstance(rollout_state, list): diff --git a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py index 203a065758..ba65707aba 100644 --- a/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py +++ b/xtuner/v1/rl/agent_loop/single_turn_agent_loop.py @@ -65,11 +65,7 @@ def __init__( enable_batch_judge=enable_batch_judge, ) - @trace_function( - "xtuner.agent_loop.generate_sample", - target="rollout_state", - result="return", - ) + @trace_function("xtuner.agent_loop.generate_sample") async def generate_sample( self, rollout_state: RolloutState, diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index bdad231c72..e634c9c71c 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -357,19 +357,13 @@ async def available_count(self) -> int: @trace_function( "xtuner.producer.sample_group", - result="return", trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs(), ) 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) - @trace_function( - "xtuner.producer.generate_group", - target="rollout_state", - result="return", - trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs(), - ) + @trace_function("xtuner.producer.generate_group", trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs()) async def generate_group( self, rollout_state: list[RolloutState], diff --git a/xtuner/v1/rl/rollout/controller.py b/xtuner/v1/rl/rollout/controller.py index de2f3b8746..ea0a1a2cc4 100644 --- a/xtuner/v1/rl/rollout/controller.py +++ b/xtuner/v1/rl/rollout/controller.py @@ -174,11 +174,7 @@ def get_generate_concurrency(self) -> int: return active_workers * concurrency_per_worker @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) - @trace_function( - "xtuner.rollout_controller.generate", - target="rollout_state", - result="return", - ) + @trace_function("xtuner.rollout_controller.generate") async def generate(self, rollout_state: RolloutState) -> RolloutState: if XTUNER_DETERMINISTIC: sample_params = rollout_state.sample_params.model_copy(deep=True) diff --git a/xtuner/v1/rl/rollout/utils.py b/xtuner/v1/rl/rollout/utils.py index 7991280db8..1b291973a8 100644 --- a/xtuner/v1/rl/rollout/utils.py +++ b/xtuner/v1/rl/rollout/utils.py @@ -13,7 +13,7 @@ from xtuner.v1.data_proto.rl_data import RolloutState, Status from xtuner.v1.rl.utils import free_object_refs from xtuner.v1.utils import get_logger - +from xtuner.v1.rl.trace import trace_function if TYPE_CHECKING: from .controller import WorkerInfo @@ -275,6 +275,7 @@ class PartialRolloutHandler: def __init__(self) -> None: self.logger = get_logger(self.__class__.__name__) + @trace_function("xtuner.partial_rollout_handler.postprocess") def preprocess(self, rollout_state: RolloutState, max_tokens: int) -> RolloutState: # Set up token and length variable response_ids = list(rollout_state.response_ids or []) @@ -291,6 +292,7 @@ def preprocess(self, rollout_state: RolloutState, max_tokens: int) -> RolloutSta ) return rollout_state + @trace_function("xtuner.partial_rollout_handler.postprocess") async def postprocess( self, rollout_state: RolloutState, diff --git a/xtuner/v1/rl/rollout/worker.py b/xtuner/v1/rl/rollout/worker.py index 8899d15f56..c1d06c38d6 100644 --- a/xtuner/v1/rl/rollout/worker.py +++ b/xtuner/v1/rl/rollout/worker.py @@ -681,12 +681,7 @@ async def _decode_routed_experts(self, routed_experts: Any) -> Any: return routed_experts @ray.method(concurrency_group=ROLLOUT_CONCURRENCY_GROUP_GENERATE) - @trace_function( - "xtuner.rollout_worker.generate", - target="rollout_state", - result="return", - trace_kwargs_getter=lambda self, *args, **kwargs: {"worker_rank": self.rank}, - ) + @trace_function("xtuner.rollout_worker.generate", trace_kwargs_getter=lambda self, *args, **kwargs: {"worker_rank": self.rank}) async def generate(self, rollout_state: RolloutState) -> RolloutState: try: # TODO(@duanyanhui): diff --git a/xtuner/v1/rl/trace.py b/xtuner/v1/rl/trace.py index f8d57af194..78356a2c09 100644 --- a/xtuner/v1/rl/trace.py +++ b/xtuner/v1/rl/trace.py @@ -636,6 +636,10 @@ def resolve( return target(*args, **kwargs) if target is not None: return target + rollout_state = bound_arguments.get("rollout_state") + rollout_states = cls.as_rollout_state_list(rollout_state) + if rollout_states: + return rollout_states if len(rollout_states) > 1 else rollout_states[0] for value in bound_arguments.values(): states = cls.as_rollout_state_list(value) if states: @@ -796,14 +800,12 @@ def __init__( target: str | RolloutState | Sequence[RolloutState] | Callable[..., Any] | None = None, target_getter: Callable[..., Any] | None = None, trace_kwargs_getter: Callable[..., dict[str, Any] | None] | None = None, - result: str | Callable[[Any], Any] | None = None, trace_kwargs: dict[str, Any] | None = None, ) -> None: self.name = name self.target = target self.target_getter = target_getter self.trace_kwargs_getter = trace_kwargs_getter - self.result = result self.trace_kwargs = trace_kwargs or {} def decorate(self, func: Callable[..., Any]): @@ -824,9 +826,7 @@ def _start_target(self, signature: inspect.Signature, args: tuple[Any, ...], kwa ) def _end_target(self, start_target: Any, return_value: Any) -> Any: - if callable(self.result): - return self.result(return_value) - if self.result == "return": + if TraceTargetResolver.as_rollout_state_list(return_value): return return_value return start_target @@ -911,15 +911,33 @@ def trace_function( target: str | RolloutState | Sequence[RolloutState] | Callable[..., Any] | None = None, target_getter: Callable[..., Any] | None = None, trace_kwargs_getter: Callable[..., dict[str, Any] | None] | None = None, - result: str | Callable[[Any], Any] | None = None, **trace_kwargs: Any, ): + """Trace a whole sync/async function as one task-level span. + + Target resolution for the `.start` event: + - If `target_getter` is provided, use its return value. + - Else if `target` is provided, resolve that explicit target. + - Else prefer the argument named `rollout_state` when it is a `RolloutState` + or `list[RolloutState]`. + - Else fall back to the first `RolloutState` / `list[RolloutState]` found + in the bound arguments. + + Target resolution for the `.end` event: + - If the function returns a `RolloutState` or `list[RolloutState]`, use the + return value so the end event reflects the latest task state. + - Otherwise reuse the start target. + + In practice this means standard XTuner functions whose task parameter is + named `rollout_state` usually do not need to pass `target=...`. Functions + with non-standard parameter names such as `group` should still pass an + explicit `target`. + """ return TraceFunctionDecorator( name, target=target, target_getter=target_getter, trace_kwargs_getter=trace_kwargs_getter, - result=result, trace_kwargs=trace_kwargs, ).decorate From 191507773ee44de9704b30f96bad82fe7dc9865a Mon Sep 17 00:00:00 2001 From: YanhuiDua Date: Tue, 9 Jun 2026 09:05:20 +0000 Subject: [PATCH 3/3] fix claude comments --- xtuner/tools/producer_trace_hotspots.py | 23 ++------------------- xtuner/v1/rl/agent_loop_manager/producer.py | 12 +++++------ xtuner/v1/rl/rollout/utils.py | 5 +++-- xtuner/v1/rl/trace.py | 2 +- 4 files changed, 12 insertions(+), 30 deletions(-) diff --git a/xtuner/tools/producer_trace_hotspots.py b/xtuner/tools/producer_trace_hotspots.py index b8418d2d0c..a5a729f4f5 100644 --- a/xtuner/tools/producer_trace_hotspots.py +++ b/xtuner/tools/producer_trace_hotspots.py @@ -14,6 +14,8 @@ TraceViewerScope, ) from xtuner.tools.producer_trace_analysis import ( + _percentile, + _split_span_stage, display_trace_stage, filter_trace_events_by_scope, load_trace_jsonl, @@ -329,27 +331,6 @@ def _build_stage_stats(records: Iterable[TraceSpanRecord]) -> list[dict[str, Any return sorted(stats, key=lambda item: (item["p95_s"], item["total_s"]), reverse=True) -def _split_span_stage(stage: str) -> tuple[str | None, str | None]: - for suffix in (".start", ".end", ".error"): - if stage.endswith(suffix): - return stage[: -len(suffix)], suffix[1:] - return None, None - - -def _percentile(sorted_values: list[float], percentile: float) -> float: - if not sorted_values: - return 0.0 - if len(sorted_values) == 1: - return sorted_values[0] - position = (len(sorted_values) - 1) * percentile - lower = int(position) - upper = min(lower + 1, len(sorted_values) - 1) - if lower == upper: - return sorted_values[lower] - fraction = position - lower - return sorted_values[lower] * (1 - fraction) + sorted_values[upper] * fraction - - _HTML_TEMPLATE = """ diff --git a/xtuner/v1/rl/agent_loop_manager/producer.py b/xtuner/v1/rl/agent_loop_manager/producer.py index e634c9c71c..b1a87830de 100644 --- a/xtuner/v1/rl/agent_loop_manager/producer.py +++ b/xtuner/v1/rl/agent_loop_manager/producer.py @@ -363,7 +363,9 @@ 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) - @trace_function("xtuner.producer.generate_group", trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs()) + @trace_function( + "xtuner.producer.generate_group", trace_kwargs_getter=lambda self, *args, **kwargs: self.trace_kwargs() + ) async def generate_group( self, rollout_state: list[RolloutState], @@ -392,11 +394,9 @@ async def generate_group( ) elapsed = time.perf_counter() - start for item in result: - extra_fields = getattr(item, "extra_fields", None) - if extra_fields is None: - extra_fields = {} - setattr(item, "extra_fields", extra_fields) - extra_fields[GROUP_GENERATE_TIME_KEY] = elapsed + item_extra_fields = dict(item.extra_fields) + item_extra_fields[GROUP_GENERATE_TIME_KEY] = elapsed + item.extra_fields = item_extra_fields return result @trace_function( diff --git a/xtuner/v1/rl/rollout/utils.py b/xtuner/v1/rl/rollout/utils.py index 1b291973a8..b822e27a2c 100644 --- a/xtuner/v1/rl/rollout/utils.py +++ b/xtuner/v1/rl/rollout/utils.py @@ -11,9 +11,10 @@ from ray import ObjectRef as RayObjectRef from xtuner.v1.data_proto.rl_data import RolloutState, Status +from xtuner.v1.rl.trace import trace_function from xtuner.v1.rl.utils import free_object_refs from xtuner.v1.utils import get_logger -from xtuner.v1.rl.trace import trace_function + if TYPE_CHECKING: from .controller import WorkerInfo @@ -275,7 +276,7 @@ class PartialRolloutHandler: def __init__(self) -> None: self.logger = get_logger(self.__class__.__name__) - @trace_function("xtuner.partial_rollout_handler.postprocess") + @trace_function("xtuner.partial_rollout_handler.preprocess") def preprocess(self, rollout_state: RolloutState, max_tokens: int) -> RolloutState: # Set up token and length variable response_ids = list(rollout_state.response_ids or []) diff --git a/xtuner/v1/rl/trace.py b/xtuner/v1/rl/trace.py index 78356a2c09..68a1b8c2c1 100644 --- a/xtuner/v1/rl/trace.py +++ b/xtuner/v1/rl/trace.py @@ -373,7 +373,7 @@ def _ensure_file(self, next_bytes: int) -> None: def _open_file(self) -> None: path = self.output_dir / f"{TRACE_JSONL_BASENAME}_{self._writer_id}_{self._shard_idx:06d}.jsonl" self._file = path.open("a", encoding="utf-8") - self._bytes_written = path.stat().st_size if path.exists() else 0 + self._bytes_written = path.stat().st_size def _flush_file(self) -> None: if self._file is not None: