Website can't handl large files

[FindHao]

When a parsed trace is too large, e.g., a 700MB ndjson.gz, the website crashes. We should add an automatic file split function to divide a large trace into multiple files.

[Aryan-Grg]

Hi I am new to open source work but I am good with file handling stuff if you want i would love to help fix this

[FindHao]

Hi I am new to open source work but I am good with file handling stuff if you want i would love to help fix this

Hi @Aryan-Grg , you are welcome to contribute! I'll share a short design doc later.

[FindHao]

Hi @Aryan-Grg , here is the design doc. You don't need to exactly follow the example code, but the high level design is supposed to be followed. Let me know if you have any questions. Thanks!

Large Trace File Splitting Design

1. Background

1.1 Problem Description

When parsed trace files are too large (e.g., 700MB ndjson.gz, which decompresses to 18GB), the website crashes. We need to add an automatic file splitting feature to divide large traces into multiple files.

1.2 Browser Memory Limits

Browser/Platform	64-bit Process Limit
Chrome/Edge (V8)	~4 GB (single ArrayBuffer limit ~2 GB)
Firefox	~4-8 GB
Safari	~4 GB

1.3 Known Test Results

• 1GB raw log: Loads successfully (both raw and parsed)
• 700MB ndjson.gz: Decompresses to 18GB, causes crash

Further testing is needed to determine appropriate thresholds.

---

2. Current Code Analysis

2.1 Write Flow

trace_processor.py                         common.py
      │                                         │
      │  parse_single_file()                    │
      │       │                                 │
      │       ▼                                 │
      │  all_output_lines = defaultdict(list)   │
      │       │                                 │
      │       ▼                                 │
      │  with open(output_file, "w") as out:    │  <-- Writes uncompressed .ndjson
      │      out.writelines(final_lines)        │
      │                                         │
      └─────────────────────────────────────────┼──► After parse_logs() completes
                                                │
                                                ▼
                                          gzip_single_file()  <-- Compresses afterwards
                                                │
                                                ▼
                                          Generates .ndjson.gz

Key Points:

• trace_processor.py writes uncompressed .ndjson files
• common.py's gzip_single_file() compresses after writing
• Splitting should track uncompressed size (i.e., the size of strings accumulated in all_output_lines)

2.2 Current Inductor Splitting Logic

# trace_processor.py lines 466-477
if split_inductor_compilations:
    pt_info = payload.get("pt_info", {})
    frame_id = pt_info.get("frame_id")
    frame_compile_id = pt_info.get("frame_compile_id")
    attempt_id = pt_info.get("attempt_id", 0)
    cai = pt_info.get("compiled_autograd_id", "-")
    if frame_id is not None or frame_compile_id is not None:
        fname = f"f{frame_id}_fc{frame_compile_id}_a{attempt_id}_cai{cai}.ndjson"
    else:
        fname = f"{file_name_without_extension}_mapped.ndjson"

Current Limitation: Only splits by Inductor compilation unit. If a compilation unit has many kernels, the output file can still be too large.

2.3 Data Relationships

Each kernel contains a group of tightly related events:

kernel (by hash)
├── compilation event (1)
├── launch events (multiple)
├── ir_analysis event (1, optional)
├── launch_diff event (1, optional)
└── autotune_analysis events (multiple, optional)

Key Constraint: All events for the same kernel must be in the same partition; they cannot be split.

---

3. Splitting Strategy Design

3.1 Core Approach

1. Split after Inductor splitting: First split by Inductor compilation unit (if enabled), then check each output file's size and split further if it exceeds the threshold
2. Split by file size: Track the actual accumulated size in memory (uncompressed)
3. Group by kernel name: Not just by hash, because autotune kernels share the same name but have different hashes
4. No new data structures: Reuse existing all_output_lines, handle splitting in the write phase
5. Determine threshold through testing: Rather than theoretical estimates

3.2 Splitting Hierarchy

Layer 1: Split by Inductor compilation unit (existing, optional)
    │
    ▼
