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

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/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/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",
 ):
@@ -213,6 +213,8 @@ def plot_strong_scaling(
     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,
@@ -241,8 +243,8 @@ def plot_strong_scaling(
         available_data_sizes,
         "gpus",
         "x",
-        "Number of GPUs",
-        "Data size",
+        xlabel,
+        title,
         figure_size,
         output,
         dark_bg,
@@ -270,6 +272,8 @@ def plot_weak_scaling(
     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,
@@ -297,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/timer.py

@@ -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)
 
@@ -174,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},"
@@ -249,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/utils.py

@@ -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",

{jax_hpc_profiler-0.2.10.dist-info → jax_hpc_profiler-0.2.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: jax_hpc_profiler
-Version: 0.2.10
+Version: 0.2.12
 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.12.dist-info/RECORD

@@ -0,0 +1,12 @@
+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,,

{jax_hpc_profiler-0.2.10.dist-info → jax_hpc_profiler-0.2.12.dist-info}/WHEEL

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

jax_hpc_profiler-0.2.10.dist-info/RECORD

@@ -1,12 +0,0 @@
-jax_hpc_profiler/__init__.py,sha256=yDWt2S0xJZeS6YLBFvyPj1p5txCgFx2kCxJzVZZcdtI,367
-jax_hpc_profiler/create_argparse.py,sha256=CSdl76LvaTVVn43dkwpVyiIkyl4lHlDCiI5jvUrIoj0,6059
-jax_hpc_profiler/main.py,sha256=2zPVTGRgFkYV75EJA1eoOqf92gCRXAtg-28cFgRy3Bw,2164
-jax_hpc_profiler/plotting.py,sha256=PdRdEIjsiiDbpr7iwvDUnC4mXz9QWE3JIvOAiWiSS3w,9382
-jax_hpc_profiler/timer.py,sha256=DlbB4O4qJZImcjFq9T6zuywVbZT0x4dBxiJTAR0LNtY,9913
-jax_hpc_profiler/utils.py,sha256=hyWldQIjlNs2laPxgc19szzPnJSHYnqiT9p3knPor8Y,14657
-jax_hpc_profiler-0.2.10.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-jax_hpc_profiler-0.2.10.dist-info/METADATA,sha256=Ca9EYifDf9zDD8nRMmLvxvBu-JUwsI3wIgqClT35egI,49271
-jax_hpc_profiler-0.2.10.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-jax_hpc_profiler-0.2.10.dist-info/entry_points.txt,sha256=_cFlxSINscX3ZyNiklfjyOOO7vNkddhoYy_v1JQHSO4,51
-jax_hpc_profiler-0.2.10.dist-info/top_level.txt,sha256=DKAhVKDwkerhth-xo7oKFSnnKE0Xm46m94b06vZksA4,17
-jax_hpc_profiler-0.2.10.dist-info/RECORD,,