feat: extend usage_report_logger to track all subcommands and API calls

Summary:
This change extends the usage logging functionality to track all tritonparse usage, including:

1. **CLI layer logging**: Added unified usage_report_logger() call in cli.py main() function to record all subcommand invocations (parse, reproduce, info)

2. **API layer logging**: Added usage_report_logger() calls to reproduce() and info_command() functions to capture direct API usage (non-CLI calls)

3. **Entry function detection**: Modified usage_report_logger() to use sys._getframe() to traverse the call stack and identify the outermost entry function (unified_parse, reproduce, info_command, fb_run, or main), storing it in extra_data["entry_function"]

4. **skip_logger parameter**: Added skip_logger parameter to unified_parse(), reproduce(), info_command(), and oss_run() functions to prevent duplicate logging when called from CLI (which already logs at the top level)

This enables better tracking of how tritonparse is being used - both through CLI commands and through direct Python API calls.

Reviewed By: wychi

Differential Revision: D88983900

fbshipit-source-id: fd193a0ff64ba01ce425899220a22deb710ebdaf

Author: FindHao

tritonparse/cli.py

@@ -73,11 +73,18 @@ def main():
 
     args = parser.parse_args()
 
+    # Log usage for CLI invocations
+    if is_fbcode():
+        from tritonparse.fb.utils import usage_report_logger
+
+        usage_report_logger()
+
     if args.func == "parse":
         parse_args = {
             k: v for k, v in vars(args).items() if k not in ["command", "func"]
         }
-        unified_parse(**parse_args)
+        # skip_logger=True because we already logged above
+        unified_parse(**parse_args, skip_logger=True)
     elif args.func == "reproduce":
         # Check mutual exclusivity between --line and --kernel/--launch-id
         if args.kernel and args.line != 0:
@@ -99,9 +106,12 @@ def main():
             launch_id=args.launch_id if args.kernel else 0,
             kernel_import=args.kernel_import,
             replacer=replacer,
+            skip_logger=True,  # Already logged above
         )
     elif args.func == "info":
-        info_command(input_path=args.input, kernel_name=args.kernel)
+        info_command(
+            input_path=args.input, kernel_name=args.kernel, skip_logger=True
+        )  # Already logged above
     else:
         raise RuntimeError(f"Unknown command: {args.func}")
 

tritonparse/info/cli.py

@@ -11,6 +11,7 @@
 import tempfile
 from typing import Optional
 
+from tritonparse.common import is_fbcode
 from tritonparse.info.kernel_query import (
     find_similar_kernels,
     list_kernels_fast,
@@ -34,14 +35,22 @@ def _add_info_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
-def info_command(input_path: str, kernel_name: Optional[str] = None) -> None:
+def info_command(
+    input_path: str, kernel_name: Optional[str] = None, skip_logger: bool = False
+) -> None:
     """
     Main function for the info command.
 
     Args:
         input_path: Path to ndjson file
         kernel_name: Optional kernel name to list launches for
+        skip_logger: Whether to skip usage logging (default: False).
     """
+    if not skip_logger and is_fbcode():
+        from tritonparse.fb.utils import usage_report_logger
+
+        usage_report_logger()
+
     # 1. Load and detect type
     events = load_ndjson(input_path)
     has_launch_diff = any(e.get("event_type") == "launch_diff" for e in events)

tritonparse/reproducer/orchestrator.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 from typing import Optional
 
+from tritonparse.common import is_fbcode
 from tritonparse.info.kernel_query import find_launch_index_by_kernel
 from tritonparse.reproducer.ingestion.ndjson import build_context_bundle
 from tritonparse.reproducer.placeholder_replacer import (
@@ -25,6 +26,7 @@ def reproduce(
     launch_id: int = 0,
     replacer: Optional[PlaceholderReplacer] = None,
     kernel_import: KernelImportMode = KernelImportMode.DEFAULT,
+    skip_logger: bool = False,
 ) -> dict[str, str]:
     """
     Generate a reproducer script from NDJSON trace file.
@@ -43,7 +45,13 @@ def reproduce(
         launch_id: 0-based launch index for the kernel (default: 0, first launch).
         replacer: Optional custom PlaceholderReplacer instance. If None, uses DefaultPlaceholderReplacer.
         kernel_import: Kernel import mode (DEFAULT or COPY).
+        skip_logger: Whether to skip usage logging (default: False).
     """
+    if not skip_logger and is_fbcode():
+        from tritonparse.fb.utils import usage_report_logger
+
+        usage_report_logger()
+
     events = load_ndjson(Path(input_path))
     logger.debug(f"Loaded {len(events)} events")
 

tritonparse/utils.py

@@ -60,6 +60,7 @@ def oss_run(
     all_ranks: bool = False,
     verbose: bool = False,
     split_inductor_compilations: bool = True,
+    skip_logger: bool = True,
 ):
     """
     Main function for tritonparse. It is for OSS only.
@@ -71,6 +72,7 @@ def oss_run(
         rank: Rank of logs to be analyzed
         all_ranks: Analyze all ranks
         verbose: Verbose logging
+        skip_logger: Unused in OSS, kept for API compatibility.
     """
     source = Source(source, verbose)
     rank_config = RankConfig.from_cli_args(rank, all_ranks, source.type)
@@ -123,6 +125,7 @@ def unified_parse(
     all_ranks: bool = False,
     verbose: bool = False,
     split_inductor_compilations: bool = True,
+    skip_logger: bool = False,
     **kwargs,
 ):
     """
@@ -135,7 +138,14 @@ def unified_parse(
         rank: Specific rank to analyze
         all_ranks: Whether to analyze all ranks
         verbose: Whether to enable verbose logging
+        skip_logger: Whether to skip usage logging (default: False).
     """
+    # Log usage for API invocations
+    if not skip_logger and is_fbcode():
+        from tritonparse.fb.utils import usage_report_logger
+
+        usage_report_logger()
+
     # Choose the appropriate parse function
     if is_fbcode():
         from tritonparse.fb.utils import fb_run as parse
@@ -150,6 +160,7 @@ def unified_parse(
         all_ranks=all_ranks,
         verbose=verbose,
         split_inductor_compilations=split_inductor_compilations,
+        skip_logger=skip_logger,
         **kwargs,
     )
     return output