change line-length to 100

Author: FindHao

pyproject.toml

@@ -18,7 +18,7 @@ test = [
 include = ["tritonparse*"]
 
 [tool.black]
-line-length = 88
+line-length = 100
 target-version = ["py310"]
 
 [tool.ufmt]

tests/test_add.py

@@ -72,6 +72,4 @@ def test_tensor_add():
 if __name__ == "__main__":
     test_tensor_add()
     # Use improved unified_parse with explicit output directory
-    tritonparse.utils.unified_parse(
-        source=log_path, out="./parsed_output", overwrite=True
-    )
+    tritonparse.utils.unified_parse(source=log_path, out="./parsed_output", overwrite=True)

tests/test_tritonparse.py

@@ -211,18 +211,14 @@ def run_test_kernel(x):
         torch.cuda.synchronize()
 
         # Check that temp_dir_logs folder has content
-        assert os.path.exists(
-            temp_dir_logs
-        ), f"Log directory {temp_dir_logs} does not exist."
+        assert os.path.exists(temp_dir_logs), f"Log directory {temp_dir_logs} does not exist."
         log_files = os.listdir(temp_dir_logs)
         assert (
             len(log_files) > 0
         ), f"No log files found in {temp_dir_logs}. Expected log files to be generated during Triton compilation."
         print(f"Found {len(log_files)} log files in {temp_dir_logs}: {log_files}")
 
-        tritonparse.utils.unified_parse(
-            source=temp_dir_logs, out=temp_dir_parsed, overwrite=True
-        )
+        tritonparse.utils.unified_parse(source=temp_dir_logs, out=temp_dir_parsed, overwrite=True)
 
         # Clean up temporary directory
         try:

tritonparse/common.py

@@ -247,9 +247,7 @@ def copy_local_to_tmpdir(local_path: str, verbose: bool = False) -> str:
 
     for item in os.listdir(local_path):
         item_path = os.path.join(local_path, item)
-        if os.path.isfile(item_path) and os.path.basename(item_path).startswith(
-            LOG_PREFIX
-        ):
+        if os.path.isfile(item_path) and os.path.basename(item_path).startswith(LOG_PREFIX):
             if verbose:
                 logger.info(f"Copying {item_path} to {temp_dir}")
             shutil.copy2(item_path, temp_dir)
@@ -309,9 +307,7 @@ def parse_logs(
     for rank, files in ranks.items():
         use_filenames = False
         if len(files) > 1:
-            logger.warning(
-                "Warning: multiple logs found for the same rank. Using filenames."
-            )
+            logger.warning("Warning: multiple logs found for the same rank. Using filenames.")
             use_filenames = True
         # Determine rank key for file mapping
         rank_key = "rank_default" if rank.is_default else f"rank_{rank.value}"
@@ -350,9 +346,7 @@ def parse_logs(
                 # Add files to the mapping (now with .gz extensions)
                 file_mapping[rank_key]["regular_files"].extend(generated_files)
                 # this is used to generate the tritonparse url
-                file_mapping[rank_key]["rank_suffix"] = rank_config.to_rank().to_string(
-                    suffix="/"
-                )
+                file_mapping[rank_key]["rank_suffix"] = rank_config.to_rank().to_string(suffix="/")
                 if mapped_file:
                     file_mapping[rank_key]["mapped_file"] = mapped_file
 

tritonparse/extract_source_mappings.py

@@ -40,16 +40,12 @@
 
 # the definition of the PTX loc directive.
 # Example: .loc 1 0 50 // abcdef.py:0:50
-PTX_LOC_PATTERN = re.compile(
-    r"^\s*\.loc\s+\d+\s+(\d+)\s+(\d+)\s+//\s*(.+?):(\d+):(\d+)"
-)
+PTX_LOC_PATTERN = re.compile(r"^\s*\.loc\s+\d+\s+(\d+)\s+(\d+)\s+//\s*(.+?):(\d+):(\d+)")
 
 # the definition of the AMDGCN loc directive.
 # Example: .loc	1 32 30                         ; abcd.py:32:30
 # .loc	1 32 46 is_stmt 0               ; abcd.py:32:46
-AMDGCN_LOC_PATTERN = re.compile(
-    r".*loc\s+(\d+)\s+(\d+)\s+(\d+)(?:\s+[^;]*)?;\s*(.+?):(\d+):(\d+)"
-)
+AMDGCN_LOC_PATTERN = re.compile(r".*loc\s+(\d+)\s+(\d+)\s+(\d+)(?:\s+[^;]*)?;\s*(.+?):(\d+):(\d+)")
 
 
 def get_file_extension(filename: str) -> str:
@@ -257,9 +253,7 @@ def extract_ptx_amdgcn_mappings(
     def get_file_path(filename: str) -> str:
         file_path = filename
         if not os.path.isabs(filename):
-            logger.debug(
-                f"Filename '{filename}' does not contain a path. Attempting to resolve."
-            )
+            logger.debug(f"Filename '{filename}' does not contain a path. Attempting to resolve.")
             # Attempt to resolve the filename to a full path using referenced_files
             if filename in referenced_files:
                 if len(referenced_files[filename]) > 1:
@@ -488,9 +482,7 @@ def parse_single_trace_content(trace_content: str) -> str:
                 create_bidirectional_mapping(
                     ir_maps[src_type], ir_maps[tgt_type], src_type, tgt_type
                 )
-                logger.debug(
-                    f"Created bidirectional mapping between {src_type} and {tgt_type}"
-                )
+                logger.debug(f"Created bidirectional mapping between {src_type} and {tgt_type}")
 
     py_map = {}
 
@@ -602,9 +594,7 @@ def parse_single_file(
                 if frame_id is not None or frame_compile_id is not None:
                     output_file_name = f"f{frame_id}_fc{frame_compile_id}_a{attempt_id}_cai{compiled_autograd_id}.ndjson"
                 else:
-                    logger.debug(
-                        "No frame_id or frame_compile_id found in the payload."
-                    )
+                    logger.debug("No frame_id or frame_compile_id found in the payload.")
                     output_file_name = f"{file_name_without_extension}_mapped.ndjson"
             else:
                 output_file_name = f"{file_name_without_extension}_mapped.ndjson"

tritonparse/shared_vars.py

@@ -3,7 +3,5 @@
 
 # The compilation information will be stored to /logs/DEFAULT_TRACE_FILE_PREFIX by default
 # unless other flags disable or set another store. Add USER to avoid permission issues in shared servers.
-DEFAULT_TRACE_FILE_PREFIX = (
-    f"dedicated_log_triton_trace_{os.getenv('USER', 'unknown')}_"
-)
+DEFAULT_TRACE_FILE_PREFIX = f"dedicated_log_triton_trace_{os.getenv('USER', 'unknown')}_"
 DEFAULT_TRACE_FILE_PREFIX_WITHOUT_USER = "dedicated_log_triton_trace_"

tritonparse/source_type.py

@@ -51,6 +51,4 @@ def _parse_source(self) -> Tuple[SourceType, str]:
         elif path.is_file():
             return SourceType.LOCAL_FILE, str(path.absolute())
         else:
-            raise ValueError(
-                f"Source '{self.source_str}' is not a valid directory or file"
-            )
+            raise ValueError(f"Source '{self.source_str}' is not a valid directory or file")

tritonparse/structured_logging.py

@@ -173,9 +173,7 @@ def convert(obj):
         return str(obj)
 
     if is_dataclass(obj):
-        return convert(
-            asdict(obj)
-        )  # Convert dataclass to dict and then process that dict
+        return convert(asdict(obj))  # Convert dataclass to dict and then process that dict
     log.warning(f"Unknown type: {type(obj)}")
     return str(obj)  # Return primitive types as-is
 
@@ -193,8 +191,7 @@ def maybe_enable_debug_logging():
 
         # Check if we already have a debug handler
         has_debug_handler = any(
-            isinstance(handler, logging.StreamHandler)
-            and handler.level <= logging.DEBUG
+            isinstance(handler, logging.StreamHandler) and handler.level <= logging.DEBUG
             for handler in log.handlers
         )
 
@@ -276,11 +273,7 @@ def extract_kernel_name(src) -> Optional[str]:
             return src.getattr("name", None)
         else:
             # For ASTSource, get the function name
-            if (
-                hasattr(src, "fn")
-                and hasattr(src.fn, "fn")
-                and hasattr(src.fn.fn, "__name__")
-            ):
+            if hasattr(src, "fn") and hasattr(src.fn, "fn") and hasattr(src.fn.fn, "__name__"):
                 return src.fn.fn.__name__
             return None
     except Exception as e:
@@ -316,9 +309,7 @@ def should_trace_kernel(
             log.debug(f"Kernel '{kernel_name}' matches pattern '{pattern}', will trace")
             return True
 
-    log.debug(
-        f"Kernel '{kernel_name}' does not match any allowlist pattern, skipping trace"
-    )
+    log.debug(f"Kernel '{kernel_name}' does not match any allowlist pattern, skipping trace")
     return False
 
 
@@ -371,18 +362,14 @@ def extract_file_content(trace_data: Dict[str, Any], metadata_group: Dict[str, s
                 # Check file size before reading to avoid memory issues
                 file_size = os.path.getsize(file_path)
                 if file_size > MAX_FILE_SIZE:
-                    trace_data["file_content"][
-                        ir_filename
-                    ] = f"<file too large: {file_size} bytes>"
+                    trace_data["file_content"][ir_filename] = f"<file too large: {file_size} bytes>"
                     continue
 
                 with open(file_path, "r") as f:
                     trace_data["file_content"][ir_filename] = f.read()
             except (UnicodeDecodeError, OSError) as e:
                 # add more specific error type
-                trace_data["file_content"][
-                    ir_filename
-                ] = f"<error reading file: {str(e)}>"
+                trace_data["file_content"][ir_filename] = f"<error reading file: {str(e)}>"
                 log.debug(f"Error reading file {file_path}: {e}")
 
 
@@ -439,9 +426,7 @@ class TritonTraceHandler(logging.StreamHandler):
     it automatically adds rank information to filenames.
     """
 
-    def __init__(
-        self, root_dir: Optional[str] = None, prefix=DEFAULT_TRACE_FILE_PREFIX
-    ):
+    def __init__(self, root_dir: Optional[str] = None, prefix=DEFAULT_TRACE_FILE_PREFIX):
         logging.Handler.__init__(self)
         self.root_dir = root_dir
         self.prefix = prefix
@@ -462,13 +447,8 @@ def get_root_dir(self):
         if TORCH_INSTALLED:
             import torch.version as torch_version
 
-            if (
-                hasattr(torch_version, "git_version")
-                and os.getenv("MAST_HPC_JOB_NAME") is None
-            ):
-                log.info(
-                    "TritonTraceHandler: disabled because not fbcode or conda on mast"
-                )
+            if hasattr(torch_version, "git_version") and os.getenv("MAST_HPC_JOB_NAME") is None:
+                log.info("TritonTraceHandler: disabled because not fbcode or conda on mast")
                 should_set_root_dir = False
             # TODO: change to tritonparse knob
             elif not torch._utils_internal.justknobs_check("pytorch/trace:enable"):
@@ -711,9 +691,7 @@ def maybe_trace_triton(
     # Add cache_hit to metadata
     trace_data["metadata"]["cache_hit"] = cache_hit
     if not metadata:
-        metadata_path = next(
-            (Path(p) for c, p in metadata_group.items() if c.endswith(".json"))
-        )
+        metadata_path = next((Path(p) for c, p in metadata_group.items() if c.endswith(".json")))
         with open(metadata_path, "r") as f:
             metadata = json.load(f)
             trace_data["metadata"].update(metadata)
@@ -919,9 +897,7 @@ def __call__(self, metadata):
         trace_data["name"] = metadata_dict["name"]
         trace_data["function"] = metadata_dict["function"]
         trace_data["stream"] = metadata_dict["stream"]
-        launch_metadata_tritonparse = metadata_dict.get(
-            "launch_metadata_tritonparse", None
-        )
+        launch_metadata_tritonparse = metadata_dict.get("launch_metadata_tritonparse", None)
         if launch_metadata_tritonparse is not None:
             trace_data["grid"] = launch_metadata_tritonparse[0]
             trace_data["metadata"] = launch_metadata_tritonparse[1]
@@ -967,9 +943,7 @@ def init_basic(trace_folder: Optional[str] = None):
     # Parse and store kernel allowlist configuration
     _KERNEL_ALLOWLIST_PATTERNS = parse_kernel_allowlist()
     if _KERNEL_ALLOWLIST_PATTERNS:
-        log.debug(
-            f"Kernel allowlist enabled with patterns: {_KERNEL_ALLOWLIST_PATTERNS}"
-        )
+        log.debug(f"Kernel allowlist enabled with patterns: {_KERNEL_ALLOWLIST_PATTERNS}")
     else:
         log.debug("Kernel allowlist not set, tracing all kernels")
 

tritonparse/tools/decompress_bin_ndjson.py

@@ -62,9 +62,7 @@ def decompress_bin_ndjson(input_file: str, output_file: str = None) -> None:
         # 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
-        )
+        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")
@@ -100,9 +98,7 @@ def main():
         help="Output .ndjson file path (default: replace .bin.ndjson with .ndjson)",
     )
 
-    parser.add_argument(
-        "-v", "--verbose", action="store_true", help="Enable verbose output"
-    )
+    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")
 
     args = parser.parse_args()
 

tritonparse/tools/prettify_ndjson.py

@@ -98,9 +98,7 @@ def parse_line_ranges(lines_arg: str) -> set[int]:
     return line_numbers
 
 
-def load_ndjson(
-    file_path: Path, save_irs: bool = False, line_filter: set[int] = None
-) -> List[Any]:
+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.
 
@@ -157,17 +155,13 @@ def load_ndjson(
                                     )  # 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 = 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"Error parsing JSON on line {line_num}: {e}", file=sys.stderr)
                     print(f"Problematic line: {line[:100]}...", file=sys.stderr)
                     raise
 
@@ -230,9 +224,7 @@ def main():
         """,
     )
 
-    parser.add_argument(
-        "ndjson_file", type=str, help="Path to the NDJSON file to convert"
-    )
+    parser.add_argument("ndjson_file", type=str, help="Path to the NDJSON file to convert")
 
     parser.add_argument(
         "--save-irs",
@@ -282,9 +274,7 @@ def main():
         if args.lines:
             try:
                 line_filter = parse_line_ranges(args.lines)
-                print(
-                    f"Line filtering enabled: will process {len(line_filter)} specified 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)
@@ -295,9 +285,7 @@ def main():
             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
-        )
+        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