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

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

@@ -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.2.13 → jax_hpc_profiler-0.3.0}/README.md

@@ -1,5 +1,11 @@
 # 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
@@ -181,9 +187,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.2.13 → jax_hpc_profiler-0.3.0}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jax_hpc_profiler"
-version = "0.2.13"
-description = "HPC Plotter and profiler for benchmarking data made for JAX"
+version = "0.3.0"
+description = "A comprehensive benchmarking and profiling tool designed for JAX in HPC environments, offering automated instrumentation, strong/weak scaling analysis, and performance visualization."
 authors = [
     { name="Wassim Kabalan" }
 ]
@@ -14,12 +14,15 @@ dependencies = [
     "pandas",
     "matplotlib",
     "seaborn",
-    "tabulate"
+    "tabulate",
+    "adjustText",
+    "jax>=0.4.0",
+    "jaxtyping",
 ]
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.8"
-keywords = ["jax", "hpc", "profiler", "plotter", "benchmarking"]
+keywords = ["jax", "hpc", "profiling", "benchmarking", "visualization", "scaling", "performance-analysis", "gpu", "distributed-computing"]
 
 # For a list of valid classifiers, see https://pypi.org/classifiers/
 classifiers = [
@@ -45,6 +48,12 @@ classifiers = [
 
 urls = { "Homepage" = "https://github.com/ASKabalan/jax-hpc-profiler" }
 
+[project.optional-dependencies]
+test = [
+    "pytest",
+    "pytest-cov",
+]
+
 [project.scripts]
 jhp = "jax_hpc_profiler.main:main"
 
@@ -70,7 +79,16 @@ ignore = [
     'E731',
     'E741',
     'F722', # conflicts with jaxtyping Array annotations
+    "E402", # module level import not at top of file
 ]
 
+[tool.ruff.lint.per-file-ignores]
+"*.ipynb" = ["F401"]
+
 [tool.ruff.format]
 quote-style = 'single'
+
+[tool.pytest.ini_options]
+addopts = "--cov=src --cov-report=term-missing"
+testpaths = ["tests"]
+python_files = "test_*.py"

{jax_hpc_profiler-0.2.13 → jax_hpc_profiler-0.3.0}/src/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-0.2.13 → jax_hpc_profiler-0.3.0}/src/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-0.2.13 → jax_hpc_profiler-0.3.0}/src/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-0.2.13 → jax_hpc_profiler-0.3.0}/src/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-0.2.13 → jax_hpc_profiler-0.3.0}/src/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 → jax_hpc_profiler-0.3.0}/src/jax_hpc_profiler.egg-info/PKG-INFO

@@ -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.2.13 → jax_hpc_profiler-0.3.0}/src/jax_hpc_profiler.egg-info/SOURCES.txt

@@ -12,4 +12,6 @@ src/jax_hpc_profiler.egg-info/SOURCES.txt
 src/jax_hpc_profiler.egg-info/dependency_links.txt
 src/jax_hpc_profiler.egg-info/entry_points.txt
 src/jax_hpc_profiler.egg-info/requires.txt
-src/jax_hpc_profiler.egg-info/top_level.txt
+src/jax_hpc_profiler.egg-info/top_level.txt
+tests/test_plotting.py
+tests/test_timer.py

jax_hpc_profiler-0.3.0/src/jax_hpc_profiler.egg-info/requires.txt

@@ -0,0 +1,12 @@
+numpy
+pandas
+matplotlib
+seaborn
+tabulate
+adjustText
+jax>=0.4.0
+jaxtyping
+
+[test]
+pytest
+pytest-cov

jax_hpc_profiler-0.3.0/tests/test_plotting.py

