jax-hpc-profiler 0.2.8__py3-none-any.whl → 0.2.9__py3-none-any.whl

jax_hpc_profiler/create_argparse.py

@@ -11,151 +11,187 @@ def create_argparser():
         Parsed and validated arguments.
     """
     parser = argparse.ArgumentParser(
-        description='HPC Plotter for benchmarking data')
+        description="HPC Plotter for benchmarking data")
 
     # Group for concatenation to ensure mutually exclusive behavior
-    subparsers = parser.add_subparsers(dest='command', required=True)
+    subparsers = parser.add_subparsers(dest="command", required=True)
 
-    concat_parser = subparsers.add_parser('concat',
-                                          help='Concatenate CSV files')
-    concat_parser.add_argument('input',
+    concat_parser = subparsers.add_parser("concat",
+                                          help="Concatenate CSV files")
+    concat_parser.add_argument("input",
                                type=str,
-                               help='Input directory for concatenation')
-    concat_parser.add_argument('output',
+                               help="Input directory for concatenation")
+    concat_parser.add_argument("output",
                                type=str,
-                               help='Output directory for concatenation')
+                               help="Output directory for concatenation")
 
     # Arguments for plotting
-    plot_parser = subparsers.add_parser('plot', help='Plot CSV data')
-    plot_parser.add_argument('-f',
-                             '--csv_files',
-                             nargs='+',
-                             help='List of CSV files to plot',
+    plot_parser = subparsers.add_parser("plot", help="Plot CSV data")
+    plot_parser.add_argument("-f",
+                             "--csv_files",
+                             nargs="+",
+                             help="List of CSV files to plot",
                              required=True)
-    plot_parser.add_argument('-g',
-                             '--gpus',
-                             nargs='*',
-                             type=int,
-                             help='List of number of GPUs to plot')
-    plot_parser.add_argument('-d',
-                             '--data_size',
-                             nargs='*',
-                             type=int,
-                             help='List of data sizes to plot')
+    plot_parser.add_argument(
+        "-g",
+        "--gpus",
+        nargs="*",
+        type=int,
+        help="List of number of GPUs to plot",
+        default=None,
+    )
+    plot_parser.add_argument(
+        "-d",
+        "--data_size",
+        nargs="*",
+        type=int,
+        help="List of data sizes to plot",
+        default=None,
+    )
 
     # pdims related arguments
-    plot_parser.add_argument('-fd',
-                             '--filter_pdims',
-                             nargs='*',
-                             help='List of pdims to filter, e.g., 1x4 2x2 4x8')
     plot_parser.add_argument(
-        '-ps',
-        '--pdim_strategy',
-        choices=['plot_all', 'plot_fastest', 'slab_yz', 'slab_xy', 'pencils'],
-        nargs='*',
-        default=['plot_fastest'],
-        help='Strategy for plotting pdims')
+        "-fd",
+        "--filter_pdims",
+        nargs="*",
+        help="List of pdims to filter, e.g., 1x4 2x2 4x8",
+        default=None,
+    )
+    plot_parser.add_argument(
+        "-ps",
+        "--pdim_strategy",
+        choices=["plot_all", "plot_fastest", "slab_yz", "slab_xy", "pencils"],
+        nargs="*",
+        default=["plot_fastest"],
+        help="Strategy for plotting pdims",
+    )
 
     # Function and precision related arguments
     plot_parser.add_argument(
-        '-pr',
-        '--precision',
-        choices=['float32', 'float64'],
-        default=['float32', 'float64'],
-        nargs='*',
-        help='Precision to filter by (float32 or float64)')
-    plot_parser.add_argument('-fn',
-                             '--function_name',
-                             nargs='+',
-                             help='Function names to filter')
+        "-pr",
+        "--precision",
+        choices=["float32", "float64"],
+        default=["float32", "float64"],
+        nargs="*",
+        help="Precision to filter by (float32 or float64)",
+    )
+    plot_parser.add_argument(
+        "-fn",
+        "--function_name",
+        nargs="+",
+        help="Function names to filter",
+        default=None,
+    )
 
     # Time or memory related arguments
     plotting_group = plot_parser.add_mutually_exclusive_group(required=True)
-    plotting_group.add_argument('-pt',
-                                '--plot_times',
-                                nargs='*',
-                                choices=[
-                                    'jit_time', 'min_time', 'max_time',
-                                    'mean_time', 'std_time', 'last_time'
-                                ],
-                                help='Time columns to plot')
-    plotting_group.add_argument('-pm',
-                                '--plot_memory',
-                                nargs='*',
-                                choices=[
-                                    'generated_code', 'argument_size',
-                                    'output_size', 'temp_size'
-                                ],
-                                help='Memory columns to plot')
-    plot_parser.add_argument('-mu',
-                             '--memory_units',
-                             default='GB',
-                             help='Memory units to plot (KB, MB, GB, TB)')
+    plotting_group.add_argument(
+        "-pt",
+        "--plot_times",
+        nargs="*",
+        choices=[
+            "jit_time",
+            "min_time",
+            "max_time",
+            "mean_time",
+            "std_time",
+            "last_time",
+        ],
+        help="Time columns to plot",
+    )
+    plotting_group.add_argument(
+        "-pm",
+        "--plot_memory",
+        nargs="*",
+        choices=[
+            "generated_code", "argument_size", "output_size", "temp_size"
+        ],
+        help="Memory columns to plot",
+    )
+    plot_parser.add_argument(
+        "-mu",
+        "--memory_units",
+        default="GB",
+        help="Memory units to plot (KB, MB, GB, TB)",
+    )
 
     # Plot customization arguments
-    plot_parser.add_argument('-fs',
-                             '--figure_size',
+    plot_parser.add_argument("-fs",
+                             "--figure_size",
                              nargs=2,
                              type=int,
-                             help='Figure size')
-    plot_parser.add_argument('-o',
-                             '--output',
-                             help='Output file (if none then only show plot)',
+                             help="Figure size",
+                             default=(10, 6))
+    plot_parser.add_argument("-o",
+                             "--output",
+                             help="Output file (if none then only show plot)",
                              default=None)
-    plot_parser.add_argument('-db',
-                             '--dark_bg',
-                             action='store_true',
-                             help='Use dark background for plotting')
-    plot_parser.add_argument('-pd',
-                             '--print_decompositions',
-                             action='store_true',
-                             help='Print decompositions on plot')
+    plot_parser.add_argument("-db",
+                             "--dark_bg",
+                             action="store_true",
+                             help="Use dark background for plotting")
+    plot_parser.add_argument(
+        "-pd",
+        "--print_decompositions",
+        action="store_true",
+        help="Print decompositions on plot",
+    )
 
     # Backend related arguments
-    plot_parser.add_argument('-b',
-                             '--backends',
-                             nargs='*',
-                             default=['MPI', 'NCCL', 'MPI4JAX'],
-                             help='List of backends to include')
+    plot_parser.add_argument(
+        "-b",
+        "--backends",
+        nargs="*",
+        default=["MPI", "NCCL", "MPI4JAX"],
+        help="List of backends to include",
+    )
 
     # Scaling type argument
-    plot_parser.add_argument('-sc',
-                             '--scaling',
-                             choices=['Weak', 'Strong'],
-                             required=True,
-                             help='Scaling type (Weak or Strong)')
+    plot_parser.add_argument(
+        "-sc",
+        "--scaling",
+        choices=["Weak", "Strong", "w", "s"],
+        required=True,
+        help="Scaling type (Weak or Strong)",
+    )
 
     # Label customization argument
     plot_parser.add_argument(
-        '-l',
-        '--label_text',
+        "-l",
+        "--label_text",
         type=str,
         help=
-        'Custom label for the plot. You can use placeholders: %decomposition% (or %p%), %precision% (or %pr%), %plot_name% (or %pn%), %backend% (or %b%), %node% (or %n%), %methodname% (or %m%)',
-        default="%m%-%f%-%pn%-%pr%-%b%-%p%-%n%")
+        ("Custom label for the plot. You can use placeholders: %%decomposition%% "
+         "(or %%p%%), %%precision%% (or %%pr%%), %%plot_name%% (or %%pn%%), "
+         "%%backend%% (or %%b%%), %%node%% (or %%n%%), %%methodname%% (or %%m%%)"
+         ),
+        default="%m%-%f%-%pn%-%pr%-%b%-%p%-%n%",
+    )
 
-    subparsers.add_parser('label_help',help='Label customization help')
+    subparsers.add_parser("label_help", help="Label customization help")
 
     args = parser.parse_args()
-    
+
     # if command was plot, then check if pdim_strategy is validat
-    if args.command == 'plot':
-      if 'plot_all' in args.pdim_strategy and len(args.pdim_strategy) > 1:
-          print(
-              "Warning: 'plot_all' strategy is combined with other strategies. Using 'plot_all' only."
-          )
-          args.pdim_strategy = ['plot_all']
-
-      if 'plot_fastest' in args.pdim_strategy and len(args.pdim_strategy) > 1:
-          print(
-              "Warning: 'plot_fastest' strategy is combined with other strategies. Using 'plot_fastest' only."
-          )
-          args.pdim_strategy = ['plot_fastest']
-      if args.plot_times is not None:
-          args.plot_columns = args.plot_times
-      elif args.plot_memory is not None:
-          args.plot_columns = args.plot_memory
-      else:
-          raise ValueError('Either plot_times or plot_memory should be provided')
+    if args.command == "plot":
+        if "plot_all" in args.pdim_strategy and len(args.pdim_strategy) > 1:
+            print(
+                "Warning: 'plot_all' strategy is combined with other strategies. Using 'plot_all' only."
+            )
+            args.pdim_strategy = ["plot_all"]
+
+        if "plot_fastest" in args.pdim_strategy and len(
+                args.pdim_strategy) > 1:
+            print(
+                "Warning: 'plot_fastest' strategy is combined with other strategies. Using 'plot_fastest' only."
+            )
+            args.pdim_strategy = ["plot_fastest"]
+        if args.plot_times is not None:
+            args.plot_columns = args.plot_times
+        elif args.plot_memory is not None:
+            args.plot_columns = args.plot_memory
+        else:
+            raise ValueError(
+                "Either plot_times or plot_memory should be provided")
 
     return args

jax_hpc_profiler/main.py

@@ -1,4 +1,5 @@
 import sys
+from typing import List, Optional
 
 from .create_argparse import create_argparser
 from .plotting import plot_strong_scaling, plot_weak_scaling
@@ -8,55 +9,56 @@ from .utils import clean_up_csv, concatenate_csvs
 def main():
     args = create_argparser()
 
-    if args.command == 'concat':
+    if args.command == "concat":
         input_dir, output_dir = args.input, args.output
         concatenate_csvs(input_dir, output_dir)
-    elif args.command == 'label_help':
+    elif args.command == "label_help":
         print(f"Customize the label text for the plot. using these commands.")
-        print(' -- %m% or %methodname%: method name')
-        print(' -- %f% or %function%: function name')
-        print(' -- %pn% or %plot_name%: plot name')
-        print(' -- %pr% or %precision%: precision')
-        print(' -- %b% or %backend%: backend')
-        print(' -- %p% or %pdims%: pdims')
-        print(' -- %n% or %node%: node')
-    elif args.command == 'plot':
-        dataframes, available_gpu_counts, available_data_sizes = clean_up_csv(
-            args.csv_files, args.precision, args.function_name, args.gpus,
-            args.data_size, args.filter_pdims, args.pdim_strategy,
-            args.backends, args.memory_units)
-        if len(dataframes) == 0:
-            print(f"No dataframes found for the given arguments. Exiting...")
-            sys.exit(1)
-        print(
-            f"requested GPUS: {args.gpus} available GPUS: {available_gpu_counts}"
-        )
-        # filter back the requested data sizes and gpus
-        
-        args.gpus = available_gpu_counts if args.gpus is None else [gpu for gpu in args.gpus if gpu in available_gpu_counts]
-        args.data_size = available_data_sizes if args.data_size is None else [data_size for data_size in args.data_size if data_size in available_data_sizes]
+        print(" -- %m% or %methodname%: method name")
+        print(" -- %f% or %function%: function name")
+        print(" -- %pn% or %plot_name%: plot name")
+        print(" -- %pr% or %precision%: precision")
+        print(" -- %b% or %backend%: backend")
+        print(" -- %p% or %pdims%: pdims")
+        print(" -- %n% or %node%: node")
+    elif args.command == "plot":
 
-        if len(args.gpus) == 0:
-            print(f"No dataframes found for the given GPUs. Exiting...")
-            sys.exit(1)
-        if len(args.data_size) == 0:
-            print(f"No dataframes found for the given data sizes. Exiting...")
-            sys.exit(1)
-
-        if args.scaling == 'Weak':
-            plot_weak_scaling(dataframes, args.gpus, args.figure_size,
-                              args.output, args.dark_bg,
-                              args.print_decompositions, args.backends,
-                              args.precision, args.function_name,
-                              args.plot_columns, args.memory_units,
-                              args.label_text, args.pdim_strategy)
-        elif args.scaling == 'Strong':
-            plot_strong_scaling(dataframes, args.data_size, args.figure_size,
-                                args.output, args.dark_bg,
-                                args.print_decompositions, args.backends,
-                                args.precision, args.function_name,
-                                args.plot_columns, args.memory_units,
-                                args.label_text, args.pdim_strategy)
+        if args.scaling.lower() == "weak" or args.scaling.lower() == "w":
+            plot_weak_scaling(
+                args.csv_files,
+                args.gpus,
+                args.data_size,
+                args.function_name,
+                args.precision,
+                args.filter_pdims,
+                args.pdim_strategy,
+                args.print_decompositions,
+                args.backends,
+                args.plot_columns,
+                args.memory_units,
+                args.label_text,
+                args.figure_size,
+                args.dark_bg,
+                args.output,
+            )
+        elif args.scaling.lower() == "strong" or args.scaling.lower() == "s":
+            plot_strong_scaling(
+                args.csv_files,
+                args.gpus,
+                args.data_size,
+                args.function_name,
+                args.precision,
+                args.filter_pdims,
+                args.pdim_strategy,
+                args.print_decompositions,
+                args.backends,
+                args.plot_columns,
+                args.memory_units,
+                args.label_text,
+                args.figure_size,
+                args.dark_bg,
+                args.output,
+            )
 
 
 if __name__ == "__main__":

jax_hpc_profiler/plotting.py

@@ -4,22 +4,24 @@ from typing import Dict, List, Optional
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import seaborn as sns
 from matplotlib.axes import Axes
 from matplotlib.patches import FancyBboxPatch
 
-from .utils import inspect_df, plot_with_pdims_strategy, inspect_data
+from .utils import clean_up_csv, inspect_df, plot_with_pdims_strategy
 
-np.seterr(divide='ignore')
-plt.rcParams.update({'font.size': 15})
+np.seterr(divide="ignore")
 
 
-def configure_axes(ax: Axes,
-                   x_values: List[int],
-                   y_values: List[float],
-                   xlabel: str,
-                   title: str,
-                   plotting_memory: bool = False,
-                   memory_units: str = 'bytes'):
+def configure_axes(
+    ax: Axes,
+    x_values: List[int],
+    y_values: List[float],
+    xlabel: str,
+    title: str,
+    plotting_memory: bool = False,
+    memory_units: str = "bytes",
+):
     """
     Configure the axes for the plot.
 
