Support recursive functions, and @benchmark / @pytest.mark.benchmark ways of using benchmark. created tests for all of them

Author: alvin-r

code_to_optimize/bubble_sort_codeflash_trace.py

@@ -9,6 +9,24 @@ def sorter(arr):
                 arr[j + 1] = temp
     return arr
 
+@codeflash_trace
+def recursive_bubble_sort(arr, n=None):
+    # Initialize n if not provided
+    if n is None:
+        n = len(arr)
+
+    # Base case: if n is 1, the array is already sorted
+    if n == 1:
+        return arr
+
+    # One pass of bubble sort - move the largest element to the end
+    for i in range(n - 1):
+        if arr[i] > arr[i + 1]:
+            arr[i], arr[i + 1] = arr[i + 1], arr[i]
+
+    # Recursively sort the remaining n-1 elements
+    return recursive_bubble_sort(arr, n - 1)
+
 class Sorter:
     @codeflash_trace
     def __init__(self, arr):

code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py

@@ -0,0 +1,6 @@
+from code_to_optimize.bubble_sort_codeflash_trace import recursive_bubble_sort
+
+
+def test_recursive_sort(benchmark):
+    result = benchmark(recursive_bubble_sort, list(reversed(range(500))))
+    assert result == list(range(500))

code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py

@@ -0,0 +1,11 @@
+import pytest
+from code_to_optimize.bubble_sort_codeflash_trace import sorter
+
+def test_benchmark_sort(benchmark):
+    @benchmark
+    def do_sort():
+        sorter(list(reversed(range(500))))
+
+@pytest.mark.benchmark(group="benchmark_decorator")
+def test_pytest_mark(benchmark):
+    benchmark(sorter, list(reversed(range(500))))

codeflash/benchmarking/codeflash_trace.py

@@ -3,6 +3,7 @@
 import pickle
 import sqlite3
 import sys
+import threading
 import time
 from typing import Callable
 
@@ -18,6 +19,8 @@ def __init__(self) -> None:
         self.pickle_count_limit = 1000
         self._connection = None
         self._trace_path = None
+        self._thread_local = threading.local()
+        self._thread_local.active_functions = set()
 
     def setup(self, trace_path: str) -> None:
         """Set up the database connection for direct writing.
@@ -98,23 +101,29 @@ def __call__(self, func: Callable) -> Callable:
             The wrapped function
 
         """
+        func_id = (func.__module__,func.__name__)
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
+            # Initialize thread-local active functions set if it doesn't exist
+            if not hasattr(self._thread_local, "active_functions"):
+                self._thread_local.active_functions = set()
+            # If it's in a recursive function, just return the result
+            if func_id in self._thread_local.active_functions:
+                return func(*args, **kwargs)
+            # Track active functions so we can detect recursive functions
+            self._thread_local.active_functions.add(func_id)
             # Measure execution time
             start_time = time.thread_time_ns()
             result = func(*args, **kwargs)
             end_time = time.thread_time_ns()
             # Calculate execution time
             execution_time = end_time - start_time
-
             self.function_call_count += 1
 
-            # Measure overhead
-            original_recursion_limit = sys.getrecursionlimit()
             # Check if currently in pytest benchmark fixture
             if os.environ.get("CODEFLASH_BENCHMARKING", "False") == "False":
+                self._thread_local.active_functions.remove(func_id)
                 return result
-
             # Get benchmark info from environment
             benchmark_function_name = os.environ.get("CODEFLASH_BENCHMARK_FUNCTION_NAME", "")
             benchmark_module_path = os.environ.get("CODEFLASH_BENCHMARK_MODULE_PATH", "")
@@ -125,32 +134,54 @@ def wrapper(*args, **kwargs):
             if "." in qualname:
                 class_name = qualname.split(".")[0]
 