Layer 2: Split by file size (new)
    │
    ▼
Ensure all hashes for the same kernel name are in the same partition

3.3 Why Group by Kernel Name?

Autotune Scenario:

kernel_name: "triton_add_kernel"
│
├── hash: "abc123" (num_warps=4, num_stages=2)
│    ├── compilation event
│    ├── launch event (benchmark)
│    └── ir_analysis event
│
├── hash: "def456" (num_warps=8, num_stages=3)
│    ├── compilation event
│    ├── launch event (benchmark)
│    └── ir_analysis event
│
└── autotune_analysis event (references abc123 and def456)

If split arbitrarily by hash:

• part_001.ndjson: contains hash abc123
• part_002.ndjson: contains hash def456
• Problem: The autotune_analysis event references both hashes — which partition should it go in?

Benefits of grouping by kernel name:

1. All autotune variants of the same-named kernel stay together
2. autotune_analysis events naturally belong to the same partition
3. When users load a partition on the website, they see complete kernel data

3.4 About autotune_summary

autotune_summary's winner_run_counts is keyed by hash, but different hashes of the same kernel name should be in the same Inductor compilation unit (same output_file).

Therefore, autotune_summary is cross-hash but not cross-kernel-name, and can naturally be placed in the partition containing kernels of that name.

---

4. Implementation

4.1 No New Data Structures

Existing data structures:

# trace_processor.py lines 377-391
kernels_by_hash = defaultdict(
    lambda: {"compilation": None, "launches": [], "output_file": None}
)
autotune_sessions = defaultdict(...)
autotune_winners = {}
session_stacks = {}
launch_by_group_hash = {}

# lines 552-553
all_output_lines = defaultdict(list)  # {output_file: [lines...]}

No new data structures needed, just:

1. Sort iteration of kernels_by_hash by kernel name
2. Check size and split in the write phase

4.2 Modify Iteration Order

# Original (line 554):
for _kernel_hash, data in kernels_by_hash.items():

# Change to:
sorted_kernels = sorted(
    kernels_by_hash.items(),
    key=lambda x: (
        x[1]["output_file"] or "",
        x[1]["compilation"].get("payload", {}).get("metadata", {}).get("name", "") if x[1]["compilation"] else "",
    )
)
for _kernel_hash, data in sorted_kernels:

This ensures that in final_lines:

1. Kernels for the same output_file are together
2. Different hashes of the same kernel name are adjacent

4.3 Write Phase Splitting Logic

# Original (lines 640-642):
for output_file, final_lines in all_output_lines.items():
    with open(output_file, "w") as out:
        out.writelines(final_lines)

# Change to:
for output_file, final_lines in all_output_lines.items():
    total_size = sum(len(line.encode('utf-8')) for line in final_lines)
    
    if total_size <= max_size_bytes:
        # No split needed, write directly
        with open(output_file, "w") as out:
            out.writelines(final_lines)
        continue
    
    # Split needed
    _write_with_split(output_file, final_lines, max_size_bytes)

4.4 Split Write Function

def _write_with_split(output_file: str, final_lines: List[str], max_size_bytes: int):
    """Split by size, using compilation events as split boundaries"""
    
    base_name = output_file.rsplit('.ndjson', 1)[0]
    part_index = 1
    current_lines = []
    current_size = 0
    
    for line in final_lines:
        line_size = len(line.encode('utf-8'))
        event = json.loads(line)
        event_type = event.get("event_type")
        
        # Check for split at compilation events
        # This ensures all events for a kernel stay together
        if event_type == "compilation" and current_lines:
            if current_size + line_size > max_size_bytes:
                # Write current partition
                part_file = f"{base_name}_part_{part_index:03d}.ndjson"
                with open(part_file, "w") as out:
                    out.writelines(current_lines)
                part_index += 1
                current_lines = []
                current_size = 0
        
        current_lines.append(line)
        current_size += line_size
    
    # Write final partition
    if current_lines:
        if part_index == 1:
            # No split actually needed (all content in first partition)
            with open(output_file, "w") as out:
                out.writelines(current_lines)
        else:
            part_file = f"{base_name}_part_{part_index:03d}.ndjson"
            with open(part_file, "w") as out:
                out.writelines(current_lines)

