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

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: jax_hpc_profiler
-Version: 0.2.9
+Version: 0.2.11
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
 License:                     GNU GENERAL PUBLIC LICENSE
@@ -698,8 +698,7 @@ Requires-Dist: pandas
 Requires-Dist: matplotlib
 Requires-Dist: seaborn
 Requires-Dist: tabulate
-
-Here's the updated README with the additional information about the timer.report and the multi-GPU setup:
+Dynamic: license-file
 
 # JAX HPC Profiler
 

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.11}/README.md

@@ -1,5 +1,3 @@
-Here's the updated README with the additional information about the timer.report and the multi-GPU setup:
-
 # JAX HPC Profiler
 
 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.

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

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

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.11}/src/jax_hpc_profiler/create_argparse.py

@@ -168,6 +168,19 @@ def create_argparser():
         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()

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.11}/src/jax_hpc_profiler/main.py

@@ -37,6 +37,8 @@ def main():
                 args.plot_columns,
                 args.memory_units,
                 args.label_text,
+                args.title,
+                args.label_text,
                 args.figure_size,
                 args.dark_bg,
                 args.output,
@@ -55,6 +57,8 @@ def main():
                 args.plot_columns,
                 args.memory_units,
                 args.label_text,
+                args.title,
+                args.label_text,
                 args.figure_size,
                 args.dark_bg,
                 args.output,

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.11}/src/jax_hpc_profiler/plotting.py

@@ -17,8 +17,8 @@ def configure_axes(
     ax: Axes,
     x_values: List[int],
     y_values: List[float],
-    xlabel: str,
     title: str,
+    xlabel: str,
     plotting_memory: bool = False,
     memory_units: str = "bytes",
 ):
