diff --git a/src/cloudai/report_generator/report_generator.py b/src/cloudai/report_generator/report_generator.py index 26f27f16f..17f834db5 100644 --- a/src/cloudai/report_generator/report_generator.py +++ b/src/cloudai/report_generator/report_generator.py @@ -71,4 +71,4 @@ def _generate_test_report(self, directory_path: str, test: Test) -> None: for subdir in os.listdir(directory_path): subdir_path = os.path.join(directory_path, subdir) if os.path.isdir(subdir_path) and test.test_template.can_handle_directory(subdir_path): - test.test_template.generate_report(subdir_path, test.sol) + test.test_template.generate_report(test.name, subdir_path, test.sol) diff --git a/src/cloudai/schema/core/strategy/report_generation_strategy.py b/src/cloudai/schema/core/strategy/report_generation_strategy.py index 2f3bab2b1..03e9067a8 100644 --- a/src/cloudai/schema/core/strategy/report_generation_strategy.py +++ b/src/cloudai/schema/core/strategy/report_generation_strategy.py @@ -38,11 +38,12 @@ def can_handle_directory(self, directory_path: str) -> bool: pass @abstractmethod - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: """ Generate a report from the directory. Args: + test_name (str): The name of the test. directory_path (str): Path to the directory. sol (Optional[float]): Speed-of-light performance for reference. """ diff --git a/src/cloudai/schema/core/test_template.py b/src/cloudai/schema/core/test_template.py index 358368a45..d5ab9c727 100644 --- a/src/cloudai/schema/core/test_template.py +++ b/src/cloudai/schema/core/test_template.py @@ -208,16 +208,17 @@ def can_handle_directory(self, directory_path: str) -> bool: else: return False - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: """ Generate a report from the directory. Args: + test_name (str): The name of the test. directory_path (str): Path to the directory. sol (Optional[float]): Speed-of-light performance for reference. """ if self.report_generation_strategy is not None: - return self.report_generation_strategy.generate_report(directory_path, sol) + return self.report_generation_strategy.generate_report(test_name, directory_path, sol) def grade(self, directory_path: str, ideal_perf: float) -> Optional[float]: """ diff --git a/src/cloudai/schema/test_template/chakra_replay/report_generation_strategy.py b/src/cloudai/schema/test_template/chakra_replay/report_generation_strategy.py index dae3380aa..9aca34a2f 100644 --- a/src/cloudai/schema/test_template/chakra_replay/report_generation_strategy.py +++ b/src/cloudai/schema/test_template/chakra_replay/report_generation_strategy.py @@ -47,7 +47,7 @@ def can_handle_directory(self, directory_path: str) -> bool: return True return False - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: stdout_path = os.path.join(directory_path, "stdout.txt") if not os.path.isfile(stdout_path): return diff --git a/src/cloudai/schema/test_template/jax_toolbox/report_generation_strategy.py b/src/cloudai/schema/test_template/jax_toolbox/report_generation_strategy.py index ea7ae94dd..ecb7eaba4 100644 --- a/src/cloudai/schema/test_template/jax_toolbox/report_generation_strategy.py +++ b/src/cloudai/schema/test_template/jax_toolbox/report_generation_strategy.py @@ -39,7 +39,7 @@ def can_handle_directory(self, directory_path: str) -> bool: return True return False - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: times = self._extract_times(directory_path) if times: stats = { diff --git a/src/cloudai/schema/test_template/nccl_test/report_generation_strategy.py b/src/cloudai/schema/test_template/nccl_test/report_generation_strategy.py index fa244398e..1d3a1d4bc 100644 --- a/src/cloudai/schema/test_template/nccl_test/report_generation_strategy.py +++ b/src/cloudai/schema/test_template/nccl_test/report_generation_strategy.py @@ -49,7 +49,7 @@ def can_handle_directory(self, directory_path: str) -> bool: return True return False - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: report_data, _ = self._parse_output(directory_path) if report_data: df = pd.DataFrame( @@ -76,7 +76,7 @@ def generate_report(self, directory_path: str, sol: Optional[float] = None) -> N df["Algbw (GB/s) In-place"] = df["Algbw (GB/s) In-place"].astype(float) df["Busbw (GB/s) In-place"] = df["Busbw (GB/s) In-place"].astype(float) df = add_human_readable_sizes(df, "Size (B)", "Size Human-readable") - self._generate_plots(df, directory_path, sol) + self._generate_plots(test_name, df, directory_path, sol) def _parse_output(self, directory_path: str) -> Tuple[List[List[str]], Optional[float]]: """ @@ -99,11 +99,12 @@ def _parse_output(self, directory_path: str) -> Tuple[List[List[str]], Optional[ return data, avg_bus_bw return [], None - def _generate_plots(self, df: pd.DataFrame, directory_path: str, sol: Optional[float]) -> None: + def _generate_plots(self, test_name: str, df: pd.DataFrame, directory_path: str, sol: Optional[float]) -> None: """ Create and saves plots to visualize NCCL test metrics. Args: + test_name (str): The name of the test. df (pd.DataFrame): DataFrame containing the NCCL test data. directory_path (str): Output directory path for saving the plots. sol (Optional[float]): Speed-of-light performance for reference. @@ -115,7 +116,7 @@ def _generate_plots(self, df: pd.DataFrame, directory_path: str, sol: Optional[f ] for col_name, color, title in line_plots: report_tool.add_log_x_linear_y_single_line_plot( - title=title, + title=f"{test_name} {title}", x_column="Size (B)", y_column=col_name, x_axis_label="Message Size", @@ -127,7 +128,7 @@ def _generate_plots(self, df: pd.DataFrame, directory_path: str, sol: Optional[f combined_columns = [("Busbw (GB/s) Out-of-place", "blue"), ("Busbw (GB/s) In-place", "green")] report_tool.add_log_x_linear_y_multi_line_plot( - title="Combined Bus Bandwidth", + title=f"{test_name} Combined Bus Bandwidth", x_column="Size (B)", y_columns=combined_columns, x_axis_label="Message Size", diff --git a/src/cloudai/schema/test_template/nemo_launcher/report_generation_strategy.py b/src/cloudai/schema/test_template/nemo_launcher/report_generation_strategy.py index 44cab60f5..8ca21443f 100644 --- a/src/cloudai/schema/test_template/nemo_launcher/report_generation_strategy.py +++ b/src/cloudai/schema/test_template/nemo_launcher/report_generation_strategy.py @@ -43,7 +43,7 @@ def can_handle_directory(self, directory_path: str) -> bool: return True return False - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: tags = ["train_step_timing"] data_reader = TensorBoardDataReader(directory_path) report_tool = BokehReportTool(directory_path) diff --git a/src/cloudai/schema/test_template/sleep/report_generation_strategy.py b/src/cloudai/schema/test_template/sleep/report_generation_strategy.py index 4a5285172..f428ef92a 100644 --- a/src/cloudai/schema/test_template/sleep/report_generation_strategy.py +++ b/src/cloudai/schema/test_template/sleep/report_generation_strategy.py @@ -27,5 +27,5 @@ class SleepReportGenerationStrategy(ReportGenerationStrategy): def can_handle_directory(self, directory_path: str) -> bool: return False - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: pass diff --git a/src/cloudai/schema/test_template/ucc_test/report_generation_strategy.py b/src/cloudai/schema/test_template/ucc_test/report_generation_strategy.py index 59420ec7f..6adafd048 100644 --- a/src/cloudai/schema/test_template/ucc_test/report_generation_strategy.py +++ b/src/cloudai/schema/test_template/ucc_test/report_generation_strategy.py @@ -46,7 +46,7 @@ def can_handle_directory(self, directory_path: str) -> bool: return True return False - def generate_report(self, directory_path: str, sol: Optional[float] = None) -> None: + def generate_report(self, test_name: str, directory_path: str, sol: Optional[float] = None) -> None: report_data = [] stdout_path = os.path.join(directory_path, "stdout.txt") if os.path.isfile(stdout_path):