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

jax_hpc_profiler/utils.py

@@ -1,13 +1,11 @@
 import os
 from typing import Dict, List, Optional, Tuple
 
-import matplotlib.pyplot as plt
-import numpy as np
 import pandas as pd
 from matplotlib.axes import Axes
 
 
-def inspect_data(dataframes: Dict[str, pd.DataFrame]):
+def inspect_data(dataframes: Dict[str, pd.DataFrame]) -> None:
     """
     Inspect the dataframes.
 
@@ -16,16 +14,16 @@ def inspect_data(dataframes: Dict[str, pd.DataFrame]):
     dataframes : Dict[str, pd.DataFrame]
         Dictionary of method names to dataframes.
     """
-    print("=" * 80)
-    print("Inspecting dataframes...")
-    print("=" * 80)
+    print('=' * 80)
+    print('Inspecting dataframes...')
+    print('=' * 80)
     for method, df in dataframes.items():
-        print(f"Method: {method}")
+        print(f'Method: {method}')
         inspect_df(df)
-    print("=" * 80)
+    print('=' * 80)
 
 
-def inspect_df(df: pd.DataFrame):
+def inspect_df(df: pd.DataFrame) -> None:
     """
     Inspect the dataframe.
 
@@ -35,24 +33,24 @@ def inspect_df(df: pd.DataFrame):
         The dataframe to inspect.
     """
     print(df.to_markdown())
-    print("-" * 80)
+    print('-' * 80)
 
 
 params_dict = {
-    "%pn%": "%plot_name%",
-    "%m%": "%method_name%",
-    "%n%": "%node%",
-    "%b%": "%backend%",
-    "%f%": "%function%",
-    "%cn%": "%column_name%",
-    "%pr%": "%precision%",
-    "%p%": "%decomposition%",
-    "%d%": "%data_size%",
-    "%g%": "%nb_gpu%"
+    '%pn%': '%plot_name%',
+    '%m%': '%method_name%',
+    '%n%': '%node%',
+    '%b%': '%backend%',
+    '%f%': '%function%',
+    '%cn%': '%column_name%',
+    '%pr%': '%precision%',
+    '%p%': '%decomposition%',
+    '%d%': '%data_size%',
+    '%g%': '%nb_gpu%',
 }
 
 
-def expand_label(label_template: str, params: dict) -> str:
+def expand_label(label_template: str, params: dict[str, str]) -> str:
     """
     Expand the label template with the provided parameters.
 
@@ -72,14 +70,20 @@ def expand_label(label_template: str, params: dict) -> str:
         label_template = label_template.replace(key, value)
 
     for key, value in params.items():
-        label_template = label_template.replace(f"%{key}%", str(value))
+        label_template = label_template.replace(f'%{key}%', str(value))
     return label_template
 
 
-def plot_with_pdims_strategy(ax: Axes, df: pd.DataFrame, method: str,
-                             pdims_strategy: List[str],
-                             print_decompositions: bool, x_col: str,
-                             y_col: str, label_template: str):
+def plot_with_pdims_strategy(
+    ax: Axes,
+    df: pd.DataFrame,
+    method: str,
+    pdims_strategy: List[str],
+    print_decompositions: bool,
+    x_col: str,
+    y_col: str,
+    label_template: str,
+) -> Optional[Tuple[List[float], List[float]]]:
     """
     Plot the data based on the pdims strategy.
 