@@ -77,9 +77,9 @@ def plot_scaling(
     output: Optional[str] = None,
     dark_bg: bool = False,
     print_decompositions: bool = False,
-    backends: List[str] = ["NCCL"],
-    precisions: List[str] = ["float32"],
-    functions: List[str] | None = None,
+    backends: Optional[List[str]] = None,
+    precisions: Optional[List[str]] = None,
+    functions: Optional[List[str]] = None,
     plot_columns: List[str] = ["mean_time"],
     memory_units: str = "bytes",
     label_text: str = "plot",
@@ -141,6 +141,11 @@ def plot_scaling(
                 by=[size_column])
             functions = (pd.unique(filtered_method_df["function"])
                          if functions is None else functions)
+            precisions = (pd.unique(filtered_method_df["precision"])
+                          if precisions is None else precisions)
+            backends = (pd.unique(filtered_method_df["backend"])
+                        if backends is None else backends)
+
             combinations = product(backends, precisions, functions,
                                    plot_columns)
 
@@ -199,15 +204,17 @@ def plot_strong_scaling(
     csv_files: List[str],
     fixed_gpu_size: Optional[List[int]] = None,
     fixed_data_size: Optional[List[int]] = None,
-    functions: List[str] | None = None,
-    precisions: List[str] = ["float32"],
+    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: List[str] = ["NCCL"],
+    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: str = "Data sizes",
     figure_size: tuple = (6, 4),
     dark_bg: bool = False,
     output: Optional[str] = None,
@@ -236,8 +243,8 @@ def plot_strong_scaling(
         available_data_sizes,
         "gpus",
         "x",
-        "Number of GPUs",
-        "Data size",
+        xlabel,
+        title,
         figure_size,
         output,
         dark_bg,
@@ -256,15 +263,17 @@ def plot_weak_scaling(
     csv_files: List[str],
     fixed_gpu_size: Optional[List[int]] = None,
     fixed_data_size: Optional[List[int]] = None,
-    functions: List[str] | None = None,
-    precisions: List[str] = ["float32"],
+    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: List[str] = ["NCCL"],
+    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 = "Data sizes",
+    title: str = "Number of GPUs",
     figure_size: tuple = (6, 4),
     dark_bg: bool = False,
     output: Optional[str] = None,
@@ -292,8 +301,8 @@ def plot_weak_scaling(
         available_gpu_counts,
         "x",
         "gpus",
-        "Data size",
-        "Number of GPUs",
+        xlabel,
+        title,
         figure_size,
         output,
         dark_bg,

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.11}/src/jax_hpc_profiler/timer.py

@@ -1,7 +1,7 @@
 import os
 import time
 from functools import partial
-from typing import Any, Callable, List, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -11,23 +11,31 @@ from jax.experimental import mesh_utils
 from jax.experimental.shard_map import shard_map
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
+from jaxtyping import Array
 from tabulate import tabulate
 
 
 class Timer:
 
-    def __init__(self, save_jaxpr=False, jax_fn=True, devices=None):
+    def __init__(self,
+                 save_jaxpr=False,
+                 compile_info=True,
+                 jax_fn=True,
+                 devices=None,
+                 static_argnums=()):
         self.jit_time = 0.0
         self.times = []
         self.profiling_data = {}
         self.compiled_code = {}
         self.save_jaxpr = save_jaxpr
+        self.compile_info = compile_info
         self.jax_fn = jax_fn
         self.devices = devices
+        self.static_argnums = static_argnums
 
     def _normalize_memory_units(self, memory_analysis) -> str:
 
-        if not self.jax_fn:
+        if not (self.jax_fn and self.compile_info):
             return memory_analysis
 
         sizes_str = ["B", "KB", "MB", "GB", "TB", "PB"]
@@ -47,23 +55,36 @@ class Timer:
             memory_analysis.temp_size_in_bytes,
         )
 
-    def chrono_jit(self, fun: Callable, *args, ndarray_arg=None) -> np.ndarray:
+    def chrono_jit(self, fun: Callable, *args, **kwargs) -> Array:
         start = time.perf_counter()
-        out = fun(*args)
+        out = fun(*args, **kwargs)
         if self.jax_fn:
-            if ndarray_arg is None:
-                out.block_until_ready()
-            else:
-                out[ndarray_arg].block_until_ready()
+
+            def _block(x):
+                if isinstance(x, Array):
+                    x.block_until_ready()
+
+            jax.tree_map(_block, out)
         end = time.perf_counter()
         self.jit_time = (end - start) * 1e3
 
+        self.compiled_code["JAXPR"] = "N/A"
+        self.compiled_code["LOWERED"] = "N/A"
+        self.compiled_code["COMPILED"] = "N/A"
+        self.profiling_data["generated_code"] = "N/A"
+        self.profiling_data["argument_size"] = "N/A"
+        self.profiling_data["output_size"] = "N/A"
+        self.profiling_data["temp_size"] = "N/A"
+
         if self.save_jaxpr:
-            jaxpr = make_jaxpr(fun)(*args)
+            jaxpr = make_jaxpr(fun,
+                               static_argnums=self.static_argnums)(*args,
+                                                                   **kwargs)
             self.compiled_code["JAXPR"] = jaxpr.pretty_print()
 
-        if self.jax_fn:
-            lowered = jax.jit(fun).lower(*args)
+        if self.jax_fn and self.compile_info:
+            lowered = jax.jit(fun, static_argnums=self.static_argnums).lower(
+                *args, **kwargs)
             compiled = lowered.compile()
             memory_analysis = self._read_memory_analysis(
                 compiled.memory_analysis())
@@ -77,19 +98,21 @@ class Timer:
 
         return out
 
-    def chrono_fun(self, fun: Callable, *args, ndarray_arg=None) -> np.ndarray:
+    def chrono_fun(self, fun: Callable, *args, **kwargs) -> Array:
         start = time.perf_counter()
-        out = fun(*args)
+        out = fun(*args, **kwargs)
         if self.jax_fn:
-            if ndarray_arg is None:
-                out.block_until_ready()
-            else:
-                out[ndarray_arg].block_until_ready()
+
+            def _block(x):
+                if isinstance(x, Array):
+                    x.block_until_ready()
+
+            jax.tree_map(_block, out)
         end = time.perf_counter()
         self.times.append((end - start) * 1e3)
         return out
 
-    def _get_mean_times(self) -> np.ndarray:
+    def _get_mean_times(self) -> Array:
         if jax.device_count() == 1 or jax.process_count() == 1:
             return np.array(self.times)
 
@@ -133,8 +156,12 @@ class Timer:
         backend: str = "NCCL",
         nodes: int = 1,
         md_filename: str | None = None,
+        npz_data: Optional[dict] = None,
         extra_info: dict = {},
     ):
+        if self.jit_time == 0.0 and len(self.times) == 0:
+            print(f"No profiling data to report for {function}")
+            return
 
         if md_filename is None:
             dirname, filename = (
@@ -147,6 +174,18 @@ class Timer:
                 f"{report_folder}/{x}_{px}_{py}_{backend}_{precision}_{function}.md"
             )
 
+        if npz_data is not None:
+            dirname, filename = (
+                os.path.dirname(csv_filename),
+                os.path.splitext(os.path.basename(csv_filename))[0],
+            )
+            report_folder = filename if dirname == "" else f"{dirname}/{filename}"
+            os.makedirs(report_folder, exist_ok=True)
+            npz_filename = (
+                f"{report_folder}/{x}_{px}_{py}_{backend}_{precision}_{function}.npz"
+            )
+            np.savez(npz_filename, **npz_data)
+
         y = x if y is None else y
         z = x if z is None else z
 
@@ -158,18 +197,10 @@ class Timer:
             mean_time = np.mean(times_array)
             std_time = np.std(times_array)
             last_time = times_array[-1]
-
-            if self.jax_fn:
-
-                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"]
-            else:
-                generated_code = "N/A"
-                argument_size = "N/A"
-                output_size = "N/A"
-                temp_size = "N/A"
+            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"]
 
             csv_line = (
                 f"{function},{precision},{x},{y},{z},{px},{py},{backend},{nodes},"
@@ -233,7 +264,7 @@ class Timer:
                         headers=["Iteration", "Time"],
                         tablefmt="github",
                     ))
-                if self.jax_fn:
+                if self.jax_fn and self.compile_info:
                     f.write("\n---\n")
                     f.write(f"## Compiled Code\n")
                     f.write(f"```hlo\n")

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.11}/src/jax_hpc_profiler/utils.py