Note: Since we've already sorted by kernel name, splitting at compilation boundaries keeps same-named kernels together. A split only occurs at a different kernel name when the same-named kernel group's total size exceeds the threshold.

---

5. Output Structure

5.1 File Naming

No new directories created. Partition files go in the original directory with this naming format:

output_dir/
├── f1_fc0_a0_cai-.ndjson.gz              # Small file, no split
├── f2_fc0_a1_cai-_part_001.ndjson.gz     # Large file, split into multiple
├── f2_fc0_a1_cai-_part_002.ndjson.gz
├── f2_fc0_a1_cai-_part_003.ndjson.gz
└── log_file_list.json

5.2 Manifest File (Optional)

If a manifest is needed, generate f2_fc0_a1_cai-_manifest.json:

{
    "version": 1,
    "original_name": "f2_fc0_a1_cai-.ndjson",
    "total_parts": 3,
    "parts": [
        {"filename": "f2_fc0_a1_cai-_part_001.ndjson.gz", "kernel_count": 50},
        {"filename": "f2_fc0_a1_cai-_part_002.ndjson.gz", "kernel_count": 50},
        {"filename": "f2_fc0_a1_cai-_part_003.ndjson.gz", "kernel_count": 45}
    ]
}

---

6. Website

No website updates required.

Current website behavior: Users select which file to load. Each partition after splitting is an independent, self-contained ndjson file that users can directly select and load.

---

7. Threshold Determination

7.1 Testing Required

Known:

• 1GB raw log loads successfully
• 700MB gzip (decompresses to 18GB) crashes

Testing needed:

• Different sizes of parsed files (100MB, 200MB, 500MB, 1GB uncompressed)
• Record load time and memory usage

7.2 CLI Parameter

tritonparse parse --max-part-size-mb 500 <input_file>

Default value can be set to a larger value (e.g., 500MB or 1GB), so only truly large files get split.

---

8. Implementation Steps

Step 1: Test to determine appropriate threshold
        │
        ▼
Step 2: Modify iteration order to sort by (output_file, kernel_name)
        │
        ▼
Step 3: Add size checking and splitting logic in write phase
        │
        ▼
Step 4: Update common.py to handle partition files (collect into log_file_list.json)
        │
        ▼
Step 5: Add CLI parameter (--max-part-size-mb)
        │
        ▼
Step 6: Add test cases

---

9. Files to Modify

File	Changes
tritonparse/parse/trace_processor.py	Add sorting and split write logic
tritonparse/parse/common.py	Handle partition files, update log_file_list.json
tritonparse/cli.py	Add --max-part-size-mb parameter

---

10. Backward Compatibility

• Don't split small files by default: Only split when threshold is exceeded
• Maintain single-file format compatibility: Each partition after splitting is still a valid ndjson file
• Can disable via parameter: --max-part-size-mb=0 disables splitting
• No website changes needed: Users directly select and load partitions

---

11. Edge Cases

11.1 Single Kernel Group Too Large

If all events for a kernel name exceed the threshold:

• Handling: Log a warning, but still place in a separate partition
• Reason: Cannot break autotune relationships for same-named kernels

if group_size > max_size_bytes:
    logger.warning(f"Kernel group '{kernel_name}' ({group_size} bytes) exceeds "
                   f"max_size ({max_size_bytes} bytes), placing in separate part")

11.2 Only One Partition

If no split is actually needed after checking:

• Write to original filename, don't generate _part_001 suffix

[Aryan-Grg]

Hi I will try to get it done with as soon as possible but just to know is there a particular day before which you would like it to be completed in?

[FindHao]

Hi I will try to get it done with as soon as possible but just to know is there a particular day before which you would like it to be completed in?

It's not a high-priority task, so no need to worry about deadlines.