diff --git a/TraceLens/Reporting/compare_perf_reports_pytorch.py b/TraceLens/Reporting/compare_perf_reports_pytorch.py index 8595dea35..f7fc00871 100644 --- a/TraceLens/Reporting/compare_perf_reports_pytorch.py +++ b/TraceLens/Reporting/compare_perf_reports_pytorch.py @@ -11,7 +11,8 @@ from typing import Dict, List, Optional, Sequence import pandas as pd -from openpyxl.utils import get_column_letter + +from TraceLens.Reporting.reporting_utils import write_report_outputs # ────────────────────────────────────────────────────────────────────────────── # Configuration @@ -588,31 +589,12 @@ def generate_compare_perf_reports_pytorch( cols_to_hide_xl[sheet_name] = cols_to_hide # ── Write workbook ──────────────────────────────────────────────────────── - if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in results.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print( - f"Wrote '{sheet_name}.csv' with {len(df)} rows × {len(df.columns)} columns" - ) - - if output is not None: - with pd.ExcelWriter(output, engine="openpyxl") as xls: - for sheet_name, df in results.items(): - if df.empty: - print(f"Sheet '{sheet_name}' is empty (no matching rows)") - df.to_excel( - xls, sheet_name=sheet_name[:31], index=False - ) # Excel 31-char limit - for col in cols_to_hide_xl.get(sheet_name, []): - col_idx = df.columns.get_loc(col) + 1 - col_letter = get_column_letter(col_idx) - worksheet = xls.sheets[sheet_name[:31]] - worksheet.column_dimensions[col_letter].hidden = True - print( - f"Wrote sheet '{sheet_name}' with {len(df)} rows × {len(df.columns)} columns" - ) + write_report_outputs( + results, + xlsx_path=output, + csvs_dir=output_csvs_dir, + hide_columns=cols_to_hide_xl, + ) return results diff --git a/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py b/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py index a861f683d..f4a867b5f 100644 --- a/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py +++ b/TraceLens/Reporting/generate_multi_rank_collective_report_pytorch.py @@ -4,7 +4,6 @@ # See LICENSE for license information. ############################################################################### -import importlib.util import os import re import argparse @@ -16,7 +15,7 @@ from TraceLens.Reporting.reporting_utils import ( add_node_span_columns, detect_gpus_per_node, - request_install, + write_report_outputs, ) DEFAULT_RANK_REGEX = r"rank[\[\-_/]?(?P\d+)" @@ -257,25 +256,9 @@ def generate_collective_report( ) # Export DataFrames - if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in report_dfs.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - - if output_xlsx_path: - if importlib.util.find_spec("openpyxl") is None: - print("Error importing openpyxl") - request_install("openpyxl") - - print(f"Writing Excel report to {output_xlsx_path}...") - with pd.ExcelWriter(output_xlsx_path, engine="openpyxl") as writer: - for sheet_name, df in report_dfs.items(): - df.to_excel( - writer, sheet_name=sheet_name[:31], index=False - ) # Excel limits sheet names to 31 chars - print(f"Excel report successfully written to {output_xlsx_path}") + write_report_outputs( + report_dfs, xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir + ) return report_dfs diff --git a/TraceLens/Reporting/generate_perf_report_genesis.py b/TraceLens/Reporting/generate_perf_report_genesis.py index fc22dbd76..6882ed55c 100644 --- a/TraceLens/Reporting/generate_perf_report_genesis.py +++ b/TraceLens/Reporting/generate_perf_report_genesis.py @@ -36,6 +36,7 @@ pftrace_to_json, resolve_profile_json, ) +from TraceLens.Reporting.reporting_utils import write_report_outputs logging.basicConfig( stream=sys.stdout, @@ -45,30 +46,13 @@ logger = logging.getLogger(__name__) -def _safe_sheet(name: str, used: set) -> str: - base = name[:31] - n = 0 - while base in used: - n += 1 - suffix = f"_{n}" - base = name[: 31 - len(suffix)] + suffix - used.add(base) - return base - - def write_excel(path: Path, sections: Dict[str, Dict[str, pd.DataFrame]]) -> None: - used: set = set() - with pd.ExcelWriter(path, engine="openpyxl") as writer: - for prefix, dfs in sections.items(): - for sheet, df in dfs.items(): - if df is None or df.empty: - continue - # rocprof sheets use short names (no prefix); pftrace keeps prefix for clarity - sheet_label = sheet if prefix == "rocprof" else f"{prefix}_{sheet}" - df.to_excel( - writer, sheet_name=_safe_sheet(sheet_label, used), index=False - ) - logger.info("Wrote %s", path) + flat: Dict[str, pd.DataFrame] = {} + for prefix, dfs in sections.items(): + for sheet, df in dfs.items(): + label = sheet if prefix == "rocprof" else f"{prefix}_{sheet}" + flat[label] = df + write_report_outputs(flat, xlsx_path=str(path), skip_empty=True) def _rocprof_sheets_for_excel( diff --git a/TraceLens/Reporting/generate_perf_report_jax.py b/TraceLens/Reporting/generate_perf_report_jax.py index a0e2b81bf..f0bafdfb8 100644 --- a/TraceLens/Reporting/generate_perf_report_jax.py +++ b/TraceLens/Reporting/generate_perf_report_jax.py @@ -5,8 +5,6 @@ ############################################################################### import argparse -import importlib.util -import os import sys from typing import Optional, Dict import pandas as pd @@ -22,8 +20,8 @@ from TraceLens.TreePerf import JaxTreePerfAnalyzer from TraceLens.Reporting.reporting_utils import ( add_gpu_arch_cli_args, - request_install, resolve_gpu_arch, + write_report_outputs, ) from TraceLens.util import TraceEventUtils @@ -165,26 +163,14 @@ def generate_perf_report_jax( ) # Write all DataFrames to separate sheets in an Excel workbook - if output_csvs_dir: - # Ensure the output directory exists - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - else: - if output_xlsx_path is None: - # 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" - 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: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - print(f"DataFrames successfully written to {output_xlsx_path}") + if not output_csvs_dir and output_xlsx_path is None: + base_path = profile_path.rsplit(".xplane.pb", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path, + csvs_dir=output_csvs_dir, + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py b/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py index d9f14c9b0..044de15d2 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_hip_activity.py @@ -11,7 +11,6 @@ API↔kernel correlation (see generate_perf_report_pftrace_hip_api for that). """ -import importlib.util import os import argparse import sys @@ -29,7 +28,11 @@ logger = logging.getLogger(__name__) from TraceLens.util import PftraceParser -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) +from TraceLens.Reporting.reporting_utils import write_report_outputs from TraceLens.Reporting.pftrace_hip_activity_analysis import ( PftraceHipActivityAnalyzer, ns_to_ms, @@ -154,33 +157,15 @@ def generate_perf_report_pftrace_hip_activity( dict_name2df["hip_summary"] = analyzer.get_df_hip_summary() logger.info(" - hip_summary (%d rows)", len(dict_name2df["hip_summary"])) - if output_csvs_dir: - logger.info("Writing CSV files to: %s", output_csvs_dir) - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(" - %s.csv (%d rows)", sheet_name, len(df)) - else: - if output_xlsx_path is None: - base = Path(trace_path).resolve() - if base.suffix.lower() == ".pftrace": - base = base.with_suffix("") - elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): - base = base.parent / base.name.replace(".json.gz", "") - else: - base = base.with_suffix("") - output_xlsx_path = str(base) + "_pftrace_activity_report.xlsx" - logger.info("Writing Excel to: %s", output_xlsx_path) - 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] - df.to_excel(writer, sheet_name=sn, index=False) - logger.info(" - Sheet '%s' (%d rows)", sn, len(df)) - logger.info("Successfully written to %s", output_xlsx_path) + if not output_csvs_dir and output_xlsx_path is None: + output_xlsx_path = derive_pftrace_output_path( + trace_path, "_pftrace_activity_report.xlsx" + ) + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path, + csvs_dir=output_csvs_dir, + ) if output_md_path: logger.info("Writing Markdown to: %s", output_md_path) diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py index 278a1bc88..d8b0cedf7 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_hip_api.py @@ -4,12 +4,10 @@ # See LICENSE for license information. ############################################################################### -import importlib.util import os import re import argparse import sys -from pathlib import Path from typing import Optional, Dict import pandas as pd @@ -24,7 +22,11 @@ logger = logging.getLogger(__name__) from TraceLens.util import PftraceParser -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) +from TraceLens.Reporting.reporting_utils import write_report_outputs from TraceLens.Reporting.pftrace_hip_api_analysis import PftraceHipApiAnalyzer @@ -99,37 +101,15 @@ def generate_perf_report_pftrace_hip_api( ) # Write output - if output_csvs_dir: - logger.info(f"Writing CSV files to: {output_csvs_dir}") - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(f" - {sheet_name}.csv ({len(df)} rows)") - else: - if output_xlsx_path is None: - base = Path(trace_path).resolve() - if base.suffix.lower() == ".pftrace": - base = base.with_suffix("") - elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): - base = base.parent / base.name.replace(".json.gz", "") - else: - base = base.with_suffix("") - output_xlsx_path = str(base) + "_pftrace_hip_api_report.xlsx" - - logger.info(f"Writing Excel file to: {output_xlsx_path}") - if importlib.util.find_spec("openpyxl") is None: - logger.error( - "Error importing openpyxl. Please install: 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 dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - logger.info(f" - Sheet '{sheet_name}' ({len(df)} rows)") - - logger.info(f"Successfully written to {output_xlsx_path}") + if not output_csvs_dir and output_xlsx_path is None: + output_xlsx_path = derive_pftrace_output_path( + trace_path, "_pftrace_hip_api_report.xlsx" + ) + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path, + csvs_dir=output_csvs_dir, + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py index e612443b0..b8dcfffa0 100644 --- a/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py +++ b/TraceLens/Reporting/generate_perf_report_pftrace_memory_copy.py @@ -11,11 +11,9 @@ Uses shared pftrace_utils (traceconv) and PftraceParser. """ -import importlib.util import os import argparse import sys -from pathlib import Path from typing import Optional, Dict, List, Any, Tuple import pandas as pd @@ -29,7 +27,11 @@ logger = logging.getLogger(__name__) from TraceLens.util import PftraceParser -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) +from TraceLens.Reporting.reporting_utils import write_report_outputs # Event name substrings for direction (ROCm/Perfetto) NAME_HOST_TO_DEVICE = "MEMORY_COPY_HOST_TO_DEVICE" @@ -146,32 +148,15 @@ def generate_perf_report_pftrace_memory_copy( count_df = build_memory_copy_count_df(events) dfs = {"memory_copy_by_copy_bytes": count_df} - if output_csvs_dir: - logger.info("Writing CSV files to: %s", output_csvs_dir) - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dfs.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(" - %s.csv (%d rows)", sheet_name, len(df)) - else: - if output_xlsx_path is None: - base = Path(trace_path).resolve() - if base.suffix.lower() == ".pftrace": - base = base.with_suffix("") - elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): - base = base.parent / base.name.replace(".json.gz", "") - else: - base = base.with_suffix("") - output_xlsx_path = str(base) + "_pftrace_memory_copy_report.xlsx" - logger.info("Writing Excel file to: %s", output_xlsx_path) - 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) - logger.info(" - Sheet '%s' (%d rows)", sheet_name, len(df)) - logger.info("Successfully written to %s", output_xlsx_path) + if not output_csvs_dir and output_xlsx_path is None: + output_xlsx_path = derive_pftrace_output_path( + trace_path, "_pftrace_memory_copy_report.xlsx" + ) + write_report_outputs( + dfs, + xlsx_path=output_xlsx_path, + csvs_dir=output_csvs_dir, + ) return dfs diff --git a/TraceLens/Reporting/generate_perf_report_pytorch.py b/TraceLens/Reporting/generate_perf_report_pytorch.py index 16bb9b644..e1e6ca4d8 100644 --- a/TraceLens/Reporting/generate_perf_report_pytorch.py +++ b/TraceLens/Reporting/generate_perf_report_pytorch.py @@ -19,8 +19,8 @@ from TraceLens.PerfModel.torch_op_mapping import build_sheet_category_to_op_names from TraceLens.Reporting.reporting_utils import ( add_gpu_arch_cli_args, - request_install, resolve_gpu_arch, + write_report_outputs, ) _WRAPPER_FILE_PATTERNS = frozenset( @@ -1029,25 +1029,12 @@ def _launcher_category(name): print(f"Added {len(additional_dfs)} additional sheets from extension") # Write CSVs and/or Excel (independent options) - if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - - if output_xlsx_path is not None or output_csvs_dir is None: - if output_xlsx_path is None: - base_path = profile_json_path.rsplit(".json", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - 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: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - print(f"DataFrames successfully written to {output_xlsx_path}") + if output_xlsx_path is None and output_csvs_dir is None: + base_path = profile_json_path.rsplit(".json", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py b/TraceLens/Reporting/generate_perf_report_pytorch_inference.py index 905193a34..4f6e1a640 100644 --- a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py +++ b/TraceLens/Reporting/generate_perf_report_pytorch_inference.py @@ -25,8 +25,8 @@ from TraceLens.Reporting.generate_perf_report_pytorch import _find_entry_point from TraceLens.Reporting.reporting_utils import ( add_gpu_arch_cli_args, - request_install, resolve_gpu_arch, + write_report_outputs, ) from TraceLens.util import TraceEventUtils from TraceLens.TraceUtils.annotation_utils import ( @@ -1168,25 +1168,12 @@ def _launcher_category(name): print(f"Added {len(additional_dfs)} additional sheets from extension") # Write CSVs and/or Excel (independent options) - if output_csvs_dir: - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - print(f"DataFrame '{sheet_name}' written to {csv_path}") - - if output_xlsx_path is not None or output_csvs_dir is None: - if output_xlsx_path is None: - base_path = profile_json_path.rsplit(".json", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - 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: - for sheet_name, df in dict_name2df.items(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - print(f"DataFrames successfully written to {output_xlsx_path}") + if output_xlsx_path is None and output_csvs_dir is None: + base_path = profile_json_path.rsplit(".json", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, xlsx_path=output_xlsx_path, csvs_dir=output_csvs_dir + ) return dict_name2df diff --git a/TraceLens/Reporting/generate_perf_report_rocprof.py b/TraceLens/Reporting/generate_perf_report_rocprof.py index d1835f341..94696bdc1 100644 --- a/TraceLens/Reporting/generate_perf_report_rocprof.py +++ b/TraceLens/Reporting/generate_perf_report_rocprof.py @@ -4,7 +4,6 @@ # See LICENSE for license information. ############################################################################### -import importlib.util import os import argparse import sys @@ -22,6 +21,7 @@ from TraceLens.util import RocprofParser from TraceLens.Reporting.rocprof_analysis import RocprofAnalyzer +from TraceLens.Reporting.reporting_utils import write_report_outputs def generate_perf_report_rocprof( @@ -132,37 +132,19 @@ def generate_perf_report_rocprof( ) # 5. Write output - if output_csvs_dir: - logger.info(f"Writing CSV files to: {output_csvs_dir}") - os.makedirs(output_csvs_dir, exist_ok=True) - for sheet_name, df in dict_name2df.items(): - csv_path = os.path.join(output_csvs_dir, f"{sheet_name}.csv") - df.to_csv(csv_path, index=False) - logger.info(f" - {sheet_name}.csv ({len(df)} rows)") - else: - if output_xlsx_path is None: - # Auto-generate output filename - if profile_json_path.endswith("_results.json"): - output_xlsx_path = profile_json_path.replace( - "_results.json", "_perf_report.xlsx" - ) - else: - base_path = profile_json_path.rsplit(".json", 1)[0] - output_xlsx_path = base_path + "_perf_report.xlsx" - - logger.info(f"Writing Excel file to: {output_xlsx_path}") - if importlib.util.find_spec("openpyxl") is None: - logger.error( - "Error importing openpyxl. Please install it with: pip install openpyxl" + if not output_csvs_dir and output_xlsx_path is None: + if profile_json_path.endswith("_results.json"): + output_xlsx_path = profile_json_path.replace( + "_results.json", "_perf_report.xlsx" ) - 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(): - df.to_excel(writer, sheet_name=sheet_name, index=False) - logger.info(f" - Sheet '{sheet_name}' ({len(df)} rows)") - - logger.info(f"Successfully written to {output_xlsx_path}") + else: + base_path = profile_json_path.rsplit(".json", 1)[0] + output_xlsx_path = base_path + "_perf_report.xlsx" + write_report_outputs( + dict_name2df, + xlsx_path=output_xlsx_path, + csvs_dir=output_csvs_dir, + ) return dict_name2df diff --git a/TraceLens/Reporting/pftrace_utils.py b/TraceLens/Reporting/pftrace_utils.py index b01cd95c7..5906bf121 100644 --- a/TraceLens/Reporting/pftrace_utils.py +++ b/TraceLens/Reporting/pftrace_utils.py @@ -86,3 +86,20 @@ def ensure_trace_json(trace_path: str, traceconv_path: Optional[str] = None) -> raise ValueError( f"Unsupported trace format: {trace_path}. Use .json, .json.gz, or .pftrace." ) + + +def derive_pftrace_output_path(trace_path: str, report_suffix: str) -> str: + """Derive a default output xlsx path from a pftrace input path. + + Strips ``.pftrace``, ``.json.gz``, or other extensions from + *trace_path* and appends *report_suffix* (e.g. + ``"_pftrace_activity_report.xlsx"``). + """ + base = Path(trace_path).resolve() + if base.suffix.lower() == ".pftrace": + base = base.with_suffix("") + elif base.suffix.lower() == ".gz" and base.name.endswith(".json.gz"): + base = base.parent / base.name.replace(".json.gz", "") + else: + base = base.with_suffix("") + return str(base) + report_suffix diff --git a/TraceLens/Reporting/reporting_utils.py b/TraceLens/Reporting/reporting_utils.py index 10ddc520a..6b06fb29a 100644 --- a/TraceLens/Reporting/reporting_utils.py +++ b/TraceLens/Reporting/reporting_utils.py @@ -8,6 +8,7 @@ import ast import json import logging +import os import re import subprocess import sys @@ -15,6 +16,7 @@ from typing import Dict, List, Optional, Union import pandas as pd +from openpyxl.utils import get_column_letter logger = logging.getLogger(__name__) @@ -124,6 +126,75 @@ def export_data_df( data_df.to_csv(output_path, index=False) +def _safe_sheet_name(name: str, used: set) -> str: + """Truncate *name* to Excel's 31-char limit with collision avoidance. + + If the truncated name already appears in *used*, a numeric suffix + (``_1``, ``_2``, ...) is appended while staying within the limit. + The final name is added to *used* before returning. + """ + base = name[:31] + n = 0 + while base in used: + n += 1 + suffix = f"_{n}" + base = name[: 31 - len(suffix)] + suffix + used.add(base) + return base + + +def write_report_outputs( + dfs: Dict[str, pd.DataFrame], + *, + xlsx_path: Optional[str] = None, + csvs_dir: Optional[str] = None, + hide_columns: Optional[Dict[str, List[str]]] = None, + skip_empty: bool = False, +) -> None: + """Write report DataFrames to CSV files and/or an Excel workbook. + + Sheet names are truncated to 31 characters (Excel limit) with + collision-safe suffixes. + + Args: + dfs: Mapping of sheet name -> DataFrame. + xlsx_path: If set, write an ``.xlsx`` workbook here. + csvs_dir: If set, write one CSV per DataFrame into this directory. + hide_columns: Optional mapping of sheet name -> column names to + hide in the Excel output. Hidden columns stay in the file; + names absent from a DataFrame are ignored. + skip_empty: If True, ``None`` or empty DataFrames are omitted from + all outputs. + """ + if skip_empty: + dfs = {name: df for name, df in dfs.items() if df is not None and not df.empty} + + if csvs_dir: + os.makedirs(csvs_dir, exist_ok=True) + for sheet_name, df in dfs.items(): + csv_path = os.path.join(csvs_dir, f"{sheet_name}.csv") + df.to_csv(csv_path, index=False) + logger.info("Wrote %s (%d rows)", csv_path, len(df)) + + if xlsx_path: + hide_columns = hide_columns or {} + used: set = set() + with pd.ExcelWriter(xlsx_path, engine="openpyxl") as writer: + for sheet_name, df in dfs.items(): + safe = _safe_sheet_name(sheet_name, used) + df.to_excel(writer, sheet_name=safe, index=False) + cols_to_hide = hide_columns.get(sheet_name, []) + if not cols_to_hide: + continue + worksheet = writer.book.worksheets[-1] + for col in cols_to_hide: + if col not in df.columns: + continue + col_letter = get_column_letter(df.columns.get_loc(col) + 1) + worksheet.column_dimensions[col_letter].hidden = True + logger.info("Wrote %s", xlsx_path) + + def request_install(package_name): """ Prompts the user to install a Python package via pip. If the user agrees, attempts installation. diff --git a/tests/test_genesis.py b/tests/test_genesis.py index 2c6ec5fc0..fc1bdd2a1 100644 --- a/tests/test_genesis.py +++ b/tests/test_genesis.py @@ -40,7 +40,6 @@ from TraceLens.Reporting.generate_perf_report_genesis import ( _resolve_steady_state_fallback_s, _rocprof_sheets_for_excel, - _safe_sheet, write_excel, write_genesis_summary_md, ) @@ -776,46 +775,6 @@ def test_raises_when_no_data(self): resolve_profile_json(capture_dict, output_dir, include_api=False) -############################################################################### -# generate_perf_report_genesis — _safe_sheet -############################################################################### - - -class TestSafeSheet: - """Validate Excel sheet name deduplication and length limits.""" - - def test_no_collision(self): - used = set() - name = _safe_sheet("gpu_timeline", used) - assert name == "gpu_timeline" - assert "gpu_timeline" in used - - def test_collision_adds_suffix(self): - used = {"gpu_timeline"} - name = _safe_sheet("gpu_timeline", used) - assert name == "gpu_timeline_1" - assert "gpu_timeline_1" in used - - def test_multiple_collisions(self): - used = {"test_sheet", "test_sheet_1", "test_sheet_2"} - name = _safe_sheet("test_sheet", used) - assert name == "test_sheet_3" - - def test_truncates_to_31_chars(self): - used = set() - long_name = "a" * 50 - name = _safe_sheet(long_name, used) - assert len(name) <= 31 - - def test_truncation_with_collision(self): - long_name = "a" * 31 - used = {long_name} - name = _safe_sheet(long_name, used) - assert len(name) <= 31 - assert name != long_name - assert name.endswith("_1") - - ############################################################################### # generate_perf_report_genesis — _rocprof_sheets_for_excel ############################################################################### diff --git a/tests/test_inference_perf_report.py b/tests/test_inference_perf_report.py index 8fe3ae9f0..2dba090e8 100644 --- a/tests/test_inference_perf_report.py +++ b/tests/test_inference_perf_report.py @@ -11,7 +11,7 @@ # - Optionally capture_traces/ (graph capture mode) # - Optionally gpu_arch.json -import os, numpy as np, pandas as pd, pytest, ast, re, gzip, json, glob +import os, shutil, numpy as np, pandas as pd, pytest, ast, re, gzip, json, glob from pandas.api.types import is_float_dtype from TraceLens.Reporting.generate_perf_report_pytorch_inference import ( classify_graph_capture_trace, @@ -318,6 +318,22 @@ def test_inference_perf_report( ) +def test_inference_perf_report_default_output_path(tmp_path): + """No output_xlsx_path/output_csvs_dir given -> auto-derive next to the + input trace, replacing the '.json' suffix.""" + cases = find_inference_test_cases() + if not cases: + pytest.skip("No inference trace fixtures found") + dirpath, trace_gz, _capture_folder, _gpu_arch_path = cases[0].values + profile_path = shutil.copy(os.path.join(dirpath, trace_gz), tmp_path) + + result = generate_perf_report_pytorch(profile_json_path=profile_path) + + expected_xlsx = profile_path.rsplit(".json", 1)[0] + "_perf_report.xlsx" + assert os.path.exists(expected_xlsx) + assert isinstance(result, dict) + + # --------------------------------------------------------------------------- # Capture merge validation: verifies that kernel timing from the replay trace # and input args/call stacks from the capture trace are correctly reflected diff --git a/tests/test_jax_perf_report.py b/tests/test_jax_perf_report.py index 707a1170b..0d16220a4 100644 --- a/tests/test_jax_perf_report.py +++ b/tests/test_jax_perf_report.py @@ -280,6 +280,20 @@ def test_jax_llama_helpers(self, tmp_path): assert d_model == 4096 +def test_jax_report_default_output_path(tmp_path): + """No output_csvs_dir/output_xlsx_path given -> auto-derive next to the + input trace, replacing the '.xplane.pb' suffix.""" + src = os.path.join( + os.path.dirname(__file__), + "traces/mi300/jax_conv_minimal_legacy/chi-mi300x-013.ord.vultr.cpe.ice.amd.com.xplane.pb", + ) + trace = shutil.copy(src, tmp_path) + dict_name2df = generate_perf_report_jax(profile_path=trace) + expected_xlsx = trace.rsplit(".xplane.pb", 1)[0] + "_perf_report.xlsx" + assert os.path.exists(expected_xlsx) + assert isinstance(dict_name2df, dict) + + def test_jax_report_main(tmp_path): trace = os.path.join( os.path.dirname(__file__), diff --git a/tests/test_pftrace_memory_copy_report.py b/tests/test_pftrace_memory_copy_report.py index 82a02e695..ce1486a19 100644 --- a/tests/test_pftrace_memory_copy_report.py +++ b/tests/test_pftrace_memory_copy_report.py @@ -4,6 +4,8 @@ # See LICENSE for license information. ############################################################################### +from TraceLens.Reporting.pftrace_utils import derive_pftrace_output_path + import json import os import tempfile @@ -184,3 +186,20 @@ def test_generate_csvs_dir(self): assert list(df.columns) == ["copy_bytes", "direction", "count"] finally: os.unlink(trace_path) + + def test_generate_default_output_path(self): + """No output_xlsx_path/output_csvs_dir given -> auto-derive next to trace_path.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump({"traceEvents": _make_memory_copy_events()}, f) + trace_path = f.name + expected_xlsx = derive_pftrace_output_path( + trace_path, "_pftrace_memory_copy_report.xlsx" + ) + try: + dfs = generate_perf_report_pftrace_memory_copy(trace_path=trace_path) + assert "memory_copy_by_copy_bytes" in dfs + assert os.path.isfile(expected_xlsx) + finally: + os.unlink(trace_path) + if os.path.isfile(expected_xlsx): + os.unlink(expected_xlsx) diff --git a/tests/test_reporting_utils.py b/tests/test_reporting_utils.py index c90a127de..69ca93690 100644 --- a/tests/test_reporting_utils.py +++ b/tests/test_reporting_utils.py @@ -11,12 +11,14 @@ from TraceLens.Reporting.reporting_utils import ( _node_span_for_pg, _parse_pg_ranks, + _safe_sheet_name, add_gpu_arch_cli_args, add_node_span_columns, detect_gpus_per_node, export_data_df, request_install, resolve_gpu_arch, + write_report_outputs, ) from TraceLens.Agent.Analysis.category_analyses import ( analysis_utils as au, @@ -120,7 +122,10 @@ _cleanup_work_dir, generate_perf_report_genesis, ) -from TraceLens.Reporting.pftrace_utils import ensure_trace_json +from TraceLens.Reporting.pftrace_utils import ( + derive_pftrace_output_path, + ensure_trace_json, +) GPU_ONLY_TRACE = os.path.join( os.path.dirname(__file__), @@ -3165,3 +3170,194 @@ def test_pytorch_report_extension_and_arch(self, tmp_path): include_call_stack=True, ) assert (tmp_path / "ext_out" / "gpu_timeline.csv").exists() + + +############################################################################### +# reporting_utils — _safe_sheet_name +############################################################################### + + +class TestSafeSheetName: + """Validate Excel sheet name deduplication and length limits.""" + + def test_no_collision(self): + used = set() + name = _safe_sheet_name("gpu_timeline", used) + assert name == "gpu_timeline" + assert "gpu_timeline" in used + + def test_collision_adds_suffix(self): + used = {"gpu_timeline"} + name = _safe_sheet_name("gpu_timeline", used) + assert name == "gpu_timeline_1" + assert "gpu_timeline_1" in used + + def test_multiple_collisions(self): + used = {"test_sheet", "test_sheet_1", "test_sheet_2"} + name = _safe_sheet_name("test_sheet", used) + assert name == "test_sheet_3" + + def test_truncates_to_31_chars(self): + used = set() + long_name = "a" * 50 + name = _safe_sheet_name(long_name, used) + assert len(name) <= 31 + + def test_truncation_with_collision(self): + long_name = "a" * 31 + used = {long_name} + name = _safe_sheet_name(long_name, used) + assert len(name) <= 31 + assert name != long_name + assert name.endswith("_1") + + +############################################################################### +# reporting_utils — write_report_outputs +############################################################################### + + +def _read_sheets(xlsx_path): + """Return {sheet_name: DataFrame} for an .xlsx file.""" + return pd.read_excel(xlsx_path, sheet_name=None) + + +class TestWriteReportOutputs: + def _dfs(self): + return { + "alpha": pd.DataFrame({"a": [1, 2], "b": [3, 4]}), + "beta": pd.DataFrame({"x": [5]}), + } + + def test_csvs_only(self, tmp_path): + out = tmp_path / "csvs" + write_report_outputs(self._dfs(), csvs_dir=str(out)) + assert (out / "alpha.csv").exists() + assert (out / "beta.csv").exists() + # No xlsx requested -> none written + assert not list(tmp_path.glob("*.xlsx")) + pd.testing.assert_frame_equal( + pd.read_csv(out / "alpha.csv"), self._dfs()["alpha"] + ) + + def test_xlsx_only(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + write_report_outputs(self._dfs(), xlsx_path=str(xlsx)) + assert xlsx.exists() + assert not list(tmp_path.glob("*.csv")) + sheets = _read_sheets(xlsx) + assert set(sheets) == {"alpha", "beta"} + + def test_both_outputs_written(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + csvs = tmp_path / "csvs" + write_report_outputs(self._dfs(), xlsx_path=str(xlsx), csvs_dir=str(csvs)) + assert xlsx.exists() + assert (csvs / "alpha.csv").exists() + assert (csvs / "beta.csv").exists() + + def test_neither_output_is_noop(self, tmp_path): + write_report_outputs(self._dfs()) + assert not list(tmp_path.iterdir()) + + def test_long_sheet_names_truncated_and_deduped(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + dfs = { + "a" * 40: pd.DataFrame({"c": [1]}), + "a" * 45: pd.DataFrame({"c": [2]}), # truncates to same 31 chars -> deduped + } + write_report_outputs(dfs, xlsx_path=str(xlsx)) + names = list(_read_sheets(xlsx)) + assert all(len(n) <= 31 for n in names) + assert len(names) == 2 # no collision -> both sheets present + + def test_skip_empty_drops_empty_and_none(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + csvs = tmp_path / "csvs" + dfs = { + "keep": pd.DataFrame({"a": [1]}), + "empty": pd.DataFrame(), + "none": None, + } + write_report_outputs( + dfs, xlsx_path=str(xlsx), csvs_dir=str(csvs), skip_empty=True + ) + assert set(_read_sheets(xlsx)) == {"keep"} + assert (csvs / "keep.csv").exists() + assert not (csvs / "empty.csv").exists() + assert not (csvs / "none.csv").exists() + + def test_skip_empty_default_keeps_empty(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + dfs = { + "keep": pd.DataFrame({"a": [1]}), + "empty": pd.DataFrame({"a": []}), + } + write_report_outputs(dfs, xlsx_path=str(xlsx)) + assert set(_read_sheets(xlsx)) == {"keep", "empty"} + + def test_hide_columns_hides_and_keeps_data(self, tmp_path): + from openpyxl import load_workbook + from openpyxl.utils import get_column_letter + + xlsx = tmp_path / "report.xlsx" + df = pd.DataFrame({"keep": [1], "hide_me": [2], "also_keep": [3]}) + write_report_outputs( + {"data": df}, xlsx_path=str(xlsx), hide_columns={"data": ["hide_me"]} + ) + ws = load_workbook(xlsx)["data"] + hidden_col = get_column_letter(df.columns.get_loc("hide_me") + 1) # "B" + assert ws.column_dimensions[hidden_col].hidden is True + # Non-hidden columns are not hidden, and the data is still present + keep_col = get_column_letter(df.columns.get_loc("keep") + 1) + assert ws.column_dimensions[keep_col].hidden is False + assert set(_read_sheets(xlsx)["data"].columns) == { + "keep", + "hide_me", + "also_keep", + } + + def test_hide_columns_ignores_missing_column(self, tmp_path): + xlsx = tmp_path / "report.xlsx" + df = pd.DataFrame({"a": [1]}) + # Should not raise even though "nonexistent" isn't a column + write_report_outputs( + {"data": df}, xlsx_path=str(xlsx), hide_columns={"data": ["nonexistent"]} + ) + assert xlsx.exists() + + def test_sheet_named_like_openpyxl_default(self, tmp_path): + # "sheet" collides case-insensitively with openpyxl's default "Sheet", + # so openpyxl renames it; the worksheet lookup must not rely on the + # requested name. hide_columns exercises that lookup path. + xlsx = tmp_path / "report.xlsx" + df = pd.DataFrame({"a": [1], "b": [2]}) + write_report_outputs( + {"sheet": df}, xlsx_path=str(xlsx), hide_columns={"sheet": ["b"]} + ) + assert xlsx.exists() + + +############################################################################### +# pftrace_utils — derive_pftrace_output_path +############################################################################### + + +class TestDerivePftraceOutputPath: + def test_pftrace_suffix(self): + assert ( + derive_pftrace_output_path("/tmp/trace.pftrace", "_activity_report.xlsx") + == "/tmp/trace_activity_report.xlsx" + ) + + def test_json_gz_suffix(self): + assert ( + derive_pftrace_output_path("/tmp/trace.json.gz", "_hip_api_report.xlsx") + == "/tmp/trace_hip_api_report.xlsx" + ) + + def test_plain_json_suffix(self): + assert ( + derive_pftrace_output_path("/tmp/trace.json", "_memory_copy_report.xlsx") + == "/tmp/trace_memory_copy_report.xlsx" + ) diff --git a/tests/test_rocprof_perf_report.py b/tests/test_rocprof_perf_report.py index fb7ed3eda..ec2b4febb 100644 --- a/tests/test_rocprof_perf_report.py +++ b/tests/test_rocprof_perf_report.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -import os, tempfile, pandas as pd, pytest, importlib, sys +import os, tempfile, gzip, shutil, pandas as pd, pytest, importlib, sys from TraceLens.Reporting.generate_perf_report_rocprof import ( generate_perf_report_rocprof, ) @@ -283,6 +283,34 @@ def test_generate_with_all_options(self, rocprof_file): assert len(dfs["kernel_details"]) <= 50 +def _copy_decompressed(src_gz, dst): + with gzip.open(src_gz, "rb") as fin, open(dst, "wb") as fout: + shutil.copyfileobj(fin, fout) + return str(dst) + + +def test_generate_default_output_path_for_results_json(tmp_path): + """No output_xlsx_path/output_csvs_dir given + filename ends with + '_results.json' -> auto-derive via string replace.""" + profile_json_path = _copy_decompressed( + ROCprof_FILE, tmp_path / "trace_results.json" + ) + expected_xlsx = profile_json_path.replace("_results.json", "_perf_report.xlsx") + dfs = generate_perf_report_rocprof(profile_json_path=profile_json_path) + assert os.path.exists(expected_xlsx) + assert isinstance(dfs, dict) + + +def test_generate_default_output_path_for_generic_json(tmp_path): + """No output_xlsx_path/output_csvs_dir given + filename doesn't end with + '_results.json' -> auto-derive via rsplit on '.json'.""" + profile_json_path = _copy_decompressed(ROCprof_FILE, tmp_path / "trace.json") + expected_xlsx = profile_json_path.rsplit(".json", 1)[0] + "_perf_report.xlsx" + dfs = generate_perf_report_rocprof(profile_json_path=profile_json_path) + assert os.path.exists(expected_xlsx) + assert isinstance(dfs, dict) + + if __name__ == "__main__": pytest.main([__file__, "-v"])