 |
VOOZH | about |
|
VOOZH | about |
This page documents AReaL's trace visualization system, which provides performance profiling and session lifecycle tracking for distributed training and inference workloads. The system captures timestamped events during execution and exports them in formats compatible with visualization tools.
For distributed system architecture, see Architecture Overview For performance optimization guidance, see Performance Optimization Guide
AReaL provides two complementary tracing systems:
| System | Purpose | Output Format | Visualization Tool |
|---|---|---|---|
| PerfTracer | Low-level performance profiling with duration scopes and instant markers | JSONL → Chrome Trace JSON | Chrome Tracing (chrome://tracing) |
| SessionTracer | High-level session lifecycle tracking with phase breakdowns and metrics | JSONL | Plotly interactive plots |
Both systems write JSONL files to disk with configurable flush intervals, enabling post-training analysis of distributed execution patterns. The tracing infrastructure is designed for minimal overhead and uses context variables to propagate metadata through async execution contexts areal/utils/perf_tracer.py30-40
Sources: areal/utils/perf_tracer.py1-118
The PerfTracer class captures fine-grained performance events across distributed workers. Each trace event follows the Chrome Trace Event Format with timestamp, duration, category, and metadata.
Trace Event Generation Flow
Trace Event Categories:
The PerfTraceCategory enum areal/utils/perf_tracer.py62-93 defines standard categories used to organize events in visualization tools:
| Category | Usage | Examples |
|---|---|---|
COMPUTE | CPU/GPU computation | forward pass, backward pass, loss calculation |
COMM | Distributed communication | all-reduce, broadcast, NCCL operations |
IO | Disk I/O operations | checkpoint save/load, data loading |
SYNC | Synchronization barriers | torch.cuda.synchronize(), barriers |
SCHEDULER | Task scheduling | queue management, worker allocation |
INSTR | Instrumentation overhead | tracing itself |
MISC | Uncategorized events | general utilities |
Sources: areal/utils/perf_tracer.py62-114 areal/utils/perf_tracer.py1602-1950
The SessionTracer tracks high-level session lifecycle events for rollout episodes. Each session represents a single data sample's journey from submission through generation, reward computation, and finalization.
Session Lifecycle and Data Flow
Session Lifecycle Events:
The SessionTraceEvent enum defines lifecycle markers used by the tracer:
FINALIZED: Terminal state reached areal/utils/perf_tracer.py288GENERATE_START/END: Generation phase markers areal/utils/perf_tracer.py265-266REWARD_START/END: Reward computation markers areal/utils/perf_tracer.py268-269TOOLCALL_START/END: Tool calling markers areal/utils/perf_tracer.py271-272Sources: areal/utils/perf_tracer.py244-293 areal/workflow/rlvr.py82-136
The SessionRecord class represents a complete session execution trace with lifecycle timestamps, phase executions, and derived metrics.
| Field | Type | Description |
|---|---|---|
task_id | int | Dataset-level task identifier areal/utils/perf_tracer.py433 |
session_id | int | Unique session identifier areal/utils/perf_tracer.py434 |
rank | int | Process rank identifier areal/utils/perf_tracer.py435 |
role | str | None | Optional role (e.g., "rollout", "train") areal/utils/perf_tracer.py436 |
submit_ts | float | Submission timestamp (wall-clock) areal/utils/perf_tracer.py437 |
finalized_ts | float | None | Finalization timestamp areal/utils/perf_tracer.py438 |
status | str | "pending", "accepted", "rejected", "failed", "dropped" areal/utils/perf_tracer.py439 |
phases | dict[str, list[PhaseSpan]] | Phase execution history areal/utils/perf_tracer.py441 |
Each phase (generate, reward, toolcall) can execute multiple times within a session. A PhaseSpan records each execution with start_ts and end_ts areal/utils/perf_tracer.py296-318
The SessionRecord.to_dict() method computes derived metrics for analysis areal/utils/perf_tracer.py652-667:
total_s: Total session duration (submit to finalized).generate_s: Total time spent in all generate phase executions.reward_s: Total time spent in all reward phase executions.toolcall_s: Total time spent in all tool call phase executions.Sources: areal/utils/perf_tracer.py296-668
Function Decorator: Applied to methods to track execution time automatically.
Defined in areal/utils/perf_tracer.py1602-1750
Context Managers: Used for fine-grained scoping within functions.
Defined in areal/utils/perf_tracer.py1752-1900
Instant Markers: Mark a point-in-time event without duration.
Defined in areal/utils/perf_tracer.py1902-1950
Sources: areal/utils/perf_tracer.py1600-1950
Workflow Usage Example:
The following pattern is implemented in RLVRWorkflow areal/workflow/rlvr.py82-136 and VisionRLVRWorkflow areal/workflow/vision_rlvr.py45-101
Key Decorators:
| Decorator | Purpose | Usage |
|---|---|---|
@session_context() | Registers a new session and propagates session_id | Applied to top-level workflow methods areal/utils/perf_tracer.py816-860 |
@trace_session(phase) | Wraps a method with phase start/end markers | Applied to phase-specific methods areal/utils/perf_tracer.py862-895 |
atrace_session_phase(phase) | Async context manager for phase tracking | Used within session context areal/utils/perf_tracer.py897-918 |
Sources: areal/workflow/rlvr.py82-136 areal/workflow/vision_rlvr.py45-101 areal/utils/perf_tracer.py816-918
The convert_jsonl_to_chrome_trace() function areal/tools/perf_trace_converter.py295-438 transforms per-rank JSONL files into a single Chrome Trace JSON file.
Trace Conversion Logic
Key Transformations:
(rank, role, original_pid) areal/tools/perf_trace_converter.py121-215Sources: areal/tools/perf_trace_converter.py69-438
The plot_session_trace.py script generates interactive Plotly visualizations from session trace JSONL files.
Session Visualization Pipeline
Generated Visualizations:
_determine_step_timepoints() function implements global step synchronization by partitioning sessions into steps based on completion times and batch sizes areal/tools/plot_session_trace.py154-215Sources: areal/tools/plot_session_trace.py20-60 areal/tools/plot_session_trace.py154-250
The tracing behavior is controlled via configuration dataclasses areal/api/cli_args.py25:
| Parameter | Type | Default | Description |
|---|---|---|---|
enable | bool | False | Enable performance tracing |
save_interval | int | 100 | Flush events to disk every N events areal/utils/perf_tracer.py153-154 |
flush_threshold | int | 128 | Flush completed sessions every N sessions areal/utils/perf_tracer.py157-166 |
Traces are written to a standardized path under fileroot/logs/user/experiment/trial/ with rank-qualified filenames areal/utils/perf_tracer.py199-230
Sources: areal/utils/perf_tracer.py153-230 areal/api/cli_args.py25
The tracing system uses Python context variables to propagate metadata through async execution:
These variables are automatically inherited by async tasks created within a traced context, enabling task-level correlation of events and global step annotation on performance events.
Sources: areal/utils/perf_tracer.py31-40
Refresh this wiki