initial attempt for codeflash_trace_decorator

codeflash/benchmarking/codeflash_trace_decorator.py

@@ -0,0 +1,73 @@
+import functools
+import pickle
+import sqlite3
+import time
+import os
+
+def codeflash_trace(output_file: str):
+ """A decorator factory that returns a decorator that measures the execution time
+ of a function and pickles its arguments using the highest protocol available.
+
+ Args:
+ output_file: Path to the SQLite database file where results will be stored
+
+ Returns:
+ The decorator function
+
+ """
+ def decorator(func):
+ @functools.wraps(func)
+ def wrapper(*args, **kwargs):
+ # Measure execution time
+ start_time = time.time()
+ result = func(*args, **kwargs)
+ end_time = time.time()
+
+ # Calculate execution time
+ execution_time = end_time - start_time
+
+ # Measure overhead
+ overhead_start_time = time.time()
+
+ try:
+ # Connect to the database
+ con = sqlite3.connect(output_file)
+ cur = con.cursor()
+ cur.execute("PRAGMA synchronous = OFF")
+
+ # Check if table exists and create it if it doesn't
+ cur.execute(
+ "CREATE TABLE IF NOT EXISTS function_calls(function_name TEXT, class_name TEXT, file_name TEXT, benchmark_function_name TEXT, benchmark_file_name TEXT,"
+ "time_ns INTEGER, args BLOB, kwargs BLOB)"
+ )
+
+ # Pickle the arguments
+ pickled_args = pickle.dumps(args, protocol=pickle.HIGHEST_PROTOCOL)
+ pickled_kwargs = pickle.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL)
+
+ # Get benchmark info from environment
+ benchmark_function_name = os.environ.get("CODEFLASH_BENCHMARK_FUNCTION_NAME")
+ benchmark_file_name = os.environ.get("CODEFLASH_BENCHMARK_FILE_NAME")
+ # Insert the data
+ cur.execute(
+ "INSERT INTO function_calls (function_name, classname, filename, benchmark_function_name, benchmark_file_name, time_ns, args, kwargs) "
+ "VALUES (?, ?, ?, ?, ?, ?)",
+ (func.__name__, func.__module__, func.__code__.co_filename,
+ execution_time, pickled_args, pickled_kwargs)
+ )
+
+ # Commit and close
+ con.commit()
+ con.close()
+
+ overhead_end_time = time.time()
+
+ print(f"Function '{func.__name__}' took {execution_time:.6f} seconds to execute")
+ print(f"Function '{func.__name__}' overhead took {overhead_end_time - overhead_start_time:.6f} seconds to execute")
+
+ except Exception as e:
+ print(f"Error in codeflash_trace: {e}")
+
+ return result
+ return wrapper
+ return decorator

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -0,0 +1 @@
+

codeflash/benchmarking/plugin/plugin.py

@@ -1,7 +1,5 @@
 import pytest
-
-from codeflash.tracer import Tracer
-from pathlib import Path
+import time
 
 class CodeFlashPlugin:
  @staticmethod
@@ -12,18 +10,6 @@ def pytest_addoption(parser):
  default=False,
  help="Enable CodeFlash tracing"
  )
- parser.addoption(
- "--functions",
- action="store",
- default="",
- help="Comma-separated list of additional functions to trace"
- )
- parser.addoption(
- "--benchmarks-root",
- action="store",
- default=".",
- help="Root directory for benchmarks"
- )
 
  @staticmethod
  def pytest_plugin_registered(plugin, manager):
@@ -49,32 +35,10 @@ def benchmark(request):
 
  class Benchmark:
  def __call__(self, func, *args, **kwargs):
- func_name = func.__name__
- test_name = request.node.name
- additional_functions = request.config.getoption("--functions").split(",")
- trace_functions = [f for f in additional_functions if f]
- print("Tracing functions: ", trace_functions)
-
- # Get benchmarks root directory from command line option
- benchmarks_root = Path(request.config.getoption("--benchmarks-root"))
-
- # Create .trace directory if it doesn't exist
- trace_dir = benchmarks_root / '.codeflash_trace'
- trace_dir.mkdir(exist_ok=True)
-
- # Set output path to the .trace directory
- output_path = trace_dir / f"{test_name}.trace"
-
- tracer = Tracer(
- output=str(output_path), # Convert Path to string for Tracer
- functions=trace_functions,
- max_function_count=256,
- benchmark=True
- )
-
- with tracer:
- result = func(*args, **kwargs)
-
+ start = time.time_ns()
+ result = func(*args, **kwargs)
+ end = time.time_ns()
+ print(f"Benchmark: {func.__name__} took {end - start} ns")
  return result
 
  return Benchmark()

