tritonparse 0.2.4.dev20250929071529__tar.gz → 0.2.4.dev20251001071529__tar.gz

{tritonparse-0.2.4.dev20250929071529/tritonparse.egg-info → tritonparse-0.2.4.dev20251001071529}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20250929071529
+Version: 0.2.4.dev20251001071529
 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.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529}/tests/test_tritonparse.py

@@ -776,7 +776,7 @@ class TestTritonparseCUDA(unittest.TestCase):
 
         from tritonparse.tools.prettify_ndjson import load_ndjson as _load_ndjson
 
-        events = _load_ndjson(_Path(ndjson_path), save_irs=True)
+        events = _load_ndjson(_Path(ndjson_path))
         launch_indices = [
             i for i, ev in enumerate(events) if ev.get("event_type") == "launch"
         ]

tritonparse-0.2.4.dev20251001071529/tritonparse/context_manager.py

@@ -0,0 +1,23 @@
+import os
+import shutil
+import tempfile
+
+from .structured_logging import clear_logging_config, init
+from .utils import unified_parse
+
+
+def createUniqueTempDirectory():
+    return tempfile.mkdtemp()
+
+
+class TritonParseManager:
+    def __enter__(self):
+        self.dir_path = createUniqueTempDirectory()
+        init(self.dir_path)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.output_link = unified_parse(source=self.dir_path, overwrite=True)
+        clear_logging_config()
+        if os.path.exists(self.dir_path):
+            shutil.rmtree(self.dir_path)

{tritonparse-0.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529}/tritonparse/reproducer/orchestrator.py

@@ -18,7 +18,7 @@ def reproduce(
     line_index: int,
     out_dir: str,
     template: str,
-):
+) -> dict[str, Path]:
     """
     Generate a reproducer script from NDJSON trace file.
 
@@ -28,7 +28,7 @@ def reproduce(
         out_dir: Output directory for reproducer files.
     """
     logger.debug(f"Building bundle from {input_path} at line {line_index}")
-    events = load_ndjson(Path(input_path), save_irs=True)
+    events = load_ndjson(Path(input_path))
     logger.debug(f"Loaded {len(events)} events")
 
     # Build context bundle from the specified launch event
@@ -55,9 +55,15 @@ def reproduce(
         "# {{KERNEL_INVOCATION_PLACEHOLDER}}", invocation_snippet
     )
     out_py_path.write_text(final_code, encoding="utf-8")
-    logger.info(
-        "REPRODUCER_OUTPUT script=%s json=%s kernel=%s",
-        str(out_py_path.resolve()),
-        str(temp_json_path.resolve()),
-        context_bundle.kernel_info.function_name,
-    )
+
+    filepath = context_bundle.kernel_info.file_path
+    filepath = "/".join(filepath.split("/")[5:])
+    ret = {
+        "kernel_src_path": filepath,
+        "kernel": context_bundle.kernel_info.function_name,
+        "repo_script": str(out_py_path.resolve()),
+        "repo_context": str(temp_json_path.resolve()),
+    }
+    logger.info("REPRODUCER_OUTPUT\n%s", ret)
+
+    return ret

{tritonparse-0.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529}/tritonparse/structured_logging.py

@@ -28,8 +28,6 @@ log = logging.getLogger(__name__)
 
 TEXT_FILE_EXTENSIONS = [".ttir", ".ttgir", ".llir", ".ptx", ".amdgcn", ".json"]
 MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB limit for file content extraction
-# Enable ndjson output. json format is only for debugging purpose.
-TRITONPARSE_NDJSON = os.getenv("TRITONPARSE_NDJSON", "1") in ["1", "true", "True"]
 # Enable gzip compression for each line in trace files
 TRITON_TRACE_GZIP = os.getenv("TRITON_TRACE_GZIP", "0") in ["1", "true", "True"]
 triton_trace_log = logging.getLogger("tritonparse_trace")
@@ -505,12 +503,9 @@ class TritonJsonFormatter(logging.Formatter):
         if payload is not None:
             log_entry["payload"] = json.loads(payload)
         clean_log_entry = convert(log_entry)
-        if not TRITONPARSE_NDJSON:
-            return json.dumps(clean_log_entry, indent=2)
-        else:
-            # NDJSON format requires a newline at the end of each line
-            json_str = json.dumps(clean_log_entry, separators=(",", ":"))
-            return json_str + "\n"
+        # NDJSON format requires a newline at the end of each line
+        json_str = json.dumps(clean_log_entry, separators=(",", ":"))
+        return json_str + "\n"
 
 
 class TritonTraceHandler(logging.StreamHandler):

{tritonparse-0.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529}/tritonparse/tools/prettify_ndjson.py

