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

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: jax_hpc_profiler
-Version: 0.2.12
-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.12 → 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.12 → 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.12"
-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,5 +48,47 @@ 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"
+
+
+
+[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
+    "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.3.0/src/jax_hpc_profiler/__init__.py

@@ -0,0 +1,15 @@
+from .create_argparse import create_argparser
+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
+
+__all__ = [
+    'create_argparser',
+    'plot_strong_scaling',
+    'plot_weak_scaling',
+    'plot_weak_fixed_scaling',
+    'Timer',
+    'clean_up_csv',
+    'concatenate_csvs',
+    'plot_with_pdims_strategy',
+]

jax_hpc_profiler-0.3.0/src/jax_hpc_profiler/create_argparse.py

@@ -0,0 +1,218 @@
+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', 'WeakFixed', 'w', 's', 'wf'],
+        required=True,
+        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
+    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')
+
+        # 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.3.0/src/jax_hpc_profiler/main.py

@@ -0,0 +1,88 @@
+from .create_argparse import create_argparser
+from .plotting import plot_strong_scaling, plot_weak_fixed_scaling, plot_weak_scaling
+from .utils import 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('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':
+        scaling = args.scaling.lower()
+        if scaling in ('weak', '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.xlabel if getattr(args, 'xlabel', None) is not None else 'Number of GPUs',
+                args.title,
+                args.figure_size,
+                args.dark_bg,
+                args.output,
+                args.weak_ideal_line,
+                args.weak_reverse_axes,
+            )
+        elif scaling in ('strong', '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.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,
+            )
+
+
+if __name__ == '__main__':
+    main()