tritonparse 0.1.1__py3-none-any.whl

tritonparse/tools/decompress_bin_ndjson.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+"""
+Script to decompress .bin.ndjson files back to regular .ndjson format.
+
+The .bin.ndjson format stores each JSON record as a separate gzip member,
+concatenated in sequence within a single binary file. This script uses
+gzip.open() which automatically handles member concatenation to read
+the compressed file and write out the original NDJSON format.
+
+Usage:
+    python decompress_bin_ndjson.py trace.bin.ndjson
+"""
+
+import argparse
+import gzip
+import sys
+from pathlib import Path
+
+
+def decompress_bin_ndjson(input_file: str, output_file: str = None) -> None:
+    """
+    Decompress a .bin.ndjson file to regular .ndjson format.
+
+    Args:
+        input_file: Path to the .bin.ndjson file
+        output_file: Path for the output .ndjson file (optional)
+    """
+    input_path = Path(input_file)
+
+    # Validate input file
+    if not input_path.exists():
+        print(f"Error: Input file '{input_file}' does not exist", file=sys.stderr)
+        return
+
+    if not input_path.suffix.endswith(".bin.ndjson"):
+        print(f"Warning: Input file '{input_file}' doesn't have .bin.ndjson extension")
+
+    # Determine output file path
+    if output_file is None:
+        if input_path.name.endswith(".bin.ndjson"):
+            # Replace .bin.ndjson with .ndjson
+            output_file = str(input_path.with_suffix("").with_suffix(".ndjson"))
+        else:
+            # Add .decompressed.ndjson suffix
+            output_file = str(input_path.with_suffix(".decompressed.ndjson"))
+
+    output_path = Path(output_file)
+
+    try:
+        line_count = 0
+        # Because we use NDJSON format, each line is a complete JSON record.
+        # It is guruanteed here https://github.com/meta-pytorch/tritonparse/blob/
+        # c8dcc2a94ac10ede4342dba7456f6ebd8409b95d/tritonparse/structured_logging.py#L320
+        with gzip.open(input_path, "rt", encoding="utf-8") as compressed_file:
+            with open(output_path, "w", encoding="utf-8") as output:
+                for line in compressed_file:
+                    # gzip.open automatically handles member concatenation
+                    # Each line is already a complete JSON record with newline
+                    output.write(line)
+                    line_count += 1
+
+        # Get file sizes for comparison
+        input_size = input_path.stat().st_size
+        output_size = output_path.stat().st_size
+        compression_ratio = (
+            (1 - input_size / output_size) * 100 if output_size > 0 else 0
+        )
+
+        print(f"Successfully decompressed '{input_file}' to '{output_file}'")
+        print(f"  Input size:  {input_size:,} bytes")
+        print(f"  Output size: {output_size:,} bytes")
+        print(f"  Compression ratio: {compression_ratio:.1f}%")
+        print(f"  Records processed: {line_count:,}")
+
+    except gzip.BadGzipFile as e:
+        print(f"Error: Invalid gzip format in '{input_file}': {e}", file=sys.stderr)
+    except UnicodeDecodeError as e:
+        print(f"Error: Unicode decode error in '{input_file}': {e}", file=sys.stderr)
+    except Exception as e:
+        print(f"Error: Failed to decompress '{input_file}': {e}", file=sys.stderr)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Decompress .bin.ndjson files to regular .ndjson format",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  %(prog)s trace.bin.ndjson
+  %(prog)s trace.bin.ndjson -o output.ndjson
+  %(prog)s /logs/dedicated_log_triton_trace_user_.bin.ndjson
+        """,
+    )
+
+    parser.add_argument("input_file", help="Input .bin.ndjson file to decompress")
+
+    parser.add_argument(
+        "-o",
+        "--output",
+        help="Output .ndjson file path (default: replace .bin.ndjson with .ndjson)",
+    )
+
+    parser.add_argument(
+        "-v", "--verbose", action="store_true", help="Enable verbose output"
+    )
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        print(f"Decompressing: {args.input_file}")
+        if args.output:
+            print(f"Output file: {args.output}")
+
+    decompress_bin_ndjson(args.input_file, args.output)
+
+
+if __name__ == "__main__":
+    main()

tritonparse/tools/format_fix.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+"""
+Format fix script for tritonparse project.
+
+This script runs all linter tools to format and fix code issues:
+- usort: Import sorting
+- ruff: Linting only
+- black: Code formatting
+
+Usage:
+    python -m tritonparse.tools.format_fix [options]
+
+Options:
+    --check-only    Only check for issues, don't fix them
+    --verbose       Verbose output
+    --help          Show this help message
+"""
+
+import argparse
+import subprocess
+import sys
+
+
+def run_command(cmd: list[str], verbose: bool = False) -> bool:
+    """Run a command and return success status."""
+    if verbose:
+        print(f"Running: {' '.join(cmd)}")
+
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
+
+        if result.returncode != 0:
+            if verbose:
+                print(f"Command failed with return code {result.returncode}")
+                if result.stdout:
+                    print("STDOUT:", result.stdout)
+                if result.stderr:
+                    print("STDERR:", result.stderr)
+            return False
+
+        if verbose and result.stdout:
+            print(result.stdout)
+
+        return True
+    except Exception as e:
+        if verbose:
+            print(f"Error running command: {e}")
+        return False
+
+
+def run_usort(check_only: bool = False, verbose: bool = False) -> bool:
+    """Run usort for import sorting."""
+    cmd = ["usort"]
+
+    if check_only:
+        cmd.extend(["check", "."])
+    else:
+        cmd.extend(["format", "."])
+
+    return run_command(cmd, verbose)
+
+
+def run_ruff_check(check_only: bool = False, verbose: bool = False) -> bool:
+    """Run ruff for linting only."""
+    cmd = ["ruff", "check", "."]
+
+    if check_only:
+        cmd.append("--diff")
+    else:
+        cmd.append("--fix")
+
+    return run_command(cmd, verbose)
+
+
+def run_black(check_only: bool = False, verbose: bool = False) -> bool:
+    """Run black for code formatting."""
+    cmd = ["black"]
+
+    if check_only:
+        cmd.extend(["--check", "--diff", "."])
+    else:
+        cmd.append(".")
+
+    return run_command(cmd, verbose)
+
+
+def main():
+    """Main function."""
+    parser = argparse.ArgumentParser(
+        description="Format fix script for tritonparse project",
+        epilog="""
+Examples:
+    # Fix all formatting issues
+    python -m tritonparse.tools.format_fix
+    
+    # Check for issues without fixing
+    python -m tritonparse.tools.format_fix --check-only
+    
+    # Verbose output
+    python -m tritonparse.tools.format_fix --verbose
+        """,
+    )
+
+    parser.add_argument(
+        "--check-only",
+        action="store_true",
+        help="Only check for issues, don't fix them",
+    )
+    parser.add_argument("--verbose", action="store_true", help="Verbose output")
+
+    args = parser.parse_args()
+
+    # Run formatters on the entire project
+    success = True
+
+    # 1. Run usort for import sorting
+    print("Running usort for import sorting...")
+    if not run_usort(args.check_only, args.verbose):
+        print("❌ usort failed")
+        success = False
+    else:
+        print("✅ usort completed")
+
+    # 2. Run ruff for linting only
+    print("Running ruff for linting...")
+    if not run_ruff_check(args.check_only, args.verbose):
+        print("❌ ruff linting failed")
+        success = False
+    else:
+        print("✅ ruff linting completed")
+
+    # 3. Run black for code formatting
+    print("Running black for code formatting...")
+    if not run_black(args.check_only, args.verbose):
+        print("❌ black failed")
+        success = False
+    else:
+        print("✅ black completed")
+
+    if success:
+        print("\n🎉 All formatting tools completed successfully!")
+        return 0
+    else:
+        print("\n❌ Some formatting tools failed. Please check the output above.")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tritonparse/tools/load_tensor.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+"""
+Simple tensor loading utility for tritonparse saved tensors.
+Usage:
+import tritonparse.tools.load_tensor as load_tensor
+tensor = load_tensor.load_tensor(tensor_file_path, device)
+"""
+
+import hashlib
+from pathlib import Path
+
+import torch
+
+
+def load_tensor(tensor_file_path: str, device: str = None) -> torch.Tensor:
+    """
+    Load a tensor from its file path and verify its integrity using the hash in the filename.
+
+    Args:
+        tensor_file_path (str): Direct path to the tensor .bin file. The filename should be
+                               the hash of the file contents followed by .bin extension.
+        device (str, optional): Device to load the tensor to (e.g., 'cuda:0', 'cpu').
+                               If None, keeps the tensor on its original device.
+
+    Returns:
+        torch.Tensor: The loaded tensor (moved to the specified device if provided)
+
+    Raises:
+        FileNotFoundError: If the tensor file doesn't exist
+        RuntimeError: If the tensor cannot be loaded
+        ValueError: If the computed hash doesn't match the filename hash
+    """
+    blob_path = Path(tensor_file_path)
+
+    if not blob_path.exists():
+        raise FileNotFoundError(f"Tensor blob not found: {blob_path}")
+
+    # Extract expected hash from filename (remove .bin extension)
+    expected_hash = blob_path.stem
+
+    # Compute actual hash of file contents
+    with open(blob_path, "rb") as f:
+        file_contents = f.read()
+        computed_hash = hashlib.blake2b(file_contents).hexdigest()
+
+    # Verify hash matches filename
+    if computed_hash != expected_hash:
+        raise ValueError(
+            f"Hash verification failed: expected '{expected_hash}' but computed '{computed_hash}'"
+        )
+
+    try:
+        # Load the tensor using torch.load (tensors are saved with torch.save)
+        # If device is None, keep tensor on its original device, otherwise move to specified device
+        tensor = torch.load(blob_path, map_location=device)
+        return tensor
+    except Exception as e:
+        raise RuntimeError(f"Failed to load tensor from {blob_path}: {str(e)}")

tritonparse/tools/prettify_ndjson.py

@@ -0,0 +1,315 @@
+#!/usr/bin/env python3
+"""
+Convert an NDJSON file to a prettified JSON file.
+
+This script takes an NDJSON (newline-delimited JSON) file and converts it to a
+standard human-readable JSON file where each line becomes an element in a JSON array, with
+pretty formatting applied.
+
+Example:
+    Input NDJSON file (data.ndjson):
+        {"name": "Alice", "age": 30}
+        {"name": "Bob", "age": 25}
+        {"name": "Charlie", "age": 35}
+
+    Output JSON file (data_prettified.json):
+        [
+          {
+            "age": 30,
+            "name": "Alice"
+          },
+          {
+            "age": 25,
+            "name": "Bob"
+          },
+          {
+            "age": 35,
+            "name": "Charlie"
+          }
+        ]
+
+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
+
+
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any, List
+
+
+def parse_line_ranges(lines_arg: str) -> set[int]:
+    """
+    Parse line ranges from string like "1,2,3,5-10" into a set of line numbers.
+
+    Line numbers use 1-based indexing (first line is line 1, not 0).
+
+    Args:
+        lines_arg: String containing comma-separated line numbers and ranges
+                  Examples: "1", "1,2,3", "5-10", "1,3,5-10,15"
+
+    Returns:
+        Set of line numbers (1-based indexing, where 1 = first line)
+
+    Raises:
+        ValueError: If the format is invalid or contains non-positive numbers
+    """
+    line_numbers = set()
+
+    if not lines_arg.strip():
+        return line_numbers
+
+    parts = lines_arg.split(",")
+    for part in parts:
+        part = part.strip()
+        if not part:
+            continue
+
+        if "-" in part:
+            # Handle range like "5-10"
+            try:
+                start, end = part.split("-", 1)
+                start_num = int(start.strip())
+                end_num = int(end.strip())
+                if start_num <= 0 or end_num <= 0:
+                    raise ValueError("Line numbers must be positive")
+                if start_num > end_num:
+                    raise ValueError(f"Invalid range: {part} (start > end)")
+                line_numbers.update(range(start_num, end_num + 1))
+            except ValueError as e:
+                if "invalid literal" in str(e):
+                    raise ValueError(f"Invalid range format: {part}")
+                raise
+        else:
+            # Handle single number like "1"
+            try:
+                line_num = int(part)
+                if line_num <= 0:
+                    raise ValueError("Line numbers must be positive")
+                line_numbers.add(line_num)
+            except ValueError:
+                raise ValueError(f"Invalid line number: {part}")
+
+    return line_numbers
+
+
+def load_ndjson(
+    file_path: Path, save_irs: bool = False, line_filter: set[int] = 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
+        line_filter: Set of line numbers to include (1-based indexing), None means include all
+
+    Returns:
+        List of parsed JSON objects
+
+    Raises:
+        FileNotFoundError: If the input file doesn't exist
+        json.JSONDecodeError: If a line contains invalid JSON
+    """
+    json_objects = []
+    filtered_compilation_events = 0
+    total_lines_processed = 0
+
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            # enumerate(f, 1) starts line numbering from 1 (1-based indexing)
+            for line_num, line in enumerate(f, 1):
+                line = line.strip()
+                if not line:  # Skip empty lines
+                    continue
+
+                # Skip line if line filtering is enabled and this line is not in the filter
+                # line_num is 1-based (first line = 1, second line = 2, etc.)
+                if line_filter is not None and line_num not in line_filter:
+                    continue
+
+                total_lines_processed += 1
+
+                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):
+                        event_type = json_obj.get("event_type")
+                        if event_type == "compilation":
+                            # Remove file_content and python_source from payload if they exist
+                            payload = json_obj.get("payload")
+                            if isinstance(payload, dict):
+                                fields_to_remove = []
+                                if "file_content" in payload:
+                                    fields_to_remove.append("file_content")
+                                if "python_source" in payload:
+                                    fields_to_remove.append("python_source")
+
+                                if fields_to_remove:
+                                    payload = (
+                                        payload.copy()
+                                    )  # Create a copy to avoid modifying original
+                                    for field in fields_to_remove:
+                                        del payload[field]
+                                    json_obj = (
+                                        json_obj.copy()
+                                    )  # Create a copy of the main object
+                                    json_obj["payload"] = payload
+                                    filtered_compilation_events += 1
+
+                    json_objects.append(json_obj)
+                except json.JSONDecodeError as e:
+                    print(
+                        f"Error parsing JSON on line {line_num}: {e}", file=sys.stderr
+                    )
+                    print(f"Problematic line: {line[:100]}...", file=sys.stderr)
+                    raise
+
+    except FileNotFoundError:
+        print(f"Error: File '{file_path}' not found.", file=sys.stderr)
+        raise
+    except Exception as e:
+        print(f"Error reading file '{file_path}': {e}", file=sys.stderr)
+        raise
+
+    # Print informational messages
+    if line_filter is not None:
+        if line_filter:
+            print(
+                f"Line filtering: processed {total_lines_processed} out of {len(line_filter)} specified lines"
+            )
+        else:
+            print("Line filtering: no valid lines specified")
+
+    # Print warning if compilation events were filtered
+    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.",
+            file=sys.stderr,
+        )
+
+    return json_objects
+
+
+def save_prettified_json(json_objects: List[Any], output_path: Path) -> None:
+    """
+    Save list of JSON objects to a prettified JSON file.
+
+    Args:
+        json_objects: List of JSON objects to save
+        output_path: Path where to save the prettified JSON file
+    """
+    try:
+        with open(output_path, "w", encoding="utf-8") as f:
+            json.dump(json_objects, f, indent=2, ensure_ascii=False, sort_keys=True)
+        print(f"Successfully converted to prettified JSON: {output_path}")
+    except Exception as e:
+        print(f"Error writing to file '{output_path}': {e}", file=sys.stderr)
+        raise
+
+
+def main():
+    """Main function to handle command line arguments and orchestrate the conversion."""
+    parser = argparse.ArgumentParser(
+        description="Convert NDJSON file to prettified JSON file",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python prettify_ndjson.py data.ndjson
+  python prettify_ndjson.py /path/to/logs.ndjson
+        """,
+    )
+
+    parser.add_argument(
+        "ndjson_file", type=str, help="Path to the NDJSON file to convert"
+    )
+
+    parser.add_argument(
+        "--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)",
+    )
+
+    parser.add_argument(
+        "--lines",
+        type=str,
+        help="Specify line numbers to include using 1-based indexing (e.g., '1,2,3,5-10'). "
+        "Line 1 is the first line of the file. Only these lines from the original NDJSON will be processed. "
+        "Supports individual lines (1,2,3) and ranges (5-10).",
+    )
+
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        help="Specify output file path (default: {input_stem}_prettified.json in the same directory as input)",
+    )
+
+    args = parser.parse_args()
+
+    # Convert to Path object and validate
+    input_path = Path(args.ndjson_file)
+
+    if not input_path.exists():
+        print(f"Error: File '{input_path}' does not exist.", file=sys.stderr)
+        sys.exit(1)
+
+    if not input_path.is_file():
+        print(f"Error: '{input_path}' is not a file.", file=sys.stderr)
+        sys.exit(1)
+
+    # Generate output filename
+    if args.output:
+        output_path = Path(args.output)
+    else:
+        # Default: original_prettified.json in same directory as input
+        output_path = input_path.parent / f"{input_path.stem}_prettified.json"
+
+    try:
+        # Parse line filter if provided
+        line_filter = None
+        if args.lines:
+            try:
+                line_filter = parse_line_ranges(args.lines)
+                print(
+                    f"Line filtering enabled: will process {len(line_filter)} specified lines"
+                )
+            except ValueError as e:
+                print(f"Error parsing --lines argument: {e}", file=sys.stderr)
+                sys.exit(1)
+
+        # Load NDJSON file
+        print(f"Loading NDJSON file: {input_path}")
+        if not args.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
+        )
+        print(f"Loaded {len(json_objects)} JSON objects")
+
+        # Save as prettified JSON
+        print(f"Saving prettified JSON to: {output_path}")
+        save_prettified_json(json_objects, output_path)
+
+        print("Conversion completed successfully!")
+
+    except Exception as e:
+        print(f"Conversion failed: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

tritonparse/tp_logger.py

@@ -0,0 +1,9 @@
+#  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import logging
+
+logger = logging.getLogger("tritonparse")
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)