@@ -0,0 +1,112 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from jax_hpc_profiler.plotting import (
+    configure_axes,
+    plot_strong_scaling,
+    plot_weak_fixed_scaling,
+    plot_weak_scaling,
+)
+
+
+@pytest.fixture
+def mock_plt():
+    with patch('jax_hpc_profiler.plotting.plt') as mock:
+        mock_fig = MagicMock()
+        mock_ax = MagicMock()
+        mock.subplots.return_value = (mock_fig, mock_ax)
+        yield mock
+
+
+@pytest.fixture
+def sample_csv(tmp_path):
+    csv_file = tmp_path / 'test_data.csv'
+    # Create data matching Timer.report format (19 columns)
+    # function,precision,x,y,z,px,py,backend,nodes,jit,min,max,mean,std,last,gen_code,arg_size,
+    # out_size,tmp_size
+    data = [
+        'fun1,float32,100,100,100,1,1,NCCL,1,0.1,1.0,1.2,1.1,0.01,1.1,1000,1000,1000,1000',
+        'fun1,float32,200,200,200,1,1,NCCL,1,0.2,2.0,2.4,2.2,0.02,2.2,2000,2000,2000,2000',
+        'fun1,float32,400,400,400,1,1,NCCL,1,0.4,4.0,4.8,4.4,0.04,4.4,4000,4000,4000,4000',
+        # Add entries for strong scaling (same x, diff nodes/gpus)
+        # nodes is used as 'gpus' in some logic?
+        # utils.py: df['gpus'] = df['px'] * df['py']
+        # So we vary px*py.
+        'fun2,float32,1000,1000,1000,1,1,NCCL,1,0.1,10.0,12.0,11.0,0.1,11.0,1000,1000,1000,1000',
+        'fun2,float32,1000,1000,1000,2,1,NCCL,2,0.1,5.0,6.0,5.5,0.05,5.5,1000,1000,1000,1000',
+        'fun2,float32,1000,1000,1000,2,2,NCCL,4,0.1,2.5,3.0,2.75,0.025,2.75,1000,1000,1000,1000',
+    ]
+    with open(csv_file, 'w') as f:
+        f.write('\n'.join(data) + '\n')
+    return str(csv_file)
+
+
+@pytest.fixture
+def mock_adjust_text():
+    with patch('jax_hpc_profiler.plotting.adjust_text') as mock:
+        yield mock
+
+
+def test_plot_weak_fixed_scaling(mock_plt, sample_csv):
+    # WeakFixed: vary data size (x), fixed GPUs (calculated from px*py).
+    # In sample_csv fun1: gpus=1, x=[100, 200, 400]
+
+    plot_weak_fixed_scaling(
+        csv_files=[sample_csv],
+        fixed_gpu_size=[1],
+        fixed_data_size=[100, 200, 400],
+        functions=['fun1'],
+        xlabel='Data Size',
+        title='Weak Fixed Scaling',
+    )
+
+    assert mock_plt.show.called or mock_plt.savefig.called
+
+
+def test_plot_strong_scaling(mock_plt, sample_csv):
+    # Strong: fixed data size (x), vary GPUs.
+    # In sample_csv fun2: x=1000, gpus=[1, 2, 4]
+
+    plot_strong_scaling(
+        csv_files=[sample_csv],
+        fixed_data_size=[1000],
+        fixed_gpu_size=[1, 2, 4],
+        functions=['fun2'],
+        xlabel='GPUs',
+        title='Strong Scaling',
+    )
+
+    assert mock_plt.show.called or mock_plt.savefig.called
+
+
+def test_plot_weak_scaling(mock_plt, mock_adjust_text, sample_csv):
+    # Weak: explicit pairs of (gpus, data_size)
+    # We have (1, 100) for fun1.
+
+    plot_weak_scaling(
+        csv_files=[sample_csv],
+        fixed_gpu_size=[1],
+        fixed_data_size=[100],
+        functions=['fun1'],
+        xlabel='GPUs',
+        title='Weak Scaling',
+    )
+
+    assert mock_plt.show.called or mock_plt.savefig.called
+    assert (
+        mock_adjust_text.called or not mock_adjust_text.called
+    )  # It's okay if called or not, just don't crash.
+
+
+def test_configure_axes():
+    # Test the helper directly
+    mock_ax = MagicMock()
+    configure_axes(
+        mock_ax, x_values=[1, 2, 4], y_values=[10, 5, 2.5], title='Test Plot', xlabel='X Label'
+    )
+
+    mock_ax.set_title.assert_called_with('Test Plot')
+    mock_ax.set_xlabel.assert_called_with('X Label')
+    mock_ax.set_xscale.assert_called()
+    mock_ax.set_yscale.assert_called_with('symlog')

