From fc43ff5176b3343f36b019c293fc8bc18cc30513 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Wed, 25 Mar 2026 15:17:02 +0100 Subject: [PATCH 1/8] add pyinstruments flamegraph benchmark script --- pyproject.toml | 1 + tests/benchmarks/README.md | 46 ++++++++ tests/benchmarks/__init__.py | 0 tests/benchmarks/graph_creation_flamegraph.py | 110 ++++++++++++++++++ 4 files changed, 157 insertions(+) create mode 100644 tests/benchmarks/README.md create mode 100644 tests/benchmarks/__init__.py create mode 100644 tests/benchmarks/graph_creation_flamegraph.py diff --git a/pyproject.toml b/pyproject.toml index 01eb298..28f773e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,4 +47,5 @@ dev = [ "nbval>=0.11.0", "ipdb>=0.13.13", "pre-commit>=4.3.0", + "pyinstrument>=5.1.2", ] diff --git a/tests/benchmarks/README.md b/tests/benchmarks/README.md new file mode 100644 index 0000000..212baab --- /dev/null +++ b/tests/benchmarks/README.md @@ -0,0 +1,46 @@ +# Graph Creation Benchmarks + +This directory contains benchmarking scripts to profile the execution time and performance bottlenecks during graph creation. + +## Requirements + +The benchmarks rely on `pyinstrument` to generate call-stack flamegraphs and timing hierarchies. +Make sure you have installed the development dependencies: + +```bash +uv sync --all-extras --dev +# or specifically +uv add --dev pyinstrument +``` + +## Running the Benchmark + +You can run the script from the root of the project to profile graph creation for a specific archetype and grid size. Because the script uses the `tests` utility module, run it via the Python module syntax. By default, it will open an interactive HTML flamegraph in your browser! + +```bash +uv run python -m tests.benchmarks.graph_creation_flamegraph +``` + +### Options + +- `--N `: Set the size of the input grid ($N \times N$). Default is `425` which produces ~180k points (a roughly 10s baseline for the `keisler` graph). +- `--archetype `: The archetype graph to create. Options are `keisler`, `oskarsson_hierarchical`, and `graphcast`. +- `--console`: Print the profiling hierarchy to the console instead of opening the HTML flamegraph in the browser. +- `--save-flamegraph [FILENAME]`: Saves the interactive HTML flamegraph to disk and exits. If no filename is provided, defaults to `pyinstrument_profile.html`. + +**Examples:** + +Profile the hierarchical archetype with $200 \times 200$ points (opens in browser): +```bash +uv run python -m tests.benchmarks.graph_creation_flamegraph --N 200 --archetype oskarsson_hierarchical +``` + +Print a fast profile stack to the console instead of the browser: +```bash +uv run python -m tests.benchmarks.graph_creation_flamegraph --N 100 --console +``` + +Save the flamegraph to a custom file without opening a server: +```bash +uv run python -m tests.benchmarks.graph_creation_flamegraph --N 425 --save-flamegraph my_profile.html +``` diff --git a/tests/benchmarks/__init__.py b/tests/benchmarks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/benchmarks/graph_creation_flamegraph.py b/tests/benchmarks/graph_creation_flamegraph.py new file mode 100644 index 0000000..74c4ccb --- /dev/null +++ b/tests/benchmarks/graph_creation_flamegraph.py @@ -0,0 +1,110 @@ +import argparse +import http.server +import os +import socketserver +import tempfile +import threading +import time +import webbrowser + +from pyinstrument import Profiler + +import tests.utils as test_utils +import weather_model_graphs as wmg + + +def main(): + parser = argparse.ArgumentParser( + description="Profile graph creation with pyinstrument." + ) + parser.add_argument( + "--N", + type=int, + default=425, + help="Size of grid (NxN points). Default is 425 (~180k points).", + ) + parser.add_argument( + "--archetype", + type=str, + default="keisler", + choices=["keisler", "oskarsson_hierarchical", "graphcast"], + help="Graph archetype to create.", + ) + parser.add_argument( + "--console", + action="store_true", + help="Print the profile to the console instead of opening a flamegraph in the browser.", + ) + parser.add_argument( + "--save-flamegraph", + type=str, + nargs="?", + const="pyinstrument_profile.html", + help="Save the HTML flamegraph to a file (default: pyinstrument_profile.html).", + ) + + args = parser.parse_args() + + print(f"Generating input coordinates for N={args.N} ({args.N**2} points)...") + xy = test_utils.create_fake_xy(N=args.N) + + # Get the graph creation function dynamically based on the argument + fn_name = f"create_{args.archetype}_graph" + create_fn = getattr(wmg.create.archetype, fn_name) + + print(f"Starting pyinstrument profiling for {fn_name}...") + profiler = Profiler(interval=0.001) # 1ms precision + + # Profile the function + profiler.start() + t0 = time.time() + graph = create_fn(coords=xy) + t1 = time.time() + profiler.stop() + + print(f"Graph creation finished in {t1 - t0:.2f} seconds.") + print(f"Graph has {len(graph.nodes)} nodes and {len(graph.edges)} edges.") + + if args.save_flamegraph: + with open(args.save_flamegraph, "w") as f: + f.write(profiler.output_html()) + print(f"Detailed report saved to '{args.save_flamegraph}'.") + + if args.console: + print("\n--- Profile Output ---") + print(profiler.output_text(unicode=True, color=True)) + elif not args.save_flamegraph: + with tempfile.TemporaryDirectory() as temp_dir: + html_path = os.path.join(temp_dir, "index.html") + with open(html_path, "w") as f: + f.write(profiler.output_html()) + + class Handler(http.server.SimpleHTTPRequestHandler): + def __init__(self, *args, **kwargs): + super().__init__(*args, directory=temp_dir, **kwargs) + + def log_message(self, format, *args): + pass # suppress noisy server logs + + # Find a free port by binding to port 0 + with socketserver.TCPServer(("127.0.0.1", 0), Handler) as httpd: + port = httpd.server_address[1] + url = f"http://127.0.0.1:{port}" + print(f"\nServing flamegraph at {url}") + print("Press Ctrl+C to shut down the server and exit.") + + # Open the browser in a separate thread so we can start serving immediately + def open_browser(): + time.sleep(0.5) + webbrowser.open(url) + + threading.Thread(target=open_browser, daemon=True).start() + + try: + httpd.serve_forever() + except KeyboardInterrupt: + print("\nShutting down server.") + + +if __name__ == "__main__": + main() From 8a081987c9e5d286567c07590bce134022a2048c Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Wed, 25 Mar 2026 15:54:02 +0100 Subject: [PATCH 2/8] add runtime scaling benchmark cli script --- tests/benchmarks/README.md | 31 ++++++-- tests/benchmarks/graph_creation_scaling.py | 91 ++++++++++++++++++++++ 2 files changed, 116 insertions(+), 6 deletions(-) create mode 100644 tests/benchmarks/graph_creation_scaling.py diff --git a/tests/benchmarks/README.md b/tests/benchmarks/README.md index 212baab..720edda 100644 --- a/tests/benchmarks/README.md +++ b/tests/benchmarks/README.md @@ -4,7 +4,7 @@ This directory contains benchmarking scripts to profile the execution time and p ## Requirements -The benchmarks rely on `pyinstrument` to generate call-stack flamegraphs and timing hierarchies. +The benchmarks rely on `pyinstrument` and `matplotlib` to generate call-stack flamegraphs and scaling plots. Make sure you have installed the development dependencies: ```bash @@ -13,7 +13,7 @@ uv sync --all-extras --dev uv add --dev pyinstrument ``` -## Running the Benchmark +## 1. Call Stack Flamegraphs (`graph_creation_flamegraph.py`) You can run the script from the root of the project to profile graph creation for a specific archetype and grid size. Because the script uses the `tests` utility module, run it via the Python module syntax. By default, it will open an interactive HTML flamegraph in your browser! @@ -35,12 +35,31 @@ Profile the hierarchical archetype with $200 \times 200$ points (opens in browse uv run python -m tests.benchmarks.graph_creation_flamegraph --N 200 --archetype oskarsson_hierarchical ``` -Print a fast profile stack to the console instead of the browser: +Save the flamegraph to a custom file without opening a server: ```bash -uv run python -m tests.benchmarks.graph_creation_flamegraph --N 100 --console +uv run python -m tests.benchmarks.graph_creation_flamegraph --N 425 --save-flamegraph my_profile.html ``` -Save the flamegraph to a custom file without opening a server: +## 2. Runtime Scaling Plot (`graph_creation_scaling.py`) + +This script runs the graph creation process across a range of different grid sizes and plots the execution time versus the number of input nodes. This helps visualize how the algorithm's runtime scales as the coordinate size increases. + ```bash -uv run python -m tests.benchmarks.graph_creation_flamegraph --N 425 --save-flamegraph my_profile.html +uv run python -m tests.benchmarks.graph_creation_scaling +``` + +### Options + +- `--min-N `: The minimum grid size N ($N \times N$ nodes). Default: 50 +- `--max-N `: The maximum grid size N ($N \times N$ nodes). Default: 400 +- `--num-steps `: Number of intermediate grid sizes to test between min and max. Default: 8 +- `--archetype `: The archetype graph to create. Options are `keisler`, `oskarsson_hierarchical`, and `graphcast`. +- `--output `: The file path to save the generated plot. Default: `scaling_plot.png` +- `--show`: Opens a matplotlib interactive window to display the plot after benchmarking. + +**Examples:** + +Test scaling from $100 \times 100$ to $500 \times 500$ and open the plot interactively: +```bash +uv run python -m tests.benchmarks.graph_creation_scaling --min-N 100 --max-N 500 --num-steps 10 --show ``` diff --git a/tests/benchmarks/graph_creation_scaling.py b/tests/benchmarks/graph_creation_scaling.py new file mode 100644 index 0000000..a521b43 --- /dev/null +++ b/tests/benchmarks/graph_creation_scaling.py @@ -0,0 +1,91 @@ +import argparse +import time + +import matplotlib.pyplot as plt +import numpy as np + +import tests.utils as test_utils +import weather_model_graphs as wmg + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark graph creation scaling.") + parser.add_argument( + "--min-N", type=int, default=50, help="Minimum grid size N (NxN nodes)." + ) + parser.add_argument( + "--max-N", type=int, default=400, help="Maximum grid size N (NxN nodes)." + ) + parser.add_argument( + "--num-steps", type=int, default=8, help="Number of intermediate steps." + ) + parser.add_argument( + "--archetype", + type=str, + default="keisler", + choices=["keisler", "oskarsson_hierarchical", "graphcast"], + help="Graph archetype to create.", + ) + parser.add_argument( + "--output", + type=str, + default="scaling_plot.png", + help="Path to save the output plot.", + ) + parser.add_argument( + "--show", action="store_true", help="Show the plot interactively." + ) + + args = parser.parse_args() + + # Generate an array of N values + Ns = np.linspace(args.min_N, args.max_N, args.num_steps, dtype=int) + + fn_name = f"create_{args.archetype}_graph" + create_fn = getattr(wmg.create.archetype, fn_name) + + num_nodes_list = [] + times = [] + + print(f"Benchmarking scaling for {fn_name}...") + for n in Ns: + num_nodes = n * n + print(f"Testing N={n:4d} ({num_nodes:7d} nodes)...", end="", flush=True) + xy = test_utils.create_fake_xy(N=n) + + t0 = time.time() + _ = create_fn(coords=xy) + t1 = time.time() + + duration = t1 - t0 + print(f" {duration:.3f} seconds.") + + num_nodes_list.append(num_nodes) + times.append(duration) + + # Create the plot + plt.figure(figsize=(10, 6)) + plt.plot(num_nodes_list, times, marker="o", linestyle="-", linewidth=2) + + # Add a reference line for linear scaling (O(N)) fitted to the first point + ref_linear = [times[0] * (nodes / num_nodes_list[0]) for nodes in num_nodes_list] + plt.plot( + num_nodes_list, ref_linear, linestyle="--", color="gray", label="O(N) Reference" + ) + + plt.title(f"Graph Creation Scaling: {args.archetype}") + plt.xlabel("Number of Input Grid Nodes (N²)") + plt.ylabel("Execution Time (seconds)") + plt.grid(True, which="both", ls="--", alpha=0.7) + plt.legend() + plt.tight_layout() + + plt.savefig(args.output) + print(f"\nPlot saved to {args.output}") + + if args.show: + plt.show() + + +if __name__ == "__main__": + main() From 098a48ca82a71fe64540e82ecd73d23e3343f8d2 Mon Sep 17 00:00:00 2001 From: Yuvraaj Date: Tue, 21 Apr 2026 22:55:40 +0530 Subject: [PATCH 3/8] feat : testing memory profiling args implemented --- tests/benchmarks/graph_creation_scaling.py | 176 ++++++++++++++------- 1 file changed, 122 insertions(+), 54 deletions(-) diff --git a/tests/benchmarks/graph_creation_scaling.py b/tests/benchmarks/graph_creation_scaling.py index a521b43..9711b8a 100644 --- a/tests/benchmarks/graph_creation_scaling.py +++ b/tests/benchmarks/graph_creation_scaling.py @@ -1,5 +1,8 @@ import argparse +import json import time +import tracemalloc +from typing import Dict, List, Optional import matplotlib.pyplot as plt import numpy as np @@ -8,84 +11,149 @@ import weather_model_graphs as wmg -def main(): - parser = argparse.ArgumentParser(description="Benchmark graph creation scaling.") - parser.add_argument( - "--min-N", type=int, default=50, help="Minimum grid size N (NxN nodes)." - ) - parser.add_argument( - "--max-N", type=int, default=400, help="Maximum grid size N (NxN nodes)." - ) - parser.add_argument( - "--num-steps", type=int, default=8, help="Number of intermediate steps." - ) - parser.add_argument( - "--archetype", - type=str, - default="keisler", - choices=["keisler", "oskarsson_hierarchical", "graphcast"], - help="Graph archetype to create.", - ) - parser.add_argument( - "--output", - type=str, - default="scaling_plot.png", - help="Path to save the output plot.", - ) - parser.add_argument( - "--show", action="store_true", help="Show the plot interactively." - ) - - args = parser.parse_args() - - # Generate an array of N values - Ns = np.linspace(args.min_N, args.max_N, args.num_steps, dtype=int) - - fn_name = f"create_{args.archetype}_graph" +def run_benchmark( + min_N: int, + max_N: int, + num_steps: int, + archetype: str, + track_memory: bool = False, +) -> List[Dict[str, float]]: + """ + Run the graph creation benchmark over a range of grid sizes. + + Returns a list of dicts with keys: + "grid_points" (int), "runtime_s" (float), "peak_memory_mb" (float, optional). + """ + Ns = np.linspace(min_N, max_N, num_steps, dtype=int) + fn_name = f"create_{archetype}_graph" create_fn = getattr(wmg.create.archetype, fn_name) - num_nodes_list = [] - times = [] + results = [] - print(f"Benchmarking scaling for {fn_name}...") for n in Ns: num_nodes = n * n print(f"Testing N={n:4d} ({num_nodes:7d} nodes)...", end="", flush=True) + xy = test_utils.create_fake_xy(N=n) + if track_memory: + tracemalloc.start() + t0 = time.time() - _ = create_fn(coords=xy) + graph = create_fn(coords=xy) t1 = time.time() - duration = t1 - t0 - print(f" {duration:.3f} seconds.") - num_nodes_list.append(num_nodes) - times.append(duration) + peak_mb = None + if track_memory: + _, peak = tracemalloc.get_traced_memory() + peak_mb = peak / (1024 * 1024) + tracemalloc.stop() + + print(f" {duration:.3f} seconds.", end="") + if peak_mb is not None: + print(f" Peak memory: {peak_mb:.1f} MB") + else: + print() + + results.append({ + "grid_points": num_nodes, + "runtime_s": duration, + "peak_memory_mb": peak_mb, + }) + + return results + + +def plot_runtime_scaling(results: List[Dict[str, float]], archetype: str, output_path: str): + """Create a scaling plot for runtime vs number of grid points.""" + grid_points = [r["grid_points"] for r in results] + times = [r["runtime_s"] for r in results] - # Create the plot plt.figure(figsize=(10, 6)) - plt.plot(num_nodes_list, times, marker="o", linestyle="-", linewidth=2) + plt.plot(grid_points, times, marker="o", linestyle="-", linewidth=2) - # Add a reference line for linear scaling (O(N)) fitted to the first point - ref_linear = [times[0] * (nodes / num_nodes_list[0]) for nodes in num_nodes_list] - plt.plot( - num_nodes_list, ref_linear, linestyle="--", color="gray", label="O(N) Reference" - ) + # Add O(N) reference line fitted to the first point + ref_linear = [times[0] * (gp / grid_points[0]) for gp in grid_points] + plt.plot(grid_points, ref_linear, linestyle="--", color="gray", label="O(N) Reference") - plt.title(f"Graph Creation Scaling: {args.archetype}") - plt.xlabel("Number of Input Grid Nodes (N²)") + plt.title(f"Graph Creation Runtime Scaling: {archetype}") + plt.xlabel("Number of Input Grid Nodes") plt.ylabel("Execution Time (seconds)") plt.grid(True, which="both", ls="--", alpha=0.7) plt.legend() plt.tight_layout() + plt.savefig(output_path) + print(f"Runtime scaling plot saved to {output_path}") + + +def plot_memory_scaling(results: List[Dict[str, float]], archetype: str, output_path: str): + """Create a scaling plot for peak memory vs number of grid points.""" + # Filter out results without memory data (should not happen if track_memory=True) + memory_results = [r for r in results if r["peak_memory_mb"] is not None] + if not memory_results: + raise ValueError("No memory data available. Run with --track-memory to collect memory profiles.") + + grid_points = [r["grid_points"] for r in memory_results] + memory = [r["peak_memory_mb"] for r in memory_results] + + plt.figure(figsize=(10, 6)) + plt.plot(grid_points, memory, marker="s", linestyle="-", linewidth=2, color="green") + + plt.title(f"Graph Creation Memory Scaling: {archetype}") + plt.xlabel("Number of Input Grid Nodes") + plt.ylabel("Peak Memory Usage (MB)") + plt.grid(True, which="both", ls="--", alpha=0.7) + plt.tight_layout() + plt.savefig(output_path) + print(f"Memory scaling plot saved to {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark graph creation scaling.") + parser.add_argument("--min-N", type=int, default=50, help="Minimum grid size N (NxN nodes)") + parser.add_argument("--max-N", type=int, default=400, help="Maximum grid size N (NxN nodes)") + parser.add_argument("--num-steps", type=int, default=8, help="Number of intermediate steps") + parser.add_argument( + "--archetype", + choices=["keisler", "oskarsson_hierarchical", "graphcast"], + default="keisler", + help="Graph archetype to create", + ) + parser.add_argument("--output-plot-runtime", type=str, default="runtime_scaling.png", help="Output file for runtime plot") + parser.add_argument("--output-plot-memory", type=str, help="Output file for memory scaling plot (requires --track-memory)") + parser.add_argument("--output-json", type=str, help="Save raw results to JSON file") + parser.add_argument("--track-memory", action="store_true", help="Profile peak memory usage") + parser.add_argument("--show", action="store_true", help="Show plots interactively") + + args = parser.parse_args() + + if args.output_plot_memory and not args.track_memory: + parser.error("--output-plot-memory requires --track-memory") + + results = run_benchmark( + min_N=args.min_N, + max_N=args.max_N, + num_steps=args.num_steps, + archetype=args.archetype, + track_memory=args.track_memory, + ) + + if args.output_json: + with open(args.output_json, "w") as f: + json.dump(results, f, indent=2) + print(f"Raw results saved to {args.output_json}") + + # Always plot runtime (if we have results) + if results: + plot_runtime_scaling(results, args.archetype, args.output_plot_runtime) - plt.savefig(args.output) - print(f"\nPlot saved to {args.output}") + if args.output_plot_memory: + plot_memory_scaling(results, args.archetype, args.output_plot_memory) if args.show: plt.show() if __name__ == "__main__": - main() + main() \ No newline at end of file From 9b708c11c4cd63b33e69dad2053949228f3ad741 Mon Sep 17 00:00:00 2001 From: Yuvraaj Date: Tue, 21 Apr 2026 22:59:34 +0530 Subject: [PATCH 4/8] chore --- tests/benchmarks/graph_creation_scaling.py | 61 ++++++++++++++++------ 1 file changed, 44 insertions(+), 17 deletions(-) diff --git a/tests/benchmarks/graph_creation_scaling.py b/tests/benchmarks/graph_creation_scaling.py index 9711b8a..ee6ff95 100644 --- a/tests/benchmarks/graph_creation_scaling.py +++ b/tests/benchmarks/graph_creation_scaling.py @@ -2,7 +2,7 @@ import json import time import tracemalloc -from typing import Dict, List, Optional +from typing import Dict, List import matplotlib.pyplot as plt import numpy as np @@ -56,16 +56,20 @@ def run_benchmark( else: print() - results.append({ - "grid_points": num_nodes, - "runtime_s": duration, - "peak_memory_mb": peak_mb, - }) + results.append( + { + "grid_points": num_nodes, + "runtime_s": duration, + "peak_memory_mb": peak_mb, + } + ) return results -def plot_runtime_scaling(results: List[Dict[str, float]], archetype: str, output_path: str): +def plot_runtime_scaling( + results: List[Dict[str, float]], archetype: str, output_path: str +): """Create a scaling plot for runtime vs number of grid points.""" grid_points = [r["grid_points"] for r in results] times = [r["runtime_s"] for r in results] @@ -75,7 +79,9 @@ def plot_runtime_scaling(results: List[Dict[str, float]], archetype: str, output # Add O(N) reference line fitted to the first point ref_linear = [times[0] * (gp / grid_points[0]) for gp in grid_points] - plt.plot(grid_points, ref_linear, linestyle="--", color="gray", label="O(N) Reference") + plt.plot( + grid_points, ref_linear, linestyle="--", color="gray", label="O(N) Reference" + ) plt.title(f"Graph Creation Runtime Scaling: {archetype}") plt.xlabel("Number of Input Grid Nodes") @@ -87,12 +93,16 @@ def plot_runtime_scaling(results: List[Dict[str, float]], archetype: str, output print(f"Runtime scaling plot saved to {output_path}") -def plot_memory_scaling(results: List[Dict[str, float]], archetype: str, output_path: str): +def plot_memory_scaling( + results: List[Dict[str, float]], archetype: str, output_path: str +): """Create a scaling plot for peak memory vs number of grid points.""" # Filter out results without memory data (should not happen if track_memory=True) memory_results = [r for r in results if r["peak_memory_mb"] is not None] if not memory_results: - raise ValueError("No memory data available. Run with --track-memory to collect memory profiles.") + raise ValueError( + "No memory data available. Run with --track-memory to collect memory profiles." + ) grid_points = [r["grid_points"] for r in memory_results] memory = [r["peak_memory_mb"] for r in memory_results] @@ -111,19 +121,36 @@ def plot_memory_scaling(results: List[Dict[str, float]], archetype: str, output_ def main(): parser = argparse.ArgumentParser(description="Benchmark graph creation scaling.") - parser.add_argument("--min-N", type=int, default=50, help="Minimum grid size N (NxN nodes)") - parser.add_argument("--max-N", type=int, default=400, help="Maximum grid size N (NxN nodes)") - parser.add_argument("--num-steps", type=int, default=8, help="Number of intermediate steps") + parser.add_argument( + "--min-N", type=int, default=50, help="Minimum grid size N (NxN nodes)" + ) + parser.add_argument( + "--max-N", type=int, default=400, help="Maximum grid size N (NxN nodes)" + ) + parser.add_argument( + "--num-steps", type=int, default=8, help="Number of intermediate steps" + ) parser.add_argument( "--archetype", choices=["keisler", "oskarsson_hierarchical", "graphcast"], default="keisler", help="Graph archetype to create", ) - parser.add_argument("--output-plot-runtime", type=str, default="runtime_scaling.png", help="Output file for runtime plot") - parser.add_argument("--output-plot-memory", type=str, help="Output file for memory scaling plot (requires --track-memory)") + parser.add_argument( + "--output-plot-runtime", + type=str, + default="runtime_scaling.png", + help="Output file for runtime plot", + ) + parser.add_argument( + "--output-plot-memory", + type=str, + help="Output file for memory scaling plot (requires --track-memory)", + ) parser.add_argument("--output-json", type=str, help="Save raw results to JSON file") - parser.add_argument("--track-memory", action="store_true", help="Profile peak memory usage") + parser.add_argument( + "--track-memory", action="store_true", help="Profile peak memory usage" + ) parser.add_argument("--show", action="store_true", help="Show plots interactively") args = parser.parse_args() @@ -156,4 +183,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() From db3b0f687c2974683bf41ba4216df28d5c8ff845 Mon Sep 17 00:00:00 2001 From: Yuvraaj Date: Wed, 22 Apr 2026 16:29:02 +0530 Subject: [PATCH 5/8] chore : int conversion fixes --- tests/benchmarks/graph_creation_scaling.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/tests/benchmarks/graph_creation_scaling.py b/tests/benchmarks/graph_creation_scaling.py index ee6ff95..211aac7 100644 --- a/tests/benchmarks/graph_creation_scaling.py +++ b/tests/benchmarks/graph_creation_scaling.py @@ -31,7 +31,7 @@ def run_benchmark( results = [] for n in Ns: - num_nodes = n * n + num_nodes = int(n * n) # convert to Python int print(f"Testing N={n:4d} ({num_nodes:7d} nodes)...", end="", flush=True) xy = test_utils.create_fake_xy(N=n) @@ -47,7 +47,7 @@ def run_benchmark( peak_mb = None if track_memory: _, peak = tracemalloc.get_traced_memory() - peak_mb = peak / (1024 * 1024) + peak_mb = float(peak) / (1024 * 1024) # convert to float tracemalloc.stop() print(f" {duration:.3f} seconds.", end="") @@ -56,13 +56,11 @@ def run_benchmark( else: print() - results.append( - { - "grid_points": num_nodes, - "runtime_s": duration, - "peak_memory_mb": peak_mb, - } - ) + results.append({ + "grid_points": num_nodes, + "runtime_s": duration, + "peak_memory_mb": peak_mb, + }) return results @@ -183,4 +181,4 @@ def main(): if __name__ == "__main__": - main() + main() \ No newline at end of file From 123eb72452dad709335c779bdb687c79d57b8bff Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Thu, 16 Jul 2026 14:23:23 +0200 Subject: [PATCH 6/8] final tweaks --- tests/benchmarks/graph_creation_scaling.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tests/benchmarks/graph_creation_scaling.py b/tests/benchmarks/graph_creation_scaling.py index 8d97fe2..0585909 100644 --- a/tests/benchmarks/graph_creation_scaling.py +++ b/tests/benchmarks/graph_creation_scaling.py @@ -6,6 +6,7 @@ import matplotlib.pyplot as plt import numpy as np +from loguru import logger import tests.utils as test_utils import weather_model_graphs as wmg @@ -32,7 +33,7 @@ def run_benchmark( for n in Ns: num_nodes = int(n * n) # convert to Python int - print(f"Testing N={n:4d} ({num_nodes:7d} nodes)...", end="", flush=True) + logger.info(f"Testing N={n:4d} ({num_nodes:7d} nodes)...") xy = test_utils.create_fake_xy(N=n) @@ -50,11 +51,9 @@ def run_benchmark( peak_mb = float(peak) / (1024 * 1024) # convert to float tracemalloc.stop() - print(f" {duration:.3f} seconds.", end="") + logger.info(f" {duration:.3f} seconds.") if peak_mb is not None: - print(f" Peak memory: {peak_mb:.1f} MB") - else: - print() + logger.info(f" Peak memory: {peak_mb:.1f} MB") results.append( { @@ -161,7 +160,7 @@ def main(): if args.output_json: with open(args.output_json, "w") as f: json.dump(results, f, indent=2) - print(f"Raw results saved to {args.output_json}") + logger.info(f"Raw results saved to {args.output_json}") # Always plot runtime (if we have results) if results: @@ -170,8 +169,8 @@ def main(): if args.output_plot_memory: plot_memory_scaling(results, args.archetype, args.output_plot_memory) - plt.savefig(args.output) - print(f"\nPlot saved to {args.output}") + plt.savefig(args.output_plot_runtime) + logger.info(f"Runtime plot saved to {args.output_plot_runtime}") if args.show: plt.show() From 7b308e01d53979356f369f6ac51741de46ac29d9 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Thu, 16 Jul 2026 14:24:58 +0200 Subject: [PATCH 7/8] changelog --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d0a8008..3794f22 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (`save.base`, `save.neural_lam.torch_tensors`, `save.neural_lam.deprecated`) with all existing entrypoints re-exported unchanged. [\#123](https://github.com/mllam/weather-model-graphs/pull/123), @prajwal-tech07 +- Add support for writing benchmarking results to json, + [\#140](https://github.com/mllam/weather-model-graphs/pull/140), + @yuvraajnarula & @leifdenby ### Deprecated From d896ee78e76bb9087f6197e6e8b1efa4031de5a4 Mon Sep 17 00:00:00 2001 From: Leif Denby Date: Thu, 16 Jul 2026 15:49:30 +0200 Subject: [PATCH 8/8] ensure benchmark scaling plots are saved --- tests/benchmarks/graph_creation_scaling.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/benchmarks/graph_creation_scaling.py b/tests/benchmarks/graph_creation_scaling.py index 0585909..b29e12e 100644 --- a/tests/benchmarks/graph_creation_scaling.py +++ b/tests/benchmarks/graph_creation_scaling.py @@ -85,6 +85,9 @@ def plot_runtime_scaling( plt.title(f"Graph Creation Runtime Scaling: {archetype}") plt.xlabel("Number of Input Grid Nodes") + plt.savefig(output_path) + logger.info(f"Runtime scaling plot saved to {output_path}") + def plot_memory_scaling( results: List[Dict[str, float]], archetype: str, output_path: str @@ -109,6 +112,9 @@ def plot_memory_scaling( plt.grid(True, which="both", ls="--", alpha=0.7) plt.tight_layout() + plt.savefig(output_path) + logger.info(f"Memory scaling plot saved to {output_path}") + def main(): parser = argparse.ArgumentParser(description="Benchmark graph creation scaling.") @@ -169,9 +175,6 @@ def main(): if args.output_plot_memory: plot_memory_scaling(results, args.archetype, args.output_plot_memory) - plt.savefig(args.output_plot_runtime) - logger.info(f"Runtime plot saved to {args.output_plot_runtime}") - if args.show: plt.show()