@@ -31,7 +31,7 @@ Example:
 Usage:
     python prettify_ndjson.py data.ndjson
     python prettify_ndjson.py --lines 1,3 data.ndjson  # Only process lines 1 and 3
-    python prettify_ndjson.py --save-irs logs.ndjson   # Keep all fields for compilation events
+    python prettify_ndjson.py --not-save-irs logs.ndjson   # Remove large fields for compilation events
 
 
 """
@@ -99,14 +99,14 @@ def parse_line_ranges(lines_arg: str) -> set[int]:
 
 
 def load_ndjson(
-    file_path: Path, save_irs: bool = False, line_filter: set[int] = None
+    file_path: Path, not_save_irs: bool = False, line_filter: set[int] | None = None
 ) -> List[Any]:
     """
     Load NDJSON file and return list of JSON objects.
 
     Args:
         file_path: Path to the NDJSON file
-        save_irs: Whether to save file_content and python_source for compilation events
+        not_save_irs: Whether to NOT save file_content and python_source for compilation events
         line_filter: Set of line numbers to include (1-based indexing), None means include all
 
     Returns:
@@ -138,8 +138,8 @@ def load_ndjson(
                 try:
                     json_obj = json.loads(line)
 
-                    # Filter out file_content and python_source for compilation events if save_irs is False
-                    if not save_irs and isinstance(json_obj, dict):
+                    # Filter out file_content and python_source for compilation events if not_save_irs is True
+                    if not_save_irs and isinstance(json_obj, dict):
                         event_type = json_obj.get("event_type")
                         if event_type == "compilation":
                             # Remove file_content and python_source from payload if they exist
@@ -188,13 +188,13 @@ def load_ndjson(
             print("Line filtering: no valid lines specified")
 
     # Print warning if compilation events were filtered
-    if not save_irs and filtered_compilation_events > 0:
+    if not_save_irs and filtered_compilation_events > 0:
         print(
             f"WARNING: Removed 'file_content' and 'python_source' fields from {filtered_compilation_events} compilation events to reduce file size.",
             file=sys.stderr,
         )
         print(
-            "         Use --save-irs flag to preserve these fields if needed.",
+            "         Re-run without --not-save-irs flag to preserve these fields if needed.",
             file=sys.stderr,
         )
 
@@ -237,10 +237,10 @@ Examples:
     )
 
     parser.add_argument(
-        "--save-irs",
+        "--not-save-irs",
         action="store_true",
         default=False,
-        help="Save file_content and python_source for compilation events (default: False, removes these fields to reduce size)",
+        help="Do not save file_content and python_source for compilation events (default: False, removes these fields to reduce size)",
     )
 
     parser.add_argument(
@@ -293,12 +293,12 @@ Examples:
 
         # Load NDJSON file
         print(f"Loading NDJSON file: {input_path}")
-        if not args.save_irs:
+        if args.not_save_irs:
             print(
                 "Filtering out file_content and python_source from compilation events to reduce size"
             )
         json_objects = load_ndjson(
-            input_path, save_irs=args.save_irs, line_filter=line_filter
+            input_path, not_save_irs=args.not_save_irs, line_filter=line_filter
         )
         print(f"Loaded {len(json_objects)} JSON objects")
 

{tritonparse-0.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529}/tritonparse/utils.py

@@ -106,6 +106,7 @@ def oss_run(
     else:
         out_dir = str(Path(parsed_log_dir).absolute())
     print_parsed_files_summary(out_dir)
+    return None
 
 
 def unified_parse(
@@ -134,7 +135,7 @@ def unified_parse(
     else:
         parse = oss_run
 
-    parse(
+    output = parse(
         source=source,
         out=out,
         overwrite=overwrite,
@@ -143,3 +144,4 @@ def unified_parse(
         verbose=verbose,
         **kwargs,
     )
+    return output

{tritonparse-0.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529/tritonparse.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20250929071529
+Version: 0.2.4.dev20251001071529
 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.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529}/tritonparse.egg-info/SOURCES.txt

@@ -36,6 +36,7 @@ tests/example_output/parsed_output_complex/log_file_list.json
 tests/example_output/repro/repro_context_20250816192455.json
 tritonparse/__init__.py
 tritonparse/common.py
+tritonparse/context_manager.py
 tritonparse/event_diff.py
 tritonparse/extract_source_mappings.py
 tritonparse/ir_parser.py

{tritonparse-0.2.4.dev20250929071529 → tritonparse-0.2.4.dev20251001071529}/website/src/index.css

@@ -7,9 +7,7 @@
   line-height: 1.5;
   font-weight: 400;
 
-  color-scheme: light dark;
-  color: rgba(255, 255, 255, 0.87);
-  background-color: #242424;
+  color-scheme: light;
 
   font-synthesis: none;
   text-rendering: optimizeLegibility;