viz_traces.py

#!/usr/bin/env python3
"""
Convert cudf-polars traces from JSONL format to Chrome Tracing format.

Chrome Tracing format can be viewed in chrome://tracing or similar tools.
"""

import json
import sys
from pathlib import Path


def convert_trace_to_chrome_format(trace: dict) -> dict:
    """
    Convert a single trace event to Chrome Tracing format.
    
    Parameters
    ----------
    trace : dict
        A trace event with 'type', 'start', 'stop', 'process', 'thread' fields.
    
    Returns
    -------
    dict
        Chrome Tracing event with 'name', 'cat', 'ph', 'ts', 'dur', 'pid', 'tid'.
    """
    # Chrome Tracing expects timestamps in microseconds
    # Our traces are in nanoseconds, so divide by 1000
    ts_us = trace["start"] // 1000
    dur_us = (trace["stop"] - trace["start"]) // 1000
    
    return {
        "name": trace["type"],
        "cat": "IR",  # category
        "ph": "X",    # Complete event (duration event)
        "ts": ts_us,
        "dur": dur_us,
        "pid": trace["process"],
        "tid": trace["thread"],
        "args": {
            # Include additional metadata as arguments
            "query_id": trace.get("query_id"),
            "iteration": trace.get("iteration"),
            "total_bytes_input": trace.get("total_bytes_input"),
            "total_bytes_output": trace.get("total_bytes_output"),
            "count_frames_input": trace.get("count_frames_input"),
            "count_frames_output": trace.get("count_frames_output"),
        }
    }


def convert_jsonl_to_chrome_trace(input_file: str, output_file: str = None):
    """
    Convert traces from JSONL file to Chrome Tracing JSON format.
    
    Parameters
    ----------
    input_file : str
        Path to input JSONL file containing traces.
    output_file : str, optional
        Path to output JSON file. If not provided, uses input filename with .json extension.
    """
    input_path = Path(input_file)
    
    if output_file is None:
        output_file = input_path.with_suffix(".chrome_trace.json")
    
    chrome_events = []
    
    # Read the JSONL file
    with open(input_path) as f:
        for line in f:
            if not line.strip():
                continue
                
            data = json.loads(line)
            
            # Extract traces from all queries and iterations
            if "records" in data:
                for query_id, iterations in data["records"].items():
                    for iteration_data in iterations:
                        if "traces" in iteration_data:
                            for trace in iteration_data["traces"]:
                                chrome_event = convert_trace_to_chrome_format(trace)
                                chrome_events.append(chrome_event)
    
    # Write Chrome Tracing format
    with open(output_file, "w") as f:
        json.dump(chrome_events, f, indent=2)
    
    print(f"Converted {len(chrome_events)} trace events")
    print(f"Output written to: {output_file}")
    print(f"\nTo view:")
    print(f"  1. Open chrome://tracing in Chrome/Chromium")
    print(f"  2. Click 'Load' and select {output_file}")


def main():
    if len(sys.argv) < 2:
        print("Usage: python viz_traces.py <input_jsonl_file> [output_json_file]")
        print("\nExample:")
        print("  python viz_traces.py pdsh_results.jsonl")
        print("  python viz_traces.py pdsh_results.jsonl traces.json")
        sys.exit(1)
    
    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else None
    
    convert_jsonl_to_chrome_trace(input_file, output_file)


if __name__ == "__main__":
    main()