jax-hpc-profiler 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

jax_hpc_profiler/__init__.py

@@ -4,6 +4,11 @@ from .timer import Timer
 from .utils import clean_up_csv, concatenate_csvs, plot_with_pdims_strategy
 
 __all__ = [
-    'create_argparser', 'plot_strong_scaling', 'plot_weak_scaling', 'Timer',
-    'clean_up_csv', 'concatenate_csvs', 'plot_with_pdims_strategy'
+    'create_argparser',
+    'plot_strong_scaling',
+    'plot_weak_scaling',
+    'Timer',
+    'clean_up_csv',
+    'concatenate_csvs',
+    'plot_with_pdims_strategy',
 ]

jax_hpc_profiler/create_argparse.py

@@ -10,201 +10,190 @@ def create_argparser():
     argparse.Namespace
         Parsed and validated arguments.
     """
-    parser = argparse.ArgumentParser(
-        description="HPC Plotter for benchmarking data")
+    parser = argparse.ArgumentParser(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",
-                               type=str,
-                               help="Input directory for concatenation")
-    concat_parser.add_argument("output",
-                               type=str,
-                               help="Output directory for concatenation")
+    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', type=str, 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",
-                             required=True)
-    plot_parser.add_argument(
-        "-g",
-        "--gpus",
-        nargs="*",
+    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",
+        help='List of number of GPUs to plot',
         default=None,
     )
     plot_parser.add_argument(
-        "-d",
-        "--data_size",
-        nargs="*",
+        '-d',
+        '--data_size',
+        nargs='*',
         type=int,
-        help="List of data sizes to plot",
+        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",
+        '-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",
+        '-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)",
+        '-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",
+        '-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="*",
+        '-pt',
+        '--plot_times',
+        nargs='*',
         choices=[
-            "jit_time",
-            "min_time",
-            "max_time",
-            "mean_time",
-            "std_time",
-            "last_time",
+            'jit_time',
+            'min_time',
+            'max_time',
+            'mean_time',
+            'std_time',
+            'last_time',
         ],
-        help="Time columns to plot",
+        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",
+        '-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)",
+        '-mu',
+        '--memory_units',
+        default='GB',
+        help='Memory units to plot (KB, MB, GB, TB)',
     )
 
     # Plot customization arguments
-    plot_parser.add_argument("-fs",
-                             "--figure_size",
-                             nargs=2,
-                             type=int,
-                             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(
+        '-fs', '--figure_size', nargs=2, type=int, 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',
     )
 
     # Backend related arguments
     plot_parser.add_argument(
-        "-b",
-        "--backends",
-        nargs="*",
-        default=["MPI", "NCCL", "MPI4JAX"],
-        help="List of backends to include",
+        '-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", "w", "s"],
+        '-sc',
+        '--scaling',
+        choices=['Weak', 'Strong', 'w', 's'],
         required=True,
-        help="Scaling type (Weak or Strong)",
+        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%",
+        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%',
     )
 
     plot_parser.add_argument(
-        "-xl",
-        "--xlabel",
+        '-xl',
+        '--xlabel',
         type=str,
-        help="X-axis label for the plot",
+        help='X-axis label for the plot',
     )
     plot_parser.add_argument(
-        "-tl",
-        "--title",
+        '-tl',
+        '--title',
         type=str,
-        help="Title for the plot",
+        help='Title for the plot',
     )
 
-    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:
+    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."
+                """
+                Warning: 'plot_all' strategy is combined with other strategies.
+                Using 'plot_all' only.
+                """
             )
-            args.pdim_strategy = ["plot_all"]
+            args.pdim_strategy = ['plot_all']
 
-        if "plot_fastest" in args.pdim_strategy and len(
-                args.pdim_strategy) > 1:
+        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."
+                """
+                Warning: 'plot_fastest' strategy is combined with other strategies.
+                Using 'plot_fastest' only.
+                """
             )
-            args.pdim_strategy = ["plot_fastest"]
+            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")
+            raise ValueError('Either plot_times or plot_memory should be provided')
 
     return args

jax_hpc_profiler/main.py

@@ -1,29 +1,25 @@
-import sys
-from typing import List, Optional
-
 from .create_argparse import create_argparser
 from .plotting import plot_strong_scaling, plot_weak_scaling
-from .utils import clean_up_csv, concatenate_csvs
+from .utils import 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":
-        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":
-
-        if args.scaling.lower() == "weak" or args.scaling.lower() == "w":
+    elif args.command == 'label_help':
+        print('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':
+        if args.scaling.lower() == 'weak' or args.scaling.lower() == 'w':
             plot_weak_scaling(
                 args.csv_files,
                 args.gpus,
@@ -43,7 +39,7 @@ def main():
                 args.dark_bg,
                 args.output,
             )
-        elif args.scaling.lower() == "strong" or args.scaling.lower() == "s":
+        elif args.scaling.lower() == 'strong' or args.scaling.lower() == 's':
             plot_strong_scaling(
                 args.csv_files,
                 args.gpus,
@@ -65,5 +61,5 @@ def main():
             )
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()

jax_hpc_profiler/plotting.py

@@ -4,23 +4,22 @@ 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 clean_up_csv, inspect_df, plot_with_pdims_strategy
+from .utils import clean_up_csv, plot_with_pdims_strategy
 
-np.seterr(divide="ignore")
+np.seterr(divide='ignore')
 
 
 def configure_axes(
     ax: Axes,
     x_values: List[int],
     y_values: List[float],
-    title: str,
+    title: Optional[str],
     xlabel: str,
     plotting_memory: bool = False,
-    memory_units: str = "bytes",
+    memory_units: str = 'bytes',
 ):
     """
     Configure the axes for the plot.
