jax-hpc-profiler 0.2.6__tar.gz → 0.2.8__tar.gz

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/PKG-INFO

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

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/pyproject.toml

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

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/src/jax_hpc_profiler/create_argparse.py

@@ -67,7 +67,6 @@ def create_argparser():
     plot_parser.add_argument('-fn',
                              '--function_name',
                              nargs='+',
-                             default=['FFT'],
                              help='Function names to filter')
 
     # Time or memory related arguments
@@ -135,6 +134,8 @@ def create_argparser():
         '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%")
 
+    subparsers.add_parser('label_help',help='Label customization help')
+
     args = parser.parse_args()
     
     # if command was plot, then check if pdim_strategy is validat

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/src/jax_hpc_profiler/main.py

@@ -11,6 +11,15 @@ def main():
     if args.command == 'concat':
         input_dir, output_dir = args.input, args.output
         concatenate_csvs(input_dir, output_dir)
+    elif args.command == 'label_help':
+        print(f"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':
         dataframes, available_gpu_counts, available_data_sizes = clean_up_csv(
             args.csv_files, args.precision, args.function_name, args.gpus,
@@ -23,11 +32,10 @@ def main():
             f"requested GPUS: {args.gpus} available GPUS: {available_gpu_counts}"
         )
         # filter back the requested data sizes and gpus
-        args.gpus = [gpu for gpu in args.gpus if gpu in available_gpu_counts]
-        args.data_size = [
-            data_size for data_size in args.data_size
-            if data_size in available_data_sizes
-        ]
+        
+        args.gpus = available_gpu_counts if args.gpus is None else [gpu for gpu in args.gpus if gpu in available_gpu_counts]
+        args.data_size = available_data_sizes if args.data_size is None else [data_size for data_size in args.data_size if data_size in available_data_sizes]
+
         if len(args.gpus) == 0:
             print(f"No dataframes found for the given GPUs. Exiting...")
             sys.exit(1)

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/src/jax_hpc_profiler/plotting.py

@@ -7,10 +7,10 @@ import pandas as pd
 from matplotlib.axes import Axes
 from matplotlib.patches import FancyBboxPatch
 
-from .utils import inspect_df, plot_with_pdims_strategy
+from .utils import inspect_df, plot_with_pdims_strategy, inspect_data
 
 np.seterr(divide='ignore')
-plt.rcParams.update({'font.size': 10})
+plt.rcParams.update({'font.size': 15})
 
 
 def configure_axes(ax: Axes,
@@ -38,7 +38,7 @@ def configure_axes(ax: Axes,
     f2 = lambda x: np.log2(x)
     g2 = lambda x: 2**x
     ax.set_xlim([min(x_values), max(x_values)])
-    y_min, y_max = min(y_values) * 0.9, max(y_values) * 1.1
+    y_min, y_max = min(y_values) * 0.6, max(y_values) * 1.1
     ax.set_title(title)
     ax.set_ylim([y_min, y_max])
     ax.set_xscale('function', functions=(f2, g2))
@@ -57,6 +57,7 @@ def configure_axes(ax: Axes,
     ax.legend(loc='lower center',
               bbox_to_anchor=(0.5, 0.05),
               ncol=4,
+              fontsize="x-large",
               prop={'size': 14})
 
 
@@ -117,13 +118,13 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
         axs = axs.flatten()
     else:
         axs = [axs]
-
+    
     for i, fixed_size in enumerate(fixed_sizes):
         ax: Axes = axs[i]
 
+        x_values = []
+        y_values = []
         for method, df in dataframes.items():
-            x_values = []
-            y_values = []
 
             filtered_method_df = df[df[fixed_column] == int(fixed_size)]
             if filtered_method_df.empty:
@@ -136,6 +137,7 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
                                    plot_columns)
 
             for backend, precision, function, plot_column in combinations:
+              
                 filtered_params_df = filtered_method_df[
                     (filtered_method_df['backend'] == backend)
                     & (filtered_method_df['precision'] == precision) &
@@ -148,10 +150,11 @@ def plot_scaling(dataframes: Dict[str, pd.DataFrame],
 
                 x_values.extend(x_vals)
                 y_values.extend(y_vals)
-
-        plotting_memory = 'time' not in plot_columns[0].lower()
-        configure_axes(ax, x_values, y_values, f"{title} {fixed_size}", xlabel,
-                       plotting_memory, memory_units)
+        
+        if len(x_values) != 0: 
+          plotting_memory = 'time' not in plot_columns[0].lower()
+          configure_axes(ax, x_values, y_values, f"{title} {fixed_size}", xlabel,
+                         plotting_memory, memory_units)
 
     for i in range(num_subplots, num_rows * num_cols):
         fig.delaxes(axs[i])

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/src/jax_hpc_profiler/timer.py

@@ -23,16 +23,11 @@ class Timer:
         self.compiled_code = {}
         self.save_jaxpr = save_jaxpr
 
-    def _read_cost_analysis(self, cost_analysis: Any) -> str | None:
-        if cost_analysis is None:
-            return None
-        return cost_analysis[0]['flops']
-
     def _normalize_memory_units(self, memory_analysis) -> str:
 
         sizes_str = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
         factors = [1, 1024, 1024**2, 1024**3, 1024**4, 1024**5]
-        factor = int(np.log10(memory_analysis) // 3)
+        factor = 0 if memory_analysis == 0 else int(np.log10(memory_analysis) // 3)
 
         return f"{memory_analysis / factors[factor]:.2f} {sizes_str[factor]}"
 
@@ -62,11 +57,9 @@ class Timer:
         compiled = lowered.compile()
         memory_analysis = self._read_memory_analysis(
             compiled.memory_analysis())
-        cost_analysis = self._read_cost_analysis(compiled.cost_analysis())
 
         self.compiled_code["LOWERED"] = lowered.as_text()
         self.compiled_code["COMPILED"] = compiled.as_text()
-        self.profiling_data["FLOPS"] = cost_analysis
         self.profiling_data["generated_code"] = memory_analysis[0]
         self.profiling_data["argument_size"] = memory_analysis[1]
         self.profiling_data["output_size"] = memory_analysis[2]
@@ -145,7 +138,6 @@ class Timer:
             std_time = np.std(times_array)
             last_time = times_array[-1]
 
-            flops = self.profiling_data["FLOPS"]
             generated_code = self.profiling_data["generated_code"]
             argument_size = self.profiling_data["argument_size"]
             output_size = self.profiling_data["output_size"]
@@ -154,7 +146,7 @@ class Timer:
             csv_line = (
                 f"{function},{precision},{x},{y},{z},{px},{py},{backend},{nodes},"
                 f"{self.jit_time:.4f},{min_time:.4f},{max_time:.4f},{mean_time:.4f},{std_time:.4f},{last_time:.4f},"
-                f"{generated_code},{argument_size},{output_size},{temp_size},{flops}\n"
+                f"{generated_code},{argument_size},{output_size},{temp_size}\n"
             )
 
             with open(csv_filename, 'a') as f:

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/src/jax_hpc_profiler/utils.py

@@ -123,7 +123,7 @@ def plot_with_pdims_strategy(ax: Axes, df: pd.DataFrame, method: str,
         # Sort all and keep fastest
         sorted_dfs = []
         for _, group in df_decomp:
-            group.sort_values(by=[y_col], inplace=True, ascending=False)
+            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({
@@ -336,12 +336,11 @@ def clean_up_csv(
         # Filter data sizes
         if data_sizes:
             df = df[df['x'].isin(data_sizes)]
-
+        
         # Filter pdims
         if pdims:
             px_list, py_list = zip(*[map(int, p.split('x')) for p in pdims])
             df = df[(df['px'].isin(px_list)) & (df['py'].isin(py_list))]
-
         # convert memory units columns to remquested memory_units
         match memory_units:
             case 'KB':
@@ -385,6 +384,7 @@ def clean_up_csv(
             df.drop(columns=['px', 'py'], inplace=True)
             if not 'plot_all' in pdims_strategy:
                 df = df[df['decomp'].isin(pdims_strategy)]
+
         # check available gpus in dataset
         available_gpu_counts.update(df['gpus'].unique())
         available_data_sizes.update(df['x'].unique())

{jax_hpc_profiler-0.2.6 → jax_hpc_profiler-0.2.8}/src/jax_hpc_profiler.egg-info/PKG-INFO

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