jax_hpc_profiler-0.3.0/tests/test_timer.py

@@ -0,0 +1,103 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from jax_hpc_profiler import Timer
+
+
+# Simple JAX function
+@jax.jit
+def simple_add(x, y):
+    return x + y
+
+
+def test_timer_initialization():
+    timer = Timer()
+    assert timer.save_jaxpr is False
+    assert timer.compile_info is True
+    assert timer.jit_time == 0.0
+    assert len(timer.times) == 0
+
+
+def test_chrono_jit():
+    timer = Timer(save_jaxpr=True, compile_info=True)
+    x = jnp.ones((10, 10))
+    y = jnp.ones((10, 10))
+
+    out = timer.chrono_jit(simple_add, x, y)
+    assert out.shape == (10, 10)
+    assert timer.jit_time > 0
+    assert timer.compiled_code['JAXPR'] != 'N/A'
+    assert timer.compiled_code['COMPILED'] != 'N/A'
+    assert timer.profiling_data['generated_code'] != 'N/A'
+
+
+def test_chrono_fun():
+    timer = Timer()
+    x = jnp.ones((10, 10))
+    y = jnp.ones((10, 10))
+
+    # Run once to compile (if using jit externally, but here we just call the function)
+    out = timer.chrono_fun(simple_add, x, y)
+    assert out.shape == (10, 10)
+    assert len(timer.times) == 1
+
+
+def test_report(tmp_path):
+    timer = Timer(save_jaxpr=False)
+    x = jnp.ones((10, 10))
+    y = jnp.ones((10, 10))
+
+    timer.chrono_jit(simple_add, x, y)
+    for _ in range(5):
+        timer.chrono_fun(simple_add, x, y)
+
+    csv_file = tmp_path / 'report.csv'
+    md_file = tmp_path / 'report.md'
+
+    # We specify nodes=1 explicitely to match what we expect in some tests, though default is 1
+    timer.report(
+        str(csv_file),
+        function='simple_add',
+        x=10,
+        y=10,
+        precision='float32',
+        md_filename=str(md_file),
+        extra_info={'custom_key': 'custom_val'},
+    )
+
+    assert csv_file.exists()
+    assert md_file.exists()
+
+    with open(csv_file) as f:
+        content = f.read()
+        assert 'simple_add' in content
+        assert 'float32' in content
+
+    with open(md_file) as f:
+        content = f.read()
+        assert '# Reporting for simple_add' in content
+        assert 'custom_key' in content
+        assert 'custom_val' in content
+
+
+def test_normalize_memory_units():
+    timer = Timer()
+    # Mocking internal state as if we had data
+    timer.jax_fn = True
+    timer.compile_info = True
+
+    assert timer._normalize_memory_units(100) == '100.00 B'
+    assert timer._normalize_memory_units(1024) == '1.00 KB'
+    assert timer._normalize_memory_units(1024**2) == '1.00 MB'
+    assert timer._normalize_memory_units(1024**3) == '1.00 GB'
+
+
+def test_get_mean_times():
+    timer = Timer()
+    timer.times = [10.0, 20.0, 30.0]
+
+    means = timer._get_mean_times()
+    assert isinstance(means, np.ndarray)
+    assert len(means) == 3
+    assert means[0] == 10.0

jax_hpc_profiler-0.2.13/src/jax_hpc_profiler.egg-info/requires.txt

@@ -1,5 +0,0 @@
-numpy
-pandas
-matplotlib
-seaborn
-tabulate