@@ -34,16 +36,17 @@ def configure_axes(ax: Axes,
     xlabel : str
         The label for the x-axis.
     """
-    ylabel = 'Time (milliseconds)' if not plotting_memory else f'Memory ({memory_units})'
+    ylabel = ("Time (milliseconds)"
+              if not plotting_memory else f"Memory ({memory_units})")
     f2 = lambda x: np.log2(x)
     g2 = lambda x: 2**x
     ax.set_xlim([min(x_values), max(x_values)])
     y_min, y_max = min(y_values) * 0.6, max(y_values) * 1.1
     ax.set_title(title)
     ax.set_ylim([y_min, y_max])
-    ax.set_xscale('function', functions=(f2, g2))
+    ax.set_xscale("function", functions=(f2, g2))
     if not plotting_memory:
-        ax.set_yscale('symlog')
+        ax.set_yscale("symlog")
         time_ticks = [
             10**t for t in range(int(np.floor(np.log10(y_min))), 1 +
                                  int(np.ceil(np.log10(y_max))))
@@ -53,31 +56,35 @@ def configure_axes(ax: Axes,
     ax.set_xlabel(xlabel)
     ax.set_ylabel(ylabel)
     for x_value in x_values:
-        ax.axvline(x=x_value, color='gray', linestyle='--', alpha=0.5)
-    ax.legend(loc='lower center',
-              bbox_to_anchor=(0.5, 0.05),
-              ncol=4,
-              fontsize="x-large",
-              prop={'size': 14})
-
-
-def plot_scaling(dataframes: Dict[str, pd.DataFrame],
-                 fixed_sizes: List[int],
-                 size_column: str,
-                 fixed_column: str,
-                 xlabel: str,
-                 title: str,
-                 figure_size: tuple = (6, 4),
-                 output: Optional[str] = None,
-                 dark_bg: bool = False,
-                 print_decompositions: bool = False,
-                 backends: List[str] = ['NCCL'],
-                 precisions: List[str] = ['float32'],
-                 functions: List[str] | None = None,
-                 plot_columns: List[str] = ['mean_time'],
-                 memory_units: str = 'bytes',
-                 label_text: str = 'plot',
-                 pdims_strategy: str = 'plot_fastest'):
+        ax.axvline(x=x_value, color="gray", linestyle="--", alpha=0.5)
+    ax.legend(
+        loc="lower center",
+        bbox_to_anchor=(0.5, 0.05),
+        ncol=4,
+        fontsize="x-large",
+        prop={"size": 14},
+    )
+
+
+def plot_scaling(
+    dataframes: Dict[str, pd.DataFrame],
+    fixed_sizes: List[int],
+    size_column: str,
+    fixed_column: str,
+    xlabel: str,
+    title: str,
+    figure_size: tuple = (6, 4),
+    output: Optional[str] = None,
+    dark_bg: bool = False,
+    print_decompositions: bool = False,
+    backends: List[str] = ["NCCL"],
+    precisions: List[str] = ["float32"],
+    functions: List[str] | None = None,
+    plot_columns: List[str] = ["mean_time"],
+    memory_units: str = "bytes",
+    label_text: str = "plot",
+    pdims_strategy: List[str] = ["plot_fastest"],
+):
     """
     General scaling plot function based on the number of GPUs or data size.
 
@@ -106,8 +113,9 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
     pdims_strategy : str, optional
         Strategy for plotting pdims ('plot_all' or 'plot_fastest'), by default 'plot_fastest'.
     """
+
     if dark_bg:
-        plt.style.use('dark_background')
+        plt.style.use("dark_background")
 
     num_subplots = len(fixed_sizes)
     num_rows = int(np.ceil(np.sqrt(num_subplots)))
@@ -118,7 +126,7 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
         axs = axs.flatten()
     else:
         axs = [axs]
-    
+
     for i, fixed_size in enumerate(fixed_sizes):
         ax: Axes = axs[i]
 
@@ -131,30 +139,44 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
                 continue
             filtered_method_df = filtered_method_df.sort_values(
                 by=[size_column])
-            functions = pd.unique(filtered_method_df['function']
-                                  ) if functions is None else functions
+            functions = (pd.unique(filtered_method_df["function"])
+                         if functions is None else functions)
             combinations = product(backends, precisions, functions,
                                    plot_columns)
 
             for backend, precision, function, plot_column in combinations:
-              
+
                 filtered_params_df = filtered_method_df[
-                    (filtered_method_df['backend'] == backend)
-                    & (filtered_method_df['precision'] == precision) &
-                    (filtered_method_df['function'] == function)]
+                    (filtered_method_df["backend"] == backend)
+                    & (filtered_method_df["precision"] == precision)
+                    & (filtered_method_df["function"] == function)]
                 if filtered_params_df.empty:
                     continue
                 x_vals, y_vals = plot_with_pdims_strategy(
-                    ax, filtered_params_df, method, pdims_strategy,
-                    print_decompositions, size_column, plot_column, label_text)
+                    ax,
+                    filtered_params_df,
+                    method,
+                    pdims_strategy,
+                    print_decompositions,
+                    size_column,
+                    plot_column,
+                    label_text,
+                )
 
                 x_values.extend(x_vals)
                 y_values.extend(y_vals)
-        
-        if len(x_values) != 0: 
-          plotting_memory = 'time' not in plot_columns[0].lower()
-          configure_axes(ax, x_values, y_values, f"{title} {fixed_size}", xlabel,
-                         plotting_memory, memory_units)
+
+        if len(x_values) != 0:
+            plotting_memory = "time" not in plot_columns[0].lower()
+            configure_axes(
+                ax,
+                x_values,
+                y_values,
+                f"{title} {fixed_size}",
+                xlabel,
+                plotting_memory,
+                memory_units,
+            )
 
     for i in range(num_subplots, num_rows * num_cols):
         fig.delaxes(axs[i])
@@ -170,48 +192,117 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
     if output is None:
         plt.show()
     else:
-        plt.savefig(output, bbox_inches='tight', transparent=False)
-
-
-def plot_strong_scaling(dataframes: Dict[str, pd.DataFrame],
-                        fixed_data_size: List[int],
-                        figure_size: tuple = (6, 4),
-                        output: Optional[str] = None,
-                        dark_bg: bool = False,
-                        print_decompositions: bool = False,
-                        backends: List[str] = ['NCCL'],
-                        precisions: List[str] = ['float32'],
-                        functions: List[str] | None = None,
-                        plot_columns: List[str] = ['mean_time'],
-                        memory_units: str = 'bytes',
-                        label_text: str = 'plot',
-                        pdims_strategy: str = 'plot_fastest'):
+        plt.savefig(output, bbox_inches="tight", transparent=True)
+
+
+def plot_strong_scaling(
+    csv_files: List[str],
+    fixed_gpu_size: Optional[List[int]] = None,
+    fixed_data_size: Optional[List[int]] = None,
+    functions: List[str] | None = None,
+    precisions: List[str] = ["float32"],
+    pdims: Optional[List[str]] = None,
+    pdims_strategy: List[str] = ["plot_fastest"],
+    print_decompositions: bool = False,
+    backends: List[str] = ["NCCL"],
+    plot_columns: List[str] = ["mean_time"],
+    memory_units: str = "bytes",
+    label_text: str = "%m%-%f%-%pn%-%pr%-%b%-%p%-%n%",
+    figure_size: tuple = (6, 4),
+    dark_bg: bool = False,
+    output: Optional[str] = None,
+):
     """
     Plot strong scaling based on the number of GPUs.
     """
-    plot_scaling(dataframes, fixed_data_size, 'gpus', 'x', 'Number of GPUs',
-                 'Data size', figure_size, output, dark_bg,
-                 print_decompositions, backends, precisions, functions,
-                 plot_columns, memory_units, label_text, pdims_strategy)
-
-
-def plot_weak_scaling(dataframes: Dict[str, pd.DataFrame],
-                      fixed_gpu_size: List[int],
-                      figure_size: tuple = (6, 4),
-                      output: Optional[str] = None,
-                      dark_bg: bool = False,
-                      print_decompositions: bool = False,
-                      backends: List[str] = ['NCCL'],
-                      precisions: List[str] = ['float32'],
-                      functions: List[str] | None = None,
-                      plot_columns: List[str] = ['mean_time'],
-                      memory_units: str = 'bytes',
-                      label_text: str = 'plot',
-                      pdims_strategy: str = 'plot_fastest'):
+
+    dataframes, _, available_data_sizes = clean_up_csv(
+        csv_files,
+        precisions,
+        functions,
+        fixed_gpu_size,
+        fixed_data_size,
+        pdims,
+        pdims_strategy,
+        backends,
+        memory_units,
+    )
+    if len(dataframes) == 0:
+        print(f"No dataframes found for the given arguments. Exiting...")
+        return
+
+    plot_scaling(
+        dataframes,
+        available_data_sizes,
+        "gpus",
+        "x",
+        "Number of GPUs",
+        "Data size",
+        figure_size,
+        output,
+        dark_bg,
+        print_decompositions,
+        backends,
+        precisions,
+        functions,
+        plot_columns,
+        memory_units,
+        label_text,
+        pdims_strategy,
+    )
+
+
+def plot_weak_scaling(
+    csv_files: List[str],
+    fixed_gpu_size: Optional[List[int]] = None,
+    fixed_data_size: Optional[List[int]] = None,
+    functions: List[str] | None = None,
+    precisions: List[str] = ["float32"],
+    pdims: Optional[List[str]] = None,
+    pdims_strategy: List[str] = ["plot_fastest"],
+    print_decompositions: bool = False,
+    backends: List[str] = ["NCCL"],
+    plot_columns: List[str] = ["mean_time"],
+    memory_units: str = "bytes",
+    label_text: str = "%m%-%f%-%pn%-%pr%-%b%-%p%-%n%",
+    figure_size: tuple = (6, 4),
+    dark_bg: bool = False,
+    output: Optional[str] = None,
+):
     """
     Plot weak scaling based on the data size.
     """
-    plot_scaling(dataframes, fixed_gpu_size, 'x', 'gpus', 'Data size',
-                 'Number of GPUs', figure_size, output, dark_bg,
-                 print_decompositions, backends, precisions, functions,
-                 plot_columns, memory_units, label_text, pdims_strategy)
+    dataframes, available_gpu_counts, _ = clean_up_csv(
+        csv_files,
+        precisions,
+        functions,
+        fixed_gpu_size,
+        fixed_data_size,
+        pdims,
+        pdims_strategy,
+        backends,
+        memory_units,
+    )
+    if len(dataframes) == 0:
+        print(f"No dataframes found for the given arguments. Exiting...")
+        return
+
+    plot_scaling(
+        dataframes,
+        available_gpu_counts,
+        "x",
+        "gpus",
+        "Data size",
+        "Number of GPUs",
+        figure_size,
+        output,
+        dark_bg,
+        print_decompositions,
+        backends,
+        precisions,
+        functions,
+        plot_columns,
+        memory_units,
+        label_text,
+        pdims_strategy,
+    )

jax_hpc_profiler/timer.py

@@ -16,36 +16,45 @@ from tabulate import tabulate
 
 class Timer:
 
-    def __init__(self, save_jaxpr=False):
-        self.jit_time = None
+    def __init__(self, save_jaxpr=False, jax_fn=True, devices=None):
+        self.jit_time = 0.0
         self.times = []
         self.profiling_data = {}
         self.compiled_code = {}
         self.save_jaxpr = save_jaxpr
+        self.jax_fn = jax_fn
+        self.devices = devices
 
     def _normalize_memory_units(self, memory_analysis) -> str:
 
-        sizes_str = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
+        if not self.jax_fn:
+            return memory_analysis
+
+        sizes_str = ["B", "KB", "MB", "GB", "TB", "PB"]
         factors = [1, 1024, 1024**2, 1024**3, 1024**4, 1024**5]
-        factor = 0 if memory_analysis == 0 else int(np.log10(memory_analysis) // 3)
+        factor = 0 if memory_analysis == 0 else int(
+            np.log10(memory_analysis) // 3)
 
         return f"{memory_analysis / factors[factor]:.2f} {sizes_str[factor]}"
 
     def _read_memory_analysis(self, memory_analysis: Any) -> Tuple:
         if memory_analysis is None:
             return None, None, None, None
-        return (memory_analysis.generated_code_size_in_bytes,
-                memory_analysis.argument_size_in_bytes,
-                memory_analysis.output_size_in_bytes,
-                memory_analysis.temp_size_in_bytes)
+        return (
+            memory_analysis.generated_code_size_in_bytes,
+            memory_analysis.argument_size_in_bytes,
+            memory_analysis.output_size_in_bytes,
+            memory_analysis.temp_size_in_bytes,
+        )
 
     def chrono_jit(self, fun: Callable, *args, ndarray_arg=None) -> np.ndarray:
         start = time.perf_counter()
         out = fun(*args)
-        if ndarray_arg is None:
-            out.block_until_ready()
-        else:
-            out[ndarray_arg].block_until_ready()
+        if self.jax_fn:
+            if ndarray_arg is None:
+                out.block_until_ready()
+            else:
+                out[ndarray_arg].block_until_ready()
         end = time.perf_counter()
         self.jit_time = (end - start) * 1e3
 
@@ -53,78 +62,90 @@ class Timer:
             jaxpr = make_jaxpr(fun)(*args)
             self.compiled_code["JAXPR"] = jaxpr.pretty_print()
 
-        lowered = jax.jit(fun).lower(*args)
-        compiled = lowered.compile()
-        memory_analysis = self._read_memory_analysis(
-            compiled.memory_analysis())
-
-        self.compiled_code["LOWERED"] = lowered.as_text()
-        self.compiled_code["COMPILED"] = compiled.as_text()
-        self.profiling_data["generated_code"] = memory_analysis[0]
-        self.profiling_data["argument_size"] = memory_analysis[1]
-        self.profiling_data["output_size"] = memory_analysis[2]
-        self.profiling_data["temp_size"] = memory_analysis[3]
+        if self.jax_fn:
+            lowered = jax.jit(fun).lower(*args)
+            compiled = lowered.compile()
+            memory_analysis = self._read_memory_analysis(
+                compiled.memory_analysis())
+
+            self.compiled_code["LOWERED"] = lowered.as_text()
+            self.compiled_code["COMPILED"] = compiled.as_text()
+            self.profiling_data["generated_code"] = memory_analysis[0]
+            self.profiling_data["argument_size"] = memory_analysis[1]
+            self.profiling_data["output_size"] = memory_analysis[2]
+            self.profiling_data["temp_size"] = memory_analysis[3]
+
         return out
 
     def chrono_fun(self, fun: Callable, *args, ndarray_arg=None) -> np.ndarray:
         start = time.perf_counter()
         out = fun(*args)
-        if ndarray_arg is None:
-            out.block_until_ready()
-        else:
-            out[ndarray_arg].block_until_ready()
+        if self.jax_fn:
+            if ndarray_arg is None:
+                out.block_until_ready()
+            else:
+                out[ndarray_arg].block_until_ready()
         end = time.perf_counter()
         self.times.append((end - start) * 1e3)
         return out
 
     def _get_mean_times(self) -> np.ndarray:
-        if jax.device_count() == 1:
+        if jax.device_count() == 1 or jax.process_count() == 1:
             return np.array(self.times)
 
-        devices = mesh_utils.create_device_mesh((jax.device_count(), ))
-        mesh = Mesh(devices, ('x', ))
-        sharding = NamedSharding(mesh, P('x'))
+        if self.devices is None:
+            self.devices = jax.devices()
+
+        mesh = jax.make_mesh((len(self.devices), ), ("x", ),
+                             devices=self.devices)
+        sharding = NamedSharding(mesh, P("x"))
 
         times_array = jnp.array(self.times)
         global_shape = (jax.device_count(), times_array.shape[0])
         global_times = jax.make_array_from_callback(
             shape=global_shape,
             sharding=sharding,
-            data_callback=lambda _: jnp.expand_dims(times_array, axis=0))
+            data_callback=lambda _: jnp.expand_dims(times_array, axis=0),
+        )
 
         @partial(shard_map,
                  mesh=mesh,
-                 in_specs=P('x'),
+                 in_specs=P("x"),
                  out_specs=P(),
                  check_rep=False)
         def get_mean_times(times):
-            return jax.lax.pmean(times, axis_name='x')
+            return jax.lax.pmean(times, axis_name="x")
 
         times_array = get_mean_times(global_times)
         times_array.block_until_ready()
         return np.array(times_array.addressable_data(0)[0])
 
-    def report(self,
-               csv_filename: str,
-               function: str,
-               x: int,
-               y: int | None = None,
-               z: int | None = None,
-               precision: str = "float32",
-               px: int = 1,
-               py: int = 1,
-               backend: str = "NCCL",
-               nodes: int = 1,
-               md_filename: str | None = None,
-               extra_info: dict = {}):
+    def report(
+        self,
+        csv_filename: str,
+        function: str,
+        x: int,
+        y: int | None = None,
+        z: int | None = None,
+        precision: str = "float32",
+        px: int = 1,
+        py: int = 1,
+        backend: str = "NCCL",
+        nodes: int = 1,
+        md_filename: str | None = None,
+        extra_info: dict = {},
+    ):
 
         if md_filename is None:
-            dirname, filename = os.path.dirname(
-                csv_filename), os.path.splitext(
-                    os.path.basename(csv_filename))[0]
+            dirname, filename = (
+                os.path.dirname(csv_filename),
+                os.path.splitext(os.path.basename(csv_filename))[0],
+            )
             report_folder = filename if dirname == "" else f"{dirname}/{filename}"
             os.makedirs(report_folder, exist_ok=True)
-            md_filename = f"{report_folder}/{x}_{px}_{py}_{backend}_{precision}_{function}.md"
+            md_filename = (
+                f"{report_folder}/{x}_{px}_{py}_{backend}_{precision}_{function}.md"
+            )
 
         y = x if y is None else y
         z = x if z is None else z
@@ -138,10 +159,17 @@ class Timer:
             std_time = np.std(times_array)
             last_time = times_array[-1]
 
-            generated_code = self.profiling_data["generated_code"]
-            argument_size = self.profiling_data["argument_size"]
-            output_size = self.profiling_data["output_size"]
-            temp_size = self.profiling_data["temp_size"]
+            if self.jax_fn:
+
+                generated_code = self.profiling_data["generated_code"]
+                argument_size = self.profiling_data["argument_size"]
+                output_size = self.profiling_data["output_size"]
+                temp_size = self.profiling_data["temp_size"]
+            else:
+                generated_code = "N/A"
+                argument_size = "N/A"
+                output_size = "N/A"
+                temp_size = "N/A"
 
             csv_line = (
                 f"{function},{precision},{x},{y},{z},{px},{py},{backend},{nodes},"
@@ -149,7 +177,7 @@ class Timer:
                 f"{generated_code},{argument_size},{output_size},{temp_size}\n"
             )
 
-            with open(csv_filename, 'a') as f:
+            with open(csv_filename, "a") as f:
                 f.write(csv_line)
 
             param_dict = {
@@ -175,44 +203,56 @@ class Timer:
                 "Argument Size": self._normalize_memory_units(argument_size),
                 "Output Size": self._normalize_memory_units(output_size),
                 "Temporary Size": self._normalize_memory_units(temp_size),
-                "FLOPS": self.profiling_data["FLOPS"]
             }
             iteration_runs = {}
             for i in range(len(times_array)):
                 iteration_runs[f"Run {i}"] = times_array[i]
 
-            with open(md_filename, 'w') as f:
+            with open(md_filename, "w") as f:
                 f.write(f"# Reporting for {function}\n")
                 f.write(f"## Parameters\n")
                 f.write(
-                    tabulate(param_dict.items(),
-                             headers=["Parameter", "Value"],
-                             tablefmt='github'))
+                    tabulate(
+                        param_dict.items(),
+                        headers=["Parameter", "Value"],
+                        tablefmt="github",
+                    ))
                 f.write("\n---\n")
                 f.write(f"## Profiling Data\n")
                 f.write(
-                    tabulate(profiling_result.items(),
-                             headers=["Parameter", "Value"],
-                             tablefmt='github'))
+                    tabulate(
+                        profiling_result.items(),
+                        headers=["Parameter", "Value"],
+                        tablefmt="github",
+                    ))
                 f.write("\n---\n")
                 f.write(f"## Iteration Runs\n")
                 f.write(
-                    tabulate(iteration_runs.items(),
-                             headers=["Iteration", "Time"],
-                             tablefmt='github'))
-                f.write("\n---\n")
-                f.write(f"## Compiled Code\n")
-                f.write(f"```hlo\n")
-                f.write(self.compiled_code["COMPILED"])
-                f.write(f"\n```\n")
-                f.write("\n---\n")
-                f.write(f"## Lowered Code\n")
-                f.write(f"```hlo\n")
-                f.write(self.compiled_code["LOWERED"])
-                f.write(f"\n```\n")
-                f.write("\n---\n")
-                if self.save_jaxpr:
-                    f.write(f"## JAXPR\n")
-                    f.write(f"```haskel\n")
-                    f.write(self.compiled_code["JAXPR"])
+                    tabulate(
+                        iteration_runs.items(),
+                        headers=["Iteration", "Time"],
+                        tablefmt="github",
+                    ))
+                if self.jax_fn:
+                    f.write("\n---\n")
+                    f.write(f"## Compiled Code\n")
+                    f.write(f"```hlo\n")
+                    f.write(self.compiled_code["COMPILED"])
                     f.write(f"\n```\n")
+                    f.write("\n---\n")
+                    f.write(f"## Lowered Code\n")
+                    f.write(f"```hlo\n")
+                    f.write(self.compiled_code["LOWERED"])
+                    f.write(f"\n```\n")
+                    f.write("\n---\n")
+                    if self.save_jaxpr:
+                        f.write(f"## JAXPR\n")
+                        f.write(f"```haskel\n")
+                        f.write(self.compiled_code["JAXPR"])
+                        f.write(f"\n```\n")
+
+        # Reset the timer
+        self.jit_time = 0.0
+        self.times = []
+        self.profiling_data = {}
+        self.compiled_code = {}

jax_hpc_profiler/utils.py

@@ -250,7 +250,7 @@ def clean_up_csv(
     pdims_strategy: List[str] = ['plot_fastest'],
     backends: List[str] = ['MPI', 'NCCL', 'MPI4JAX'],
     memory_units: str = 'KB',
-) -> Tuple[Dict[str, pd.DataFrame], set, set]:
+) -> Tuple[Dict[str, pd.DataFrame], List[int], List[int]]:
     """
     Clean up and aggregate data from CSV files.
 
@@ -336,7 +336,7 @@ def clean_up_csv(
         # Filter data sizes
         if data_sizes:
             df = df[df['x'].isin(data_sizes)]
-        
+
         # Filter pdims
         if pdims:
             px_list, py_list = zip(*[map(int, p.split('x')) for p in pdims])
@@ -394,4 +394,17 @@ def clean_up_csv(
         else:
             dataframes[file_name] = pd.concat([dataframes[file_name], df])
 
+    print(f"requested GPUS: {gpus} available GPUS: {available_gpu_counts}")
+    print(
+        f"requested data sizes: {data_sizes} available data sizes: {available_data_sizes}"
+    )
+
+    available_gpu_counts = (available_gpu_counts if gpus is None else [
+        gpu for gpu in gpus if gpu in available_gpu_counts
+    ])
+    available_data_sizes = (available_data_sizes if data_sizes is None else [
+        data_size for data_size in data_sizes
+        if data_size in available_data_sizes
+    ])
+
     return dataframes, available_gpu_counts, available_data_sizes

{jax_hpc_profiler-0.2.8.dist-info → jax_hpc_profiler-0.2.9.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: jax_hpc_profiler
-Version: 0.2.8
+Version: 0.2.9
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
-License: GNU GENERAL PUBLIC LICENSE
+License:                     GNU GENERAL PUBLIC LICENSE
                                Version 3, 29 June 2007
         
          Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>

jax_hpc_profiler-0.2.9.dist-info/RECORD

@@ -0,0 +1,12 @@
+jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,367
+jax_hpc_profiler/create_argparse.py,sha256=CSdl76LvaTVVn43dkwpVyiIkyl4lHlDCiI5jvUrIoj0,6059
+jax_hpc_profiler/main.py,sha256=2zPVTGRgFkYV75EJA1eoOqf92gCRXAtg-28cFgRy3Bw,2164
+jax_hpc_profiler/plotting.py,sha256=8ELOB_Yv_AdSVWtS-jrRNm0HtK5FgKwf_ljeNRfdp14,9087
+jax_hpc_profiler/timer.py,sha256=p7MUcbd2H4_tRAevhG9T4jJ8XL-liComvJn2sis4psM,9209
+jax_hpc_profiler/utils.py,sha256=hSsS34i46WdCR9XRW1-02fI_k0RUty78imnI-xAc-tY,14644
+jax_hpc_profiler-0.2.9.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jax_hpc_profiler-0.2.9.dist-info/METADATA,sha256=CelrNVm13lK7L1ZkdOqD8Tm7qLBIF1oCyaghzDdrLRg,49270
+jax_hpc_profiler-0.2.9.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+jax_hpc_profiler-0.2.9.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
+jax_hpc_profiler-0.2.9.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
+jax_hpc_profiler-0.2.9.dist-info/RECORD,,

{jax_hpc_profiler-0.2.8.dist-info → jax_hpc_profiler-0.2.9.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (74.0.0)
+Generator: setuptools (75.6.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

jax_hpc_profiler-0.2.8.dist-info/RECORD

@@ -1,12 +0,0 @@
-jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,367
-jax_hpc_profiler/create_argparse.py,sha256=m9_lg9HHxq2JDMITiHXQW1Ximua0ClwsEq1Zd9Y0hvo,6511
-jax_hpc_profiler/main.py,sha256=VJKvVc4m2XGJI2yp9ZF9tmmBmnTDpZ7-6LGo8ZIrWLc,2906
-jax_hpc_profiler/plotting.py,sha256=jnWpyfZOB5w5L4BNIU5euKbUt__wfGPyCjUITrEwScM,8431
-jax_hpc_profiler/timer.py,sha256=qPp3NcCJlMM-Cmw2mEWn63BlvPqmj_k7E8P9m0-Fy7k,8294
-jax_hpc_profiler/utils.py,sha256=tUXnNHwQSSCqA6XxLd4MoV2gyFIC7ncB-uvVc2INhms,14140
-jax_hpc_profiler-0.2.8.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jax_hpc_profiler-0.2.8.dist-info/METADATA,sha256=NDOKx7RzKr4Z4mJnY457-4uxaP2TDhwWYH5rZRxfukQ,49250
-jax_hpc_profiler-0.2.8.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
-jax_hpc_profiler-0.2.8.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
-jax_hpc_profiler-0.2.8.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
-jax_hpc_profiler-0.2.8.dist-info/RECORD,,