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

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: jax_hpc_profiler
-Version: 0.2.9
+Version: 0.2.10
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
 License:                     GNU GENERAL PUBLIC LICENSE

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "jax_hpc_profiler"
-version = "0.2.9"
+version = "0.2.10"
 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.10}/src/jax_hpc_profiler/plotting.py

@@ -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,12 +204,12 @@ 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%",
@@ -256,12 +261,12 @@ 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%",

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.10}/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
@@ -133,8 +133,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 +151,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
 

{jax_hpc_profiler-0.2.9 → jax_hpc_profiler-0.2.10}/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]]:
     """
@@ -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.10}/src/jax_hpc_profiler.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: jax_hpc_profiler
-Version: 0.2.9
+Version: 0.2.10
 Summary: HPC Plotter and profiler for benchmarking data made for JAX
 Author: Wassim Kabalan
 License:                     GNU GENERAL PUBLIC LICENSE