-            if self.function_call_count <= self.pickle_count_limit:
+            # Limit pickle count so memory does not explode
+            if self.function_call_count > self.pickle_count_limit:
+                print("Pickle limit reached")
+                self._thread_local.active_functions.remove(func_id)
+                overhead_time = time.thread_time_ns() - end_time
+                self.function_calls_data.append(
+                    (func.__name__, class_name, func.__module__, func.__code__.co_filename,
+                     benchmark_function_name, benchmark_module_path, benchmark_line_number, execution_time,
+                     overhead_time, None, None)
+                )
+                return result
+
+            try:
+                original_recursion_limit = sys.getrecursionlimit()
+                sys.setrecursionlimit(10000)
+                # args = dict(args.items())
+                # if class_name and func.__name__ == "__init__" and "self" in args:
+                #     del args["self"]
+                # Pickle the arguments
+                pickled_args = pickle.dumps(args, protocol=pickle.HIGHEST_PROTOCOL)
+                pickled_kwargs = pickle.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL)
+                sys.setrecursionlimit(original_recursion_limit)
+            except (TypeError, pickle.PicklingError, AttributeError, RecursionError, OSError):
+                # Retry with dill if pickle fails. It's slower but more comprehensive
                 try:
-                    sys.setrecursionlimit(1000000)
-                    args = dict(args.items())
-                    if class_name and func.__name__ == "__init__" and "self" in args:
-                        del args["self"]
-                    # Pickle the arguments
-                    pickled_args = pickle.dumps(args, protocol=pickle.HIGHEST_PROTOCOL)
-                    pickled_kwargs = pickle.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL)
+                    pickled_args = dill.dumps(args, protocol=pickle.HIGHEST_PROTOCOL)
+                    pickled_kwargs = dill.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL)
                     sys.setrecursionlimit(original_recursion_limit)
-                except (TypeError, pickle.PicklingError, AttributeError, RecursionError, OSError):
-                    # we retry with dill if pickle fails. It's slower but more comprehensive
-                    try:
-                        pickled_args = dill.dumps(args, protocol=pickle.HIGHEST_PROTOCOL)
-                        pickled_kwargs = dill.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL)
-                        sys.setrecursionlimit(original_recursion_limit)
-
-                    except (TypeError, dill.PicklingError, AttributeError, RecursionError, OSError) as e:
-                        print(f"Error pickling arguments for function {func.__name__}: {e}")
-                        return result
 
+                except (TypeError, dill.PicklingError, AttributeError, RecursionError, OSError) as e:
+                    print(f"Error pickling arguments for function {func.__name__}: {e}")
+                    # Add to the list of function calls without pickled args. Used for timing info only
+                    self._thread_local.active_functions.remove(func_id)
+                    overhead_time = time.thread_time_ns() - end_time
+                    self.function_calls_data.append(
+                        (func.__name__, class_name, func.__module__, func.__code__.co_filename,
+                         benchmark_function_name, benchmark_module_path, benchmark_line_number, execution_time,
+                         overhead_time, None, None)
+                    )
+                    return result
+
+            # Flush to database every 1000 calls
             if len(self.function_calls_data) > 1000:
                 self.write_function_timings()
-            # Calculate overhead time
-            overhead_time = time.thread_time_ns() - end_time
 
