jax-hpc-profiler 0.2.8__tar.gz → 0.2.9__tar.gz

{jax_hpc_profiler-0.2.8 → jax_hpc_profiler-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 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

{jax_hpc_profiler-0.2.8 → jax_hpc_profiler-0.2.9}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jax_hpc_profiler"
-version = "0.2.8"
+version = "0.2.9"
 description = "HPC Plotter and profiler for benchmarking data made for JAX"
 authors = [
     { name="Wassim Kabalan" }

jax_hpc_profiler-0.2.9/src/jax_hpc_profiler/create_argparse.py

@@ -0,0 +1,197 @@
+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%",
+    )
+
+    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.9/src/jax_hpc_profiler/main.py

@@ -0,0 +1,65 @@
+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
+
+
+def main():
+    args = create_argparser()
+
+    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":
+            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__":
+    main()

jax_hpc_profiler-0.2.9/src/jax_hpc_profiler/plotting.py

@@ -0,0 +1,308 @@
+from itertools import product
+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
+
+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",
+):
+    """
+    Configure the axes for the plot.
+
+    Parameters
+    ----------
+    ax : Axes
+        The axes to configure.
+    x_values : List[int]
+        The x-axis values.
+    y_values : List[float]
+        The y-axis values.
+    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
+    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))
+    if not plotting_memory:
+        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))))
+        ]
+        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.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.
+
+    Parameters
+    ----------
+    dataframes : Dict[str, pd.DataFrame]
+        Dictionary of method names to dataframes.
+    fixed_sizes : List[int]
+        List of fixed sizes (data or GPUs) to plot.
+    size_column : str
+        Column name for the size axis ('x' for weak scaling, 'gpus' for strong scaling).
+    fixed_column : str
+        Column name for the fixed axis ('gpus' for weak scaling, 'x' for strong scaling).
+    xlabel : str
+        Label for the x-axis.
+    figure_size : tuple, optional
+        Size of the figure, by default (6, 4).
+    output : Optional[str], optional
+        Output file to save the plot, by default None.
+    dark_bg : bool, optional
+        Whether to use dark background for the plot, by default False.
+    print_decompositions : bool, optional
+        Whether to print decompositions on the plot, by default False.
+    backends : Optional[List[str]], optional
+        List of backends to include, by default None.
+    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")
+
+    num_subplots = len(fixed_sizes)
+    num_rows = int(np.ceil(np.sqrt(num_subplots)))
+    num_cols = int(np.ceil(num_subplots / num_rows))
+
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=figure_size)
+    if num_subplots > 1:
+        axs = axs.flatten()
+    else:
+        axs = [axs]
+
+    for i, fixed_size in enumerate(fixed_sizes):
+        ax: Axes = axs[i]
+
+        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)
+            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)]
+                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,
+                )
+
+                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,
+            )
+
+    for i in range(num_subplots, num_rows * num_cols):
+        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")
+    fig.patches.append(rect)
+    if output is None:
+        plt.show()
+    else:
+        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.
+    """
+
+    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.
+    """
+    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,
+    )