benchmark flow is working. changed paths to use module_path instead of file_path for Benchmarkkey

Author: alvin-r

codeflash/benchmarking/codeflash_trace.py

@@ -1,13 +1,13 @@
 import functools
 import os
+import pickle
 import sqlite3
 import sys
+import time
+from typing import Callable
 
-import pickle
 import dill
 
-import time
-from typing import Callable, Optional
 
 class CodeflashTrace:
     """Decorator class that traces and profiles function execution."""
@@ -35,7 +35,7 @@ def setup(self, trace_path: str) -> None:
             cur.execute(
                 "CREATE TABLE IF NOT EXISTS benchmark_function_timings("
                 "function_name TEXT, class_name TEXT, module_name TEXT, file_path TEXT,"
-                "benchmark_function_name TEXT, benchmark_file_path TEXT, benchmark_line_number INTEGER,"
+                "benchmark_function_name TEXT, benchmark_module_path TEXT, benchmark_line_number INTEGER,"
                 "function_time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)"
             )
             self._connection.commit()
@@ -51,6 +51,7 @@ def write_function_timings(self) -> None:
 
         Args:
             data: List of function call data tuples to write
+
         """
         if not self.function_calls_data:
             return  # No data to write
@@ -64,7 +65,7 @@ def write_function_timings(self) -> None:
             cur.executemany(
                 "INSERT INTO benchmark_function_timings"
                 "(function_name, class_name, module_name, file_path, benchmark_function_name, "
-                "benchmark_file_path, benchmark_line_number, function_time_ns, overhead_time_ns, args, kwargs) "
+                "benchmark_module_path, benchmark_line_number, function_time_ns, overhead_time_ns, args, kwargs) "
                 "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                 self.function_calls_data
             )
@@ -116,7 +117,7 @@ def wrapper(*args, **kwargs):
 
             # Get benchmark info from environment
             benchmark_function_name = os.environ.get("CODEFLASH_BENCHMARK_FUNCTION_NAME", "")
-            benchmark_file_path = os.environ.get("CODEFLASH_BENCHMARK_FILE_PATH", "")
+            benchmark_module_path = os.environ.get("CODEFLASH_BENCHMARK_MODULE_PATH", "")
             benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "")
             # Get class name
             class_name = ""
@@ -143,7 +144,7 @@ def wrapper(*args, **kwargs):
 
                     except (TypeError, dill.PicklingError, AttributeError, RecursionError, OSError) as e:
                         print(f"Error pickling arguments for function {func.__name__}: {e}")
-                        return
+                        return None
 
             if len(self.function_calls_data) > 1000:
                 self.write_function_timings()
@@ -152,7 +153,7 @@ def wrapper(*args, **kwargs):
 
             self.function_calls_data.append(
                 (func.__name__, class_name, func.__module__, func.__code__.co_filename,
-                 benchmark_function_name, benchmark_file_path, benchmark_line_number, execution_time,
+                 benchmark_function_name, benchmark_module_path, benchmark_line_number, execution_time,
                  overhead_time, pickled_args, pickled_kwargs)
             )
             return result

codeflash/benchmarking/plugin/plugin.py

@@ -31,7 +31,7 @@ def setup(self, trace_path:str, project_root:str) -> None:
             cur.execute("PRAGMA journal_mode = MEMORY")
             cur.execute(
                 "CREATE TABLE IF NOT EXISTS benchmark_timings("
-                "benchmark_file_path TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER,"
+                "benchmark_module_path TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER,"
                 "benchmark_time_ns INTEGER)"
             )
             self._connection.commit()
@@ -54,7 +54,7 @@ def write_benchmark_timings(self) -> None:
             cur = self._connection.cursor()
             # Insert data into the benchmark_timings table
             cur.executemany(
-                "INSERT INTO benchmark_timings (benchmark_file_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)",
+                "INSERT INTO benchmark_timings (benchmark_module_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)",
                 self.benchmark_timings
             )
             self._connection.commit()
@@ -93,7 +93,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[Benchmark
             # Query the function_calls table for all function calls
             cursor.execute(
                 "SELECT module_name, class_name, function_name, "
-                "benchmark_file_path, benchmark_function_name, benchmark_line_number, function_time_ns "
+                "benchmark_module_path, benchmark_function_name, benchmark_line_number, function_time_ns "
                 "FROM benchmark_function_timings"
             )
 
@@ -108,7 +108,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[Benchmark
                     qualified_name = f"{module_name}.{function_name}"
 
                 # Create the benchmark key (file::function::line)
-                benchmark_key = BenchmarkKey(file_path=benchmark_file, function_name=benchmark_func)
+                benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
                 # Initialize the inner dictionary if needed
                 if qualified_name not in result:
                     result[qualified_name] = {}
@@ -150,20 +150,20 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
         try:
             # Query the benchmark_function_timings table to get total overhead for each benchmark
             cursor.execute(
-                "SELECT benchmark_file_path, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) "
+                "SELECT benchmark_module_path, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) "
                 "FROM benchmark_function_timings "
-                "GROUP BY benchmark_file_path, benchmark_function_name, benchmark_line_number"
+                "GROUP BY benchmark_module_path, benchmark_function_name, benchmark_line_number"
             )
 
             # Process overhead information
             for row in cursor.fetchall():
                 benchmark_file, benchmark_func, benchmark_line, total_overhead_ns = row
-                benchmark_key = BenchmarkKey(file_path=benchmark_file, function_name=benchmark_func)
+                benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
                 overhead_by_benchmark[benchmark_key] = total_overhead_ns or 0  # Handle NULL sum case
 
             # Query the benchmark_timings table for total times
             cursor.execute(
-                "SELECT benchmark_file_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns "
+                "SELECT benchmark_module_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns "
                 "FROM benchmark_timings"
             )
 
@@ -172,7 +172,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
                 benchmark_file, benchmark_func, benchmark_line, time_ns = row
 
                 # Create the benchmark key (file::function::line)
-                benchmark_key = BenchmarkKey(file_path=benchmark_file, function_name=benchmark_func)
+                benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
                 # Subtract overhead from total time
                 overhead = overhead_by_benchmark.get(benchmark_key, 0)
                 result[benchmark_key] = time_ns - overhead
@@ -244,13 +244,13 @@ def test_something(benchmark):
             a
 
             """