+            # Add to the list of function calls with pickled args, to be used for replay tests
+            self._thread_local.active_functions.remove(func_id)
+            overhead_time = time.thread_time_ns() - end_time
             self.function_calls_data.append(
                 (func.__name__, class_name, func.__module__, func.__code__.co_filename,
                  benchmark_function_name, benchmark_module_path, benchmark_line_number, execution_time,

codeflash/benchmarking/plugin/plugin.py

@@ -175,6 +175,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]:
                 benchmark_key = BenchmarkKey(module_path=benchmark_file, function_name=benchmark_func)
                 # Subtract overhead from total time
                 overhead = overhead_by_benchmark.get(benchmark_key, 0)
+                print("benchmark_func:", benchmark_func, "Total time:", time_ns, "Overhead:", overhead, "Result:", time_ns - overhead)
                 result[benchmark_key] = time_ns - overhead
 
         finally:
@@ -210,61 +211,71 @@ def pytest_plugin_registered(plugin, manager):
             manager.unregister(plugin)
 
     @staticmethod
+    def pytest_configure(config):
+        """Register the benchmark marker."""
+        config.addinivalue_line(
+            "markers",
+            "benchmark: mark test as a benchmark that should be run with codeflash tracing"
+        )
+    @staticmethod
     def pytest_collection_modifyitems(config, items):
         # Skip tests that don't have the benchmark fixture
         if not config.getoption("--codeflash-trace"):
             return
 
         skip_no_benchmark = pytest.mark.skip(reason="Test requires benchmark fixture")
         for item in items:
-            if hasattr(item, "fixturenames") and "benchmark" in item.fixturenames:
-                continue
-            item.add_marker(skip_no_benchmark)
+            # Check for direct benchmark fixture usage
+            has_fixture = hasattr(item, "fixturenames") and "benchmark" in item.fixturenames
+
+            # Check for @pytest.mark.benchmark marker
+            has_marker = False
+            if hasattr(item, "get_closest_marker"):
+                marker = item.get_closest_marker("benchmark")
+                if marker is not None:
+                    has_marker = True
+
+            # Skip if neither fixture nor marker is present
+            if not (has_fixture or has_marker):
+                item.add_marker(skip_no_benchmark)
 
     # Benchmark fixture
     class Benchmark:
         def __init__(self, request):
             self.request = request
 
         def __call__(self, func, *args, **kwargs):
-            """Handle behaviour for the benchmark fixture in pytest.
-
-            For example,
-
-            def test_something(benchmark):
-                benchmark(sorter, [3,2,1])
-
-            Args:
-                func: The function to benchmark (e.g. sorter)
-                args: The arguments to pass to the function (e.g. [3,2,1])
-                kwargs: The keyword arguments to pass to the function
-
-            Returns:
-                The return value of the function
-            a
-
-            """
-            benchmark_module_path = module_name_from_file_path(Path(str(self.request.node.fspath)), Path(codeflash_benchmark_plugin.project_root))
+            """Handle both direct function calls and decorator usage."""
+            if args or kwargs:
+                # Used as benchmark(func, *args, **kwargs)
+                return self._run_benchmark(func, *args, **kwargs)
+            # Used as @benchmark decorator
+            def wrapped_func(*args, **kwargs):
+                return func(*args, **kwargs)
+            result = self._run_benchmark(func)
+            return wrapped_func
+
+        def _run_benchmark(self, func, *args, **kwargs):
+            """Actual benchmark implementation."""
+            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
+            line_number = int(str(sys._getframe(2).f_lineno))  # 2 frames up in the call stack
+            # Set env vars
             os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name
             os.environ["CODEFLASH_BENCHMARK_MODULE_PATH"] = benchmark_module_path
             os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number)
             os.environ["CODEFLASH_BENCHMARKING"] = "True"
-
-        # Run the function
-            start = time.perf_counter_ns()
+            # Run the function
+            start = time.thread_time_ns()
             result = func(*args, **kwargs)
-            end = time.perf_counter_ns()
-
+            end = time.thread_time_ns()
             # Reset the environment variable
             os.environ["CODEFLASH_BENCHMARKING"] = "False"
 
             # Write function calls
             codeflash_trace.write_function_timings()
