tritonparse 0.3.1.dev20251030071508__py3-none-any.whl → 0.3.1.dev20251101071506__py3-none-any.whl

tritonparse/cli.py

@@ -88,5 +88,4 @@ def main():
 
 
 if __name__ == "__main__":
-    # Do not add code here, it won't be run. Add them to the function called below.
     main()  # pragma: no cover

tritonparse/structured_logging.py

@@ -48,6 +48,14 @@ TRITON_TRACE_LAUNCH = os.getenv("TRITON_TRACE_LAUNCH", None) in ["1", "true", "T
 TRITONPARSE_MORE_TENSOR_INFORMATION = os.getenv(
     "TRITONPARSE_MORE_TENSOR_INFORMATION", None
 ) in ["1", "true", "True"]
+# Enable full Python source file extraction instead of just the function definition
+TRITON_FULL_PYTHON_SOURCE = os.getenv("TRITON_FULL_PYTHON_SOURCE", "0") in [
+    "1",
+    "true",
+    "True",
+]
+# Maximum file size for full source extraction (default 10MB)
+TRITON_MAX_SOURCE_SIZE = int(os.getenv("TRITON_MAX_SOURCE_SIZE", str(10 * 1024 * 1024)))
 # Inductor compiled kernel's launch tracing needs this flag to be set.
 # If TRITON_TRACE_LAUNCH is enabled, also enable TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK
 TORCHINDUCTOR_RUN_JIT_POST_COMPILE_HOOK = (
@@ -727,6 +735,17 @@ def extract_python_source_info(trace_data: Dict[str, Any], source):
     from the provided source object (typically an ASTSource or IRSource instance).
     It adds file path, line numbers, and the actual source code to the trace_data.
 
+    By default, only the function definition is extracted. Set TRITON_FULL_PYTHON_SOURCE=1
+    to extract the entire Python source file.
+    @TODO: we should enable it by default in next diff and track the compilation time regression
+
+    Environment Variables:
+        TRITON_FULL_PYTHON_SOURCE: If set to "1", extract the full Python file
+                                   instead of just the function definition.
+        TRITON_MAX_SOURCE_SIZE: Maximum file size in bytes for full source extraction
+                               (default: 10MB). Files larger than this will fall back
+                               to function-only mode.
+
     Args:
         trace_data (Dict[str, Any]): Dictionary to store extracted information
         source (Union[ASTSource, IRSource]): Source object containing kernel function information
@@ -738,23 +757,77 @@ def extract_python_source_info(trace_data: Dict[str, Any], source):
     if isinstance(source, IRSource):
         return
 
-    # Get the original Python source code for the kernel
+    # Get the function reference
+    if isinstance(fn := source.fn, JITFunction):
+        fn_ref = fn.fn
+    else:
+        fn_ref = source.fn
+
+    python_source_file = inspect.getfile(fn_ref)
+
+    # Get function range information
     if (
         isinstance(fn := source.fn, JITFunction)
         and hasattr(fn, "starting_line_number")
         and hasattr(fn, "raw_src")
     ):
-        start_line_number = fn.starting_line_number
+        function_start_line = fn.starting_line_number
         source_lines = fn.raw_src
     else:
-        source_lines, start_line_number = inspect.getsourcelines(fn.fn)
+        source_lines, function_start_line = inspect.getsourcelines(fn_ref)
+
+    function_end_line = function_start_line + len(source_lines) - 1
+
+    if TRITON_FULL_PYTHON_SOURCE:
+        # Full file mode: read the entire Python file
+        try:
+            # Check file size before reading
+            file_size = os.path.getsize(python_source_file)
+        except OSError as e:
+            log.warning(
+                f"Failed to check file size for {python_source_file}: {e}. "
+                f"Falling back to function-only mode."
+            )
+            use_full_source = False
+        else:
+            if file_size > TRITON_MAX_SOURCE_SIZE:
+                log.warning(
+                    f"Source file {python_source_file} is too large ({file_size} bytes, "
+                    f"limit: {TRITON_MAX_SOURCE_SIZE} bytes). Falling back to function-only mode."
+                )
+                use_full_source = False
+            else:
+                use_full_source = True
+
+        if use_full_source:
+            try:
+                with open(python_source_file, "r", encoding="utf-8") as f:
+                    file_content = f.read()
+
+                # Calculate total lines
+                total_lines = len(file_content.split("\n"))
+
+                trace_data["python_source"] = {
+                    "file_path": python_source_file,
+                    "start_line": 1,
+                    "end_line": total_lines,
+                    "code": file_content,
+                    # Add function range for frontend highlighting and scrolling
+                    "function_start_line": function_start_line,
+                    "function_end_line": function_end_line,
+                }
+                return
+            except (OSError, UnicodeDecodeError) as e:
+                log.warning(
+                    f"Failed to read full source file {python_source_file}: {e}. "
+                    f"Falling back to function-only mode."
+                )
 
-    python_source_file = inspect.getfile(fn.fn)
-    end_line_number = start_line_number + len(source_lines)
+    # Default behavior: only extract function definition
     trace_data["python_source"] = {
         "file_path": python_source_file,
-        "start_line": start_line_number,
-        "end_line": end_line_number,
+        "start_line": function_start_line,
+        "end_line": function_end_line,
         "code": "".join(source_lines),
     }
 
@@ -910,7 +983,7 @@ class TritonTraceHandler(logging.StreamHandler):
                 )
             elif not os.access(TRACE_LOG_DIR, os.W_OK):
                 log.info(
-                    "TritonTraceHandler: disabled because %s is not writeable",
+                    "TritonTraceHandler: disabled because %s is not writable",
                     TRACE_LOG_DIR,
                 )
             else:

{tritonparse-0.3.1.dev20251030071508.dist-info → tritonparse-0.3.1.dev20251101071506.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.3.1.dev20251030071508
+Version: 0.3.1.dev20251101071506
 Summary: TritonParse: A Compiler Tracer, Visualizer, and mini-Reproducer Generator for Triton Kernels
 Author-email: Yueming Hao <yhao@meta.com>
 License-Expression: BSD-3-Clause

{tritonparse-0.3.1.dev20251030071508.dist-info → tritonparse-0.3.1.dev20251101071506.dist-info}/RECORD

@@ -1,6 +1,6 @@
 tritonparse/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tritonparse/__main__.py,sha256=RXbkALBewcb1xlJBnsQl9IaBRUNln7U8NuRZKT8UdIk,117
-tritonparse/cli.py,sha256=JqBwzpxiFKb0TFdhovDXnz3gwkjeASYgIe311GRBy0o,2876
+tritonparse/cli.py,sha256=CBGwJWGMkctenid-DUoqmguO5gCGd0xWw5D9RlBdqNQ,2792
 tritonparse/common.py,sha256=MJo9bVCgSKkwXpEoUkUczPo_5jOYpJgXLq4UsWYqN3c,13924
 tritonparse/context_manager.py,sha256=OdMn11qbApYL2c9IlbUpcT27r04ZSa4DfvrY2mLA958,2243
 tritonparse/event_diff.py,sha256=USCjfjYr-7Ie-EfZgtCFMZMA1KRzFRDe7yDFy98zYI4,4962
@@ -11,7 +11,7 @@ tritonparse/mapper.py,sha256=QBCUMHM9pu3x3ahFp0wyXJbmv9TFGVPdkcLULok1E-k,4205
 tritonparse/shared_vars.py,sha256=RifXq55KisHgspYAmGcaCWY6ZHX8iejFHvwIewvcWZE,707
 tritonparse/source_type.py,sha256=nmYEQS8rfkIN9BhNhQbkmEvKnvS-3zAxRGLY4TaZdi8,1676
 tritonparse/sourcemap_utils.py,sha256=uI02n5Sgnlx7Nc15QAX5N6_tZZMips0PyJuo1n3eouY,2654
-tritonparse/structured_logging.py,sha256=L1xkkCx8Jr9YQbM0Kgtf2g6L3aWMkYOEeFFEOSo8Lkk,60306
+tritonparse/structured_logging.py,sha256=aBTfRVuZzd7YJV-Fgg513SPlbEceyWosONLqaTZS69k,63319
 tritonparse/tp_logger.py,sha256=vXzY7hMDmVnRBGBhIjFZe3nHZzG5NKKPONGUszJhGgU,242
 tritonparse/trace_processor.py,sha256=AW4YDrPDayURtmePkFi5m5p6P7OTi1UlTPbbrPzujwY,14418
 tritonparse/utils.py,sha256=Jnlptcd79llSDev-_1XyyOnv2izUqv0PEL74A8GF2tc,4565
@@ -34,9 +34,9 @@ tritonparse/tools/format_fix.py,sha256=ISalg_N_L7Xktag3mLr-G9T6Opxv793s1WG6A9wUt
 tritonparse/tools/load_tensor.py,sha256=7-LbpboKDNJFBLNhiKS3enoqRlVAb55OjPc70PwHXAw,2789
 tritonparse/tools/prettify_ndjson.py,sha256=kR8hmBCv-iJeuzpi2_6CZv9T4_edRQbBOSOPpMm6wrw,11117
 tritonparse/tools/readme.md,sha256=w6PWYfYnRgoPArLjxG9rVrpcLUkoVMGuRlbpF-o0IQM,110
-tritonparse-0.3.1.dev20251030071508.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
-tritonparse-0.3.1.dev20251030071508.dist-info/METADATA,sha256=GgewTGhlYm2vvn7S39qg5UiXAzJ07JzkETek6YJnwSw,8282
-tritonparse-0.3.1.dev20251030071508.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tritonparse-0.3.1.dev20251030071508.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
-tritonparse-0.3.1.dev20251030071508.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
-tritonparse-0.3.1.dev20251030071508.dist-info/RECORD,,
+tritonparse-0.3.1.dev20251101071506.dist-info/licenses/LICENSE,sha256=4ZciugpyN7wcM4L-9pyDh_etvMUeIfBhDTyH1zeZlQM,1515
+tritonparse-0.3.1.dev20251101071506.dist-info/METADATA,sha256=XavQSRG05Q5AlaCXfrfPN0BWcD8aUg2yIhOwmrs8e1M,8282
+tritonparse-0.3.1.dev20251101071506.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tritonparse-0.3.1.dev20251101071506.dist-info/entry_points.txt,sha256=wEXdaieDoRRCCdhEv2p_C68iytnaXU_2pwt5CqjfbWY,56
+tritonparse-0.3.1.dev20251101071506.dist-info/top_level.txt,sha256=ITcTKgp3vf_bXV9vixuQU9IrZa3L1EfDSZwvRzRaoJU,12
+tritonparse-0.3.1.dev20251101071506.dist-info/RECORD,,