-            benchmark_file_path = module_name_from_file_path(Path(str(self.request.node.fspath)), Path(codeflash_benchmark_plugin.project_root))
+            benchmark_module_path = module_name_from_file_path(Path(str(self.request.node.fspath)), Path(codeflash_benchmark_plugin.project_root))
             benchmark_function_name = self.request.node.name
             line_number = int(str(sys._getframe(1).f_lineno))  # 1 frame up in the call stack
 
             # Set env vars so codeflash decorator can identify what benchmark its being run in
             os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name
-            os.environ["CODEFLASH_BENCHMARK_FILE_PATH"] = benchmark_file_path
+            os.environ["CODEFLASH_BENCHMARK_MODULE_PATH"] = benchmark_module_path
             os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number)
             os.environ["CODEFLASH_BENCHMARKING"] = "True"
 
@@ -268,7 +268,7 @@ def test_something(benchmark):
             codeflash_trace.function_call_count = 0
             # Add to the benchmark timings buffer
             codeflash_benchmark_plugin.benchmark_timings.append(
-                (benchmark_file_path, benchmark_function_name, line_number, end - start))
+                (benchmark_module_path, benchmark_function_name, line_number, end - start))
 
             return result
 

codeflash/benchmarking/replay_test.py

@@ -227,18 +227,18 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework
 
         # Get distinct benchmark file paths
         cursor.execute(
-            "SELECT DISTINCT benchmark_file_path FROM benchmark_function_timings"
+            "SELECT DISTINCT benchmark_module_path FROM benchmark_function_timings"
         )
         benchmark_files = cursor.fetchall()
 
         # Generate a test for each benchmark file
         for benchmark_file in benchmark_files:
-            benchmark_file_path = benchmark_file[0]
+            benchmark_module_path = benchmark_file[0]
             # Get all benchmarks and functions associated with this file path
             cursor.execute(
             "SELECT DISTINCT benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings "
-                "WHERE benchmark_file_path = ?",
-                (benchmark_file_path,)
+                "WHERE benchmark_module_path = ?",
+                (benchmark_module_path,)
             )
 
             functions_data = []
@@ -251,7 +251,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework
                     "file_path": file_path,
                     "module_name": module_name,
                     "benchmark_function_name": benchmark_function_name,
-                    "benchmark_file_path": benchmark_file_path,
+                    "benchmark_module_path": benchmark_module_path,
                     "benchmark_line_number": benchmark_line_number,
                     "function_properties": inspect_top_level_functions_or_methods(
                         file_name=Path(file_path),
@@ -261,7 +261,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework
                 })
 
             if not functions_data:
