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

jax_hpc_profiler/__init__.py

@@ -1,9 +1,15 @@
 from .create_argparse import create_argparser
-from .plotting import plot_strong_scaling, plot_weak_scaling
+from .plotting import plot_strong_scaling, plot_weak_fixed_scaling, plot_weak_scaling
 from .timer import Timer
 from .utils import clean_up_csv, concatenate_csvs, plot_with_pdims_strategy
 
 __all__ = [
-    'create_argparser', 'plot_strong_scaling', 'plot_weak_scaling', 'Timer',
-    'clean_up_csv', 'concatenate_csvs', 'plot_with_pdims_strategy'
+    'create_argparser',
+    'plot_strong_scaling',
+    'plot_weak_scaling',
+    'plot_weak_fixed_scaling',
+    'Timer',
+    'clean_up_csv',
+    'concatenate_csvs',
+    'plot_with_pdims_strategy',
 ]

jax_hpc_profiler/create_argparse.py

@@ -10,201 +10,209 @@ def create_argparser():
     argparse.Namespace
         Parsed and validated arguments.
     """
-    parser = argparse.ArgumentParser(
-        description="HPC Plotter for benchmarking data")
+    parser = argparse.ArgumentParser(description='HPC Plotter for benchmarking data')
 
     # Group for concatenation to ensure mutually exclusive behavior
-    subparsers = parser.add_subparsers(dest="command", required=True)
+    subparsers = parser.add_subparsers(dest='command', required=True)
 
-    concat_parser = subparsers.add_parser("concat",
-                                          help="Concatenate CSV files")
-    concat_parser.add_argument("input",
-                               type=str,
-                               help="Input directory for concatenation")
-    concat_parser.add_argument("output",
-                               type=str,
-                               help="Output directory for concatenation")
+    concat_parser = subparsers.add_parser('concat', help='Concatenate CSV files')
+    concat_parser.add_argument('input', type=str, help='Input directory for concatenation')
+    concat_parser.add_argument('output', type=str, help='Output directory for concatenation')
 
     # Arguments for plotting
-    plot_parser = subparsers.add_parser("plot", help="Plot CSV data")
-    plot_parser.add_argument("-f",
-                             "--csv_files",
-                             nargs="+",
-                             help="List of CSV files to plot",
-                             required=True)
-    plot_parser.add_argument(
-        "-g",
-        "--gpus",
-        nargs="*",
+    plot_parser = subparsers.add_parser('plot', help='Plot CSV data')
+    plot_parser.add_argument(
+        '-f', '--csv_files', nargs='+', help='List of CSV files to plot', required=True
+    )
+    plot_parser.add_argument(
+        '-g',
+        '--gpus',
+        nargs='*',
         type=int,
-        help="List of number of GPUs to plot",
+        help='List of number of GPUs to plot',
         default=None,
     )
     plot_parser.add_argument(
-        "-d",
-        "--data_size",
-        nargs="*",
+        '-d',
+        '--data_size',
+        nargs='*',
         type=int,
-        help="List of data sizes to plot",
+        help='List of data sizes to plot',
         default=None,
     )
 
     # pdims related arguments
     plot_parser.add_argument(
-        "-fd",
-        "--filter_pdims",
-        nargs="*",
-        help="List of pdims to filter, e.g., 1x4 2x2 4x8",
+        '-fd',
+        '--filter_pdims',
+        nargs='*',
+        help='List of pdims to filter, e.g., 1x4 2x2 4x8',
         default=None,
     )
     plot_parser.add_argument(
-        "-ps",
-        "--pdim_strategy",
-        choices=["plot_all", "plot_fastest", "slab_yz", "slab_xy", "pencils"],
-        nargs="*",
-        default=["plot_fastest"],
-        help="Strategy for plotting pdims",
+        '-ps',
+        '--pdim_strategy',
+        choices=['plot_all', 'plot_fastest', 'slab_yz', 'slab_xy', 'pencils'],
+        nargs='*',
+        default=['plot_fastest'],
+        help='Strategy for plotting pdims',
     )
 
     # Function and precision related arguments
     plot_parser.add_argument(
-        "-pr",
-        "--precision",
-        choices=["float32", "float64"],
-        default=["float32", "float64"],
-        nargs="*",
-        help="Precision to filter by (float32 or float64)",
+        '-pr',
+        '--precision',
+        choices=['float32', 'float64'],
+        default=['float32', 'float64'],
+        nargs='*',
+        help='Precision to filter by (float32 or float64)',
     )
     plot_parser.add_argument(
-        "-fn",
-        "--function_name",
-        nargs="+",
-        help="Function names to filter",
+        '-fn',
+        '--function_name',
+        nargs='+',
+        help='Function names to filter',
         default=None,
     )
 
     # Time or memory related arguments
     plotting_group = plot_parser.add_mutually_exclusive_group(required=True)
     plotting_group.add_argument(
-        "-pt",
-        "--plot_times",
-        nargs="*",
+        '-pt',
+        '--plot_times',
+        nargs='*',
         choices=[
-            "jit_time",
-            "min_time",
-            "max_time",
-            "mean_time",
-            "std_time",
-            "last_time",
+            'jit_time',
+            'min_time',
+            'max_time',
+            'mean_time',
+            'std_time',
+            'last_time',
         ],
-        help="Time columns to plot",
+        help='Time columns to plot',
     )
     plotting_group.add_argument(
-        "-pm",
-        "--plot_memory",
-        nargs="*",
-        choices=[
-            "generated_code", "argument_size", "output_size", "temp_size"
-        ],
-        help="Memory columns to plot",
+        '-pm',
+        '--plot_memory',
+        nargs='*',
+        choices=['generated_code', 'argument_size', 'output_size', 'temp_size'],
+        help='Memory columns to plot',
     )
     plot_parser.add_argument(
-        "-mu",
-        "--memory_units",
-        default="GB",
-        help="Memory units to plot (KB, MB, GB, TB)",
+        '-mu',
+        '--memory_units',
+        default='GB',
+        help='Memory units to plot (KB, MB, GB, TB)',
     )
 
     # Plot customization arguments
-    plot_parser.add_argument("-fs",
-                             "--figure_size",
-                             nargs=2,
-                             type=int,
-                             help="Figure size",
-                             default=(10, 6))
-    plot_parser.add_argument("-o",
-                             "--output",
-                             help="Output file (if none then only show plot)",
-                             default=None)
-    plot_parser.add_argument("-db",
-                             "--dark_bg",
-                             action="store_true",
-                             help="Use dark background for plotting")
-    plot_parser.add_argument(
-        "-pd",
-        "--print_decompositions",
-        action="store_true",
-        help="Print decompositions on plot",
+    plot_parser.add_argument(
+        '-fs', '--figure_size', nargs=2, type=int, help='Figure size', default=(10, 6)
+    )
+    plot_parser.add_argument(
+        '-o', '--output', help='Output file (if none then only show plot)', default=None
+    )
+    plot_parser.add_argument(
+        '-db', '--dark_bg', action='store_true', help='Use dark background for plotting'
+    )
+    plot_parser.add_argument(
+        '-pd',
+        '--print_decompositions',
+        action='store_true',
+        help='Print decompositions on plot',
     )
 
     # Backend related arguments
     plot_parser.add_argument(
-        "-b",
-        "--backends",
-        nargs="*",
-        default=["MPI", "NCCL", "MPI4JAX"],
-        help="List of backends to include",
+        '-b',
+        '--backends',
+        nargs='*',
+        default=['MPI', 'NCCL', 'MPI4JAX'],
+        help='List of backends to include',
     )
 
     # Scaling type argument
     plot_parser.add_argument(
-        "-sc",
-        "--scaling",
-        choices=["Weak", "Strong", "w", "s"],
+        '-sc',
+        '--scaling',
+        choices=['Weak', 'Strong', 'WeakFixed', 'w', 's', 'wf'],
         required=True,
-        help="Scaling type (Weak or Strong)",
+        help='Scaling type (Strong, Weak, or WeakFixed)',
+    )
+
+    # Weak-scaling specific options
+    plot_parser.add_argument(
+        '--weak_ideal_line',
+        action='store_true',
+        help='Overlay an ideal flat line for weak scaling (Weak mode only)',
+    )
+    plot_parser.add_argument(
+        '--weak_reverse_axes',
+        action='store_true',
+        help=(
+            'Weak mode only: put data size on the x-axis and annotate each point with GPUs instead '
+            'of data size. Requires --gpus and --data_size with equal lengths.'
+        ),
     )
 
     # Label customization argument
     plot_parser.add_argument(
-        "-l",
-        "--label_text",
+        '-l',
+        '--label_text',
         type=str,
-        help=
-        ("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%",
+        help=(
+            '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%',
     )
 
     plot_parser.add_argument(
-        "-xl",
-        "--xlabel",
+        '-xl',
+        '--xlabel',
         type=str,
-        help="X-axis label for the plot",
+        help='X-axis label for the plot',
     )
     plot_parser.add_argument(
-        "-tl",
-        "--title",
+        '-tl',
+        '--title',
         type=str,
-        help="Title for the plot",
+        help='Title for the plot',
     )
 
-    subparsers.add_parser("label_help", help="Label customization help")
+    subparsers.add_parser('label_help', help='Label customization help')
 
     args = parser.parse_args()
 
     # if command was plot, then check if pdim_strategy is validat
-    if args.command == "plot":
-        if "plot_all" in args.pdim_strategy and len(args.pdim_strategy) > 1:
+    if args.command == 'plot':
+        if 'plot_all' in args.pdim_strategy and len(args.pdim_strategy) > 1:
             print(
-                "Warning: 'plot_all' strategy is combined with other strategies. Using 'plot_all' only."
+                """
+                Warning: 'plot_all' strategy is combined with other strategies.
+                Using 'plot_all' only.
+                """
             )
-            args.pdim_strategy = ["plot_all"]
+            args.pdim_strategy = ['plot_all']
 
-        if "plot_fastest" in args.pdim_strategy and len(
-                args.pdim_strategy) > 1:
+        if 'plot_fastest' in args.pdim_strategy and len(args.pdim_strategy) > 1:
             print(
-                "Warning: 'plot_fastest' strategy is combined with other strategies. Using 'plot_fastest' only."
+                """
+                Warning: 'plot_fastest' strategy is combined with other strategies.
+                Using 'plot_fastest' only.
+                """
             )
-            args.pdim_strategy = ["plot_fastest"]
+            args.pdim_strategy = ['plot_fastest']
         if args.plot_times is not None:
             args.plot_columns = args.plot_times
         elif args.plot_memory is not None:
             args.plot_columns = args.plot_memory
         else:
-            raise ValueError(
-                "Either plot_times or plot_memory should be provided")
+            raise ValueError('Either plot_times or plot_memory should be provided')
+
+        # Note: for Weak scaling, plot_weak_scaling enforces that both gpus and
+        # data_size are provided and have matching lengths. For Strong and
+        # WeakFixed, gpus/data_size remain optional as before.
 
     return args

jax_hpc_profiler/main.py

@@ -1,29 +1,26 @@
-import sys
-from typing import List, Optional
-
 from .create_argparse import create_argparser
-from .plotting import plot_strong_scaling, plot_weak_scaling
-from .utils import clean_up_csv, concatenate_csvs
+from .plotting import plot_strong_scaling, plot_weak_fixed_scaling, plot_weak_scaling
+from .utils import concatenate_csvs
 
 
 def main():
     args = create_argparser()
 
-    if args.command == "concat":
+    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":
-
-        if args.scaling.lower() == "weak" or args.scaling.lower() == "w":
+    elif args.command == 'label_help':
+        print('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':
+        scaling = args.scaling.lower()
+        if scaling in ('weak', 'w'):
             plot_weak_scaling(
                 args.csv_files,
                 args.gpus,
@@ -37,13 +34,15 @@ def main():
                 args.plot_columns,
                 args.memory_units,
                 args.label_text,
+                args.xlabel if getattr(args, 'xlabel', None) is not None else 'Number of GPUs',
                 args.title,
-                args.label_text,
                 args.figure_size,
                 args.dark_bg,
                 args.output,
+                args.weak_ideal_line,
+                args.weak_reverse_axes,
             )
-        elif args.scaling.lower() == "strong" or args.scaling.lower() == "s":
+        elif scaling in ('strong', 's'):
             plot_strong_scaling(
                 args.csv_files,
                 args.gpus,
@@ -57,13 +56,33 @@ def main():
                 args.plot_columns,
                 args.memory_units,
                 args.label_text,
-                args.title,
+                args.xlabel if getattr(args, 'xlabel', None) is not None else 'Number of GPUs',
+                args.title if getattr(args, 'title', None) is not None else 'Data sizes',
+                args.figure_size,
+                args.dark_bg,
+                args.output,
+            )
+        elif scaling in ('weakfixed', 'wf'):
+            plot_weak_fixed_scaling(
+                args.csv_files,
+                args.gpus,
+                args.data_size,
+                args.function_name,
+                args.precision,
+                args.filter_pdims,
+                args.pdim_strategy,
+                args.print_decompositions,
+                args.backends,
+                args.plot_columns,
+                args.memory_units,
                 args.label_text,
+                args.xlabel if getattr(args, 'xlabel', None) is not None else 'Data sizes',
+                args.title if getattr(args, 'title', None) is not None else 'Number of GPUs',
                 args.figure_size,
                 args.dark_bg,
                 args.output,
             )
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()