Wire torch trace dir through parse pipeline and CLI (#355)

Summary:
Pull Request resolved: #355

Add --torch-trace-dir CLI parameter and wire it through the full pipeline:
unified_parse() -> oss_run() -> parse_logs() -> parse_single_file().

Includes auto-discovery logic: when --torch-trace-dir is not specified,
torch trace log files are automatically searched in the same directory as
tritonparse logs. This enables kernel compile attribution in multi-process
scenarios without requiring explicit user configuration.

Also exports CompileInfo, discover_torch_trace_files, and
parse_torch_trace_logs from the parse module's public API.

Reviewed By: wychi

Differential Revision: D95080073

fbshipit-source-id: 09e37ac2628f00dc6cbf596f4f72b0bc718b695b

Author: FindHao

tests/cpu/test_pipeline_integration.py

@@ -0,0 +1,167 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""Integration tests for the full parse pipeline with torch trace log support."""
+
+import json
+import os
+import tempfile
+import unittest
+
+from tritonparse.parse.common import (
+    _build_kernel_compile_mapping,
+    parse_logs,
+    RankConfig,
+)
+
+
+def _make_glog_line(metadata_dict: dict) -> str:
+    """Helper to create a glog-formatted line with JSON metadata."""
+    return f"V0302 14:30:00.123456 12345 torch/_logging/_internal.py:1489] {json.dumps(metadata_dict)}"
+
+
+def _make_torch_trace_log(frame_id, frame_compile_id, kernel_paths):
+    """Create content for a torch trace log file with inductor_output_code event."""
+    metadata = {
+        "inductor_output_code": {
+            "filename": "output.py",
+            "file_path": "/tmp/output.py",
+        },
+        "frame_id": frame_id,
+        "frame_compile_id": frame_compile_id,
+        "attempt": 0,
+        "has_payload": "abc123",
+    }
+    lines = [_make_glog_line(metadata)]
+    for kp in kernel_paths:
+        lines.append(f"\t# kernel path: {kp}")
+        lines.append("\ttriton_kernel = async_compile.triton('kernel', '''...''')")
+    return "\n".join(lines) + "\n"
+
+
+def _make_tritonparse_trace(events):
+    """Create content for a tritonparse trace NDJSON file."""
+    return "\n".join(json.dumps(e) for e in events) + "\n"
+
+
+class TestBuildKernelCompileMapping(unittest.TestCase):
+    """Tests for _build_kernel_compile_mapping."""
+
+    def test_auto_discover_in_raw_log_dir(self):
+        """Test that torch trace files are auto-discovered in the raw log directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create a torch trace log file in the same directory
+            torch_log = os.path.join(tmpdir, "dedicated_log_torch_trace_rank_0_abc.log")
+            content = _make_torch_trace_log(
+                frame_id=0,
+                frame_compile_id=0,
+                kernel_paths=["/tmp/torchinductor_user/ab/kernel.py"],
+            )
+            with open(torch_log, "w") as f:
+                f.write(content)
+
+            mapping = _build_kernel_compile_mapping(tmpdir)
+            self.assertIsNotNone(mapping)
+            self.assertIn("/tmp/torchinductor_user/ab/kernel.py", mapping)
+            self.assertEqual(
+                mapping["/tmp/torchinductor_user/ab/kernel.py"].frame_id, 0
+            )
+
+    def test_explicit_torch_trace_dir(self):
+        """Test using an explicit torch_trace_dir."""
+        with tempfile.TemporaryDirectory() as log_dir:
+            with tempfile.TemporaryDirectory() as torch_dir:
+                torch_log = os.path.join(
+                    torch_dir, "dedicated_log_torch_trace_rank_0_abc.log"
+                )
+                content = _make_torch_trace_log(
+                    frame_id=1,
+                    frame_compile_id=0,
+                    kernel_paths=["/tmp/torchinductor_user/cd/kernel.py"],
+                )
+                with open(torch_log, "w") as f:
+                    f.write(content)
+
+                mapping = _build_kernel_compile_mapping(log_dir, torch_dir)
+                self.assertIsNotNone(mapping)
+                self.assertEqual(len(mapping), 1)
+
+    def test_no_torch_trace_files(self):
+        """Test that None is returned when no torch trace files exist."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mapping = _build_kernel_compile_mapping(tmpdir)
+            self.assertIsNone(mapping)
+
+
+class TestParseLogsWithTorchTrace(unittest.TestCase):
+    """End-to-end test for parse_logs with torch trace integration."""
+
+    def test_end_to_end_mapping(self):
+        """Test that kernels without pt_info are correctly attributed via torch trace logs."""
+        kernel_path = "/tmp/torchinductor_user/ab/cabcdef.py"
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create torch trace log
+            torch_log_path = os.path.join(
+                tmpdir, "dedicated_log_torch_trace_rank_0_test.log"
+            )
+            torch_content = _make_torch_trace_log(
+                frame_id=3,
+                frame_compile_id=1,
+                kernel_paths=[kernel_path],
+            )
+            with open(torch_log_path, "w") as f:
+                f.write(torch_content)
+
+            # Create tritonparse trace log (compilation without pt_info)
+            triton_events = [
+                {
+                    "event_type": "compilation",
+                    "pid": 1000,
+                    "stack": [],
+                    "payload": {
+                        "metadata": {"hash": "test_hash", "name": "test_kernel"},
+                        "file_content": {},
+                        "file_path": {},
+                        "python_source": {"file_path": kernel_path},
+                        # No pt_info — multi-process scenario
+                    },
+                },
+                {
+                    "event_type": "launch",
+                    "name": "test_kernel",
+                    "pid": 1000,
+                    "stack": [],
+                    "compilation_metadata": {"hash": "test_hash"},
+                },
+            ]
+            triton_log_path = os.path.join(
+                tmpdir, "dedicated_log_triton_trace_user_.ndjson"
+            )
+            with open(triton_log_path, "w") as f:
+                f.write(_make_tritonparse_trace(triton_events))
+
+            # Run parse_logs (use all_ranks=True to pick up no-rank files)
+            rank_config = RankConfig(all_ranks=True)
+            parsed_dir, file_mapping = parse_logs(
+                tmpdir,
+                rank_config,
+                verbose=False,
+                split_inductor_compilations=True,
+            )
+
+            # Check that the output was split into a frame-specific file
+            # Walk the output directory to find all generated files
+            all_files = []
+            for root, dirs, files in os.walk(parsed_dir):
+                for f in files:
+                    all_files.append(f)
+
+            frame_files = [f for f in all_files if f.startswith("f")]
+            # Should have f3_fc1_a0_cai-.ndjson.gz (attributed via mapping)
+            self.assertTrue(
+                any("f3_fc1" in f for f in frame_files),
+                f"Expected frame file with f3_fc1 but got: {all_files}",
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

tritonparse/parse/__init__.py

@@ -28,6 +28,11 @@
     create_python_mapping,
 )
 from .source_type import Source, SourceType
+from .torch_trace_parser import (
+    CompileInfo,
+    discover_torch_trace_files,
+    parse_torch_trace_logs,
+)
 from .trace_processor import (
     generate_source_mappings,
     parse_single_file,
@@ -55,6 +60,10 @@
     "_add_parse_args",
     "oss_run",
     "unified_parse",
+    # Torch trace parser
+    "CompileInfo",
+    "discover_torch_trace_files",
+    "parse_torch_trace_logs",
     # IR parsing
     "extract_code_locations",
     "extract_loc_definitions",

tritonparse/parse/common.py

@@ -8,7 +8,7 @@
 import tempfile
 from collections import defaultdict
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from tritonparse.shared_vars import (
     DEFAULT_TRACE_FILE_PREFIX_WITHOUT_USER as LOG_PREFIX,
@@ -312,12 +312,66 @@ def copy_local_to_tmpdir(local_path: str, verbose: bool = False) -> str:
     return temp_dir
 
 
+def _build_kernel_compile_mapping(
+    raw_log_dir: str,
+    torch_trace_dir: Optional[str] = None,
+) -> Optional[Dict[str, Any]]:
+    """
+    Build kernel compile mapping from inductor's torch trace logs.
+
+    Searches for torch trace log files and parses them to extract
+    kernel_source_path -> CompileInfo mappings. These mappings allow
+    attribution of Triton kernels to their originating compilation frame
+    when pt_info is missing (multi-process Triton JIT scenarios).
+
+    Args:
+        raw_log_dir: Directory containing tritonparse logs (used for auto-discovery).
+        torch_trace_dir: Explicit directory containing torch trace logs.
+            If None, auto-discovers in raw_log_dir.
+
+    Returns:
+        Dict mapping kernel source paths to CompileInfo, or None if no logs found.
+    """
+    from .torch_trace_parser import discover_torch_trace_files, parse_torch_trace_logs
+
+    # Determine where to look for torch trace logs
+    search_dirs = []
+    if torch_trace_dir:
+        search_dirs.append(torch_trace_dir)
+    # Also check the raw log directory (torch trace logs may coexist)
+    search_dirs.append(raw_log_dir)
+
+    all_log_paths: List[str] = []
+    seen_paths: set = set()
+    for search_dir in search_dirs:
+        if not os.path.isdir(search_dir):
+            continue
+        torch_files = discover_torch_trace_files(search_dir)
+        for rank_files in torch_files.values():
+            for path in rank_files:
+                if path not in seen_paths:
+                    all_log_paths.append(path)
+                    seen_paths.add(path)
+
+    if not all_log_paths:
+        return None
+
+    mapping = parse_torch_trace_logs(all_log_paths)
+    if mapping:
+        logger.info(
+            f"Built kernel compile mapping with {len(mapping)} entries "
+            f"from {len(all_log_paths)} torch trace log(s)"
+        )
+    return mapping if mapping else None
+
+
 def parse_logs(
     logs_to_parse: str,
     rank_config: RankConfig,
     verbose: bool = False,
     tritonparse_url_prefix: str = "",
     split_inductor_compilations: bool = True,
+    torch_trace_dir: Optional[str] = None,
 ) -> Tuple[str, dict]:
     """
     Parse logs.
@@ -330,6 +384,10 @@ def parse_logs(
         split_inductor_compilations: Whether to split
             output files by frame_id, compile_id, attempt_id, and compiled_autograd_id.
             Defaults to True. This rule follows tlparse's behavior.
+        torch_trace_dir: Optional path to directory containing inductor torch trace
+            logs. When provided, kernel compilation attribution will use these logs to
+            recover frame_id/compile_id for kernels compiled in multi-process scenarios.
+            If None, auto-discovers torch trace files in the same directory as tritonparse logs.
     Returns:
         Tuple of (parsed log directory, file mapping)
     """
@@ -372,6 +430,10 @@ def parse_logs(
                 ranks[Rank(Rank.NO_RANK)].append(path)
     if not ranks:
         raise RuntimeError(f"No eligible structured trace logs found in {raw_log_dir}")
+
+    # Build kernel compile mapping from torch trace logs (if available)
+    kernel_compile_mapping = _build_kernel_compile_mapping(raw_log_dir, torch_trace_dir)
+
     file_mapping = {"tritonparse_url_prefix": tritonparse_url_prefix}
     # Parse each eligible log
     for rank, files in ranks.items():
@@ -406,7 +468,12 @@ def parse_logs(
                 relative_path = "" if rank.is_no_rank else rank.to_string("")
             output_dir = os.path.join(parsed_log_dir, relative_path)
             # Parse the file
-            parse_single_file(input_file, output_dir, split_inductor_compilations)
+            parse_single_file(
+                input_file,
+                output_dir,
+                split_inductor_compilations,
+                kernel_compile_mapping=kernel_compile_mapping,
+            )
             # Collect generated files after parsing and gzip them immediately
             if os.path.exists(output_dir):
                 generated_files = []

tritonparse/parse/utils.py

@@ -76,6 +76,16 @@ def _add_parse_args(parser: argparse.ArgumentParser) -> None:
         action="store_true",
     )
     parser.add_argument("-v", "--verbose", help="Verbose logging", action="store_true")
+    parser.add_argument(
+        "--torch-trace-dir",
+        type=str,
+        default=None,
+        help=(
+            "Path to directory containing inductor torch trace logs. "
+            "Used to recover kernel compilation attribution in multi-process scenarios. "
+            "If not specified, auto-discovers torch trace files alongside tritonparse logs."
+        ),
+    )
     if is_fbcode():
         from tritonparse.fb.utils import append_parser
 
@@ -91,6 +101,7 @@ def oss_run(
     verbose: bool = False,
     split_inductor_compilations: bool = True,
     skip_logger: bool = True,
+    torch_trace_dir: Optional[str] = None,
 ):
     """
     Main function for tritonparse. It is for OSS only.
@@ -103,6 +114,7 @@ def oss_run(
         all_ranks: Analyze all ranks
         verbose: Verbose logging
         skip_logger: Unused in OSS, kept for API compatibility.
+        torch_trace_dir: Path to directory containing inductor torch trace logs.
     """
     source = Source(source, verbose)
     rank_config = RankConfig.from_cli_args(rank, all_ranks, source.type)
@@ -137,6 +149,7 @@ def oss_run(
             rank_config,
             verbose,
             split_inductor_compilations=split_inductor_compilations,
+            torch_trace_dir=torch_trace_dir,
         )
     else:
         parsed_log_dir = source.value
@@ -161,6 +174,7 @@ def unified_parse(
     verbose: bool = False,
     split_inductor_compilations: bool = True,
     skip_logger: bool = False,
+    torch_trace_dir: Optional[str] = None,
     **kwargs,
 ):
     """
@@ -174,6 +188,7 @@ def unified_parse(
         all_ranks: Whether to analyze all ranks
         verbose: Whether to enable verbose logging
         skip_logger: Whether to skip usage logging (default: False).
+        torch_trace_dir: Path to directory containing inductor torch trace logs.
     """
     # Log usage for API invocations
     if not skip_logger and is_fbcode():
@@ -196,6 +211,7 @@ def unified_parse(
         verbose=verbose,
         split_inductor_compilations=split_inductor_compilations,
         skip_logger=skip_logger,
+        torch_trace_dir=torch_trace_dir,
         **kwargs,
     )
     return output