Merge branch 'NVIDIA:main' into refactor-path

src/cloudai/report_generator/tool/bokeh_report_tool.py

@@ -48,6 +48,7 @@ def create_figure(
         height: int = 308,
         x_axis_type: str = "linear",
         tools: str = "pan,wheel_zoom,box_zoom,reset,save",
+        x_range: Optional[Range1d] = None,
     ) -> figure:
         """
         Create a configured Bokeh figure with common settings.
@@ -61,6 +62,7 @@ def create_figure(
             height (int): Height of the plot.
             x_axis_type (str): Type of the x-axis ('linear' or 'log').
             tools (str): Tools to include in the plot.
+            x_range (Range1d): Range for the x-axis, optional.
 
         Returns:
             figure: A Bokeh figure configured with the specified parameters.
@@ -76,6 +78,10 @@ def create_figure(
             y_range=y_range,
             align="center",
         )
+
+        if x_range is not None:
+            plot.x_range = x_range
+
         return plot
 
     def add_sol_line(
@@ -164,60 +170,6 @@ def add_linear_xy_line_plot(
 
         self.plots.append(p)
 
-    def add_log_x_linear_y_single_line_plot(
-        self,
-        title: str,
-        x_column: str,
-        y_column: str,
-        x_axis_label: str,
-        y_axis_label: str,
-        df: pd.DataFrame,
-        sol: Optional[float] = None,
-        color: str = "black",
-    ):
-        """
-        Create a single line plot with a logarithmic x-axis and linear y-axis.
-
-        Args:
-            title (str): Title of the plot.
-            x_column (str): The column used for the x-axis values.
-            y_column (str): The column used for the y-axis values.
-            x_axis_label (str): Label for the x-axis.
-            y_axis_label (str): Label for the y-axis.
-            df (pd.DataFrame): DataFrame containing the data.
-            sol (Optional[float]): Speed-of-light performance reference line.
-            color (str): Color of the line in the plot.
-
-        This function sets up a Bokeh figure and plots a single line of data. It also
-        optionally adds a reference line (SOL) if provided. The x-axis uses a logarithmic
-        scale, and custom JavaScript is used for tick formatting to enhance readability.
-        """
-        x_min, x_max = self.find_min_max(df, x_column)
-        y_min, y_max = self.find_min_max(df, y_column, sol)
-
-        # Create a Bokeh figure with logarithmic x-axis
-        p = self.create_figure(
-            title="CloudAI " + title,
-            x_axis_label=x_axis_label,
-            y_axis_label=y_axis_label,
-            x_axis_type="log",
-            y_range=Range1d(start=0, end=y_max * 1.1),
-        )
-
-        # Add main line plot
-        p.line(x=x_column, y=y_column, source=ColumnDataSource(df), line_width=2, color=color, legend_label=y_column)
-
-        self.add_sol_line(p, df, x_column, y_column, sol)
-
-        p.legend.location = "bottom_right"
-
-        p.xaxis.ticker = calculate_power_of_two_ticks(x_min, x_max)
-        p.xaxis.formatter = CustomJSTickFormatter(code=bokeh_size_unit_js_tick_formatter)
-        p.xaxis.major_label_orientation = pi / 4
-
-        # Append plot to internal list for future rendering
-        self.plots.append(p)
-
     def add_log_x_linear_y_multi_line_plot(
         self,
         title: str,
@@ -246,12 +198,25 @@ def add_log_x_linear_y_multi_line_plot(
             _, col_max = self.find_min_max(df, y_column, sol)
             y_max = max(y_max, col_max)
 
+        x_axis_type = "log"
+        x_range = None
+
+        # Check if x_min equals x_max - constant message size
+        if x_min == x_max:
+            # Use iteration number as x-axis
+            df["iteration"] = range(1, len(df) + 1)
+            x_column = "iteration"
+            x_axis_label = "Iteration"
+            x_axis_type = "linear"
+            x_range = Range1d(start=1, end=len(df))
+
         p = self.create_figure(
             title="CloudAI " + title,
             x_axis_label=x_axis_label,
             y_axis_label=y_axis_label,
-            x_axis_type="log",
+            x_axis_type=x_axis_type,
             y_range=Range1d(start=0, end=y_max * 1.1),
+            x_range=x_range,
         )
 
         # Adding lines for each data type specified
@@ -265,10 +230,11 @@ def add_log_x_linear_y_multi_line_plot(
 
         p.legend.location = "bottom_right"
 
-        # Setting up custom tick formatter for log scale readability
-        p.xaxis.ticker = calculate_power_of_two_ticks(x_min, x_max)
-        p.xaxis.formatter = CustomJSTickFormatter(code=bokeh_size_unit_js_tick_formatter)
-        p.xaxis.major_label_orientation = pi / 4
+        if x_axis_type == "log":
+            # Setting up custom tick formatter for log scale readability
+            p.xaxis.ticker = calculate_power_of_two_ticks(x_min, x_max)
+            p.xaxis.formatter = CustomJSTickFormatter(code=bokeh_size_unit_js_tick_formatter)
+            p.xaxis.major_label_orientation = pi / 4
 
         self.plots.append(p)
 

src/cloudai/schema/test_template/nccl_test/report_generation_strategy.py

@@ -120,15 +120,14 @@ def _generate_bokeh_report(
             ("Busbw (GB/s) In-place", "green", "In-place Bus Bandwidth"),
         ]
         for col_name, color, title in line_plots:
-            report_tool.add_log_x_linear_y_single_line_plot(
+            report_tool.add_log_x_linear_y_multi_line_plot(
                 title=f"{test_name} {title}",
                 x_column="Size (B)",
-                y_column=col_name,
+                y_columns=[(col_name, color)],
                 x_axis_label="Message Size",
                 y_axis_label="Bandwidth (GB/s)",
                 df=df,
                 sol=sol,
-                color=color,
             )
 
         combined_columns = [("Busbw (GB/s) Out-of-place", "blue"), ("Busbw (GB/s) In-place", "green")]

src/cloudai/schema/test_template/ucc_test/report_generation_strategy.py

@@ -95,15 +95,14 @@ def _generate_plots(self, df: pd.DataFrame, directory_path: Path, sol: Optional[
         report_tool = BokehReportTool(directory_path)
         line_plots = [("Bandwidth (GB/s) avg", "black", "Average Bandwidth")]
         for col_name, color, title in line_plots:
-            report_tool.add_log_x_linear_y_single_line_plot(
+            report_tool.add_log_x_linear_y_multi_line_plot(
                 title=title,
                 x_column="Size (B)",
-                y_column=col_name,
+                y_columns=[(col_name, color)],
                 x_axis_label="Message Size",
                 y_axis_label="Bandwidth (GB/s)",
                 df=df,
                 sol=sol,
-                color=color,
             )
 
         combined_columns = [