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

jax_hpc_profiler/__init__.py

@@ -1,5 +1,5 @@
 from .create_argparse import create_argparser
-from .plotting import plot_strong_scaling, plot_weak_scaling
+from .plotting import plot_strong_scaling, plot_weak_fixed_scaling, plot_weak_scaling
 from .timer import Timer
 from .utils import clean_up_csv, concatenate_csvs, plot_with_pdims_strategy
 
@@ -7,6 +7,7 @@ __all__ = [
     'create_argparser',
     'plot_strong_scaling',
     'plot_weak_scaling',
+    'plot_weak_fixed_scaling',
     'Timer',
     'clean_up_csv',
     'concatenate_csvs',

jax_hpc_profiler/create_argparse.py

@@ -135,9 +135,24 @@ def create_argparser():
     plot_parser.add_argument(
         '-sc',
         '--scaling',
-        choices=['Weak', 'Strong', 'w', 's'],
+        choices=['Weak', 'Strong', 'WeakFixed', 'w', 's', 'wf'],
         required=True,
-        help='Scaling type (Weak or Strong)',
+        help='Scaling type (Strong, Weak, or WeakFixed)',
+    )
+
+    # Weak-scaling specific options
+    plot_parser.add_argument(
+        '--weak_ideal_line',
+        action='store_true',
+        help='Overlay an ideal flat line for weak scaling (Weak mode only)',
+    )
+    plot_parser.add_argument(
+        '--weak_reverse_axes',
+        action='store_true',
+        help=(
+            'Weak mode only: put data size on the x-axis and annotate each point with GPUs instead '
+            'of data size. Requires --gpus and --data_size with equal lengths.'
+        ),
     )
 
     # Label customization argument
@@ -196,4 +211,8 @@ def create_argparser():
         else:
             raise ValueError('Either plot_times or plot_memory should be provided')
 
+        # Note: for Weak scaling, plot_weak_scaling enforces that both gpus and
+        # data_size are provided and have matching lengths. For Strong and
+        # WeakFixed, gpus/data_size remain optional as before.
+
     return args

jax_hpc_profiler/main.py

@@ -1,5 +1,5 @@
 from .create_argparse import create_argparser
-from .plotting import plot_strong_scaling, plot_weak_scaling
+from .plotting import plot_strong_scaling, plot_weak_fixed_scaling, plot_weak_scaling
 from .utils import concatenate_csvs
 
 
@@ -19,7 +19,8 @@ def main():
         print(' -- %p% or %pdims%: pdims')
         print(' -- %n% or %node%: node')
     elif args.command == 'plot':
-        if args.scaling.lower() == 'weak' or args.scaling.lower() == 'w':
+        scaling = args.scaling.lower()
+        if scaling in ('weak', 'w'):
             plot_weak_scaling(
                 args.csv_files,
                 args.gpus,
@@ -33,13 +34,15 @@ def main():
                 args.plot_columns,
                 args.memory_units,
                 args.label_text,
+                args.xlabel if getattr(args, 'xlabel', None) is not None else 'Number of GPUs',
                 args.title,
-                args.label_text,
                 args.figure_size,
                 args.dark_bg,
                 args.output,
+                args.weak_ideal_line,
+                args.weak_reverse_axes,
             )
-        elif args.scaling.lower() == 'strong' or args.scaling.lower() == 's':
+        elif scaling in ('strong', 's'):
             plot_strong_scaling(
                 args.csv_files,
                 args.gpus,
@@ -53,8 +56,28 @@ def main():
                 args.plot_columns,
                 args.memory_units,
                 args.label_text,
-                args.title,
+                args.xlabel if getattr(args, 'xlabel', None) is not None else 'Number of GPUs',
+                args.title if getattr(args, 'title', None) is not None else 'Data sizes',
+                args.figure_size,
+                args.dark_bg,
+                args.output,
+            )
+        elif scaling in ('weakfixed', 'wf'):
+            plot_weak_fixed_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.xlabel if getattr(args, 'xlabel', None) is not None else 'Data sizes',
+                args.title if getattr(args, 'title', None) is not None else 'Number of GPUs',
                 args.figure_size,
                 args.dark_bg,
                 args.output,

jax_hpc_profiler/plotting.py

@@ -4,6 +4,7 @@ from typing import Dict, List, Optional
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+from adjustText import adjust_text
 from matplotlib.axes import Axes
 from matplotlib.patches import FancyBboxPatch
 
@@ -252,6 +253,195 @@ def plot_strong_scaling(
 
 
 def plot_weak_scaling(
+    csv_files: List[str],
+    fixed_gpu_size: Optional[List[int]] = None,
+    fixed_data_size: Optional[List[int]] = None,
+    functions: Optional[List[str]] = None,
+    precisions: Optional[List[str]] = None,
+    pdims: Optional[List[str]] = None,
+    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: Optional[str] = None,
+    figure_size: tuple = (6, 4),
+    dark_bg: bool = False,
+    output: Optional[str] = None,
+    ideal_line: bool = False,
+    reverse_axes: bool = False,
+):
+    """
+    Plot true weak scaling: runtime vs GPUs for explicit (gpus, data size) sequences.
+
+    Both ``fixed_gpu_size`` and ``fixed_data_size`` must be provided and have the same length,
+    representing explicit weak-scaling pairs (gpus[i], data_size[i]).
+
+    reverse_axes:
+        - False (default): x-axis is GPUs, y-axis is time; points are annotated with
+        ``N=<data_size>``.
+        - True: x-axis is data size, y-axis is time; points are annotated with ``GPUs=<gpu_count>``.
+    """
+    if fixed_gpu_size is None or fixed_data_size is None:
+        raise ValueError(
+            'Weak scaling requires both fixed_gpu_size (gpus) and fixed_data_size (problem sizes).'
+        )
+    if len(fixed_gpu_size) != len(fixed_data_size):
+        raise ValueError(
+            'Weak scaling requires fixed_gpu_size and fixed_data_size lists of equal length.'
+        )
+
+    gpu_to_data = {int(g): int(d) for g, d in zip(fixed_gpu_size, fixed_data_size)}
+    data_to_gpu = {int(d): int(g) for g, d in zip(fixed_gpu_size, fixed_data_size)}
+    x_col = 'x' if reverse_axes else 'gpus'
+
+    dataframes, _, _ = clean_up_csv(
+        csv_files,
+        precisions,
+        functions,
+        fixed_gpu_size,
+        fixed_data_size,
+        pdims,
+        pdims_strategy,
+        backends,
+        memory_units,
+    )
+    if len(dataframes) == 0:
+        print('No dataframes found for the given arguments. Exiting...')
+        return
+
+    if dark_bg:
+        plt.style.use('dark_background')
+
+    fig, ax = plt.subplots(figsize=figure_size)
+
+    x_values: List[float] = []
+    y_values: List[float] = []
+    annotations: List = []
+    ideal_line_plotted = False
+
+    for method, df in dataframes.items():
+        # Determine parameter sets from the filtered dataframe if not provided
+        local_functions = pd.unique(df['function']) if functions is None else functions
+        local_precisions = pd.unique(df['precision']) if precisions is None else precisions
+        local_backends = pd.unique(df['backend']) if backends is None else backends
+
+        combinations = product(local_backends, local_precisions, local_functions, plot_columns)
+
+        for backend, precision, function, plot_column in combinations:
+            base_df = df[
+                (df['backend'] == backend)
+                & (df['precision'] == precision)
+                & (df['function'] == function)
+            ]
+            if base_df.empty:
+                continue
+
+            # Keep only rows matching any of the (gpus, x) pairs
+            mask = pd.Series(False, index=base_df.index)
+            for g, d in zip(fixed_gpu_size, fixed_data_size):
+                mask |= (base_df['gpus'] == int(g)) & (base_df['x'] == int(d))
+
+            filtered_params_df = base_df[mask]
+            if filtered_params_df.empty:
+                continue
+
+            x_vals, y_vals = plot_with_pdims_strategy(
+                ax,
+                filtered_params_df,
+                method,
+                pdims_strategy,
+                print_decompositions,
+                x_col,
+                plot_column,
+                label_text,
+            )
+            if x_vals is None or len(x_vals) == 0:
+                continue
+
+            x_arr = np.asarray(x_vals).reshape(-1)
+            y_arr = np.asarray(y_vals).reshape(-1)
+
+            # Annotate every point with data size or GPU count depending on axis choice.
+            # Use plain data coordinates for the text; adjust_text will then only move
+            # the labels slightly (mostly vertically) to avoid overlap.
+            for xv, yv in zip(x_arr, y_arr):
+                if reverse_axes:
+                    gpu = data_to_gpu.get(int(xv))
+                    if gpu is None:
+                        continue
+                    label = f'GPUs={gpu}'
+                else:
+                    data_size = gpu_to_data.get(int(xv))
+                    if data_size is None:
+                        continue
+                    label = f'N={data_size}'
+
+                text_obj = ax.text(
+                    float(xv),
+                    float(yv),
+                    label,
+                    ha='center',
+                    va='bottom',
+                    fontsize='small',
+                    clip_on=True,
+                )
+                annotations.append(text_obj)
+
+            x_values.extend(x_arr.tolist())
+            y_values.extend(y_arr.tolist())
+
+            if ideal_line and not ideal_line_plotted:
+                # Use the smallest x value in this curve as baseline
+                baseline_index = np.argmin(x_arr)
+                baseline_y = y_arr[baseline_index]
+                ax.hlines(
+                    baseline_y,
+                    xmin=float(np.min(x_arr)),
+                    xmax=float(np.max(x_arr)),
+                    colors='gray',
+                    linestyles='dashed',
+                    label='Ideal weak scaling',
+                )
+                ideal_line_plotted = True
+                y_values.append(float(baseline_y))
+
+    if x_values:
+        plotting_memory = 'time' not in plot_columns[0].lower()
+        figure_title = title if title is not None else 'Weak scaling'
+        configure_axes(
+            ax,
+            x_values,
+            y_values,
+            figure_title,
+            xlabel,
+            plotting_memory,
+            memory_units,
+        )
+        if annotations:
+            ax.figure.canvas.draw()
+            adjust_text(
+                annotations,
+                ax=ax,
+                # keep points aligned in x, only allow vertical motion
+                only_move={'text': 'y', 'static': 'y'},
+                expand=(1.02, 1.05),
+                force_text=(0.08, 0.2),
+                max_move=(0, 30),
+            )
+
+    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)
+
+
+def plot_weak_fixed_scaling(
     csv_files: List[str],
     fixed_gpu_size: Optional[List[int]] = None,
     fixed_data_size: Optional[List[int]] = None,
@@ -271,7 +461,7 @@ def plot_weak_scaling(
     output: Optional[str] = None,
 ):
     """
-    Plot weak scaling based on the data size.
+    Plot size scaling at fixed GPU count (previous weak-scaling behavior).
     """
     dataframes, available_gpu_counts, _ = clean_up_csv(
         csv_files,

jax_hpc_profiler/timer.py

@@ -6,8 +6,7 @@ from typing import Any, Callable, Optional, Tuple
 import jax
 import jax.numpy as jnp
 import numpy as np
-from jax import make_jaxpr
-from jax.experimental.shard_map import shard_map
+from jax import make_jaxpr, shard_map
 from jax.sharding import NamedSharding
 from jax.sharding import PartitionSpec as P
 from jaxtyping import Array
@@ -25,8 +24,13 @@ class Timer:
     ):
         self.jit_time = 0.0
         self.times = []
-        self.profiling_data = {}
-        self.compiled_code = {}
+        self.profiling_data = {
+            'generated_code': 'N/A',
+            'argument_size': 'N/A',
+            'output_size': 'N/A',
+            'temp_size': 'N/A',
+        }
+        self.compiled_code = {'JAXPR': 'N/A', 'LOWERED': 'N/A', 'COMPILED': 'N/A'}
         self.save_jaxpr = save_jaxpr
         self.compile_info = compile_info
         self.jax_fn = jax_fn
@@ -181,10 +185,10 @@ class Timer:
             mean_time = np.mean(times_array)
             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']
+            generated_code = self.profiling_data.get('generated_code', 'N/A')
+            argument_size = self.profiling_data.get('argument_size', 'N/A')
+            output_size = self.profiling_data.get('output_size', 'N/A')
+            temp_size = self.profiling_data.get('temp_size', 'N/A')
 
             csv_line = (
                 f'{function},{precision},{x},{y},{z},{px},{py},{backend},{nodes},'

{jax_hpc_profiler-0.2.13.dist-info → jax_hpc_profiler-0.3.0.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: jax_hpc_profiler
-Version: 0.2.13
-Summary: HPC Plotter and profiler for benchmarking data made for JAX
+Version: 0.3.0
+Summary: A comprehensive benchmarking and profiling tool designed for JAX in HPC environments, offering automated instrumentation, strong/weak scaling analysis, and performance visualization.
 Author: Wassim Kabalan
 License:                     GNU GENERAL PUBLIC LICENSE
                                Version 3, 29 June 2007
@@ -679,7 +679,7 @@ License:                     GNU GENERAL PUBLIC LICENSE
         <https://www.gnu.org/licenses/why-not-lgpl.html>.
         
 Project-URL: Homepage, https://github.com/ASKabalan/jax-hpc-profiler
-Keywords: jax,hpc,profiler,plotter,benchmarking
+Keywords: jax,hpc,profiling,benchmarking,visualization,scaling,performance-analysis,gpu,distributed-computing
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
@@ -698,10 +698,22 @@ Requires-Dist: pandas
 Requires-Dist: matplotlib
 Requires-Dist: seaborn
 Requires-Dist: tabulate
+Requires-Dist: adjustText
+Requires-Dist: jax>=0.4.0
+Requires-Dist: jaxtyping
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
 Dynamic: license-file
 
 # JAX HPC Profiler
 
+[![Build](https://github.com/ASKabalan/jax-hpc-profiler/actions/workflows/python-publish.yml/badge.svg)](https://github.com/ASKabalan/jax-hpc-profiler/actions/workflows/python-publish.yml)
+[![Code Formatting](https://github.com/ASKabalan/jax-hpc-profiler/actions/workflows/formatting.yml/badge.svg)](https://github.com/ASKabalan/jax-hpc-profiler/actions/workflows/formatting.yml)
+[![Tests](https://github.com/ASKabalan/jax-hpc-profiler/actions/workflows/tests.yml/badge.svg)](https://github.com/ASKabalan/jax-hpc-profiler/actions/workflows/tests.yml)
+[![Notebooks](https://img.shields.io/github/actions/workflow/status/ASKabalan/jax-hpc-profiler/notebooks.yml?logo=jupyter&label=notebooks)](https://github.com/ASKabalan/jax-hpc-profiler/actions/workflows/notebooks.yml)
+[![GPLv3 License](https://img.shields.io/badge/License-GPL%20v3-yellow.svg)](https://www.gnu.org/licenses/gpl-3.0)
+
 JAX HPC Profiler is a tool designed for benchmarking and visualizing performance data in high-performance computing (HPC) environments. It provides functionalities to generate, concatenate, and plot CSV data from various runs.
 
 ## Table of Contents
@@ -883,9 +895,29 @@ jax-hpc-profiler plot -f <csv_files> [options]
 - `-db, --dark_bg`: Use dark background for plotting.
 - `-pd, --print_decompositions`: Print decompositions on plot (experimental).
 - `-b, --backends`: List of backends to include. This argument can be multiple ones.
-- `-sc, --scaling`: Scaling type (`Weak`, `Strong`).
+- `-sc, --scaling`: Scaling type (`Strong`, `Weak`, `WeakFixed`).
+  - `Strong`: strong scaling with fixed global problem size(s), plotting runtime (or memory) versus number of GPUs.
+  - `Weak`: true weak scaling with explicit `(gpus, data_size)` sequences; requires that `-g/--gpus` and `-d/--data_size` are both provided and have the same length, and plots runtime (or memory) versus number of GPUs on a single figure.
+  - `WeakFixed`: size scaling at fixed GPU count (previous weak behavior); plots runtime (or memory) versus data size, grouped by number of GPUs.
+- `--weak_ideal_line`: When using `-sc Weak`, overlay an ideal flat line based on the smallest-GPU runtime for the first plotted weak-scaling curve.
 - `-l, --label_text`: 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%`).
 
+### Weak scaling CLI example
+
+For a weak-scaling run where work per GPU is kept approximately constant, you might provide matching GPU and data-size sequences, for example:
+
+```bash
+jax-hpc-profiler plot \
+  -f MYDATA.csv \
+  -pt mean_time \
+  -sc Weak \
+  -g 1 2 4 8 \
+  -d 32 64 128 256 \
+  --weak_ideal_line
+```
+
+This will produce a single weak-scaling plot of runtime versus number of GPUs, using the points `(gpus, data_size) = (1, 32), (2, 64), (4, 128), (8, 256)` and overlay an ideal weak-scaling reference line.
+
 ## Examples
 
 The repository includes examples for both profiling and plotting.

jax_hpc_profiler-0.3.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+jax_hpc_profiler/__init__.py,sha256=GIYY_D3CqhTzEf9Bh8ihar_-ntDYSQpp2utFIgYRbYg,444
+jax_hpc_profiler/create_argparse.py,sha256=ff2e8PvHmbbyF13OH2FTLlpnIUGp9xP8kS--XuJuhZ4,6582
+jax_hpc_profiler/main.py,sha256=ehqU6HwqhjKLs_34tmzWFQU2G-kSiVmhKJ1HIAw-6Lg,3262
+jax_hpc_profiler/plotting.py,sha256=vQsykw4JJNZn6Z6IR5_VABXEHKBhESQdAoAAN4dOaPk,15998
+jax_hpc_profiler/timer.py,sha256=5coHheE6eaviLCZsPuXodbl7pYW9ora-GU9M6PJqRNQ,10442
+jax_hpc_profiler/utils.py,sha256=IfGDbKldJXiDhxb02IxmQV51SFIBYLDUL7Se_OtEOkc,14963
+jax_hpc_profiler-0.3.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+jax_hpc_profiler-0.3.0.dist-info/METADATA,sha256=EI-Qb9STk1q_mJC2WMaBGLaqdj8Adtp0FpNeRRx6NNQ,51620
+jax_hpc_profiler-0.3.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+jax_hpc_profiler-0.3.0.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
+jax_hpc_profiler-0.3.0.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
+jax_hpc_profiler-0.3.0.dist-info/RECORD,,

{jax_hpc_profiler-0.2.13.dist-info → jax_hpc_profiler-0.3.0.dist-info}/WHEEL

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

jax_hpc_profiler-0.2.13.dist-info/RECORD

@@ -1,12 +0,0 @@
-jax_hpc_profiler/__init__.py,sha256=c2n33ZXVgUS8vo5xAEW-TcSi_UzJp616KdGEb3iO6p4,388
-jax_hpc_profiler/create_argparse.py,sha256=J1RF4n2e85QReoI_fqXxK5BMAUgzueHmObKOh4YHopE,5821
-jax_hpc_profiler/main.py,sha256=YPLkZCmtjzNoDrzTA4CWL8y39Spz3qbCS91eP2pqP5Y,2224
-jax_hpc_profiler/plotting.py,sha256=Lg157H3mrF3zHc4BIplddKu9f0viQkaQhtCCAQBxinE,9167
-jax_hpc_profiler/timer.py,sha256=0lbJgNh3GT1dFOpNOA4Fwvsm9JNp-J1xDdLFaaQ6jaY,10237
-jax_hpc_profiler/utils.py,sha256=IfGDbKldJXiDhxb02IxmQV51SFIBYLDUL7Se_OtEOkc,14963
-jax_hpc_profiler-0.2.13.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jax_hpc_profiler-0.2.13.dist-info/METADATA,sha256=YyHfP98Vz8ya23YsRPV2rehbRoFsO3pziOgnoX5DitE,49186
-jax_hpc_profiler-0.2.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jax_hpc_profiler-0.2.13.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
-jax_hpc_profiler-0.2.13.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
-jax_hpc_profiler-0.2.13.dist-info/RECORD,,