@@ -36,33 +35,32 @@ def configure_axes(
     xlabel : str
         The label for the x-axis.
     """
-    ylabel = ("Time (milliseconds)"
-              if not plotting_memory else f"Memory ({memory_units})")
-    f2 = lambda x: np.log2(x)
-    g2 = lambda x: 2**x
+    ylabel = 'Time (milliseconds)' if not plotting_memory else f'Memory ({memory_units})'
+
+    def f2(x):
+        return np.log2(x)
+
+    def g2(x):
+        return 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))))
+            10**t for t in range(int(np.floor(np.log10(y_min))), 1 + int(np.ceil(np.log10(y_max))))
         ]
         ax.set_yticks(time_ticks)
     ax.set_xticks(x_values)
     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.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},
+        loc='best',
     )
 
 
@@ -80,10 +78,10 @@ def plot_scaling(
     backends: Optional[List[str]] = None,
     precisions: Optional[List[str]] = None,
     functions: Optional[List[str]] = None,
-    plot_columns: List[str] = ["mean_time"],
-    memory_units: str = "bytes",
-    label_text: str = "plot",
-    pdims_strategy: List[str] = ["plot_fastest"],
+    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.
@@ -115,7 +113,7 @@ def plot_scaling(
     """
 
     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)))
@@ -133,28 +131,26 @@ def plot_scaling(
         x_values = []
         y_values = []
         for method, df in dataframes.items():
-
             filtered_method_df = df[df[fixed_column] == int(fixed_size)]
             if filtered_method_df.empty:
                 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)
-            precisions = (pd.unique(filtered_method_df["precision"])
-                          if precisions is None else precisions)
-            backends = (pd.unique(filtered_method_df["backend"])
-                        if backends is None else backends)
+            filtered_method_df = filtered_method_df.sort_values(by=[size_column])
+            functions = (
+                pd.unique(filtered_method_df['function']) if functions is None else functions
+            )
+            precisions = (
+                pd.unique(filtered_method_df['precision']) if precisions is None else precisions
+            )
+            backends = pd.unique(filtered_method_df['backend']) if backends is None else backends
 