-            # Reset function call count after a benchmark is run
+            # Reset function call count
             codeflash_trace.function_call_count = 0
             # Add to the benchmark timings buffer
             codeflash_benchmark_plugin.benchmark_timings.append(

codeflash/benchmarking/pytest_new_process_trace_benchmarks.py

@@ -16,7 +16,7 @@
         codeflash_benchmark_plugin.setup(trace_file, project_root)
         codeflash_trace.setup(trace_file)
         exitcode = pytest.main(
-            [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[codeflash_benchmark_plugin]
+            [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark","-p", "no:codspeed","-p", "no:cov-s", "-o", "addopts="], plugins=[codeflash_benchmark_plugin]
         ) # Errors will be printed to stdout, not stderr
 
     except Exception as e:

codeflash/benchmarking/replay_test.py

@@ -34,7 +34,7 @@ def get_next_arg_and_return(
         )
 
     while (val := cursor.fetchone()) is not None:
-        yield val[9], val[10]  # args and kwargs are at indices 7 and 8
+        yield val[9], val[10]  # pickled_args, pickled_kwargs
 
 
 def get_function_alias(module: str, function_name: str) -> str:

tests/test_trace_benchmarks.py

@@ -31,7 +31,7 @@ def test_trace_benchmarks():
         function_calls = cursor.fetchall()
 
         # Assert the length of function calls
-        assert len(function_calls) == 7, f"Expected 6 function calls, but got {len(function_calls)}"
+        assert len(function_calls) == 8, f"Expected 8 function calls, but got {len(function_calls)}"
 
         bubble_sort_path = (project_root / "bubble_sort_codeflash_trace.py").as_posix()
         process_and_bubble_sort_path = (project_root / "process_and_bubble_sort_codeflash_trace.py").as_posix()
@@ -64,6 +64,10 @@ def test_trace_benchmarks():
             ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace",
              f"{bubble_sort_path}",
              "test_no_func", "tests.pytest.benchmarks_test.test_process_and_sort_example", 8),
+
+            ("recursive_bubble_sort", "", "code_to_optimize.bubble_sort_codeflash_trace",
+             f"{bubble_sort_path}",
+             "test_recursive_sort", "tests.pytest.benchmarks_test.test_recursive_example", 5),
         ]
         for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)):
             assert actual[0] == expected[0], f"Mismatch at index {idx} for function_name"
@@ -222,6 +226,62 @@ def test_trace_multithreaded_benchmark() -> None:
         # Close connection
         conn.close()
 
+    finally:
+        # cleanup
+        output_file.unlink(missing_ok=True)
+
+def test_trace_benchmark_decorator() -> None:
+    project_root = Path(__file__).parent.parent / "code_to_optimize"
+    benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_test_decorator"
+    tests_root = project_root / "tests"
+    output_file = (benchmarks_root / Path("test_trace_benchmarks.trace")).resolve()
+    trace_benchmarks_pytest(benchmarks_root, tests_root, project_root, output_file)
+    assert output_file.exists()
+    try:
+        # check contents of trace file
+        # connect to database
+        conn = sqlite3.connect(output_file.as_posix())
+        cursor = conn.cursor()
+
+        # Get the count of records
+        # Get all records
+        cursor.execute(
+            "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_module_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_module_path, benchmark_function_name, function_name")
+        function_calls = cursor.fetchall()
+
+        # Assert the length of function calls
+        assert len(function_calls) == 2, f"Expected 2 function calls, but got {len(function_calls)}"
+        function_benchmark_timings = codeflash_benchmark_plugin.get_function_benchmark_timings(output_file)
+        total_benchmark_timings = codeflash_benchmark_plugin.get_benchmark_timings(output_file)
+        function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings)
+        assert "code_to_optimize.bubble_sort_codeflash_trace.sorter" in function_to_results
+
+        test_name, total_time, function_time, percent = function_to_results["code_to_optimize.bubble_sort_codeflash_trace.sorter"][0]
+        assert total_time > 0.0
+        assert function_time > 0.0
+        assert percent > 0.0
+
+        bubble_sort_path = (project_root / "bubble_sort_codeflash_trace.py").as_posix()
+        # Expected function calls
+        expected_calls = [
+            ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace",
+             f"{bubble_sort_path}",
+             "test_benchmark_sort", "tests.pytest.benchmarks_test_decorator.test_benchmark_decorator", 5),
+            ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace",
+             f"{bubble_sort_path}",
+             "test_pytest_mark", "tests.pytest.benchmarks_test_decorator.test_benchmark_decorator", 11),
+        ]
+        for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)):
+            assert actual[0] == expected[0], f"Mismatch at index {idx} for function_name"
+            assert actual[1] == expected[1], f"Mismatch at index {idx} for class_name"
+            assert actual[2] == expected[2], f"Mismatch at index {idx} for module_name"
+            assert Path(actual[3]).name == Path(expected[3]).name, f"Mismatch at index {idx} for file_path"
+            assert actual[4] == expected[4], f"Mismatch at index {idx} for benchmark_function_name"
+            assert actual[5] == expected[5], f"Mismatch at index {idx} for benchmark_module_path"
+            assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number"
+        # Close connection
+        conn.close()
+
     finally:
         # cleanup
         output_file.unlink(missing_ok=True)