Add summary printing for parsed files (#20)

Summary:
- Introduced a new function `print_parsed_files_summary` to display a formatted summary of all parsed files, including their paths and sizes.
- Updated `oss_parse` to call this new function after parsing logs, enhancing user feedback on parsing results.

Pull Request resolved: #20

Reviewed By: davidberard98

Differential Revision: D77236623

Pulled By: FindHao

fbshipit-source-id: 2fbc6485040100855de8a8eac436fb254421287f

Author: FindHao

tests/test_add.py

@@ -71,4 +71,6 @@ def test_tensor_add():
 if __name__ == "__main__":
     test_tensor_add()
     # Use improved unified_parse with explicit output directory
-    tritonparse.utils.unified_parse(source=log_path, out="./parsed_output")
+    tritonparse.utils.unified_parse(
+        source=log_path, out="./parsed_output", overwrite=True
+    )

tritonparse/common.py

@@ -135,6 +135,57 @@ def to_rank(self) -> Rank:
         return Rank()
 
 
+def print_parsed_files_summary(parsed_log_dir: str) -> None:
+    """
+    Print a beautiful summary of all parsed files.
+
+    Args:
+        parsed_log_dir: Directory containing parsed files
+    """
+    # Collect all parsed files
+    all_parsed_files = []
+    for root, _, files in os.walk(parsed_log_dir):
+        for file in files:
+            file_path = os.path.join(root, file)
+            all_parsed_files.append(file_path)
+
+    # Sort files for consistent output
+    all_parsed_files.sort()
+
+    # Print beautiful summary
+    print("\n" + "=" * 80)
+    print("📁 TRITONPARSE PARSING RESULTS")
+    print("=" * 80)
+
+    # Print log file list (required for integration)
+    print(f"📂 Parsed files directory: {parsed_log_dir}")
+    print(f"📊 Total files generated: {len(all_parsed_files)}")
+
+    if all_parsed_files:
+        print("\n📄 Generated files:")
+        print("-" * 50)
+        for i, file_path in enumerate(all_parsed_files, 1):
+            # Get relative path for cleaner display
+            rel_path = os.path.relpath(file_path, parsed_log_dir)
+            file_size = "N/A"
+            try:
+                size_bytes = os.path.getsize(file_path)
+                if size_bytes < 1024:
+                    file_size = f"{size_bytes}B"
+                elif size_bytes < 1024 * 1024:
+                    file_size = f"{size_bytes/1024:.1f}KB"
+                else:
+                    file_size = f"{size_bytes/(1024*1024):.1f}MB"
+            except OSError:
+                pass
+
+            print(f"  {i:2d}. 📝 {rel_path} ({file_size})")
+
+    print("=" * 80)
+    print("✅ Parsing completed successfully!")
+    print("=" * 80 + "\n")
+
+
 def gzip_single_file(file_path: str, verbose: bool = False) -> str:
     """
     Gzip a single file and delete the original file.
@@ -318,7 +369,9 @@ def parse_logs(
     log_file_list_path = os.path.join(parsed_log_dir, "log_file_list.json")
     with open(log_file_list_path, "w") as f:
         json.dump(file_mapping, f, indent=2)
+
     # NOTICE: this print is required for tlparser-tritonparse integration
+    # DON'T REMOVE THIS PRINT
     print(f"tritonparse log file list: {log_file_list_path}")
     return parsed_log_dir, file_mapping
 

tritonparse/structured_logging.py

@@ -851,10 +851,18 @@ def __call__(
         Returns:
             True to continue with compilation, None/False to skip
         """
-        launch_metadata_fn = fn.jit_function.launch_metadata
-        if launch_metadata_fn is not None:
+        # Check kernel allowlist early to avoid unnecessary work
+        if _KERNEL_ALLOWLIST_PATTERNS is not None:
+            kernel_name = fn.name
+            if not should_trace_kernel(kernel_name, _KERNEL_ALLOWLIST_PATTERNS):
+                # Skip overriding launch_metadata if kernel is not in allowlist
+                return True
+
+        # Get the current launch_metadata function if it exists
+        current_launch_metadata = getattr(fn.jit_function, "launch_metadata", None)
+        if current_launch_metadata is not None:
             log.warning(
-                f"fn {fn} launch_metadata_fn is not None: {launch_metadata_fn}. It will be overridden by tritonparse."
+                f"fn {fn} launch_metadata is not None: {current_launch_metadata}. It will be overridden by tritonparse."
             )
         fn.jit_function.launch_metadata = add_launch_metadata
         return True
@@ -898,8 +906,16 @@ def __call__(self, metadata):
                  https://github.com/triton-lang/triton/blob/7ce287dc24b43476cdeb30529089ac361564505d/
                  python/triton/compiler/compiler.py#L512.
         """
-        trace_data = defaultdict(dict)
         metadata_dict = metadata.get()
+        # Check kernel allowlist early to avoid unnecessary work
+        if _KERNEL_ALLOWLIST_PATTERNS is not None:
+            kernel_name = metadata_dict.get("name")
+
+            if not should_trace_kernel(kernel_name, _KERNEL_ALLOWLIST_PATTERNS):
+                # Skip tracing if kernel is not in allowlist
+                return
+
+        trace_data = defaultdict(dict)
         trace_data["name"] = metadata_dict["name"]
         trace_data["function"] = metadata_dict["function"]
         trace_data["stream"] = metadata_dict["stream"]

tritonparse/utils.py

@@ -9,7 +9,14 @@
 # argument parser for OSS
 parser = None
 
-from .common import copy_local_to_tmpdir, is_fbcode, parse_logs, RankConfig, save_logs
+from .common import (
+    copy_local_to_tmpdir,
+    is_fbcode,
+    parse_logs,
+    print_parsed_files_summary,
+    RankConfig,
+    save_logs,
+)
 from .source_type import Source, SourceType
 
 
@@ -53,7 +60,7 @@ def init_parser():
 def oss_run(
     source: str,
     out: Optional[str] = None,
-    overwrite: bool = True,
+    overwrite: Optional[bool] = False,
     rank: Optional[int] = None,
     all_ranks: bool = False,
     verbose: bool = False,
@@ -98,6 +105,12 @@ def oss_run(
     parsed_log_dir, _ = parse_logs(logs, rank_config, verbose)
     if out is not None:
         save_logs(Path(out), parsed_log_dir, overwrite, verbose)
+    # Print beautiful summary of all parsed files
+    if out is not None:
+        out_dir = str(Path(out).absolute())
+    else:
+        out_dir = str(Path(parsed_log_dir).absolute())
+    print_parsed_files_summary(out_dir)
 
 
 def unified_parse_from_cli():
@@ -109,7 +122,7 @@ def unified_parse_from_cli():
 def unified_parse(
     source: str,
     out: Optional[str] = None,
-    overwrite: bool = True,
+    overwrite: Optional[bool] = False,
     rank: Optional[int] = None,
     all_ranks: bool = False,
     verbose: bool = False,