-                logger.info(f"No benchmark test functions found in {benchmark_file_path}")
+                logger.info(f"No benchmark test functions found in {benchmark_module_path}")
                 continue
             # Generate the test code for this benchmark
             test_code = create_trace_replay_test_code(
@@ -272,7 +272,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework
             )
             test_code = isort.code(test_code)
             output_file = get_test_file_path(
-                test_dir=Path(output_dir), function_name=benchmark_file_path, test_type="replay"
+                test_dir=Path(output_dir), function_name=benchmark_module_path, test_type="replay"
             )
             # Write test code to file, parents = true
             output_dir.mkdir(parents=True, exist_ok=True)

codeflash/benchmarking/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import shutil
 from typing import Optional
 
 from rich.console import Console
@@ -35,32 +36,45 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di
         function_to_result[func_path] = sorted_tests
     return function_to_result
 
+
 def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey, float, float, float]]]) -> None:
-    console = Console()
+
+    try:
+        terminal_width = int(shutil.get_terminal_size().columns * 0.8)
+    except Exception:
+        terminal_width = 200  # Fallback width
+    console = Console(width = terminal_width)
     for func_path, sorted_tests in function_to_results.items():
         console.print()
         function_name = func_path.split(":")[-1]
 
         # Create a table for this function
         table = Table(title=f"Function: {function_name}", border_style="blue")
 
-        # Add columns
-        table.add_column("Benchmark Test", style="cyan", no_wrap=True)
+        # Add columns - split the benchmark test into two columns
+        table.add_column("Benchmark Module Path", style="cyan", no_wrap=True)
+        table.add_column("Test Function", style="magenta", no_wrap=True)
         table.add_column("Total Time (ms)", justify="right", style="green")
         table.add_column("Function Time (ms)", justify="right", style="yellow")
         table.add_column("Percentage (%)", justify="right", style="red")
 
         for benchmark_key, total_time, func_time, percentage in sorted_tests:
+            # Split the benchmark test into module path and function name
+            module_path = benchmark_key.module_path
+            test_function = benchmark_key.function_name
+
             if total_time == 0.0:
                 table.add_row(
-                    f"{benchmark_key.file_path}::{benchmark_key.function_name}",
+                    module_path,
+                    test_function,
                     "N/A",
                     "N/A",
                     "N/A"
                 )
             else:
                 table.add_row(
-                    f"{benchmark_key.file_path}::{benchmark_key.function_name}",
+                    module_path,
+                    test_function,
                     f"{total_time:.3f}",
                     f"{func_time:.3f}",
                     f"{percentage:.2f}"
@@ -108,7 +122,7 @@ def process_benchmark_data(
 
         benchmark_details.append(
             BenchmarkDetail(
-                benchmark_name=benchmark_key.file_path,
+                benchmark_name=benchmark_key.module_path,
                 test_function=benchmark_key.function_name,
                 original_timing=humanize_runtime(int(total_benchmark_timing)),
                 expected_new_timing=humanize_runtime(int(expected_new_benchmark_timing)),

codeflash/models/models.py

@@ -84,11 +84,11 @@ class BestOptimization(BaseModel):
 
 @dataclass(frozen=True)
 class BenchmarkKey:
-    file_path: str
+    module_path: str
     function_name: str
 
     def __str__(self) -> str:
-        return f"{self.file_path}::{self.function_name}"
+        return f"{self.module_path}::{self.function_name}"
 
 @dataclass
 class BenchmarkDetail:
@@ -484,7 +484,7 @@ def group_by_benchmarks(self, benchmark_keys:list[BenchmarkKey], benchmark_repla
         test_results_by_benchmark = defaultdict(TestResults)
         benchmark_module_path = {}
         for benchmark_key in benchmark_keys:
-            benchmark_module_path[benchmark_key] = module_name_from_file_path(benchmark_replay_test_dir.resolve() / f"test_{Path(benchmark_key.file_path).name.split('.')[0][5:]}__replay_test_", project_root)
+            benchmark_module_path[benchmark_key] = module_name_from_file_path(benchmark_replay_test_dir.resolve() / f"test_{benchmark_key.module_path.replace(".", "_")}__replay_test_", project_root)
         for test_result in self.test_results:
             if (test_result.test_type == TestType.REPLAY_TEST):
                 for benchmark_key, module_path in benchmark_module_path.items():