Remove unused imports across the codebase

TraceLens/Agent/Analysis/category_analyses/cpu_idle_analysis.py

@@ -15,7 +15,6 @@
 """
 
 import argparse
-import json
 import os
 import sys
 import pandas as pd

TraceLens/Agent/Analysis/category_analyses/gemm_analysis.py

@@ -45,7 +45,6 @@ def extract_category_specific(ops_df, metadata) -> dict:
 
     missing_perf_model = 0
     if "TFLOPS/s_mean" in ops_df.columns:
-        import pandas as pd
 
         missing_perf_model = ops_df["TFLOPS/s_mean"].isna().sum()
 

TraceLens/Agent/Analysis/category_analyses/other_analysis.py

@@ -15,8 +15,6 @@
 import sys
 import os
 
-import pandas as pd
-
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from analysis_utils import (

TraceLens/Agent/Analysis/utils/orchestrator_prepare.py

@@ -18,7 +18,6 @@
 import sys
 import traceback
 from collections import defaultdict
-from typing import Any
 import pandas as pd
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

TraceLens/EventReplay/batched_replay.py

@@ -7,7 +7,6 @@
 # run_repro.py
 import json
 import argparse
-import warnings
 import torch
 from utils import TensorCfg, build_tensor, benchmark_func
 

TraceLens/EventReplay/event_replay.py

@@ -11,7 +11,6 @@
 from typing import Dict, Any, List, Optional, Tuple
 import re
 import warnings
-import time
 
 from .utils import (
     _get_torch_or_raise,

TraceLens/EventReplay/utils.py

@@ -4,7 +4,7 @@
 # See LICENSE for license information.
 ###############################################################################
 
-from typing import List, Dict, Tuple, Any
+from typing import List, Any
 import time
 
 _torch_module = None

TraceLens/NcclAnalyser/nccl_analyser.py

@@ -5,17 +5,14 @@
 ###############################################################################
 
 import ast
-import gzip
 import os
-import json
 import logging
 import warnings
 
 import pandas as pd
 from concurrent.futures import ProcessPoolExecutor, as_completed
 
 from ..util import DataLoader
-from ..util import DEFAULT_CUSTOM_COLLECTIVE_PATTERNS
 from ..util import TraceEventUtils
 
 

TraceLens/NcclAnalyser/util/node_rank_to_protobuf_file_mapping.py

@@ -6,7 +6,6 @@
 
 from pathlib import Path
 import re
-from collections import defaultdict
 
 
 def get_node_rank_protobuf_mapping(traces_folder, pattern="*.xplane.pb"):

TraceLens/PerfModel/__init__.py

@@ -5,9 +5,5 @@
 ###############################################################################
 
 from .perf_model import *  # Import everything from perf_model
-from .torch_op_mapping import (
-    op_to_perf_model_class_map,
-    resolve_perf_model_class,
-)
 
 __all__ = [name for name in dir() if not name.startswith("_")]

TraceLens/PerfModel/extensions/custom_collectives_perf_model_extensions.py

@@ -10,7 +10,6 @@
 
 from math import prod
 from TraceLens.PerfModel.utils import name2bpe
-from TraceLens.PerfModel.perf_model import RMSNorm
 
 
 class CustomCollective:

TraceLens/PerfModel/extensions/perf_model_extensions.py

@@ -9,7 +9,6 @@
 """
 
 from TraceLens.PerfModel.utils import torch_dtype_map, name2bpe