@@ -248,7 +248,7 @@ def clean_up_csv(
     data_sizes: Optional[List[int]] = None,
     pdims: Optional[List[str]] = None,
     pdims_strategy: List[str] = ['plot_fastest'],
-    backends: List[str] = ['MPI', 'NCCL', 'MPI4JAX'],
+    backends: Optional[List[str]] = None,
     memory_units: str = 'KB',
 ) -> Tuple[Dict[str, pd.DataFrame], List[int], List[int]]:
     """
@@ -292,7 +292,7 @@ def clean_up_csv(
 
         df = pd.read_csv(csv_file,
                          header=None,
-                         skiprows=1,
+                         skiprows=0,
                          names=[
                              "function", "precision", "x", "y", "z", "px",
                              "py", "backend", "nodes", "jit_time", "min_time",
@@ -331,7 +331,8 @@ def clean_up_csv(
         if function_names:
             df = df[df['function'].isin(function_names)]
         # Filter backends
-        df = df[df['backend'].isin(backends)]
+        if backends:
+            df = df[df['backend'].isin(backends)]
 
         # Filter data sizes
         if data_sizes:

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.11}/src/jax_hpc_profiler.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: jax_hpc_profiler
-Version: 0.2.9
+Version: 0.2.11
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
 License:                     GNU GENERAL PUBLIC LICENSE
@@ -698,8 +698,7 @@ Requires-Dist: pandas
 Requires-Dist: matplotlib
 Requires-Dist: seaborn
 Requires-Dist: tabulate
-
-Here's the updated README with the additional information about the timer.report and the multi-GPU setup:
+Dynamic: license-file
 
 # JAX HPC Profiler