-            combinations = product(backends, precisions, functions,
-                                   plot_columns)
+            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(
@@ -172,12 +168,13 @@ def plot_scaling(
                 y_values.extend(y_vals)
 
         if len(x_values) != 0:
-            plotting_memory = "time" not in plot_columns[0].lower()
+            plotting_memory = 'time' not in plot_columns[0].lower()
+            figure_title = f'{title} {fixed_size}' if title is not None else None
             configure_axes(
                 ax,
                 x_values,
                 y_values,
-                f"{title} {fixed_size}",
+                figure_title,
                 xlabel,
                 plotting_memory,
                 memory_units,
@@ -187,17 +184,12 @@ def plot_scaling(
         fig.delaxes(axs[i])
 
     fig.tight_layout()
-    rect = FancyBboxPatch((0.1, 0.1),
-                          0.8,
-                          0.8,
-                          boxstyle="round,pad=0.02",
-                          ec="black",
-                          fc="none")
+    rect = FancyBboxPatch((0.1, 0.1), 0.8, 0.8, boxstyle='round,pad=0.02', ec='black', fc='none')
     fig.patches.append(rect)
     if output is None:
         plt.show()
     else:
-        plt.savefig(output, bbox_inches="tight", transparent=True)
+        plt.savefig(output)
 
 
 def plot_strong_scaling(
@@ -207,14 +199,14 @@ def plot_strong_scaling(
     functions: Optional[List[str]] = None,
     precisions: Optional[List[str]] = None,
     pdims: Optional[List[str]] = None,
-    pdims_strategy: List[str] = ["plot_fastest"],
+    pdims_strategy: List[str] = ['plot_fastest'],
     print_decompositions: bool = False,
     backends: Optional[List[str]] = None,
-    plot_columns: List[str] = ["mean_time"],
-    memory_units: str = "bytes",
-    label_text: str = "%m%-%f%-%pn%-%pr%-%b%-%p%-%n%",
-    xlabel: str = "Number of GPUs",
-    title: str = "Data sizes",
+    plot_columns: List[str] = ['mean_time'],
+    memory_units: str = 'bytes',
+    label_text: str = '%m%-%f%-%pn%-%pr%-%b%-%p%-%n%',
+    xlabel: str = 'Number of GPUs',
+    title: str = 'Data sizes',
     figure_size: tuple = (6, 4),
     dark_bg: bool = False,
     output: Optional[str] = None,
@@ -235,14 +227,14 @@ def plot_strong_scaling(
         memory_units,
     )
     if len(dataframes) == 0:
-        print(f"No dataframes found for the given arguments. Exiting...")
+        print('No dataframes found for the given arguments. Exiting...')
         return
 
     plot_scaling(
         dataframes,
         available_data_sizes,
-        "gpus",
-        "x",
+        'gpus',
+        'x',
         xlabel,
         title,
         figure_size,
@@ -266,14 +258,14 @@ def plot_weak_scaling(
     functions: Optional[List[str]] = None,
     precisions: Optional[List[str]] = None,
     pdims: Optional[List[str]] = None,
-    pdims_strategy: List[str] = ["plot_fastest"],
+    pdims_strategy: List[str] = ['plot_fastest'],
     print_decompositions: bool = False,
     backends: Optional[List[str]] = None,
-    plot_columns: List[str] = ["mean_time"],
-    memory_units: str = "bytes",
-    label_text: str = "%m%-%f%-%pn%-%pr%-%b%-%p%-%n%",
-    xlabel: str = "Data sizes",
-    title: str = "Number of GPUs",
+    plot_columns: List[str] = ['mean_time'],
+    memory_units: str = 'bytes',
+    label_text: str = '%m%-%f%-%pn%-%pr%-%b%-%p%-%n%',
+    xlabel: str = 'Data sizes',
+    title: str = 'Number of GPUs',
     figure_size: tuple = (6, 4),
     dark_bg: bool = False,
     output: Optional[str] = None,
@@ -293,14 +285,14 @@ def plot_weak_scaling(
         memory_units,
     )
     if len(dataframes) == 0:
-        print(f"No dataframes found for the given arguments. Exiting...")
+        print('No dataframes found for the given arguments. Exiting...')
         return
 
     plot_scaling(
         dataframes,
         available_gpu_counts,
-        "x",
-        "gpus",
+        'x',
+        'gpus',
         xlabel,
         title,
         figure_size,