-import re
 from TraceLens.PerfModel.perf_model import (
     GEMM,
     BinaryElementwise,

TraceLens/PerfModel/perf_model.py

@@ -15,7 +15,7 @@
 
 from .kernel_name_parser import gemm_name_parser
 
-from .utils import name2bpe, parse_bool, simulation_dtype_map, torch_dtype_map
+from .utils import name2bpe, parse_bool, torch_dtype_map
 
 
 # 1. GEMM

TraceLens/PerfModel/run_perf_model.py

@@ -5,8 +5,6 @@
 ###############################################################################
 
 import argparse
-import os
-import sys
 
 from TraceLens.PerfModel.perf_model import GEMM, SDPA, simulation_dtype_map
 

TraceLens/PerfModel/utils.py

@@ -8,8 +8,6 @@
 Utils. for perf. model.
 """
 
-import os
-
 
 def add_simulation_time_columns(
     dict_metrics,

TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py

@@ -4,6 +4,7 @@
 # See LICENSE for license information.
 ###############################################################################
 
+import importlib.util
 import os
 import re
 import argparse
@@ -264,9 +265,7 @@ def generate_collective_report(
             print(f"DataFrame '{sheet_name}' written to {csv_path}")
 
     if output_xlsx_path:
-        try:
-            import openpyxl
-        except (ImportError, ModuleNotFoundError):
+        if importlib.util.find_spec("openpyxl") is None:
             print("Error importing openpyxl")
             request_install("openpyxl")
 

TraceLens/Reporting/generate_perf_report_jax.py

@@ -4,7 +4,10 @@
 # See LICENSE for license information.
 ###############################################################################
 
-import argparse, os, sys
+import argparse
+import importlib.util
+import os
+import sys
 from typing import Optional, Dict
 import pandas as pd
 import logging
@@ -16,8 +19,7 @@
     format="[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
 )
 
-from TraceLens.PerfModel import jax_op_mapping
-from TraceLens.TreePerf import TreePerfAnalyzer, JaxTreePerfAnalyzer
+from TraceLens.TreePerf import JaxTreePerfAnalyzer
 from TraceLens.Reporting.reporting_utils import (
     add_gpu_arch_cli_args,
     request_install,
@@ -175,10 +177,8 @@ def generate_perf_report_jax(
             # split input path at 'xplane.pb' and take the first part and append '.xlsx'
             base_path = profile_path.rsplit(".xplane.pb", 1)[0]
             output_xlsx_path = base_path + "_perf_report.xlsx"
-        try:
-            import openpyxl
-        except (ImportError, ModuleNotFoundError) as e:
-            print(f"Error importing openpyxl: {e}")
+        if importlib.util.find_spec("openpyxl") is None:
+            print("Error importing openpyxl")
             request_install("openpyxl")
 
         with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer:

TraceLens/Reporting/generate_perf_report_jax_analysis.py

@@ -5,10 +5,7 @@
 ###############################################################################
 
 import argparse
-import json
-import os
 import sys
-import pandas as pd
 
 from pathlib import Path
 

TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py

@@ -11,6 +11,7 @@
 API↔kernel correlation (see generate_perf_report_pftrace_hip_api for that).
 """
 
+import importlib.util
 import os
 import argparse
 import sys
@@ -171,11 +172,9 @@ def generate_perf_report_pftrace_hip_activity(
                 base = base.with_suffix("")
             output_xlsx_path = str(base) + "_pftrace_activity_report.xlsx"
         logger.info("Writing Excel to: %s", output_xlsx_path)
-        try:
-            import openpyxl
-        except (ImportError, ModuleNotFoundError) as e:
-            logger.error("openpyxl required: %s", e)
-            raise
+        if importlib.util.find_spec("openpyxl") is None:
+            logger.error("openpyxl required for Excel output")
+            raise ImportError("openpyxl is required for Excel output")
         with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer:
             for sheet_name, df in dict_name2df.items():
                 sn = sheet_name[:31]

TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py

@@ -4,6 +4,7 @@
 # See LICENSE for license information.
 ###############################################################################
 
+import importlib.util
 import os
 import re
 import argparse
@@ -117,13 +118,11 @@ def generate_perf_report_pftrace_hip_api(
             output_xlsx_path = str(base) + "_pftrace_hip_api_report.xlsx"
 
         logger.info(f"Writing Excel file to: {output_xlsx_path}")
-        try:
-            import openpyxl
-        except (ImportError, ModuleNotFoundError) as e:
+        if importlib.util.find_spec("openpyxl") is None:
             logger.error(
-                f"Error importing openpyxl: {e}. Please install: pip install openpyxl"
+                "Error importing openpyxl. Please install: pip install openpyxl"
             )
-            raise
+            raise ImportError("openpyxl is required for Excel output")
 
         with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer:
             for sheet_name, df in dict_name2df.items():

TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py

@@ -11,6 +11,7 @@
 Uses shared pftrace_utils (traceconv) and PftraceParser.
 """
 
+import importlib.util
 import os
 import argparse
 import sys
@@ -163,13 +164,9 @@ def generate_perf_report_pftrace_memory_copy(
                 base = base.with_suffix("")
             output_xlsx_path = str(base) + "_pftrace_memory_copy_report.xlsx"
         logger.info("Writing Excel file to: %s", output_xlsx_path)
-        try:
-            import openpyxl  # noqa: F401
-        except (ImportError, ModuleNotFoundError) as e:
-            logger.error(
-                "openpyxl required for Excel output: %s. pip install openpyxl", e
-            )
-            raise
+        if importlib.util.find_spec("openpyxl") is None:
+            logger.error("openpyxl required for Excel output. pip install openpyxl")
+            raise ImportError("openpyxl is required for Excel output")
         with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer:
             for sheet_name, df in dfs.items():
                 df.to_excel(writer, sheet_name=sheet_name, index=False)

TraceLens/Reporting/generate_perf_report_pytorch.py

@@ -7,18 +7,15 @@
 import argparse
 import ast
 import importlib.util
-import json
 import os
 import re
-import subprocess
-import sys
 import warnings
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional
 
 import numpy as np
 import pandas as pd
 
-from TraceLens import NcclAnalyser, TraceToTree, TraceDiff, TreePerfAnalyzer
+from TraceLens import NcclAnalyser, TraceDiff, TreePerfAnalyzer
 from TraceLens.PerfModel.torch_op_mapping import build_sheet_category_to_op_names
 from TraceLens.Reporting.reporting_utils import (
     add_gpu_arch_cli_args,
@@ -1028,10 +1025,8 @@ def _launcher_category(name):
         if output_xlsx_path is None:
             base_path = profile_json_path.rsplit(".json", 1)[0]
             output_xlsx_path = base_path + "_perf_report.xlsx"
-        try:
-            import openpyxl
-        except (ImportError, ModuleNotFoundError) as e:
-            print(f"Error importing openpyxl: {e}")
+        if importlib.util.find_spec("openpyxl") is None:
+            print("Error importing openpyxl")
             request_install("openpyxl")
 
         with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer:

TraceLens/Reporting/generate_perf_report_pytorch_inference.py

@@ -8,12 +8,10 @@
 import importlib.util
 import json
 import os
-import subprocess
 import sys
 import warnings
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional
 
-from tqdm import tqdm
 
 import numpy as np
 import pandas as pd
@@ -40,8 +38,6 @@
     merge_capture_trace_into_graph,
 )
 
-import TraceLens
-
 
 def perf_report_sanity_check(
     events,
@@ -1180,10 +1176,8 @@ def _launcher_category(name):
         if output_xlsx_path is None:
             base_path = profile_json_path.rsplit(".json", 1)[0]
             output_xlsx_path = base_path + "_perf_report.xlsx"
-        try:
-            import openpyxl
-        except (ImportError, ModuleNotFoundError) as e:
-            print(f"Error importing openpyxl: {e}")
+        if importlib.util.find_spec("openpyxl") is None:
+            print("Error importing openpyxl")
             request_install("openpyxl")
 
         with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer:

TraceLens/Reporting/generate_perf_report_rocprof.py

@@ -4,6 +4,7 @@
 # See LICENSE for license information.
 ###############################################################################
 
+import importlib.util
 import os
 import argparse
 import sys
@@ -150,13 +151,11 @@ def generate_perf_report_rocprof(
                 output_xlsx_path = base_path + "_perf_report.xlsx"
 
         logger.info(f"Writing Excel file to: {output_xlsx_path}")
-        try:
-            import openpyxl
-        except (ImportError, ModuleNotFoundError) as e:
+        if importlib.util.find_spec("openpyxl") is None:
             logger.error(
-                f"Error importing openpyxl: {e}. Please install it with: pip install openpyxl"
+                "Error importing openpyxl. Please install it with: pip install openpyxl"
             )
-            raise
+            raise ImportError("openpyxl is required for Excel output")
 
         with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer:
             for sheet_name, df in dict_name2df.items():

TraceLens/Reporting/pftrace_hip_activity_analysis.py

@@ -15,7 +15,6 @@
 import re
 from collections import defaultdict
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 from ..util import TraceEventUtils

TraceLens/Reporting/rocprof_analysis.py

@@ -7,7 +7,7 @@
 from ..util import TraceEventUtils
 import pandas as pd
 import numpy as np
-from typing import List, Dict, Optional
+from typing import List, Optional
 import logging
 
 logger = logging.getLogger(__name__)

TraceLens/Reporting/tracediff_comparison_extension.py

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Dict, List, Optional, Set
 
 import numpy as np
 import pandas as pd