@@ -109,12 +113,12 @@ def plot_with_pdims_strategy(ax: Axes, df: pd.DataFrame, method: str,
         Template for plot labels with placeholders.
     """
     label_params = {
-        "plot_name": y_col,
-        "method_name": method,
-        "backend": df['backend'].values[0],
-        "node": df['nodes'].values[0],
-        "precision": df['precision'].values[0],
-        "function": df['function'].values[0],
+        'plot_name': y_col,
+        'method_name': method,
+        'backend': df['backend'].values[0],
+        'node': df['nodes'].values[0],
+        'precision': df['precision'].values[0],
+        'function': df['function'].values[0],
     }
 
     if 'plot_fastest' in pdims_strategy:
@@ -126,51 +130,50 @@ def plot_with_pdims_strategy(ax: Axes, df: pd.DataFrame, method: str,
             group.sort_values(by=[y_col], inplace=True, ascending=True)
             sorted_dfs.append(group.iloc[0])
         sorted_df = pd.DataFrame(sorted_dfs)
-        label_params.update({
-            "decomposition":
-            f"{group['px'].values[0]}x{group['py'].values[0]}"
-        })
+        label_params.update({'decomposition': f'{group["px"].values[0]}x{group["py"].values[0]}'})
         label = expand_label(label_template, label_params)
-        ax.plot(sorted_df[x_col].values,
-                sorted_df[y_col],
-                marker='o',
-                linestyle='-',
-                label=label)
+        ax.plot(
+            sorted_df[x_col].values,
+            sorted_df[y_col],
+            marker='o',
+            linestyle='-',
+            label=label,
+        )
         # TODO(wassim) : this is not working very well
         if print_decompositions:
-            for j, (px, py) in enumerate(zip(sorted_df['px'],
-                                             sorted_df['py'])):
+            for j, (px, py) in enumerate(zip(sorted_df['px'], sorted_df['py'])):
                 ax.annotate(
-                    f"{px}x{py}",
+                    f'{px}x{py}',
                     (sorted_df[x_col].values[j], sorted_df[y_col].values[j]),
-                    textcoords="offset points",
+                    textcoords='offset points',
                     xytext=(0, 10),
                     ha='center',
-                    color='red' if j == 0 else 'white')
+                    color='red' if j == 0 else 'white',
+                )
         return sorted_df[x_col].values, sorted_df[y_col].values
 
-    elif any(strategy in pdims_strategy
-             for strategy in ['plot_all', 'slab_yz', 'slab_xy', 'pencils']):
+    elif any(
+        strategy in pdims_strategy for strategy in ['plot_all', 'slab_yz', 'slab_xy', 'pencils']
+    ):
         df_decomp = df.groupby(['decomp'])
         x_values = []
         y_values = []
         for _, group in df_decomp:
-            group.drop_duplicates(subset=[x_col, 'decomp'],
-                                  keep='last',
-                                  inplace=True)
+            group.drop_duplicates(subset=[x_col, 'decomp'], keep='last', inplace=True)
             group.sort_values(by=[x_col], inplace=True, ascending=False)
             # filter decomp based on pdims_strategy
-            if 'plot_all' not in pdims_strategy and group['decomp'].values[
-                    0] not in pdims_strategy:
+            if 'plot_all' not in pdims_strategy and group['decomp'].values[0] not in pdims_strategy:
                 continue
 
-            label_params.update({"decomposition": group['decomp'].values[0]})
+            label_params.update({'decomposition': group['decomp'].values[0]})
             label = expand_label(label_template, label_params)
-            ax.plot(group[x_col].values,
-                    group[y_col],
-                    marker='o',
-                    linestyle='-',
-                    label=label)
+            ax.plot(
+                group[x_col].values,
+                group[y_col],
+                marker='o',
+                linestyle='-',
+                label=label,
+            )
             x_values.extend(group[x_col].values)
             y_values.extend(group[y_col].values)
         return x_values, y_values
@@ -204,39 +207,63 @@ def concatenate_csvs(root_dir: str, output_dir: str):
                 if file.endswith('.csv'):
                     csv_file_path = os.path.join(root, file)
                     print(f'Concatenating {csv_file_path}...')
-                    df = pd.read_csv(csv_file_path,
-                                     header=None,
-                                     names=[
-                                         "function", "precision", "x", "y",
-                                         "z", "px", "py", "backend", "nodes",
-                                         "jit_time", "min_time", "max_time",
-                                         "mean_time", "std_div", "last_time",
-                                         "generated_code", "argument_size",
-                                         "output_size", "temp_size", "flops"
-                                     ],
-                                     index_col=False)
+                    df = pd.read_csv(
+                        csv_file_path,
+                        header=None,
+                        names=[
+                            'function',
+                            'precision',
+                            'x',
+                            'y',
+                            'z',
+                            'px',
+                            'py',
+                            'backend',
+                            'nodes',
+                            'jit_time',
+                            'min_time',
+                            'max_time',
+                            'mean_time',
+                            'std_div',
+                            'last_time',
+                            'generated_code',
+                            'argument_size',
+                            'output_size',
+                            'temp_size',
+                            'flops',
+                        ],
+                        index_col=False,
+                    )
                     if file not in combined_dfs:
                         combined_dfs[file] = df
                     else:
-                        combined_dfs[file] = pd.concat(
-                            [combined_dfs[file], df], ignore_index=True)
+                        combined_dfs[file] = pd.concat([combined_dfs[file], df], ignore_index=True)
 
         # Remove duplicates based on specified columns and save
         for file_name, combined_df in combined_dfs.items():
-            combined_df.drop_duplicates(subset=[
-                "function", "precision", "x", "y", "z", "px", "py", "backend",
-                "nodes"
-            ],
-                                        keep='last',
-                                        inplace=True)
+            combined_df.drop_duplicates(
+                subset=[
+                    'function',
+                    'precision',
+                    'x',
+                    'y',
+                    'z',
+                    'px',
+                    'py',
+                    'backend',
+                    'nodes',
+                ],
+                keep='last',
+                inplace=True,
+            )
 
             gpu_output_dir = os.path.join(output_dir, gpu)
             if not os.path.exists(gpu_output_dir):
-                print(f"Creating directory {gpu_output_dir}")
+                print(f'Creating directory {gpu_output_dir}')
                 os.makedirs(gpu_output_dir)
 
             output_file = os.path.join(gpu_output_dir, file_name)
-            print(f"Writing file to {output_file}...")
+            print(f'Writing file to {output_file}...')
             combined_df.to_csv(output_file, index=False)
 
 
@@ -287,42 +314,59 @@ def clean_up_csv(
         file_name = os.path.splitext(os.path.basename(csv_file))[0]
         ext = os.path.splitext(os.path.basename(csv_file))[1]
         if ext != '.csv':
-            print(f"Ignoring {csv_file} as it is not a CSV file")
+            print(f'Ignoring {csv_file} as it is not a CSV file')
             continue
 
-        df = pd.read_csv(csv_file,
-                         header=None,
-                         skiprows=0,
-                         names=[
-                             "function", "precision", "x", "y", "z", "px",
-                             "py", "backend", "nodes", "jit_time", "min_time",
-                             "max_time", "mean_time", "std_div", "last_time",
-                             "generated_code", "argument_size", "output_size",
-                             "temp_size", "flops"
-                         ],
-                         dtype={
-                             "function": str,
-                             "precision": str,
-                             "x": int,
-                             "y": int,
-                             "z": int,
-                             "px": int,
-                             "py": int,
-                             "backend": str,
-                             "nodes": int,
-                             "jit_time": float,
-                             "min_time": float,
-                             "max_time": float,
-                             "mean_time": float,
-                             "std_div": float,
-                             "last_time": float,
-                             "generated_code": float,
-                             "argument_size": float,
-                             "output_size": float,
-                             "temp_size": float,
-                             "flops": float
-                         },
-                         index_col=False)
+        df = pd.read_csv(
+            csv_file,
+            header=None,
+            skiprows=0,
+            names=[
+                'function',
+                'precision',
+                'x',
+                'y',
+                'z',
+                'px',
+                'py',
+                'backend',
+                'nodes',
+                'jit_time',
+                'min_time',
+                'max_time',
+                'mean_time',
+                'std_div',
+                'last_time',
+                'generated_code',
+                'argument_size',
+                'output_size',
+                'temp_size',
+                'flops',
+            ],
+            dtype={
+                'function': str,
+                'precision': str,
+                'x': int,
+                'y': int,
+                'z': int,
+                'px': int,
+                'py': int,
+                'backend': str,
+                'nodes': int,
+                'jit_time': float,
+                'min_time': float,
+                'max_time': float,
+                'mean_time': float,
+                'std_div': float,
+                'last_time': float,
+                'generated_code': float,
+                'argument_size': float,
+                'output_size': float,
+                'temp_size': float,
+                'flops': float,
+            },
+            index_col=False,
+        )
 
         # Filter precisions
         if precisions:
@@ -360,18 +404,32 @@ def clean_up_csv(
         df['output_size'] = df['output_size'] / factor
         df['temp_size'] = df['temp_size'] / factor
         # in case of the same test is run multiple times, keep the last one
-        df = df.drop_duplicates(subset=[
-            "function", "precision", "x", "y", "z", "px", "py", "backend",
-            "nodes"
-        ],
-                                keep='last')
+        df = df.drop_duplicates(
+            subset=[
+                'function',
+                'precision',
+                'x',
+                'y',
+                'z',
+                'px',
+                'py',
+                'backend',
+                'nodes',
+            ],
+            keep='last',
+        )
 
         df['gpus'] = df['px'] * df['py']
 
         if gpus:
             df = df[df['gpus'].isin(gpus)]
 
-        if 'plot_all' in pdims_strategy or 'slab_yz' in pdims_strategy or 'slab_xy' in pdims_strategy or 'pencils' in pdims_strategy:
+        if (
+            'plot_all' in pdims_strategy
+            or 'slab_yz' in pdims_strategy
+            or 'slab_xy' in pdims_strategy
+            or 'pencils' in pdims_strategy
+        ):
 
             def get_decomp_from_px_py(row):
                 if row['px'] > 1 and row['py'] == 1:
@@ -383,7 +441,7 @@ def clean_up_csv(
 
             df['decomp'] = df.apply(get_decomp_from_px_py, axis=1)
             df.drop(columns=['px', 'py'], inplace=True)
-            if not 'plot_all' in pdims_strategy:
+            if 'plot_all' not in pdims_strategy:
                 df = df[df['decomp'].isin(pdims_strategy)]
 
         # check available gpus in dataset
@@ -395,17 +453,18 @@ def clean_up_csv(
         else:
             dataframes[file_name] = pd.concat([dataframes[file_name], df])
 
-    print(f"requested GPUS: {gpus} available GPUS: {available_gpu_counts}")
-    print(
-        f"requested data sizes: {data_sizes} available data sizes: {available_data_sizes}"
-    )
+    print(f'requested GPUS: {gpus} available GPUS: {available_gpu_counts}')
+    print(f'requested data sizes: {data_sizes} available data sizes: {available_data_sizes}')
 
-    available_gpu_counts = (available_gpu_counts if gpus is None else [
-        gpu for gpu in gpus if gpu in available_gpu_counts
-    ])
-    available_data_sizes = (available_data_sizes if data_sizes is None else [
-        data_size for data_size in data_sizes
-        if data_size in available_data_sizes
-    ])
+    available_gpu_counts = (
+        available_gpu_counts
+        if gpus is None
+        else [gpu for gpu in gpus if gpu in available_gpu_counts]
+    )
+    available_data_sizes = (
+        available_data_sizes
+        if data_sizes is None
+        else [data_size for data_size in data_sizes if data_size in available_data_sizes]
+    )
 
     return dataframes, available_gpu_counts, available_data_sizes

{jax_hpc_profiler-0.2.12.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.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.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.12.dist-info → jax_hpc_profiler-0.3.0.dist-info}/WHEEL

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

jax_hpc_profiler-0.2.12.dist-info/RECORD

@@ -1,12 +0,0 @@
-jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,367
-jax_hpc_profiler/create_argparse.py,sha256=6DVpYuj908L01vk09-l7BLU6dW4OGgTepFR123fsawM,6314
-jax_hpc_profiler/main.py,sha256=dwOik2rJw5YV6ocQ-EE32iFOPlq2_3CHHAAuJJFt65Q,2286
-jax_hpc_profiler/plotting.py,sha256=R0mjUhV_Q-qi02mlxWiR241sxr58USBSykwFdjBa-oM,9484
-jax_hpc_profiler/timer.py,sha256=4zc5HlJwepMK633BDz0iLTLWcLsvPdd6M1SL0-qs4js,10554
-jax_hpc_profiler/utils.py,sha256=7i8qPfKogp8nGaGdyJ2-fbQomhIZqn73PQ14qldpFTc,14657
-jax_hpc_profiler-0.2.12.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jax_hpc_profiler-0.2.12.dist-info/METADATA,sha256=bIriXXGdfJ7yTIdsuIdZgig3WpshBibBYKtGDjukr8A,49186
-jax_hpc_profiler-0.2.12.dist-info/WHEEL,sha256=tTnHoFhvKQHCh4jz3yCn0WPTYIy7wXx3CJtJ7SJGV7c,91
-jax_hpc_profiler-0.2.12.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
-jax_hpc_profiler-0.2.12.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
-jax_hpc_profiler-0.2.12.dist-info/RECORD,,