jax-hpc-profiler 0.2.11__tar.gz → 0.2.13__tar.gz

{jax_hpc_profiler-0.2.11 → jax_hpc_profiler-0.2.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax_hpc_profiler
-Version: 0.2.11
+Version: 0.2.13
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
 License:                     GNU GENERAL PUBLIC LICENSE

{jax_hpc_profiler-0.2.11 → jax_hpc_profiler-0.2.13}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jax_hpc_profiler"
-version = "0.2.11"
+version = "0.2.13"
 description = "HPC Plotter and profiler for benchmarking data made for JAX"
 authors = [
     { name="Wassim Kabalan" }
@@ -47,3 +47,30 @@ urls = { "Homepage" = "https://github.com/ASKabalan/jax-hpc-profiler" }
 
 [project.scripts]
 jhp = "jax_hpc_profiler.main:main"
+
+
+
+[tool.ruff]
+line-length = 100
+fix = true           # autofix issues
+force-exclude = true # useful with ruff-pre-commit plugin
+src = ["src"]
+
+[tool.ruff.lint]
+select = [
+    'ARG001', # flake8-unused-function-arguments
+    'E',      # pycodestyle-errors
+    'F',      # pyflakes
+    'I',      # isort
+    'UP',     # pyupgrade
+    'T10',    # flake8-debugger
+]
+ignore = [
+    'E203',
+    'E731',
+    'E741',
+    'F722', # conflicts with jaxtyping Array annotations
+]
+
+[tool.ruff.format]
+quote-style = 'single'

{jax_hpc_profiler-0.2.11 → jax_hpc_profiler-0.2.13}/src/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-0.2.13/src/jax_hpc_profiler/create_argparse.py

@@ -0,0 +1,199 @@
+import argparse
+
+
+def create_argparser():
+    """
+    Create argument parser for the HPC Plotter package.
+
+    Returns
+    -------
+    argparse.Namespace
+        Parsed and validated arguments.
+    """
+    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)
+
+    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='*',
+        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',
+        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',
+        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)',
+    )
+
+    # 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',
+    )
+
+    # Backend related arguments
+    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', 'w', 's'],
+        required=True,
+        help='Scaling type (Weak or Strong)',
+    )
+
+    # Label customization argument
+    plot_parser.add_argument(
+        '-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%',
+    )
+
+    plot_parser.add_argument(
+        '-xl',
+        '--xlabel',
+        type=str,
+        help='X-axis label for the plot',
+    )
+    plot_parser.add_argument(
+        '-tl',
+        '--title',
+        type=str,
+        help='Title for the plot',
+    )
+
+    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')
+
+    return args

{jax_hpc_profiler-0.2.11 → jax_hpc_profiler-0.2.13}/src/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-0.2.11 → jax_hpc_profiler-0.2.13}/src/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,