TraceLens Linting Automation

.github/workflows/lint.yml

@@ -0,0 +1,43 @@
+name: Lint
+
+on:
+  pull_request:
+  push:
+    branches: [ main ]
+
+jobs:
+  lint:
+    name: Black + usort
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Cache pip
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-pip-${{ hashFiles('**/pyproject.toml', '**/requirements*.txt', '**/setup.cfg') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-
+
+      - name: Install formatters
+        run: |
+          python -m pip install --upgrade pip
+          pip install black usort==1.0.7
+
+      - name: Check formatting with Black
+        # This step will fail if files need reformatting
+        run: |
+          black --check .
+
+      - name: Check import order with usort
+        # This step will fail if imports need sorting
+        run: |
+          usort check .

TraceLens/EventReplay/batched_replay.py

@@ -1,55 +1,91 @@
 # run_repro.py
-import json
 import argparse
+import json
 import warnings
+
 import torch
-from utils import TensorCfg, build_tensor, benchmark_func
+from utils import benchmark_func, build_tensor, TensorCfg
 
-def _get_args_kwargs_from_ir(event_replay_IR: dict[str, any], device: str = 'cuda') -> tuple[list[any], dict[str, any]]:
+
+def _get_args_kwargs_from_ir(
+    event_replay_IR: dict[str, any], device: str = "cuda"
+) -> tuple[list[any], dict[str, any]]:
     # (Copy the implementation of _get_args_kwargs_from_ir from Step 1 here)
     # ...
     pos_args = []
-    for arg in event_replay_IR['list_pos_args']:
-        val = arg['value']
+    for arg in event_replay_IR["list_pos_args"]:
+        val = arg["value"]
         # Check if the dict looks like a TensorCfg before attempting to create
-        if arg['arg_type'].startswith('Tensor') and val:
-             cfg = TensorCfg(**val) # Reconstruct from dict
-             pos_args.append(build_tensor(cfg, device=device))
+        if arg["arg_type"].startswith("Tensor") and val:
+            cfg = TensorCfg(**val)  # Reconstruct from dict
+            pos_args.append(build_tensor(cfg, device=device))
         else:
             pos_args.append(val)
 
     kwargs = {}
-    for arg in event_replay_IR['list_kwargs']:
-        val = arg['value']
-        key = arg['arg_name']
-        if arg['arg_type'].startswith('Tensor'):
-             cfg = TensorCfg(**val) # Reconstruct from dict
-             kwargs[key] = build_tensor(cfg, device=device)
+    for arg in event_replay_IR["list_kwargs"]:
+        val = arg["value"]
+        key = arg["arg_name"]
+        if arg["arg_type"].startswith("Tensor"):
+            cfg = TensorCfg(**val)  # Reconstruct from dict
+            kwargs[key] = build_tensor(cfg, device=device)
         else:
             kwargs[key] = val
     return pos_args, kwargs
 
+
 if __name__ == "__main__":
     # --- Argument Parsing ---
-    parser = argparse.ArgumentParser(description="Replay PyTorch operations from a repro file.")
-    parser.add_argument('repro_file', type=str, help="Path to the JSON repro file generated by extract_repro.py")
-    parser.add_argument('--device', type=str, default='cuda', choices=['cuda', 'cpu'], help="Device to run the operations on (default: cuda)")
-    parser.add_argument('--verbose', '-v', action='store_true', help="Enable verbose output during replay")
-    parser.add_argument('--stop-on-error', action='store_true', help="Stop execution immediately if any operation fails")
-    parser.add_argument('--op-limit', type=int, default=None, help="Limit the number of operations to replay")
-    parser.add_argument('--op-filter', type=str, default=None, help="Only replay operations whose name contains this string (e.g., 'aten::add')")
+    parser = argparse.ArgumentParser(
+        description="Replay PyTorch operations from a repro file."
+    )
+    parser.add_argument(
+        "repro_file",
+        type=str,
+        help="Path to the JSON repro file generated by extract_repro.py",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        choices=["cuda", "cpu"],
+        help="Device to run the operations on (default: cuda)",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        help="Enable verbose output during replay",
+    )
+    parser.add_argument(
+        "--stop-on-error",
+        action="store_true",
+        help="Stop execution immediately if any operation fails",
+    )
+    parser.add_argument(
+        "--op-limit",
+        type=int,
+        default=None,
+        help="Limit the number of operations to replay",
+    )
+    parser.add_argument(
+        "--op-filter",
+        type=str,
+        default=None,
+        help="Only replay operations whose name contains this string (e.g., 'aten::add')",
+    )
 
     args = parser.parse_args()
 
     # --- Device Check ---
-    if args.device == 'cuda' and not torch.cuda.is_available():
+    if args.device == "cuda" and not torch.cuda.is_available():
         print("Warning: CUDA device requested but not available. Falling back to CPU.")
-        args.device = 'cpu'
+        args.device = "cpu"
 
     print(f"Running repro from '{args.repro_file}' on device '{args.device}'")
 
     # --- Load Repro Data ---
-    with open(args.repro_file, 'r') as f:
+    with open(args.repro_file, "r") as f:
         # Use the custom decoder if you implemented one, otherwise standard load
         # repro_data_list = json.load(f, object_hook=decode_repro_info)
         repro_data_list = json.load(f)
@@ -60,68 +96,84 @@ def _get_args_kwargs_from_ir(event_replay_IR: dict[str, any], device: str = 'cud
 
     for i, repro_info in enumerate(repro_data_list):
 
-        op_name = repro_info['op_name']
-        replay_ir = repro_info['replay_ir']
+        op_name = repro_info["op_name"]
+        replay_ir = repro_info["replay_ir"]
         print(f"\n[{replayed_count + 1}/{len(repro_data_list)}] Replaying: {op_name}")
 
         # Get the PyTorch operation function
         try:
             func, _ = torch._C._jit_get_operation(op_name)
         except Exception as e:
-            print(f"  Error: Could not find PyTorch operation '{op_name}'. Is the PyTorch version compatible? Error: {e}")
-            if args.stop_on_error: raise
+            print(
+                f"  Error: Could not find PyTorch operation '{op_name}'. Is the PyTorch version compatible? Error: {e}"
+            )
+            if args.stop_on_error:
+                raise
             errors += 1
-            continue # Skip to next operation
+            continue  # Skip to next operation
 
         # Reconstruct arguments and keyword arguments
         try:
             if args.verbose:
                 print(f"  Reconstructing arguments for '{op_name}'...")
                 print(f"  Positional Args:")
-                for arg in replay_ir['list_pos_args']:
+                for arg in replay_ir["list_pos_args"]:
                     print(f"  {arg['arg_name']} {arg['arg_type']}: {arg['value']}")
                 print(f"  Keyword Args:")
-                for arg in replay_ir['list_kwargs']:
+                for arg in replay_ir["list_kwargs"]:
                     print(f"  {arg['arg_name']} {arg['arg_type']}: {arg['value']}")
             pos_args, kwargs = _get_args_kwargs_from_ir(replay_ir, device=args.device)
         except Exception as e:
-            print(f"  Error: Failed to reconstruct arguments for '{op_name}'. Error: {e}")
-            if args.stop_on_error: raise
+            print(
+                f"  Error: Failed to reconstruct arguments for '{op_name}'. Error: {e}"
+            )
+            if args.stop_on_error:
+                raise
             errors += 1
-            continue # Skip to next operation
-
+            continue  # Skip to next operation
 
         # Execute the operation
         # Optionally add synchronization for accurate timing if needed (cuda specific)
-        if args.device == 'cuda': torch.cuda.synchronize()
+        if args.device == "cuda":
+            torch.cuda.synchronize()
         # --- Call the function ---
         try:
             result = func(*pos_args, **kwargs)
         except Exception as e:
             print(f"  Error: Failed to execute '{op_name}'. Error: {e}")
-            if args.stop_on_error: raise
+            if args.stop_on_error:
+                raise
             errors += 1
             continue
         # --- Benchmark the function ---
-        mean_time_us = benchmark_func(lambda: func(*pos_args, **kwargs), args.device, warmup=50, avg_steps=100)
+        mean_time_us = benchmark_func(
+            lambda: func(*pos_args, **kwargs), args.device, warmup=50, avg_steps=100
+        )
         print(f"  Average time taken: {mean_time_us:.2f} microseconds")
-        if 'count' in repro_info:
-            count_workload = repro_info['count']
+        if "count" in repro_info:
+            count_workload = repro_info["count"]
             total_time_us = mean_time_us * count_workload
             print(f"  Count in workload: {count_workload}")
             print(f"  Est time in workload: {total_time_us:.2f} microseconds")
         # --- Optionally sync again ---
-        if args.device == 'cuda': torch.cuda.synchronize()
+        if args.device == "cuda":
+            torch.cuda.synchronize()
 
         if args.verbose:
             print(f"  Successfully executed {op_name}.")
             # Optionally print result info (be careful with large tensors)
             if isinstance(result, torch.Tensor):
-                print(f"  Result: Tensor(shape={result.shape}, dtype={result.dtype}, device={result.device})")
-            elif isinstance(result, (list, tuple)) and any(isinstance(r, torch.Tensor) for r in result):
+                print(
+                    f"  Result: Tensor(shape={result.shape}, dtype={result.dtype}, device={result.device})"
+                )
+            elif isinstance(result, (list, tuple)) and any(
+                isinstance(r, torch.Tensor) for r in result
+            ):
                 for r in result:
                     if isinstance(r, torch.Tensor):
-                        print(f"  Result: Tensor(shape={r.shape}, dtype={r.dtype}, device={r.device})")
+                        print(
+                            f"  Result: Tensor(shape={r.shape}, dtype={r.dtype}, device={r.device})"
+                        )
                     else:
                         print(f"  Result: {r}")
             else:
@@ -132,13 +184,14 @@ def _get_args_kwargs_from_ir(event_replay_IR: dict[str, any], device: str = 'cud
     # --- Final Summary ---
     print("\n--- Replay Summary ---")
     print(f"Total operations in file: {len(repro_data_list)}")
-    if args.op_filter: print(f"Filter applied: '{args.op_filter}'")
+    if args.op_filter:
+        print(f"Filter applied: '{args.op_filter}'")
     print(f"Attempted replays: {replayed_count}")
     print(f"Successful replays: {replayed_count - errors}")
     print(f"Errors encountered: {errors}")
     print("----------------------")
 
     if errors > 0:
-        exit(1) # Exit with error code if any operation failed
+        exit(1)  # Exit with error code if any operation failed
     else:
-        exit(0)
+        exit(0)