codeflash/benchmarking/pytest_new_process_trace_benchmarks.py

@@ -2,13 +2,12 @@
 from plugin.plugin import CodeFlashPlugin
 
 benchmarks_root = sys.argv[1]
-function_list = sys.argv[2]
 if __name__ == "__main__":
  import pytest
 
  try:
  exitcode = pytest.main(
- [benchmarks_root, "--benchmarks-root", benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "--functions", function_list,"-o", "addopts="], plugins=[CodeFlashPlugin()]
+ [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashPlugin()]
  )
  except Exception as e:
  print(f"Failed to collect tests: {e!s}")

codeflash/benchmarking/trace_benchmarks.py

@@ -3,13 +3,12 @@
 from pathlib import Path
 import subprocess
 
-def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path, function_list: list[str] = []) -> None:
+def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path) -> None:
  result = subprocess.run(
  [
  SAFE_SYS_EXECUTABLE,
  Path(__file__).parent / "pytest_new_process_trace_benchmarks.py",
- str(benchmarks_root),
- ",".join(function_list)
+ benchmarks_root,
  ],
  cwd=project_root,
  check=False,

codeflash/optimization/optimizer.py

@@ -25,6 +25,7 @@
 from codeflash.verification.verification_utils import TestConfig
 from codeflash.benchmarking.get_trace_info import get_function_benchmark_timings, get_benchmark_timings
 from codeflash.benchmarking.utils import print_benchmark_table
+from collections import defaultdict
 if TYPE_CHECKING:
  from argparse import Namespace
 
@@ -91,18 +92,28 @@ def run(self) -> None:
  module_root=self.args.module_root,
  )
  if self.args.benchmark:
- all_functions_to_optimize = [
- function
- for functions_list in file_to_funcs_to_optimize.values()
- for function in functions_list
- ]
- logger.info(f"Tracing existing benchmarks for {len(all_functions_to_optimize)} functions")
- trace_benchmarks_pytest(self.args.benchmarks_root, self.args.project_root, [fto.qualified_name_with_file_name for fto in all_functions_to_optimize])
- logger.info("Finished tracing existing benchmarks")
- trace_dir = Path(self.args.benchmarks_root) / ".codeflash_trace"
- function_benchmark_timings = get_function_benchmark_timings(trace_dir, all_functions_to_optimize)
- total_benchmark_timings = get_benchmark_timings(trace_dir)
- print_benchmark_table(function_benchmark_timings, total_benchmark_timings)
+ # Insert decorator
+ file_path_to_source_code = defaultdict(str)
+ for file in file_to_funcs_to_optimize:
+ with file.open("r", encoding="utf8") as f:
+ file_path_to_source_code[file] = f.read()
+ try:
+ for functions_to_optimize in file_to_funcs_to_optimize.values():
+ for fto in functions_to_optimize:
+ pass
+ #instrument_codeflash_trace_decorator(fto)
+ trace_benchmarks_pytest(self.args.project_root) # Simply run all tests that use pytest-benchmark
+ logger.info("Finished tracing existing benchmarks")
+ finally:
+ # Restore original source code
+ for file in file_path_to_source_code:
+ with file.open("w", encoding="utf8") as f:
+ f.write(file_path_to_source_code[file])
+
+ # trace_dir = Path(self.args.benchmarks_root) / ".codeflash_trace"
+ # function_benchmark_timings = get_function_benchmark_timings(trace_dir, all_functions_to_optimize)
+ # total_benchmark_timings = get_benchmark_timings(trace_dir)
+ # print_benchmark_table(function_benchmark_timings, total_benchmark_timings)
 
 
  optimizations_found: int = 0

tests/test_trace_benchmarks.py

@@ -5,4 +5,4 @@ def test_trace_benchmarks():
  # Test the trace_benchmarks function
  project_root = Path(__file__).parent.parent / "code_to_optimize"
  benchmarks_root = project_root / "tests" / "pytest" / "benchmarks"
- trace_benchmarks_pytest(benchmarks_root, project_root, ["sorter"])
+ trace_benchmarks_pytest(benchmarks_root, project_root)