tritonparse 0.2.4.dev20250928071447__tar.gz → 0.2.4.dev20250930071536__tar.gz

{tritonparse-0.2.4.dev20250928071447/tritonparse.egg-info → tritonparse-0.2.4.dev20250930071536}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20250928071447
+Version: 0.2.4.dev20250930071536
 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.dev20250928071447 → tritonparse-0.2.4.dev20250930071536}/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.dev20250930071536/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.dev20250928071447 → tritonparse-0.2.4.dev20250930071536}/tritonparse/reproducer/orchestrator.py

@@ -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

{tritonparse-0.2.4.dev20250928071447 → tritonparse-0.2.4.dev20250930071536}/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.dev20250928071447 → tritonparse-0.2.4.dev20250930071536}/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.dev20250928071447 → tritonparse-0.2.4.dev20250930071536}/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.dev20250928071447 → tritonparse-0.2.4.dev20250930071536/tritonparse.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tritonparse
-Version: 0.2.4.dev20250928071447
+Version: 0.2.4.dev20250930071536
 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.dev20250928071447 → tritonparse-0.2.4.dev20250930071536}/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